fbbrecon/retryAccessToken

3 files changed, 27 insertions, 24 deletions

diff --git a/src/SignalR/clients/csharp/Client/test/UnitTests/HttpConnectionTests.Transport.cs b/src/SignalR/clients/csharp/Client/test/UnitTests/HttpConnectionTests.Transport.cs
index a97ae90f7f..ae7906c459 100644
--- a/src/SignalR/clients/csharp/Client/test/UnitTests/HttpConnectionTests.Transport.cs
+++ b/src/SignalR/clients/csharp/Client/test/UnitTests/HttpConnectionTests.Transport.cs
@@ -341,10 +341,8 @@ public partial class HttpConnectionTests
             Assert.Equal(1, accessTokenCallCount);
         }
 
-        [Theory]
-        [InlineData(HttpStatusCode.Unauthorized)]
-        [InlineData(HttpStatusCode.Forbidden)]
-        public async Task HttpConnectionRetriesAccessTokenProviderWhenAuthFailsLongPolling(HttpStatusCode httpStatusCode)
+        [Fact]
+        public async Task HttpConnectionRetriesAccessTokenProviderWhenAuthFailsLongPolling()
         {
             var testHttpHandler = new TestHttpMessageHandler(autoNegotiate: false);
             var requestsExecuted = false;
@@ -365,7 +363,7 @@ public partial class HttpConnectionTests
                 if (pollCount % 2 == 0)
                 {
                     pollCount++;
-                    return ResponseUtils.CreateResponse(httpStatusCode);
+                    return ResponseUtils.CreateResponse(HttpStatusCode.Unauthorized);
                 }
                 if (pollCount / 2 >= messageFragments.Length)
                 {
@@ -385,7 +383,7 @@ public partial class HttpConnectionTests
                 if (!requestsExecuted)
                 {
                     requestsExecuted = true;
-                    return Task.FromResult(ResponseUtils.CreateResponse(httpStatusCode));
+                    return Task.FromResult(ResponseUtils.CreateResponse(HttpStatusCode.Unauthorized));
                 }
 
                 Assert.Equal("Bearer", request.Headers.Authorization.Scheme);
@@ -419,10 +417,8 @@ public partial class HttpConnectionTests
             Assert.Equal(7, accessTokenCallCount);
         }
 
-        [Theory]
-        [InlineData(HttpStatusCode.Unauthorized)]
-        [InlineData(HttpStatusCode.Forbidden)]
-        public async Task HttpConnectionFailsAfterFirstRetryFailsLongPolling(HttpStatusCode httpStatusCode)
+        [Fact]
+        public async Task HttpConnectionFailsAfterFirstRetryFailsLongPolling()
         {
             var testHttpHandler = new TestHttpMessageHandler(autoNegotiate: false);
             var accessTokenCallCount = 0;
@@ -434,7 +430,7 @@ public partial class HttpConnectionTests
 
             testHttpHandler.OnLongPoll(_ =>
             {
-                return ResponseUtils.CreateResponse(httpStatusCode);
+                return ResponseUtils.CreateResponse(HttpStatusCode.Unauthorized);
             });
 
             Task<string> AccessTokenProvider()
@@ -455,10 +451,8 @@ public partial class HttpConnectionTests
             Assert.Equal(2, accessTokenCallCount);
         }
 
-        [Theory]
-        [InlineData(HttpStatusCode.Unauthorized)]
-        [InlineData(HttpStatusCode.Forbidden)]
-        public async Task HttpConnectionRetriesAccessTokenProviderWhenAuthFailsServerSentEvents(HttpStatusCode httpStatusCode)
+        [Fact]
+        public async Task HttpConnectionRetriesAccessTokenProviderWhenAuthFailsServerSentEvents()
         {
             var testHttpHandler = new TestHttpMessageHandler(autoNegotiate: false);
             var requestsExecuted = false;
@@ -476,7 +470,7 @@ public partial class HttpConnectionTests
                 if (!sendRequestExecuted)
                 {
                     sendRequestExecuted = true;
-                    return ResponseUtils.CreateResponse(httpStatusCode);
+                    return ResponseUtils.CreateResponse(HttpStatusCode.Unauthorized);
                 }
                 sendFinishedTcs.SetResult();
                 return ResponseUtils.CreateResponse(HttpStatusCode.OK);
@@ -489,7 +483,7 @@ public partial class HttpConnectionTests
                 if (!requestsExecuted)
                 {
                     requestsExecuted = true;
-                    return Task.FromResult(ResponseUtils.CreateResponse(httpStatusCode));
+                    return Task.FromResult(ResponseUtils.CreateResponse(HttpStatusCode.Unauthorized));
                 }
 
                 Assert.Equal("Bearer", request.Headers.Authorization.Scheme);
@@ -519,10 +513,8 @@ public partial class HttpConnectionTests
             Assert.Equal(3, accessTokenCallCount);
         }
 
-        [Theory]
-        [InlineData(HttpStatusCode.Unauthorized)]
-        [InlineData(HttpStatusCode.Forbidden)]
-        public async Task HttpConnectionFailsAfterFirstRetryFailsServerSentEvents(HttpStatusCode httpStatusCode)
+        [Fact]
+        public async Task HttpConnectionFailsAfterFirstRetryFailsServerSentEvents()
         {
             var testHttpHandler = new TestHttpMessageHandler(autoNegotiate: false);
             var accessTokenCallCount = 0;
@@ -534,7 +526,7 @@ public partial class HttpConnectionTests
 
             testHttpHandler.OnSocketSend((_, _) =>
             {
-                return ResponseUtils.CreateResponse(httpStatusCode);
+                return ResponseUtils.CreateResponse(HttpStatusCode.Unauthorized);
             });
 
             var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
diff --git a/src/SignalR/clients/csharp/Http.Connections.Client/src/HttpConnection.cs b/src/SignalR/clients/csharp/Http.Connections.Client/src/HttpConnection.cs
index c5671e2450..293ac92d71 100644
--- a/src/SignalR/clients/csharp/Http.Connections.Client/src/HttpConnection.cs
+++ b/src/SignalR/clients/csharp/Http.Connections.Client/src/HttpConnection.cs
@@ -464,6 +464,12 @@ public partial class HttpConnection : ConnectionContext, IConnectionInherentKeep
                 // Corefx changed the default version and High Sierra curlhandler tries to upgrade request
                 request.Version = new Version(1, 1);
 
+#if NET5_0_OR_GREATER
+                request.Options.Set(new HttpRequestOptionsKey<bool>("IsNegotiate"), true);
+#else
+                request.Properties.Add("IsNegotiate", true);
+#endif
+
                 // ResponseHeadersRead instructs SendAsync to return once headers are read
                 // rather than buffer the entire response. This gives a small perf boost.
                 // Note that it is important to dispose of the response when doing this to
diff --git a/src/SignalR/clients/csharp/Http.Connections.Client/src/Internal/AccessTokenHttpMessageHandler.cs b/src/SignalR/clients/csharp/Http.Connections.Client/src/Internal/AccessTokenHttpMessageHandler.cs
index 090e0629f9..eca0aed789 100644
--- a/src/SignalR/clients/csharp/Http.Connections.Client/src/Internal/AccessTokenHttpMessageHandler.cs
+++ b/src/SignalR/clients/csharp/Http.Connections.Client/src/Internal/AccessTokenHttpMessageHandler.cs
@@ -24,7 +24,12 @@ internal sealed class AccessTokenHttpMessageHandler : DelegatingHandler
         var isNegotiate = false;
         if (string.IsNullOrEmpty(_accessToken) ||
             // Negotiate redirects likely will have a new access token so let's always grab a (potentially) new access token on negotiate
-            request.RequestUri!.OriginalString.Contains("/negotiate?"))
+#if NET5_0_OR_GREATER
+            request.Options.TryGetValue(new HttpRequestOptionsKey<bool>("IsNegotiate"), out var value) && value == true
+#else
+            request.Properties.TryGetValue("IsNegotiate", out var value) && value is true
+#endif
+            )
         {
             isNegotiate = true;
             _accessToken = await _httpConnection.GetAccessTokenAsync().ConfigureAwait(false);
@@ -37,7 +42,7 @@ internal sealed class AccessTokenHttpMessageHandler : DelegatingHandler
 
         var result = await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
         // retry once with a new token on auth failure
-        if (!isNegotiate && result.StatusCode is HttpStatusCode.Forbidden or HttpStatusCode.Unauthorized)
+        if (!isNegotiate && result.StatusCode is HttpStatusCode.Unauthorized)
         {
             HttpConnection.Log.RetryAccessToken(_httpConnection._logger, result.StatusCode);
             result.Dispose();