Support async access token factory (#1911)

veloek · davidfowl · commit 31dfe9196263 · 2018-04-10T05:34:10.000-07:00
diff --git a/clients/ts/FunctionalTests/ts/HubConnectionTests.ts b/clients/ts/FunctionalTests/ts/HubConnectionTests.ts
@@ -532,6 +532,33 @@ describe("hubConnection", () => {
                 }
             });
 
+            it("can connect to hub with authorization using async token factory", async (done) => {
+                const message = "你好，世界！";
+
+                try {
+                    const hubConnection = new HubConnection("/authorizedhub", {
+                        accessTokenFactory: () => getJwtToken("http://" + document.location.host + "/generateJwtToken"),
+                        ...commonOptions,
+                        transport: transportType,
+                    });
+                    hubConnection.onclose((error) => {
+                        expect(error).toBe(undefined);
+                        done();
+                    });
+                    await hubConnection.start();
+                    const response = await hubConnection.invoke("Echo", message);
+
+                    expect(response).toEqual(message);
+
+                    await hubConnection.stop();
+
+                    done();
+                } catch (err) {
+                    fail(err);
+                    done();
+                }
+            });
+
             if (transportType !== TransportType.LongPolling) {
                 it("terminates if no messages received within timeout interval", (done) => {
                     const hubConnection = new HubConnection(TESTHUBENDPOINT_URL, {
diff --git a/clients/ts/signalr/src/HttpConnection.ts b/clients/ts/signalr/src/HttpConnection.ts
@@ -13,7 +13,7 @@ export interface IHttpConnectionOptions {
     httpClient?: HttpClient;
     transport?: TransportType | ITransport;
     logger?: ILogger | LogLevel;
-    accessTokenFactory?: () => string;
+    accessTokenFactory?: () => string | Promise<string>;
     logMessageContent?: boolean;
 }
 
@@ -87,7 +87,7 @@ export class HttpConnection implements IConnection {
                 // No fallback or negotiate in this case.
                 await this.transport.connect(this.url, transferFormat, this);
             } else {
-                const token = this.options.accessTokenFactory();
+                const token = await this.options.accessTokenFactory();
                 let headers;
                 if (token) {
                     headers = {
diff --git a/clients/ts/signalr/src/Transports.ts b/clients/ts/signalr/src/Transports.ts
@@ -30,17 +30,17 @@ export interface ITransport {
 
 export class WebSocketTransport implements ITransport {
     private readonly logger: ILogger;
-    private readonly accessTokenFactory: () => string;
+    private readonly accessTokenFactory: () => string | Promise<string>;
     private readonly logMessageContent: boolean;
     private webSocket: WebSocket;
 
-    constructor(accessTokenFactory: () => string, logger: ILogger, logMessageContent: boolean) {
+    constructor(accessTokenFactory: () => string | Promise<string>, logger: ILogger, logMessageContent: boolean) {
         this.logger = logger;
         this.accessTokenFactory = accessTokenFactory || (() => null);
         this.logMessageContent = logMessageContent;
     }
 
-    public connect(url: string, transferFormat: TransferFormat, connection: IConnection): Promise<void> {
+    public async connect(url: string, transferFormat: TransferFormat, connection: IConnection): Promise<void> {
         Arg.isRequired(url, "url");
         Arg.isRequired(transferFormat, "transferFormat");
         Arg.isIn(transferFormat, TransferFormat, "transferFormat");
@@ -52,9 +52,9 @@ export class WebSocketTransport implements ITransport {
 
         this.logger.log(LogLevel.Trace, "(WebSockets transport) Connecting");
 
+        const token = await this.accessTokenFactory();
         return new Promise<void>((resolve, reject) => {
             url = url.replace(/^http/, "ws");
-            const token = this.accessTokenFactory();
             if (token) {
                 url += (url.indexOf("?") < 0 ? "?" : "&") + `access_token=${encodeURIComponent(token)}`;
             }
@@ -118,20 +118,20 @@ export class WebSocketTransport implements ITransport {
 
 export class ServerSentEventsTransport implements ITransport {
     private readonly httpClient: HttpClient;
-    private readonly accessTokenFactory: () => string;
+    private readonly accessTokenFactory: () => string | Promise<string>;
     private readonly logger: ILogger;
     private readonly logMessageContent: boolean;
     private eventSource: EventSource;
     private url: string;
 
-    constructor(httpClient: HttpClient, accessTokenFactory: () => string, logger: ILogger, logMessageContent: boolean) {
+    constructor(httpClient: HttpClient, accessTokenFactory: () => string | Promise<string>, logger: ILogger, logMessageContent: boolean) {
         this.httpClient = httpClient;
         this.accessTokenFactory = accessTokenFactory || (() => null);
         this.logger = logger;
         this.logMessageContent = logMessageContent;
     }
 
-    public connect(url: string, transferFormat: TransferFormat, connection: IConnection): Promise<void> {
+    public async connect(url: string, transferFormat: TransferFormat, connection: IConnection): Promise<void> {
         Arg.isRequired(url, "url");
         Arg.isRequired(transferFormat, "transferFormat");
         Arg.isIn(transferFormat, TransferFormat, "transferFormat");
@@ -144,12 +144,12 @@ export class ServerSentEventsTransport implements ITransport {
         this.logger.log(LogLevel.Trace, "(SSE transport) Connecting");
 
         this.url = url;
+        const token = await this.accessTokenFactory();
         return new Promise<void>((resolve, reject) => {
             if (transferFormat !== TransferFormat.Text) {
                 reject(new Error("The Server-Sent Events transport only supports the 'Text' transfer format"));
             }
 
-            const token = this.accessTokenFactory();
             if (token) {
                 url += (url.indexOf("?") < 0 ? "?" : "&") + `access_token=${encodeURIComponent(token)}`;
             }
@@ -210,15 +210,15 @@ export class ServerSentEventsTransport implements ITransport {
 
 export class LongPollingTransport implements ITransport {
     private readonly httpClient: HttpClient;
-    private readonly accessTokenFactory: () => string;
+    private readonly accessTokenFactory: () => string | Promise<string>;
     private readonly logger: ILogger;
     private readonly logMessageContent: boolean;
 
     private url: string;
     private pollXhr: XMLHttpRequest;
     private pollAbort: AbortController;
 
-    constructor(httpClient: HttpClient, accessTokenFactory: () => string, logger: ILogger, logMessageContent: boolean) {
+    constructor(httpClient: HttpClient, accessTokenFactory: () => string | Promise<string>, logger: ILogger, logMessageContent: boolean) {
         this.httpClient = httpClient;
         this.accessTokenFactory = accessTokenFactory || (() => null);
         this.logger = logger;
@@ -259,7 +259,7 @@ export class LongPollingTransport implements ITransport {
             pollOptions.responseType = "arraybuffer";
         }
 
-        const token = this.accessTokenFactory();
+        const token = await this.accessTokenFactory();
         if (token) {
             // tslint:disable-next-line:no-string-literal
             pollOptions.headers["Authorization"] = `Bearer ${token}`;
@@ -356,12 +356,12 @@ function formatArrayBuffer(data: ArrayBuffer): string {
     return str.substr(0, str.length - 1);
 }
 
-async function send(logger: ILogger, transportName: string, httpClient: HttpClient, url: string, accessTokenFactory: () => string, content: string | ArrayBuffer, logMessageContent: boolean): Promise<void> {
+async function send(logger: ILogger, transportName: string, httpClient: HttpClient, url: string, accessTokenFactory: () => string | Promise<string>, content: string | ArrayBuffer, logMessageContent: boolean): Promise<void> {
     let headers;
-    const token = accessTokenFactory();
+    const token = await accessTokenFactory();
     if (token) {
         headers = {
-            ["Authorization"]: `Bearer ${accessTokenFactory()}`,
+            ["Authorization"]: `Bearer ${token}`,
         };
     }
 
diff --git a/samples/JwtClientSample/Program.cs b/samples/JwtClientSample/Program.cs
@@ -23,13 +23,13 @@ await Task.WhenAll(
 
         private const string ServerUrl = "http://localhost:54543";
 
-        private readonly ConcurrentDictionary<string, string> _tokens = new ConcurrentDictionary<string, string>(StringComparer.Ordinal);
+        private readonly ConcurrentDictionary<string, Task<string>> _tokens = new ConcurrentDictionary<string, Task<string>>(StringComparer.Ordinal);
         private readonly Random _random = new Random();
 
         private async Task RunConnection(HttpTransportType transportType)
         {
             var userId = "C#" + transportType;
-            _tokens[userId] = await GetJwtToken(userId);
+            _tokens[userId] = GetJwtToken(userId);
 
             var hubConnection = new HubConnectionBuilder()
                 .WithUrl(ServerUrl + "/broadcast", options =>
@@ -60,7 +60,7 @@ private async Task RunConnection(HttpTransportType transportType)
                         // no need to refresh the token for websockets
                         if (transportType != HttpTransportType.WebSockets)
                         {
-                            _tokens[userId] = await GetJwtToken(userId);
+                            _tokens[userId] = GetJwtToken(userId);
                             Console.WriteLine($"[{userId}] Token refreshed");
                         }
                     }
diff --git a/src/Microsoft.AspNetCore.Http.Connections.Client/HttpOptions.cs b/src/Microsoft.AspNetCore.Http.Connections.Client/HttpOptions.cs
@@ -7,6 +7,7 @@
 using System.Net.Http;
 using System.Net.WebSockets;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading.Tasks;
 
 namespace Microsoft.AspNetCore.Http.Connections.Client
 {
@@ -19,7 +20,7 @@ public class HttpOptions
         public Func<HttpMessageHandler, HttpMessageHandler> HttpMessageHandlerFactory { get; set; }
 
         public IReadOnlyCollection<KeyValuePair<string, string>> Headers { get; set; }
-        public Func<string> AccessTokenFactory { get; set; }
+        public Func<Task<string>> AccessTokenFactory { get; set; }
         public TimeSpan CloseTimeout { get; set; } = TimeSpan.FromSeconds(5);
         public ICredentials Credentials { get; set; }
         public X509CertificateCollection ClientCertificates { get; set; } = new X509CertificateCollection();
diff --git a/src/Microsoft.AspNetCore.Http.Connections.Client/Internal/AccessTokenHttpMessageHandler.cs b/src/Microsoft.AspNetCore.Http.Connections.Client/Internal/AccessTokenHttpMessageHandler.cs
@@ -11,18 +11,19 @@ namespace Microsoft.AspNetCore.Http.Connections.Client.Internal
 {
     public class AccessTokenHttpMessageHandler : DelegatingHandler
     {
-        private readonly Func<string> _accessTokenFactory;
+        private readonly Func<Task<string>> _accessTokenFactory;
 
-        public AccessTokenHttpMessageHandler(HttpMessageHandler inner, Func<string> accessTokenFactory) : base(inner)
+        public AccessTokenHttpMessageHandler(HttpMessageHandler inner, Func<Task<string>> accessTokenFactory) : base(inner)
         {
             _accessTokenFactory = accessTokenFactory ?? throw new ArgumentNullException(nameof(accessTokenFactory));
         }
 
-        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
         {
-            request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _accessTokenFactory());
+            var accessToken = await _accessTokenFactory();
+            request.Headers.Authorization =  new AuthenticationHeaderValue("Bearer", accessToken);
 
-            return base.SendAsync(request, cancellationToken);
+            return await base.SendAsync(request, cancellationToken);
         }
     }
 }
diff --git a/src/Microsoft.AspNetCore.Http.Connections.Client/Internal/WebSocketsTransport.cs b/src/Microsoft.AspNetCore.Http.Connections.Client/Internal/WebSocketsTransport.cs
@@ -17,6 +17,7 @@ namespace Microsoft.AspNetCore.Http.Connections.Client.Internal
     public partial class WebSocketsTransport : ITransport
     {
         private readonly ClientWebSocket _webSocket;
+        private readonly Func<Task<string>> _accessTokenFactory;
         private IDuplexPipe _application;
         private WebSocketMessageType _webSocketMessageType;
         private readonly ILogger _logger;
@@ -80,7 +81,7 @@ public WebSocketsTransport(HttpOptions httpOptions, ILoggerFactory loggerFactory
 
                 if (httpOptions.AccessTokenFactory != null)
                 {
-                    _webSocket.Options.SetRequestHeader("Authorization", $"Bearer {httpOptions.AccessTokenFactory()}");
+                    _accessTokenFactory = httpOptions.AccessTokenFactory;
                 }
 
                 httpOptions.WebSocketOptions?.Invoke(_webSocket.Options);
@@ -115,6 +116,12 @@ public async Task StartAsync(Uri url, TransferFormat transferFormat)
 
             Log.StartTransport(_logger, transferFormat, resolvedUrl);
 
+            if (_accessTokenFactory != null)
+            {
+                var accessToken = await _accessTokenFactory();
+                _webSocket.Options.SetRequestHeader("Authorization", $"Bearer {accessToken}");
+            }
+
             await _webSocket.ConnectAsync(resolvedUrl, CancellationToken.None);
 
             // Create the pipe pair (Application's writer is connected to Transport's reader, and vice versa)
diff --git a/src/Microsoft.AspNetCore.SignalR.Client/HttpConnectionOptions.cs b/src/Microsoft.AspNetCore.SignalR.Client/HttpConnectionOptions.cs
@@ -7,6 +7,7 @@
 using System.Net.Http;
 using System.Net.WebSockets;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http.Connections;
 using Microsoft.AspNetCore.Http.Connections.Internal;
 
@@ -29,7 +30,7 @@ public HttpConnectionOptions()
         public bool? UseDefaultCredentials { get; set; }
         public ICredentials Credentials { get; set; }
         public IWebProxy Proxy { get; set; }
-        public Func<string> AccessTokenFactory { get; set; }
+        public Func<Task<string>> AccessTokenFactory { get; set; }
         public Action<ClientWebSocketOptions> WebSocketOptions { get; set; }
 
         public X509CertificateCollection ClientCertificates
diff --git a/test/Microsoft.AspNetCore.SignalR.Client.FunctionalTests/HubConnectionTests.cs b/test/Microsoft.AspNetCore.SignalR.Client.FunctionalTests/HubConnectionTests.cs
@@ -706,15 +706,18 @@ public async Task ClientCanUseJwtBearerTokenForAuthentication(HttpTransportType
         {
             using (StartLog(out var loggerFactory, $"{nameof(ClientCanUseJwtBearerTokenForAuthentication)}_{transportType}"))
             {
-                var httpResponse = await new HttpClient().GetAsync(_serverFixture.Url + "/generateJwtToken");
-                httpResponse.EnsureSuccessStatusCode();
-                var token = await httpResponse.Content.ReadAsStringAsync();
+                async Task<string> AccessTokenFactory()
+                {
+                    var httpResponse = await new HttpClient().GetAsync(_serverFixture.Url + "/generateJwtToken");
+                    httpResponse.EnsureSuccessStatusCode();
+                    return await httpResponse.Content.ReadAsStringAsync();
+                };
 
                 var hubConnection = new HubConnectionBuilder()
                     .WithLoggerFactory(loggerFactory)
                     .WithUrl(_serverFixture.Url + "/authorizedhub", transportType, options =>
                     {
-                        options.AccessTokenFactory = () => token;
+                        options.AccessTokenFactory = AccessTokenFactory;
                     })
                     .Build();
                 try
diff --git a/test/Microsoft.AspNetCore.SignalR.Client.Tests/HttpConnectionTests.Helpers.cs b/test/Microsoft.AspNetCore.SignalR.Client.Tests/HttpConnectionTests.Helpers.cs
@@ -21,7 +21,7 @@ private static HttpConnection CreateConnection(
             ITransport transport = null,
             ITransportFactory transportFactory = null,
             HttpTransportType transportType = HttpTransportType.LongPolling,
-            Func<string> accessTokenFactory = null)
+            Func<Task<string>> accessTokenFactory = null)
         {
             var httpOptions = new HttpOptions
             {
diff --git a/test/Microsoft.AspNetCore.SignalR.Client.Tests/HttpConnectionTests.Transport.cs b/test/Microsoft.AspNetCore.SignalR.Client.Tests/HttpConnectionTests.Transport.cs
@@ -50,10 +50,10 @@ public async Task HttpConnectionSetsAccessTokenOnAllRequests(HttpTransportType t
                     return await next();
                 });
 
-                string AccessTokenFactory()
+                Task<string> AccessTokenFactory()
                 {
                     callCount++;
-                    return callCount.ToString();
+                    return Task.FromResult(callCount.ToString());
                 }
 
                 await WithConnectionAsync(

Original file line number	Diff line number	Diff line change
`@@ -23,13 +23,13 @@ await Task.WhenAll(`
`23`	`23`
`24`	`24`	`private const string ServerUrl = "http://localhost:54543";`
`25`	`25`
`26`		`- private readonly ConcurrentDictionary<string, string> _tokens = new ConcurrentDictionary<string, string>(StringComparer.Ordinal);`
	`26`	`+ private readonly ConcurrentDictionary<string, Task<string>> _tokens = new ConcurrentDictionary<string, Task<string>>(StringComparer.Ordinal);`
`27`	`27`	`private readonly Random _random = new Random();`
`28`	`28`
`29`	`29`	`private async Task RunConnection(HttpTransportType transportType)`
`30`	`30`	`{`
`31`	`31`	`var userId = "C#" + transportType;`
`32`		`- _tokens[userId] = await GetJwtToken(userId);`
	`32`	`+ _tokens[userId] = GetJwtToken(userId);`
`33`	`33`
`34`	`34`	`var hubConnection = new HubConnectionBuilder()`
`35`	`35`	`.WithUrl(ServerUrl + "/broadcast", options =>`
`@@ -60,7 +60,7 @@ private async Task RunConnection(HttpTransportType transportType)`
`60`	`60`	`// no need to refresh the token for websockets`
`61`	`61`	`if (transportType != HttpTransportType.WebSockets)`
`62`	`62`	`{`
`63`		`- _tokens[userId] = await GetJwtToken(userId);`
	`63`	`+ _tokens[userId] = GetJwtToken(userId);`
`64`	`64`	`Console.WriteLine($"[{userId}] Token refreshed");`
`65`	`65`	`}`
`66`	`66`	`}`
Original file line number	Diff line number	Diff line change
`@@ -11,18 +11,19 @@ namespace Microsoft.AspNetCore.Http.Connections.Client.Internal`
`11`	`11`	`{`
`12`	`12`	`public class AccessTokenHttpMessageHandler : DelegatingHandler`
`13`	`13`	`{`
`14`		`- private readonly Func<string> _accessTokenFactory;`
	`14`	`+ private readonly Func<Task<string>> _accessTokenFactory;`
`15`	`15`
`16`		`- public AccessTokenHttpMessageHandler(HttpMessageHandler inner, Func<string> accessTokenFactory) : base(inner)`
	`16`	`+ public AccessTokenHttpMessageHandler(HttpMessageHandler inner, Func<Task<string>> accessTokenFactory) : base(inner)`
`17`	`17`	`{`
`18`	`18`	`_accessTokenFactory = accessTokenFactory ?? throw new ArgumentNullException(nameof(accessTokenFactory));`
`19`	`19`	`}`
`20`	`20`
`21`		`- protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)`
	`21`	`+ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)`
`22`	`22`	`{`
`23`		`- request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _accessTokenFactory());`
	`23`	`+ var accessToken = await _accessTokenFactory();`
	`24`	`+ request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", accessToken);`
`24`	`25`
`25`		`- return base.SendAsync(request, cancellationToken);`
	`26`	`+ return await base.SendAsync(request, cancellationToken);`
`26`	`27`	`}`
`27`	`28`	`}`
`28`	`29`	`}`
Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ private static HttpConnection CreateConnection(`
`21`	`21`	`ITransport transport = null,`
`22`	`22`	`ITransportFactory transportFactory = null,`
`23`	`23`	`HttpTransportType transportType = HttpTransportType.LongPolling,`
`24`		`- Func<string> accessTokenFactory = null)`
	`24`	`+ Func<Task<string>> accessTokenFactory = null)`
`25`	`25`	`{`
`26`	`26`	`var httpOptions = new HttpOptions`
`27`	`27`	`{`
Original file line number	Diff line number	Diff line change
`@@ -50,10 +50,10 @@ public async Task HttpConnectionSetsAccessTokenOnAllRequests(HttpTransportType t`
`50`	`50`	`return await next();`
`51`	`51`	`});`
`52`	`52`
`53`		`- string AccessTokenFactory()`
	`53`	`+ Task<string> AccessTokenFactory()`
`54`	`54`	`{`
`55`	`55`	`callCount++;`
`56`		`- return callCount.ToString();`
	`56`	`+ return Task.FromResult(callCount.ToString());`
`57`	`57`	`}`
`58`	`58`
`59`	`59`	`await WithConnectionAsync(`