Handle negotiate changing the connection URL

JamesNK · JamesNK · commit a698667a9478 · 2024-08-06T12:02:30.000+08:00
diff --git a/src/SignalR/clients/csharp/Client.Core/src/HubConnection.cs b/src/SignalR/clients/csharp/Client.Core/src/HubConnection.cs
@@ -729,8 +729,7 @@ async Task OnStreamCanceled(InvocationRequest irq)
 
         CheckDisposed();
 
-        var activity = StartActivity(methodName);
-        var connectionState = await WaitForActiveConnectionWithActivityAsync(nameof(StreamAsChannelCoreAsync), activity, token: cancellationToken).ConfigureAwait(false);
+        var (connectionState, activity) = await WaitForActiveConnectionWithActivityAsync(nameof(StreamAsChannelCoreAsync), methodName, token: cancellationToken).ConfigureAwait(false);
 
         ChannelReader<object?> channel;
         try
@@ -1013,11 +1012,40 @@ private async Task CommonStreaming(ConnectionState connectionState, string strea
         }
     }
 
-    private async Task<ConnectionState> WaitForActiveConnectionWithActivityAsync(string methodName, Activity? activity, CancellationToken token)
+    private async Task<(ConnectionState, Activity?)> WaitForActiveConnectionWithActivityAsync(string sendingMethodName, string invokedMethodName, CancellationToken token)
     {
+        // Start the activity before waiting on the connection.
+        // Starting the activity here means time to connect or reconnect is included in the invoke.
+        var activity = CreateActivity(invokedMethodName);
+
         try
         {
-            return await _state.WaitForActiveConnectionAsync(methodName, token).ConfigureAwait(false);
+            ConnectionState connectionState;
+            var connectionStateTask = _state.WaitForActiveConnectionAsync(sendingMethodName, token);
+            if (connectionStateTask.Status == TaskStatus.RanToCompletion)
+            {
+                // Attempt to get already connected connection and set server tags using it.
+                connectionState = connectionStateTask.Result;
+                SetServerTags(activity, connectionState.ConnectionUrl);
+                activity?.Start();
+            }
+            else
+            {
+                // Fallback to using configured endpoint.
+                var initialUri = (_endPoint as UriEndPoint)?.Uri;
+                SetServerTags(activity, initialUri);
+                activity?.Start();
+
+                connectionState = await connectionStateTask.ConfigureAwait(false);
+
+                // After connection is returned, check if URL is different. If so, update activity server tags.
+                if (connectionState.ConnectionUrl != null && connectionState.ConnectionUrl != initialUri)
+                {
+                    SetServerTags(activity, connectionState.ConnectionUrl);
+                }
+            }
+
+            return (connectionState, activity);
         }
         catch (Exception ex)
         {
@@ -1031,6 +1059,15 @@ private async Task<ConnectionState> WaitForActiveConnectionWithActivityAsync(str
 
             throw;
         }
+
+        static void SetServerTags(Activity? activity, Uri? uri)
+        {
+            if (activity != null && uri != null)
+            {
+                activity.SetTag("server.address", uri.Host);
+                activity.SetTag("server.port", uri.Port);
+            }
+        }
     }
 
     private async Task<object?> InvokeCoreAsyncCore(string methodName, Type returnType, object?[] args, CancellationToken cancellationToken)
@@ -1039,8 +1076,7 @@ private async Task<ConnectionState> WaitForActiveConnectionWithActivityAsync(str
 
         CheckDisposed();
 
-        var activity = StartActivity(methodName);
-        var connectionState = await WaitForActiveConnectionWithActivityAsync(nameof(InvokeCoreAsync), activity, token: cancellationToken).ConfigureAwait(false);
+        var (connectionState, activity) = await WaitForActiveConnectionWithActivityAsync(nameof(InvokeCoreAsync), methodName, token: cancellationToken).ConfigureAwait(false);
 
         Task<object?> invocationTask;
         try
@@ -1063,7 +1099,7 @@ private async Task<ConnectionState> WaitForActiveConnectionWithActivityAsync(str
         return await invocationTask.ConfigureAwait(false);
     }
 
-    private Activity? StartActivity(string methodName)
+    private Activity? CreateActivity(string methodName)
     {
         var activity = _activitySource.CreateActivity(ActivityName, ActivityKind.Client);
         if (activity is null && Activity.Current is not null && _logger.IsEnabled(LogLevel.Critical))
@@ -1085,14 +1121,6 @@ private async Task<ConnectionState> WaitForActiveConnectionWithActivityAsync(str
 
             activity.SetTag("rpc.system", "signalr");
             activity.SetTag("rpc.method", methodName);
-
-            if (_endPoint is UriEndPoint e)
-            {
-                activity.SetTag("server.address", e.Uri.Host);
-                activity.SetTag("server.port", e.Uri.Port);
-            }
-
-            activity.Start();
         }
 
         return activity;
@@ -1198,8 +1226,7 @@ private async Task SendCoreAsyncCore(string methodName, object?[] args, Cancella
 
         CheckDisposed();
 
-        var activity = StartActivity(methodName);
-        var connectionState = await WaitForActiveConnectionWithActivityAsync(nameof(SendCoreAsync), activity, token: cancellationToken).ConfigureAwait(false);
+        var (connectionState, activity) = await WaitForActiveConnectionWithActivityAsync(nameof(SendCoreAsync), methodName, token: cancellationToken).ConfigureAwait(false);
         try
         {
             CheckDisposed();
@@ -2101,6 +2128,7 @@ private sealed class ConnectionState : IInvocationBinder
         private long _nextActivationSendPing;
 
         public ConnectionContext Connection { get; }
+        public Uri? ConnectionUrl { get; }
         public Task? ReceiveTask { get; set; }
         public Exception? CloseException { get; set; }
         public CancellationToken UploadStreamToken { get; set; }
@@ -2119,6 +2147,7 @@ public bool Stopping
         public ConnectionState(ConnectionContext connection, HubConnection hubConnection)
         {
             Connection = connection;
+            ConnectionUrl = (connection.RemoteEndPoint is UriEndPoint ep) ? ep.Uri : null;
 
             _hubConnection = hubConnection;
             _hubConnection._logScope.ConnectionId = connection.ConnectionId;
diff --git a/src/SignalR/clients/csharp/Client/test/UnitTests/HttpConnectionTests.Negotiate.cs b/src/SignalR/clients/csharp/Client/test/UnitTests/HttpConnectionTests.Negotiate.cs
@@ -456,13 +456,15 @@ public async Task NegotiateThatReturnsRedirectUrlDoesNotAddAnotherNegotiateVersi
 
             testHttpHandler.OnLongPollDelete((token) => ResponseUtils.CreateResponse(HttpStatusCode.Accepted));
 
+            EndPoint connectedEndpoint = null;
             using (var noErrorScope = new VerifyNoErrorsScope())
             {
                 await WithConnectionAsync(
                     CreateConnection(testHttpHandler, loggerFactory: noErrorScope.LoggerFactory),
                     async (connection) =>
                     {
                         await connection.StartAsync().DefaultTimeout();
+                        connectedEndpoint = connection.RemoteEndPoint;
                     });
             }
 
@@ -471,6 +473,7 @@ await WithConnectionAsync(
             Assert.Equal("https://another.domain.url/chat?negotiateVersion=1&id=0rge0d00-0040-0030-0r00-000q00r00e00", testHttpHandler.ReceivedRequests[2].RequestUri.ToString());
             Assert.Equal("https://another.domain.url/chat?negotiateVersion=1&id=0rge0d00-0040-0030-0r00-000q00r00e00", testHttpHandler.ReceivedRequests[3].RequestUri.ToString());
             Assert.Equal(5, testHttpHandler.ReceivedRequests.Count);
+            Assert.Equal("https://another.domain.url/chat", connectedEndpoint.ToString());
         }
 
         [Fact]
diff --git a/src/SignalR/clients/csharp/Client/test/UnitTests/HubConnectionTests.Helpers.cs b/src/SignalR/clients/csharp/Client/test/UnitTests/HubConnectionTests.Helpers.cs
@@ -21,7 +21,11 @@ private static HubConnection CreateHubConnection(
         var builder = new HubConnectionBuilder().WithUrl("http://example.com");
 
         var delegateConnectionFactory = new DelegateConnectionFactory(
-            endPoint => connection.StartAsync());
+            async endPoint =>
+            {
+                connection.RemoteEndPoint = endPoint;
+                return await connection.StartAsync();
+            });
 
         builder.Services.AddSingleton<IConnectionFactory>(delegateConnectionFactory);
 
diff --git a/src/SignalR/clients/csharp/Client/test/UnitTests/HubConnectionTests.Tracing.cs b/src/SignalR/clients/csharp/Client/test/UnitTests/HubConnectionTests.Tracing.cs
@@ -4,6 +4,8 @@
 using System;
 using System.Diagnostics;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Internal;
 using Microsoft.AspNetCore.InternalTesting;
 using Microsoft.AspNetCore.SignalR.Client.Internal;
 using Microsoft.AspNetCore.SignalR.Tests;
@@ -42,6 +44,7 @@ public async Task InvokeSendsAnInvocationMessage_SendTraceHeaders()
                 var traceParent = (string)invokeMessage["headers"]["traceparent"];
 
                 Assert.Equal(clientActivity.Id, traceParent);
+                Assert.Equal("example.com", clientActivity.TagObjects.Single(t => t.Key == "server.address").Value);
 
                 Assert.Equal(TaskStatus.WaitingForActivation, invokeTask.Status);
             }
@@ -124,5 +127,91 @@ public async Task SendAnInvocationMessage_SendTraceHeaders()
             }
         }
 
+        [Fact]
+        public async Task InvokeSendsAnInvocationMessage_ConnectionRemoteEndPointChanged_UseRemoteEndpointUrl()
+        {
+            var clientSourceContainer = new SignalRClientActivitySource();
+            Activity clientActivity = null;
+
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = activitySource => ReferenceEquals(activitySource, clientSourceContainer.ActivitySource),
+                Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData,
+                ActivityStarted = activity => clientActivity = activity
+            };
+            ActivitySource.AddActivityListener(listener);
+
+            TestConnection connection = null;
+            connection = new TestConnection(onStart: () =>
+            {
+                connection.RemoteEndPoint = new UriEndPoint(new Uri("http://example.net"));
+                return Task.CompletedTask;
+            });
+            var hubConnection = CreateHubConnection(connection, clientActivitySource: clientSourceContainer);
+            try
+            {
+                await hubConnection.StartAsync().DefaultTimeout();
+
+                _ = hubConnection.InvokeAsync("Foo");
+
+                await connection.ReadSentJsonAsync().DefaultTimeout();
+
+                Assert.Equal("example.net", clientActivity.TagObjects.Single(t => t.Key == "server.address").Value);
+            }
+            finally
+            {
+                await hubConnection.DisposeAsync().DefaultTimeout();
+                await connection.DisposeAsync().DefaultTimeout();
+            }
+        }
+
+        [Fact]
+        public async Task InvokeSendsAnInvocationMessage_ConnectionRemoteEndPointChangedDuringConnect_UseRemoteEndpointUrl()
+        {
+            var clientSourceContainer = new SignalRClientActivitySource();
+            var clientActivityTcs = new TaskCompletionSource<Activity>(TaskCreationOptions.RunContinuationsAsynchronously); ;
+
+            using var listener = new ActivityListener
+            {
+                ShouldListenTo = activitySource => ReferenceEquals(activitySource, clientSourceContainer.ActivitySource),
+                Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData,
+                ActivityStarted = clientActivityTcs.SetResult
+            };
+            ActivitySource.AddActivityListener(listener);
+
+            var syncPoint = new SyncPoint();
+            TestConnection connection = null;
+            connection = new TestConnection(onStart: async () =>
+            {
+                await syncPoint.WaitToContinue();
+                connection.RemoteEndPoint = new UriEndPoint(new Uri("http://example.net"));
+            });
+            var hubConnection = CreateHubConnection(connection, clientActivitySource: clientSourceContainer);
+            try
+            {
+                var startTask = hubConnection.StartAsync();
+
+                _ = hubConnection.InvokeAsync("Foo");
+
+                var clientActivity = await clientActivityTcs.Task.DefaultTimeout();
+
+                // Initial server.address uses configured HubConnection URL.
+                Assert.Equal("example.com", clientActivity.TagObjects.Single(t => t.Key == "server.address").Value);
+
+                syncPoint.Continue();
+
+                await startTask.DefaultTimeout();
+
+                await connection.ReadSentJsonAsync().DefaultTimeout();
+
+                // After connection is started, server.address is updated to the connection's remote endpoint.
+                Assert.Equal("example.net", clientActivity.TagObjects.Single(t => t.Key == "server.address").Value);
+            }
+            finally
+            {
+                await hubConnection.DisposeAsync().DefaultTimeout();
+                await connection.DisposeAsync().DefaultTimeout();
+            }
+        }
     }
 }
diff --git a/src/SignalR/clients/csharp/Http.Connections.Client/src/HttpConnection.cs b/src/SignalR/clients/csharp/Http.Connections.Client/src/HttpConnection.cs
@@ -148,8 +148,6 @@ public HttpConnection(HttpConnectionOptions httpConnectionOptions, ILoggerFactor
 
         _url = _httpConnectionOptions.Url;
 
-        RemoteEndPoint = new UriEndPoint(_url);
-
         if (!httpConnectionOptions.SkipNegotiation || httpConnectionOptions.Transports != HttpTransportType.WebSockets)
         {
             _httpClient = CreateHttpClient();
@@ -353,6 +351,9 @@ private async Task SelectAndStartTransport(TransferFormat transferFormat, Cancel
                 throw new InvalidOperationException("Negotiate redirection limit exceeded.");
             }
 
+            // Set the final negotiated URI as the endpoint.
+            RemoteEndPoint = new UriEndPoint(Utils.CreateEndPointUri(uri));
+
             // This should only need to happen once
             var connectUrl = CreateConnectUrl(uri, negotiationResponse.ConnectionToken);
 
@@ -405,6 +406,9 @@ private async Task SelectAndStartTransport(TransferFormat transferFormat, Cancel
                             _httpConnectionOptions.UseStatefulReconnect = transportType == HttpTransportType.WebSockets ? _httpConnectionOptions.UseStatefulReconnect : false;
                             negotiationResponse = await GetNegotiationResponseAsync(uri, cancellationToken).ConfigureAwait(false);
                             connectUrl = CreateConnectUrl(uri, negotiationResponse.ConnectionToken);
+
+                            // Set the final negotiated URI as the endpoint.
+                            RemoteEndPoint = new UriEndPoint(Utils.CreateEndPointUri(uri));
                         }
 
                         Log.StartingTransport(_logger, transportType, uri);
diff --git a/src/SignalR/clients/csharp/Http.Connections.Client/src/Internal/Utils.cs b/src/SignalR/clients/csharp/Http.Connections.Client/src/Internal/Utils.cs
@@ -7,6 +7,19 @@ namespace Microsoft.AspNetCore.Http.Connections.Client.Internal;
 
 internal static class Utils
 {
+    public static Uri CreateEndPointUri(Uri url)
+    {
+        // The EndPoint URI shouldn't have querystring or target.
+        var uriBuilder = new UriBuilder
+        {
+            Scheme = url.Scheme,
+            Host = url.Host,
+            Port = url.Port,
+            Path = url.AbsolutePath
+        };
+        return uriBuilder.Uri;
+    }
+
     public static Uri AppendPath(Uri url, string path)
     {
         var builder = new UriBuilder(url);

Original file line number	Diff line number	Diff line change
`@@ -456,13 +456,15 @@ public async Task NegotiateThatReturnsRedirectUrlDoesNotAddAnotherNegotiateVersi`
`456`	`456`
`457`	`457`	`testHttpHandler.OnLongPollDelete((token) => ResponseUtils.CreateResponse(HttpStatusCode.Accepted));`
`458`	`458`
	`459`	`+ EndPoint connectedEndpoint = null;`
`459`	`460`	`using (var noErrorScope = new VerifyNoErrorsScope())`
`460`	`461`	`{`
`461`	`462`	`await WithConnectionAsync(`
`462`	`463`	`CreateConnection(testHttpHandler, loggerFactory: noErrorScope.LoggerFactory),`
`463`	`464`	`async (connection) =>`
`464`	`465`	`{`
`465`	`466`	`await connection.StartAsync().DefaultTimeout();`
	`467`	`+ connectedEndpoint = connection.RemoteEndPoint;`
`466`	`468`	`});`
`467`	`469`	`}`
`468`	`470`
`@@ -471,6 +473,7 @@ await WithConnectionAsync(`
`471`	`473`	`Assert.Equal("https://another.domain.url/chat?negotiateVersion=1&id=0rge0d00-0040-0030-0r00-000q00r00e00", testHttpHandler.ReceivedRequests[2].RequestUri.ToString());`
`472`	`474`	`Assert.Equal("https://another.domain.url/chat?negotiateVersion=1&id=0rge0d00-0040-0030-0r00-000q00r00e00", testHttpHandler.ReceivedRequests[3].RequestUri.ToString());`
`473`	`475`	`Assert.Equal(5, testHttpHandler.ReceivedRequests.Count);`
	`476`	`+ Assert.Equal("https://another.domain.url/chat", connectedEndpoint.ToString());`
`474`	`477`	`}`
`475`	`478`
`476`	`479`	`[Fact]`