Merge pull request #235 from rainersigwald/auth-caching

Auth caching, attempt 2

Author: rainersigwald

Microsoft.NET.Build.Containers/AuthHandshakeMessageHandler.cs

@@ -77,19 +77,19 @@ private record TokenResponse(string? token, string? access_token, int? expires_i
     /// Credentials for the request are retrieved from the credential provider, then used to acquire a token.
     /// That token is cached for some duration on a per-host basis.
     /// </summary>
-    /// <param name="realm"></param>
+    /// <param name="uri"></param>
     /// <param name="service"></param>
     /// <param name="scope"></param>
     /// <param name="cancellationToken"></param>
     /// <returns></returns>
-    private async Task<AuthenticationHeaderValue?> GetAuthenticationAsync(string scheme, Uri realm, string service, string? scope, CancellationToken cancellationToken)
+    private async Task<AuthenticationHeaderValue?> GetAuthenticationAsync(string scheme, Uri uri, string service, string? scope, CancellationToken cancellationToken)
     {
         // Allow overrides for auth via environment variables
         string? credU = Environment.GetEnvironmentVariable(ContainerHelpers.HostObjectUser);
         string? credP = Environment.GetEnvironmentVariable(ContainerHelpers.HostObjectPass);
 
         // fetch creds for the host
-      DockerCredentials? privateRepoCreds;
+        DockerCredentials? privateRepoCreds;
 
         if (!string.IsNullOrEmpty(credU) && !string.IsNullOrEmpty(credP))
         {
@@ -99,26 +99,26 @@ private record TokenResponse(string? token, string? access_token, int? expires_i
         {
             try
             {
-                privateRepoCreds = await CredsProvider.GetCredentialsAsync(realm.Host);
+                privateRepoCreds = await CredsProvider.GetCredentialsAsync(uri.Host);
             }
             catch (Exception e)
             {
-                throw new CredentialRetrievalException(realm.Host, e);
+                throw new CredentialRetrievalException(uri.Host, e);
             }
         }
 
         if (scheme is "Basic")
         {
             var basicAuth = new AuthenticationHeaderValue("Basic", Convert.ToBase64String(Encoding.ASCII.GetBytes($"{privateRepoCreds.Username}:{privateRepoCreds.Password}")));
-            return basicAuth;
+            return AuthHeaderCache.AddOrUpdate(uri, basicAuth);
         }
         else if (scheme is "Bearer")
         {
             // use those creds when calling the token provider
             var header = privateRepoCreds.Username == "<token>"
                             ? new AuthenticationHeaderValue("Bearer", privateRepoCreds.Password)
                             : new AuthenticationHeaderValue("Basic", Convert.ToBase64String(Encoding.ASCII.GetBytes($"{privateRepoCreds.Username}:{privateRepoCreds.Password}")));
-            var builder = new UriBuilder(realm);
+            var builder = new UriBuilder(uri);
             var queryDict = System.Web.HttpUtility.ParseQueryString("");
             queryDict["service"] = service;
             if (scope is string s)
@@ -140,7 +140,7 @@ private record TokenResponse(string? token, string? access_token, int? expires_i
 
             // save the retrieved token in the cache
             var bearerAuth = new AuthenticationHeaderValue("Bearer", token.ResolvedToken);
-            return bearerAuth;
+            return AuthHeaderCache.AddOrUpdate(uri, bearerAuth);
         }
         else
         {
@@ -155,7 +155,11 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
             throw new ArgumentException("No RequestUri specified", nameof(request));
         }
 
-        // TODO: attempt to use cached token for the request if available
+        // attempt to use cached token for the request if available
+        if (AuthHeaderCache.TryGet(request.RequestUri, out AuthenticationHeaderValue? cachedAuthentication))
+        {
+            request.Headers.Authorization = cachedAuthentication;
+        }
 
         var response = await base.SendAsync(request, cancellationToken);
         if (response is { StatusCode: HttpStatusCode.OK })
@@ -166,7 +170,7 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
         {
             if (await GetAuthenticationAsync(scheme, authInfo.Realm, authInfo.Service, authInfo.Scope, cancellationToken) is AuthenticationHeaderValue authentication)
             {
-                request.Headers.Authorization = authentication;
+                request.Headers.Authorization = AuthHeaderCache.AddOrUpdate(request.RequestUri, authentication);
                 return await base.SendAsync(request, cancellationToken);
             }
             return response;

Microsoft.NET.Build.Containers/AuthHeaderCache.cs

@@ -0,0 +1,42 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
+using System.Net.Http.Headers;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace Microsoft.NET.Build.Containers;
+
+internal static class AuthHeaderCache
+{
+
+    private static ConcurrentDictionary<string, AuthenticationHeaderValue> HostAuthenticationCache = new();
+
+    public static bool TryGet(Uri uri, [NotNullWhen(true)] out AuthenticationHeaderValue? header)
+    {
+        header = null;
+
+        // observed quirk in Azure Container Registry: if you present a token to blobs/uploads and it's wrong,
+        // it won't give back a www-authenticate header for the reauth mechanism to work. So never return
+        // a cache for that URI pattern
+        string[] segments = uri.Segments;
+        if (segments is [.., "blobs/", "uploads/"])
+        {
+            return false;
+        }
+
+        return HostAuthenticationCache.TryGetValue(GetCacheKey(uri), out header);
+    }
+
+    public static AuthenticationHeaderValue AddOrUpdate(Uri uri, AuthenticationHeaderValue header)
+    {
+        return HostAuthenticationCache.AddOrUpdate(GetCacheKey(uri), header, (_, _) => header);
+    }
+
+    private static string GetCacheKey(Uri uri)
+    {
+        return uri.Host + uri.AbsolutePath;
+    }
+}