Implement IConnectionEndPointFeature in Kestrel's TransportConnection

src/Servers/Kestrel/Core/src/Internal/Infrastructure/KestrelMetrics.cs

@@ -4,8 +4,6 @@
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Diagnostics.Metrics;
-using System.Net;
-using System.Net.Sockets;
 using System.Runtime.CompilerServices;
 using System.Security.Authentication;
 using Microsoft.AspNetCore.Connections;
@@ -321,43 +319,7 @@ private void TlsHandshakeStopCore(ConnectionMetricsContext metricsContext, long
 
     private static void InitializeConnectionTags(ref TagList tags, in ConnectionMetricsContext metricsContext)
     {
-        var localEndpoint = metricsContext.ConnectionContext.LocalEndPoint;
-        if (localEndpoint is IPEndPoint localIPEndPoint)
-        {
-            tags.Add("server.address", localIPEndPoint.Address.ToString());
-            tags.Add("server.port", localIPEndPoint.Port);
-
-            switch (localIPEndPoint.Address.AddressFamily)
-            {
-                case AddressFamily.InterNetwork:
-                    tags.Add("network.type", "ipv4");
-                    break;
-                case AddressFamily.InterNetworkV6:
-                    tags.Add("network.type", "ipv6");
-                    break;
-            }
-
-            // There isn't an easy way to detect whether QUIC is the underlying transport.
-            // This code assumes that a multiplexed connection is QUIC.
-            // Improve in the future if there are additional multiplexed connection types.
-            var transport = metricsContext.ConnectionContext is not MultiplexedConnectionContext ? "tcp" : "udp";
-            tags.Add("network.transport", transport);
-        }
-        else if (localEndpoint is UnixDomainSocketEndPoint udsEndPoint)
-        {
-            tags.Add("server.address", udsEndPoint.ToString());
-            tags.Add("network.transport", "unix");
-        }
-        else if (localEndpoint is NamedPipeEndPoint namedPipeEndPoint)
-        {
-            tags.Add("server.address", namedPipeEndPoint.ToString());
-            tags.Add("network.transport", "pipe");
-        }
-        else if (localEndpoint != null)
-        {
-            tags.Add("server.address", localEndpoint.ToString());
-            tags.Add("network.transport", localEndpoint.AddressFamily.ToString());
-        }
+        ConnectionEndpointTags.AddConnectionEndpointTags(ref tags, metricsContext.ConnectionContext);
     }
 
     public ConnectionMetricsContext CreateContext(BaseConnectionContext connection)

src/Servers/Kestrel/Core/src/Microsoft.AspNetCore.Server.Kestrel.Core.csproj

@@ -37,6 +37,7 @@
     <Compile Include="$(SharedSourceRoot)Obsoletions.cs" LinkBase="Shared" />
     <Compile Include="$(RepoRoot)src\Shared\TaskToApm.cs" Link="Internal\TaskToApm.cs" />
     <Compile Include="$(SharedSourceRoot)Metrics\MetricsExtensions.cs" />
+    <Compile Include="$(SharedSourceRoot)ConnectionEndpointTags.cs" />
   </ItemGroup>
 
   <ItemGroup>

src/Servers/Kestrel/Core/test/TransportConnectionFeatureCollectionTests.cs

@@ -0,0 +1,103 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Net;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Connections.Features;
+using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
+using Microsoft.AspNetCore.InternalTesting;
+using Moq;
+using Xunit;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests;
+
+public class TransportConnectionFeatureCollectionTests
+{
+    [Fact]
+    public void IConnectionEndPointFeature_IsAvailableInFeatureCollection()
+    {
+        var serviceContext = new TestServiceContext();
+        var connection = new Mock<DefaultConnectionContext> { CallBase = true }.Object;
+        var transportConnectionManager = new TransportConnectionManager(serviceContext.ConnectionManager);
+        var kestrelConnection = CreateKestrelConnection(serviceContext, connection, transportConnectionManager);
+
+        var endpointFeature = kestrelConnection.TransportConnection.Features.Get<IConnectionEndPointFeature>();
+
+        Assert.NotNull(endpointFeature);
+    }
+
+    [Fact]
+    public void IConnectionEndPointFeature_ReturnsCorrectLocalEndPoint()
+    {
+        var serviceContext = new TestServiceContext();
+        var connection = new Mock<DefaultConnectionContext> { CallBase = true }.Object;
+        var expectedLocalEndPoint = new IPEndPoint(IPAddress.Parse("127.0.0.1"), 8080);
+        connection.LocalEndPoint = expectedLocalEndPoint;
+        var transportConnectionManager = new TransportConnectionManager(serviceContext.ConnectionManager);
+        var kestrelConnection = CreateKestrelConnection(serviceContext, connection, transportConnectionManager);
+
+        var endpointFeature = kestrelConnection.TransportConnection.Features.Get<IConnectionEndPointFeature>();
+
+        Assert.NotNull(endpointFeature);
+        Assert.Equal(expectedLocalEndPoint, endpointFeature.LocalEndPoint);
+    }
+
+    [Fact]
+    public void IConnectionEndPointFeature_ReturnsCorrectRemoteEndPoint()
+    {
+        var serviceContext = new TestServiceContext();
+        var connection = new Mock<DefaultConnectionContext> { CallBase = true }.Object;
+        var expectedRemoteEndPoint = new IPEndPoint(IPAddress.Parse("192.168.1.100"), 54321);
+        connection.RemoteEndPoint = expectedRemoteEndPoint;
+        var transportConnectionManager = new TransportConnectionManager(serviceContext.ConnectionManager);
+        var kestrelConnection = CreateKestrelConnection(serviceContext, connection, transportConnectionManager);
+
+        var endpointFeature = kestrelConnection.TransportConnection.Features.Get<IConnectionEndPointFeature>();
+
+        Assert.NotNull(endpointFeature);
+        Assert.Equal(expectedRemoteEndPoint, endpointFeature.RemoteEndPoint);
+    }
+
+    [Fact]
+    public void IConnectionEndPointFeature_AllowsSettingLocalEndPoint()
+    {
+        var serviceContext = new TestServiceContext();
+        var connection = new Mock<DefaultConnectionContext> { CallBase = true }.Object;
+        var transportConnectionManager = new TransportConnectionManager(serviceContext.ConnectionManager);
+        var kestrelConnection = CreateKestrelConnection(serviceContext, connection, transportConnectionManager);
+        var newLocalEndPoint = new IPEndPoint(IPAddress.Parse("127.0.0.1"), 9090);
+
+        var endpointFeature = kestrelConnection.TransportConnection.Features.Get<IConnectionEndPointFeature>();
+        endpointFeature.LocalEndPoint = newLocalEndPoint;
+
+        Assert.Equal(newLocalEndPoint, kestrelConnection.TransportConnection.LocalEndPoint);
+        Assert.Equal(newLocalEndPoint, endpointFeature.LocalEndPoint);
+    }
+
+    [Fact]
+    public void IConnectionEndPointFeature_AllowsSettingRemoteEndPoint()
+    {
+        var serviceContext = new TestServiceContext();
+        var connection = new Mock<DefaultConnectionContext> { CallBase = true }.Object;
+        var transportConnectionManager = new TransportConnectionManager(serviceContext.ConnectionManager);
+        var kestrelConnection = CreateKestrelConnection(serviceContext, connection, transportConnectionManager);
+        var newRemoteEndPoint = new IPEndPoint(IPAddress.Parse("10.0.0.1"), 12345);
+
+        var endpointFeature = kestrelConnection.TransportConnection.Features.Get<IConnectionEndPointFeature>();
+        endpointFeature.RemoteEndPoint = newRemoteEndPoint;
+
+        Assert.Equal(newRemoteEndPoint, kestrelConnection.TransportConnection.RemoteEndPoint);
+        Assert.Equal(newRemoteEndPoint, endpointFeature.RemoteEndPoint);
+    }
+
+    private static KestrelConnection<ConnectionContext> CreateKestrelConnection(TestServiceContext serviceContext, DefaultConnectionContext connection, TransportConnectionManager transportConnectionManager, Func<ConnectionContext, Task> connectionDelegate = null)
+    {
+        connectionDelegate ??= _ => Task.CompletedTask;
+
+        return new KestrelConnection<ConnectionContext>(
+            id: 0, serviceContext, transportConnectionManager, connectionDelegate, connection, serviceContext.Log, TestContextFactory.CreateMetricsContext(connection));
+    }
+}

src/Servers/Kestrel/shared/TransportConnection.FeatureCollection.cs

@@ -3,6 +3,7 @@
 
 using System.Buffers;
 using System.IO.Pipelines;
+using System.Net;
 using Microsoft.AspNetCore.Connections.Features;
 
 #nullable enable
@@ -37,4 +38,16 @@ CancellationToken IConnectionLifetimeFeature.ConnectionClosed
     }
 
     void IConnectionLifetimeFeature.Abort() => Abort(new ConnectionAbortedException("The connection was aborted by the application via IConnectionLifetimeFeature.Abort()."));
+
+    EndPoint? IConnectionEndPointFeature.LocalEndPoint
+    {
+        get => LocalEndPoint;
+        set => LocalEndPoint = value;
+    }
+
+    EndPoint? IConnectionEndPointFeature.RemoteEndPoint
+    {
+        get => RemoteEndPoint;
+        set => RemoteEndPoint = value;
+    }
 }

src/Servers/Kestrel/shared/TransportConnection.Generated.cs

@@ -18,14 +18,16 @@ internal partial class TransportConnection : IFeatureCollection,
                                                  IConnectionTransportFeature,
                                                  IConnectionItemsFeature,
                                                  IMemoryPoolFeature,
-                                                 IConnectionLifetimeFeature
+                                                 IConnectionLifetimeFeature,
+                                                 IConnectionEndPointFeature
     {
         // Implemented features
         internal protected IConnectionIdFeature? _currentIConnectionIdFeature;
         internal protected IConnectionTransportFeature? _currentIConnectionTransportFeature;
         internal protected IConnectionItemsFeature? _currentIConnectionItemsFeature;
         internal protected IMemoryPoolFeature? _currentIMemoryPoolFeature;
         internal protected IConnectionLifetimeFeature? _currentIConnectionLifetimeFeature;
+        internal protected IConnectionEndPointFeature? _currentIConnectionEndPointFeature;
 
         // Other reserved feature slots
         internal protected IPersistentStateFeature? _currentIPersistentStateFeature;
@@ -48,6 +50,7 @@ private void FastReset()
             _currentIConnectionItemsFeature = this;
             _currentIMemoryPoolFeature = this;
             _currentIConnectionLifetimeFeature = this;
+            _currentIConnectionEndPointFeature = this;
 
             _currentIPersistentStateFeature = null;
             _currentIConnectionSocketFeature = null;
@@ -180,6 +183,10 @@ private void ExtraFeatureSet(Type key, object? value)
                 {
                     feature = _currentIConnectionMetricsTagsFeature;
                 }
+                else if (key == typeof(IConnectionEndPointFeature))
+                {
+                    feature = _currentIConnectionEndPointFeature;
+                }
                 else if (MaybeExtra != null)
                 {
                     feature = ExtraFeatureGet(key);
@@ -244,6 +251,10 @@ private void ExtraFeatureSet(Type key, object? value)
                 {
                     _currentIConnectionMetricsTagsFeature = (IConnectionMetricsTagsFeature?)value;
                 }
+                else if (key == typeof(IConnectionEndPointFeature))
+                {
+                    _currentIConnectionEndPointFeature = (IConnectionEndPointFeature?)value;
+                }
                 else
                 {
                     ExtraFeatureSet(key, value);
@@ -310,6 +321,10 @@ private void ExtraFeatureSet(Type key, object? value)
             {
                 feature = Unsafe.As<IConnectionMetricsTagsFeature?, TFeature?>(ref _currentIConnectionMetricsTagsFeature);
             }
+            else if (typeof(TFeature) == typeof(IConnectionEndPointFeature))
+            {
+                feature = Unsafe.As<IConnectionEndPointFeature?, TFeature?>(ref _currentIConnectionEndPointFeature);
+            }
             else if (MaybeExtra != null)
             {
                 feature = (TFeature?)(ExtraFeatureGet(typeof(TFeature)));
@@ -382,6 +397,10 @@ private void ExtraFeatureSet(Type key, object? value)
             {
                 _currentIConnectionMetricsTagsFeature = Unsafe.As<TFeature?, IConnectionMetricsTagsFeature?>(ref feature);
             }
+            else if (typeof(TFeature) == typeof(IConnectionEndPointFeature))
+            {
+                _currentIConnectionEndPointFeature = Unsafe.As<TFeature?, IConnectionEndPointFeature?>(ref feature);
+            }
             else
             {
                 ExtraFeatureSet(typeof(TFeature), feature);
@@ -442,6 +461,10 @@ private IEnumerable<KeyValuePair<Type, object>> FastEnumerable()
             {
                 yield return new KeyValuePair<Type, object>(typeof(IConnectionMetricsTagsFeature), _currentIConnectionMetricsTagsFeature);
             }
+            if (_currentIConnectionEndPointFeature != null)
+            {
+                yield return new KeyValuePair<Type, object>(typeof(IConnectionEndPointFeature), _currentIConnectionEndPointFeature);
+            }
 
             if (MaybeExtra != null)
             {

src/Servers/Kestrel/tools/CodeGenerator/TransportConnectionFeatureCollection.cs

@@ -24,7 +24,8 @@ public static string GenerateFile()
             "IStreamIdFeature",
             "IStreamAbortFeature",
             "IStreamClosedFeature",
-            "IConnectionMetricsTagsFeature"
+            "IConnectionMetricsTagsFeature",
+            "IConnectionEndPointFeature"
         };
 
         var implementedFeatures = new[]
@@ -33,7 +34,8 @@ public static string GenerateFile()
             "IConnectionTransportFeature",
             "IConnectionItemsFeature",
             "IMemoryPoolFeature",
-            "IConnectionLifetimeFeature"
+            "IConnectionLifetimeFeature",
+            "IConnectionEndPointFeature"
         };
 
         var usings = $@"