Fix client with custom ConnectCallback (#1542)

Author: JamesNK

src/Grpc.Net.Client/Balancer/Internal/BalancerHttpHandler.cs

@@ -35,15 +35,17 @@ internal class BalancerHttpHandler : DelegatingHandler
 
         private readonly ConnectionManager _manager;
 
-        public BalancerHttpHandler(HttpMessageHandler innerHandler, ConnectionManager manager)
+        public BalancerHttpHandler(HttpMessageHandler innerHandler, HttpHandlerType httpHandlerType, ConnectionManager manager)
             : base(innerHandler)
         {
             _manager = manager;
 
 #if NET5_0_OR_GREATER
-            var socketsHttpHandler = HttpRequestHelpers.GetHttpHandlerType<SocketsHttpHandler>(innerHandler);
-            if (socketsHttpHandler != null)
+            if (httpHandlerType == HttpHandlerType.SocketsHttpHandler)
             {
+                var socketsHttpHandler = HttpRequestHelpers.GetHttpHandlerType<SocketsHttpHandler>(innerHandler);
+                CompatibilityHelpers.Assert(socketsHttpHandler != null, "Should have handler with this handler type.");
+
                 socketsHttpHandler.ConnectCallback = OnConnect;
             }
 #endif

src/Grpc.Net.Client/GrpcChannel.cs

@@ -339,7 +339,7 @@ private HttpMessageInvoker CreateInternalHttpInvoker(HttpMessageHandler? handler
 #endif
 
 #if SUPPORT_LOAD_BALANCING
-            handler = new BalancerHttpHandler(handler, ConnectionManager);
+            handler = new BalancerHttpHandler(handler, HttpHandlerType, ConnectionManager);
 #endif
 
             // Use HttpMessageInvoker instead of HttpClient because it is faster

test/FunctionalTests/Client/ConnectionTests.cs

@@ -16,11 +16,14 @@
 
 #endregion
 
+using System.Net;
+using System.Net.Sockets;
 using Greet;
 using Grpc.AspNetCore.FunctionalTests.Infrastructure;
 using Grpc.Core;
 using Grpc.Net.Client;
 using Grpc.Tests.Shared;
+using Microsoft.AspNetCore.Connections.Features;
 using NUnit.Framework;
 
 namespace Grpc.AspNetCore.FunctionalTests.Client
@@ -60,7 +63,42 @@ public async Task ALPN_ProtocolDowngradedToHttp1_ThrowErrorFromServer()
 #endif
         }
 
-#if NET6_0
+#if NET5_0_OR_GREATER
+        [Test]
+        public async Task UnixDomainSockets()
+        {
+            Task<HelloReply> UnaryUds(HelloRequest request, ServerCallContext context)
+            {
+#if NET6_0_OR_GREATER
+                var endPoint = (UnixDomainSocketEndPoint)context.GetHttpContext().Features.Get<IConnectionSocketFeature>()!.Socket.LocalEndPoint!;
+                Assert.NotNull(endPoint);
+#endif
+
+                return Task.FromResult(new HelloReply { Message = "Hello " + request.Name });
+            }
+
+            // Arrange
+            var method = Fixture.DynamicGrpc.AddUnaryMethod<HelloRequest, HelloReply>(UnaryUds);
+
+            var http = Fixture.CreateHandler(TestServerEndpointName.UnixDomainSocket);
+
+            var channel = GrpcChannel.ForAddress(http.address, new GrpcChannelOptions
+            {
+                LoggerFactory = LoggerFactory,
+                HttpHandler = http.handler
+            });
+
+            var client = TestClientFactory.Create(channel, method);
+
+            // Act
+            var response = await client.UnaryCall(new HelloRequest { Name = "John" }).ResponseAsync.DefaultTimeout();
+
+            // Assert
+            Assert.AreEqual("Hello John", response.Message);
+        }
+#endif
+
+#if NET6_0_OR_GREATER
         [Test]
         [RequireHttp3]
         public async Task Http3()
@@ -79,6 +117,7 @@ public async Task Http3()
             // Act
             var response = await client.SayHelloAsync(new HelloRequest { Name = "John" }).ResponseAsync.DefaultTimeout();
 
+            // Assert
             Assert.AreEqual("Hello John", response.Message);
         }
 #endif

test/FunctionalTests/Infrastructure/GrpcTestFixture.cs

@@ -16,6 +16,7 @@
 
 #endregion
 
+using System.Net.Sockets;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Server.Kestrel.Core;
 using Microsoft.Extensions.DependencyInjection;
@@ -25,6 +26,7 @@ namespace Grpc.AspNetCore.FunctionalTests.Infrastructure
 {
     public class GrpcTestFixture<TStartup> : IDisposable where TStartup : class
     {
+        private readonly string _socketPath = Path.Combine(Path.GetTempPath(), "grpc-transporter.tmp");
         private readonly InProcessTestServer _server;
 
         public GrpcTestFixture(
@@ -86,6 +88,19 @@ public GrpcTestFixture(
                         listenOptions.UseHttps(certPath, "1111");
                     });
 
+#if NET5_0_OR_GREATER
+                    if (File.Exists(_socketPath))
+                    {
+                        File.Delete(_socketPath);
+                    }
+
+                    urls[TestServerEndpointName.UnixDomainSocket] = _socketPath;
+                    options.ListenUnixSocket(_socketPath, listenOptions =>
+                    {
+                        listenOptions.Protocols = HttpProtocols.Http2;
+                    });
+#endif
+
 #if NET6_0_OR_GREATER
                     if (RequireHttp3Attribute.IsSupported(out _))
                     {
@@ -145,6 +160,16 @@ public HttpClient CreateClient(TestServerEndpointName? endpointName = null, Dele
                 RemoteCertificateValidationCallback = (_, __, ___, ____) => true
             };
 
+#if NET5_0_OR_GREATER
+            if (endpointName == TestServerEndpointName.UnixDomainSocket)
+            {
+                var udsEndPoint = new UnixDomainSocketEndPoint(_server.GetUrl(endpointName.Value));
+                var connectionFactory = new UnixDomainSocketConnectionFactory(udsEndPoint);
+
+                socketsHttpHandler.ConnectCallback = connectionFactory.ConnectAsync;
+            }
+#endif
+
             HttpClient client;
             HttpMessageHandler handler;
             if (messageHandler != null)
@@ -175,11 +200,24 @@ public HttpClient CreateClient(TestServerEndpointName? endpointName = null, Dele
                 client.DefaultVersionPolicy = HttpVersionPolicy.RequestVersionOrHigher;
 #endif
             }
-            client.BaseAddress = new Uri(_server.GetUrl(endpointName.Value));
+
+            client.BaseAddress = CalculateBaseAddress(endpointName.Value);
 
             return (client, handler);
         }
 
+        private Uri CalculateBaseAddress(TestServerEndpointName endpointName)
+        {
+#if NET5_0_OR_GREATER
+            if (endpointName == TestServerEndpointName.UnixDomainSocket)
+            {
+                return new Uri("http://localhost");
+            }
+#endif
+
+            return new Uri(_server.GetUrl(endpointName));
+        }
+
         public Uri GetUrl(TestServerEndpointName endpointName)
         {
             switch (endpointName)
@@ -192,6 +230,10 @@ public Uri GetUrl(TestServerEndpointName endpointName)
                 case TestServerEndpointName.Http3WithTls:
 #endif
                     return new Uri(_server.GetUrl(endpointName));
+#if NET5_0_OR_GREATER
+                case TestServerEndpointName.UnixDomainSocket:
+                    return new Uri("http://localhost");
+#endif
                 default:
                     throw new ArgumentException("Unexpected value: " + endpointName, nameof(endpointName));
             }

test/FunctionalTests/Infrastructure/TestServerEndpointName.cs

@@ -24,6 +24,9 @@ public enum TestServerEndpointName
         Http1,
         Http2WithTls,
         Http1WithTls,
+#if NET5_0_OR_GREATER
+        UnixDomainSocket,
+#endif
 #if NET6_0_OR_GREATER
         Http3WithTls,
 #endif

test/FunctionalTests/Infrastructure/UnixDomainSocketConnectionFactory.cs

@@ -0,0 +1,56 @@
+#region Copyright notice and license
+
+// Copyright 2019 The gRPC Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#endregion
+
+#if NET5_0_OR_GREATER
+using System;
+using System.IO;
+using System.Net;
+using System.Net.Http;
+using System.Net.Sockets;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Grpc.AspNetCore.FunctionalTests.Infrastructure
+{
+    public class UnixDomainSocketConnectionFactory
+    {
+        private readonly EndPoint _endPoint;
+
+        public UnixDomainSocketConnectionFactory(EndPoint endPoint)
+        {
+            _endPoint = endPoint;
+        }
+
+        public async ValueTask<Stream> ConnectAsync(SocketsHttpConnectionContext _, CancellationToken cancellationToken = default)
+        {
+            var socket = new Socket(AddressFamily.Unix, SocketType.Stream, ProtocolType.Unspecified);
+
+            try
+            {
+                await socket.ConnectAsync(_endPoint, cancellationToken).ConfigureAwait(false);
+                return new NetworkStream(socket, true);
+            }
+            catch (Exception ex)
+            {
+                socket.Dispose();
+                throw new HttpRequestException($"Error connecting to '{_endPoint}'.", ex);
+            }
+        }
+    }
+}
+#endif