Added IConnectionSocketFeature (#31588)

- Removed the nullable from the Socket property as approved in API review
- Make the IConnectionSocketFeature an additional feature instead of one implemented on the TransportConnection.
- Added a test to verify it works

Author: davidfowl

src/Servers/Connections.Abstractions/src/Features/IConnectionSocketFeature.cs

@@ -0,0 +1,18 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Net.Sockets;
+
+namespace Microsoft.AspNetCore.Connections.Features
+{
+    /// <summary>
+    /// Provides access to the connection's underlying <see cref="Socket"/>.
+    /// </summary>
+    public interface IConnectionSocketFeature
+    {
+        /// <summary>
+        /// Gets the underlying <see cref="Socket"/>.
+        /// </summary>
+        Socket Socket { get; }
+    }
+}

src/Servers/Connections.Abstractions/src/PublicAPI.Unshipped.txt

@@ -1,5 +1,7 @@
 #nullable enable
 *REMOVED*Microsoft.AspNetCore.Connections.IConnectionListener.AcceptAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask<Microsoft.AspNetCore.Connections.ConnectionContext!>
+Microsoft.AspNetCore.Connections.Features.IConnectionSocketFeature
+Microsoft.AspNetCore.Connections.Features.IConnectionSocketFeature.Socket.get -> System.Net.Sockets.Socket!
 Microsoft.AspNetCore.Connections.IConnectionListener.AcceptAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.ValueTask<Microsoft.AspNetCore.Connections.ConnectionContext?>
 Microsoft.AspNetCore.Connections.Experimental.IMultiplexedConnectionBuilder
 Microsoft.AspNetCore.Connections.Experimental.IMultiplexedConnectionBuilder.ApplicationServices.get -> System.IServiceProvider!

src/Servers/Kestrel/Transport.Sockets/src/Internal/SocketConnection.FeatureCollection.cs

@@ -0,0 +1,18 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Net.Sockets;
+using Microsoft.AspNetCore.Connections.Features;
+
+namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
+{
+    internal sealed partial class SocketConnection : IConnectionSocketFeature
+    {
+        public Socket Socket => _socket;
+
+        private void InitiaizeFeatures()
+        {
+            _currentIConnectionSocketFeature = this;
+        }
+    }
+}

src/Servers/Kestrel/Transport.Sockets/src/Internal/SocketConnection.cs

@@ -13,7 +13,7 @@
 
 namespace Microsoft.AspNetCore.Server.Kestrel.Transport.Sockets.Internal
 {
-    internal sealed class SocketConnection : TransportConnection
+    internal sealed partial class SocketConnection : TransportConnection
     {
         private static readonly int MinAllocBufferSize = SlabMemoryPool.BlockSize / 2;
 
@@ -70,6 +70,8 @@ internal SocketConnection(Socket socket,
             // Set the transport and connection id
             Transport = _originalTransport = pair.Transport;
             Application = pair.Application;
+
+            InitiaizeFeatures();
         }
 
         public PipeWriter Input => Application.Output;

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

@@ -27,6 +27,9 @@ internal partial class TransportConnection : IFeatureCollection,
         internal protected IMemoryPoolFeature? _currentIMemoryPoolFeature;
         internal protected IConnectionLifetimeFeature? _currentIConnectionLifetimeFeature;
 
+        // Other reserved feature slots
+        internal protected IConnectionSocketFeature? _currentIConnectionSocketFeature;
+
         private int _featureRevision;
 
         private List<KeyValuePair<Type, object>>? MaybeExtra;
@@ -39,6 +42,7 @@ private void FastReset()
             _currentIMemoryPoolFeature = this;
             _currentIConnectionLifetimeFeature = this;
 
+            _currentIConnectionSocketFeature = null;
         }
 
         // Internal for testing
@@ -130,6 +134,10 @@ private void ExtraFeatureSet(Type key, object? value)
                 {
                     feature = _currentIConnectionLifetimeFeature;
                 }
+                else if (key == typeof(IConnectionSocketFeature))
+                {
+                    feature = _currentIConnectionSocketFeature;
+                }
                 else if (MaybeExtra != null)
                 {
                     feature = ExtraFeatureGet(key);
@@ -162,6 +170,10 @@ private void ExtraFeatureSet(Type key, object? value)
                 {
                     _currentIConnectionLifetimeFeature = (IConnectionLifetimeFeature?)value;
                 }
+                else if (key == typeof(IConnectionSocketFeature))
+                {
+                    _currentIConnectionSocketFeature = (IConnectionSocketFeature?)value;
+                }
                 else
                 {
                     ExtraFeatureSet(key, value);
@@ -196,6 +208,10 @@ private void ExtraFeatureSet(Type key, object? value)
             {
                 feature = Unsafe.As<IConnectionLifetimeFeature?, TFeature?>(ref _currentIConnectionLifetimeFeature);
             }
+            else if (typeof(TFeature) == typeof(IConnectionSocketFeature))
+            {
+                feature = Unsafe.As<IConnectionSocketFeature?, TFeature?>(ref _currentIConnectionSocketFeature);
+            }
             else if (MaybeExtra != null)
             {
                 feature = (TFeature?)(ExtraFeatureGet(typeof(TFeature)));
@@ -231,6 +247,10 @@ private void ExtraFeatureSet(Type key, object? value)
             {
                 _currentIConnectionLifetimeFeature = Unsafe.As<TFeature?, IConnectionLifetimeFeature?>(ref feature);
             }
+            else if (typeof(TFeature) == typeof(IConnectionSocketFeature))
+            {
+                _currentIConnectionSocketFeature = Unsafe.As<TFeature?, IConnectionSocketFeature?>(ref feature);
+            }
             else
             {
                 ExtraFeatureSet(typeof(TFeature), feature);
@@ -259,6 +279,10 @@ private IEnumerable<KeyValuePair<Type, object>> FastEnumerable()
             {
                 yield return new KeyValuePair<Type, object>(typeof(IConnectionLifetimeFeature), _currentIConnectionLifetimeFeature);
             }
+            if (_currentIConnectionSocketFeature != null)
+            {
+                yield return new KeyValuePair<Type, object>(typeof(IConnectionSocketFeature), _currentIConnectionSocketFeature);
+            }
 
             if (MaybeExtra != null)
             {

src/Servers/Kestrel/test/Sockets.FunctionalTests/SocketTranspotTests.cs

@@ -0,0 +1,54 @@
+using System.Net;
+using System.Net.Http;
+using System.Net.Sockets;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Builder;
+using Microsoft.AspNetCore.Connections.Features;
+using Microsoft.AspNetCore.Hosting;
+using Microsoft.AspNetCore.Server.Kestrel.FunctionalTests;
+using Microsoft.AspNetCore.Testing;
+using Microsoft.Extensions.Hosting;
+using Xunit;
+
+namespace Sockets.FunctionalTests
+{
+    public class SocketTranspotTests : LoggedTestBase
+    {
+        [Fact]
+        public async Task SocketTransportExposesSocketsFeature()
+        {
+            var builder = TransportSelector.GetHostBuilder()
+                .ConfigureWebHost(webHostBuilder =>
+                {
+                    webHostBuilder
+                        .UseKestrel()
+                        .UseUrls("http://127.0.0.1:0")
+                        .Configure(app =>
+                        {
+                            app.Run(context =>
+                            {
+                                var socket = context.Features.Get<IConnectionSocketFeature>().Socket;
+                                Assert.NotNull(socket);
+                                Assert.Equal(ProtocolType.Tcp, socket.ProtocolType);
+                                var ip = (IPEndPoint)socket.RemoteEndPoint;
+                                Assert.Equal(ip.Address, context.Connection.RemoteIpAddress);
+                                Assert.Equal(ip.Port, context.Connection.RemotePort);
+
+                                return Task.CompletedTask;
+                            });
+                        });
+                })
+                .ConfigureServices(AddTestLogging);
+
+            using var host = builder.Build();
+            using var client = new HttpClient();
+
+            await host.StartAsync();
+
+            var response = await client.GetAsync($"http://127.0.0.1:{host.GetPort()}/");
+            response.EnsureSuccessStatusCode();
+
+            await host.StopAsync();
+        }
+    }
+}

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

@@ -11,7 +11,18 @@ public static string GenerateFile()
         {
             // NOTE: This list MUST always match the set of feature interfaces implemented by TransportConnection.
             // See also: shared/TransportConnection.FeatureCollection.cs
-            var features = new[]
+
+            var allFeatures = new[]
+            {
+                "IConnectionIdFeature",
+                "IConnectionTransportFeature",
+                "IConnectionItemsFeature",
+                "IMemoryPoolFeature",
+                "IConnectionLifetimeFeature",
+                "IConnectionSocketFeature"
+            };
+            
+            var implementedFeatures = new[]
             {
                 "IConnectionIdFeature",
                 "IConnectionTransportFeature",
@@ -27,8 +38,8 @@ public static string GenerateFile()
             return FeatureCollectionGenerator.GenerateFile(
                 namespaceName: "Microsoft.AspNetCore.Connections",
                 className: "TransportConnection",
-                allFeatures: features,
-                implementedFeatures: features,
+                allFeatures: allFeatures,
+                implementedFeatures: implementedFeatures,
                 extraUsings: usings,
                 fallbackFeatures: null);
         }