[release/5.0-rc2] Implement SocketsHttpHandler.ConnectCallback (#42075)

* Implement SocketsHttpHandler.ConnectCallback

src/libraries/System.Net.Http/ref/System.Net.Http.cs

@@ -360,7 +360,15 @@ protected override void Dispose(bool disposing) { }
         protected internal override System.Net.Http.HttpResponseMessage Send(System.Net.Http.HttpRequestMessage request, System.Threading.CancellationToken cancellationToken) { throw null; }
         protected internal override System.Threading.Tasks.Task<System.Net.Http.HttpResponseMessage> SendAsync(System.Net.Http.HttpRequestMessage request, System.Threading.CancellationToken cancellationToken) { throw null; }
         public bool EnableMultipleHttp2Connections { get { throw null; } set { } }
+        public Func<SocketsHttpConnectionContext, System.Threading.CancellationToken, System.Threading.Tasks.ValueTask<System.IO.Stream>>? ConnectCallback { get { throw null; } set { } }
     }
+    public sealed class SocketsHttpConnectionContext
+    {
+        internal SocketsHttpConnectionContext() { }
+        public DnsEndPoint DnsEndPoint { get { throw null; } }
+        public HttpRequestMessage RequestMessage { get { throw null; } }
+    }
+
     public enum HttpKeepAlivePingPolicy
     {
         WithActiveRequests,

src/libraries/System.Net.Http/src/Resources/Strings.resx

@@ -585,4 +585,7 @@
   <data name="net_http_requested_version_server_refused" xml:space="preserve">
     <value>Requesting HTTP version {0} with version policy {1} while server offers only version fallback.</value>
   </data>
+  <data name="net_http_sync_operations_not_allowed_with_connect_callback" xml:space="preserve">
+    <value>Synchronous operation is not supported when a ConnectCallback is specified on the SocketsHttpHandler instance.</value>
+  </data>
 </root>

src/libraries/System.Net.Http/src/System.Net.Http.csproj

@@ -173,7 +173,7 @@
     <Compile Include="System\Net\Http\SocketsHttpHandler\MultiProxy.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\RawConnectionStream.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\RedirectHandler.cs" />
-    <Compile Include="System\Net\Http\SocketsHttpHandler\SocketsConnectionFactory.cs" />
+    <Compile Include="System\Net\Http\SocketsHttpHandler\SocketsHttpConnectionContext.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\SocketsHttpHandler.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\SystemProxyInfo.cs" />
     <Compile Include="$(CommonPath)System\Net\NTAuthentication.Common.cs"
@@ -671,6 +671,7 @@
              Link="System\System\Threading\Tasks\TaskToApm.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\HttpKeepAlivePingPolicy.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\HttpNoProxy.cs" />
+    <Compile Include="System\Net\Http\SocketsHttpHandler\SocketsHttpConnectionContext.cs" />
     <Compile Include="System\Net\Http\BrowserHttpHandler\SystemProxyInfo.Browser.cs" />
     <Compile Include="System\Net\Http\BrowserHttpHandler\SocketsHttpHandler.cs" />
     <Compile Include="System\Net\Http\BrowserHttpHandler\BrowserHttpHandler.cs" />

src/libraries/System.Net.Http/src/System/Net/Http/BrowserHttpHandler/SocketsHttpHandler.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Generic;
+using System.IO;
 using System.Net.Security;
 using System.Threading;
 using System.Threading.Tasks;
@@ -170,5 +171,11 @@ public bool EnableMultipleHttp2Connections
             get => throw new PlatformNotSupportedException();
             set => throw new PlatformNotSupportedException();
         }
+
+        public Func<SocketsHttpConnectionContext, CancellationToken, ValueTask<Stream>>? ConnectCallback
+        {
+            get => throw new PlatformNotSupportedException();
+            set => throw new PlatformNotSupportedException();
+        }
     }
 }

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ConnectHelper.cs

@@ -31,11 +31,11 @@ public CertificateCallbackMapper(Func<HttpRequestMessage, X509Certificate2?, X50
             }
         }
 
-        public static async ValueTask<Stream> ConnectAsync(SocketsConnectionFactory factory, DnsEndPoint endPoint, CancellationToken cancellationToken)
+        public static async ValueTask<Stream> ConnectAsync(Func<SocketsHttpConnectionContext, CancellationToken, ValueTask<Stream>> callback, DnsEndPoint endPoint, HttpRequestMessage requestMessage, CancellationToken cancellationToken)
         {
             try
             {
-                return await factory.ConnectAsync(endPoint, cancellationToken).ConfigureAwait(false);
+                return await callback(new SocketsHttpConnectionContext(endPoint, requestMessage), cancellationToken).ConfigureAwait(false);
             }
             catch (OperationCanceledException ex) when (ex.CancellationToken == cancellationToken)
             {

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/HttpConnectionPool.cs

@@ -1286,20 +1286,42 @@ public ValueTask<HttpResponseMessage> SendAsync(HttpRequestMessage request, bool
             }
         }
 
-        private static readonly SocketsConnectionFactory s_defaultConnectionFactory = new SocketsConnectionFactory(SocketType.Stream, ProtocolType.Tcp);
+        private static async ValueTask<Stream> DefaultConnectAsync(SocketsHttpConnectionContext context, CancellationToken cancellationToken)
+        {
+            Socket socket = new Socket(SocketType.Stream, ProtocolType.Tcp);
+            socket.NoDelay = true;
+
+            try
+            {
+                await socket.ConnectAsync(context.DnsEndPoint, cancellationToken).ConfigureAwait(false);
+                return new NetworkStream(socket, ownsSocket: true);
+            }
+            catch
+            {
+                socket.Dispose();
+                throw;
+            }
+        }
+
+        private static readonly Func<SocketsHttpConnectionContext, CancellationToken, ValueTask<Stream>> s_defaultConnectCallback = DefaultConnectAsync;
 
         private ValueTask<Stream> ConnectToTcpHostAsync(string host, int port, HttpRequestMessage initialRequest, bool async, CancellationToken cancellationToken)
         {
             if (async)
             {
-                SocketsConnectionFactory connectionFactory = s_defaultConnectionFactory;
+                Func<SocketsHttpConnectionContext, CancellationToken, ValueTask<Stream>> connectCallback = Settings._connectCallback ?? s_defaultConnectCallback;
 
                 var endPoint = new DnsEndPoint(host, port);
-                return ConnectHelper.ConnectAsync(connectionFactory, endPoint, cancellationToken);
+                return ConnectHelper.ConnectAsync(connectCallback, endPoint, initialRequest, cancellationToken);
             }
 
             // Synchronous path.
 
+            if (Settings._connectCallback is not null)
+            {
+                throw new NotSupportedException(SR.net_http_sync_operations_not_allowed_with_connect_callback);
+            }
+
             try
             {
                 return new ValueTask<Stream>(ConnectHelper.Connect(host, port, cancellationToken));

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/HttpConnectionSettings.cs

@@ -3,6 +3,7 @@
 
 using System.Collections.Generic;
 using System.Net.Security;
+using System.IO;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -55,6 +56,8 @@ internal sealed class HttpConnectionSettings
 
         internal bool _enableMultipleHttp2Connections;
 
+        internal Func<SocketsHttpConnectionContext, CancellationToken, ValueTask<Stream>>? _connectCallback;
+
         internal IDictionary<string, object?>? _properties;
 
         public HttpConnectionSettings()
@@ -108,6 +111,7 @@ public HttpConnectionSettings CloneAndNormalize()
                 _requestHeaderEncodingSelector = _requestHeaderEncodingSelector,
                 _responseHeaderEncodingSelector = _responseHeaderEncodingSelector,
                 _enableMultipleHttp2Connections = _enableMultipleHttp2Connections,
+                _connectCallback = _connectCallback,
             };
         }
 

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/SocketsHttpConnectionContext.cs

@@ -0,0 +1,30 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace System.Net.Http
+{
+    /// <summary>
+    /// Represents the context passed to the ConnectCallback for a SocketsHttpHandler instance.
+    /// </summary>
+    public sealed class SocketsHttpConnectionContext
+    {
+        private readonly DnsEndPoint _dnsEndPoint;
+        private readonly HttpRequestMessage _requestMessage;
+
+        internal SocketsHttpConnectionContext(DnsEndPoint dnsEndPoint, HttpRequestMessage requestMessage)
+        {
+            _dnsEndPoint = dnsEndPoint;
+            _requestMessage = requestMessage;
+        }
+
+        /// <summary>
+        /// The DnsEndPoint to be used by the ConnectCallback to establish the connection.
+        /// </summary>
+        public DnsEndPoint DnsEndPoint => _dnsEndPoint;
+
+        /// <summary>
+        /// The initial HttpRequestMessage that is causing the connection to be created.
+        /// </summary>
+        public HttpRequestMessage RequestMessage => _requestMessage;
+    }
+}

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/SocketsHttpHandler.cs

@@ -3,6 +3,7 @@
 
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.IO;
 using System.Net.Security;
 using System.Threading;
 using System.Threading.Tasks;
@@ -362,6 +363,19 @@ public bool EnableMultipleHttp2Connections
         internal bool SupportsProxy => true;
         internal bool SupportsRedirectConfiguration => true;
 
+        /// <summary>
+        /// When non-null, a custom callback used to open new connections.
+        /// </summary>
+        public Func<SocketsHttpConnectionContext, CancellationToken, ValueTask<Stream>>? ConnectCallback
+        {
+            get => _settings._connectCallback;
+            set
+            {
+                CheckDisposedOrStarted();
+                _settings._connectCallback = value;
+            }
+        }
+
         public IDictionary<string, object?> Properties =>
             _settings._properties ?? (_settings._properties = new Dictionary<string, object?>());