quorate/packages/fsharp-runner/oci/Transport.cs

using System.Collections.Concurrent;
using System.Net;
using System.Net.Http.Headers;
using System.Net.Security;
using System.Security.Cryptography.X509Certificates;
using System.Text.Json;

namespace Quorate.Oci;

internal sealed class RegistryTransport : IDisposable
{
    private readonly ConcurrentDictionary<RepositoryName, AuthorizationState> authorization = [];
    private readonly AuthenticationHeaderValue basic;
    private readonly HttpClient client;
    private readonly Uri endpoint;
    private readonly HttpMessageHandler? ownedHandler;

    public RegistryTransport(Uri endpoint, RegistryAuth auth, X509Certificate2 root)
        : this(endpoint, auth, CreateHandler(root), true)
    {
    }

    public RegistryTransport(Uri endpoint, RegistryAuth auth, HttpMessageHandler handler)
        : this(endpoint, auth, handler, false)
    {
    }

    private RegistryTransport(
        Uri endpoint,
        RegistryAuth auth,
        HttpMessageHandler handler,
        bool ownsHandler)
    {
        ValidateArguments(endpoint, auth);
        ArgumentNullException.ThrowIfNull(handler);
        this.endpoint = endpoint;
        basic = auth.Basic();
        ownedHandler = ownsHandler ? handler : null;
        client = new HttpClient(handler, false) { Timeout = TimeSpan.FromMinutes(15) };
    }

    public Uri Api(RepositoryName repository, string path) => new(endpoint, $"/v2/{repository}/{path}");

    public static Uri ResolveLocation(HttpResponseMessage response)
    {
        Uri location = response.Headers.Location
            ?? throw new InvalidDataException("Registry upload location is unavailable.");
        Uri request = response.RequestMessage?.RequestUri
            ?? throw new InvalidDataException("Registry request URI is unavailable.");
        Uri resolved = location.IsAbsoluteUri ? location : new Uri(request, location);

        if (resolved.Scheme != request.Scheme || resolved.IdnHost != request.IdnHost || resolved.Port != request.Port)
        {
            throw new InvalidDataException("Registry upload location changes authority.");
        }

        return resolved;
    }

    public async Task<HttpResponseMessage> SendAsync(
        RepositoryName repository,
        Func<HttpRequestMessage> create,
        CancellationToken cancellationToken,
        HttpStatusCode? allowed = null)
    {
        AuthorizationState state = authorization.GetOrAdd(repository, static _ => new AuthorizationState());

        for (int attempt = 0; attempt < 3; attempt++)
        {
            cancellationToken.ThrowIfCancellationRequested();
            int generation = state.Generation;
            AuthenticationHeaderValue current = state.Value ?? basic;
            using HttpRequestMessage request = create();
            request.Headers.Authorization = current;
            HttpResponseMessage response = await client.SendAsync(request, cancellationToken).ConfigureAwait(false);

            if (response.StatusCode != HttpStatusCode.Unauthorized)
            {
                if (!response.IsSuccessStatusCode && response.StatusCode != allowed)
                {
                    string body = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false);
                    HttpStatusCode status = response.StatusCode;
                    response.Dispose();
                    throw new HttpRequestException($"Registry returned {(int)status}: {body}", null, status);
                }

                return response;
            }

            AuthenticationHeaderValue challenge;

            using (response)
            {
                challenge = response.Headers.WwwAuthenticate.Single(header =>
                    header.Scheme.Equals("Bearer", StringComparison.OrdinalIgnoreCase));
            }

            await state.Gate.WaitAsync(cancellationToken).ConfigureAwait(false);

            try
            {
                if (state.Generation == generation)
                {
                    AuthenticationHeaderValue value = await GetBearerAsync(
                        BearerChallenge.Parse(challenge),
                        cancellationToken).ConfigureAwait(false);
                    state.Update(value);
                }
            }
            finally
            {
                state.Gate.Release();
            }
        }

        throw new InvalidDataException($"Registry authentication did not converge for {repository}.");
    }

    public void Dispose()
    {
        foreach (AuthorizationState state in authorization.Values)
        {
            state.Dispose();
        }

        client.Dispose();
        ownedHandler?.Dispose();
    }

    public static HttpRequestMessage Request(HttpMethod method, Uri uri) => new(method, uri)
    {
        Version = HttpVersion.Version20,
        VersionPolicy = HttpVersionPolicy.RequestVersionOrLower
    };

    private async Task<AuthenticationHeaderValue> GetBearerAsync(
        BearerChallenge challenge,
        CancellationToken cancellationToken)
    {
        var query = new List<string>();

        if (challenge.Service is not null)
        {
            query.Add($"service={Uri.EscapeDataString(challenge.Service)}");
        }

        if (challenge.Scope is not null)
        {
            query.Add($"scope={Uri.EscapeDataString(challenge.Scope)}");
        }

        var builder = new UriBuilder(challenge.Realm) { Query = string.Join('&', query) };
        using HttpRequestMessage request = Request(HttpMethod.Get, builder.Uri);
        request.Headers.Authorization = basic;
        using HttpResponseMessage response = await client.SendAsync(request, cancellationToken).ConfigureAwait(false);
        response.EnsureSuccessStatusCode();
        using Stream content = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
        BearerTokenModel token = await JsonSerializer.DeserializeAsync(
            content,
            OciJsonContext.Default.BearerTokenModel,
            cancellationToken).ConfigureAwait(false)
            ?? throw new InvalidDataException("Registry token response is empty.");
        string? value = token.Token ?? token.AccessToken;

        if (string.IsNullOrWhiteSpace(value))
        {
            throw new InvalidDataException("Registry token is unavailable.");
        }

        return new AuthenticationHeaderValue("Bearer", value);
    }

    private static SocketsHttpHandler CreateHandler(X509Certificate2 root)
    {
        ArgumentNullException.ThrowIfNull(root);
        var policy = new X509ChainPolicy
        {
            TrustMode = X509ChainTrustMode.CustomRootTrust,
            RevocationMode = X509RevocationMode.NoCheck
        };
        policy.CustomTrustStore.Add(root);
        return new SocketsHttpHandler
        {
            AllowAutoRedirect = false,
            SslOptions = new SslClientAuthenticationOptions
            {
                CertificateChainPolicy = policy
            }
        };
    }

    private static void ValidateArguments(Uri endpoint, RegistryAuth auth)
    {
        ArgumentNullException.ThrowIfNull(endpoint);
        ArgumentNullException.ThrowIfNull(auth);

        if (!endpoint.IsAbsoluteUri || endpoint.Scheme != Uri.UriSchemeHttps)
        {
            throw new ArgumentException("Registry endpoint must use HTTPS.", nameof(endpoint));
        }
    }

    private sealed class AuthorizationState : IDisposable
    {
        private int generation;
        private AuthenticationHeaderValue? value;

        public SemaphoreSlim Gate { get; } = new(1, 1);
        public AuthenticationHeaderValue? Value => Volatile.Read(ref value);
        public int Generation => Volatile.Read(ref generation);

        public void Update(AuthenticationHeaderValue authorization)
        {
            Volatile.Write(ref value, authorization);
            Interlocked.Increment(ref generation);
        }

        public void Dispose() => Gate.Dispose();
    }
}