git.stella-ops.org/src/StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest/Fetch/OciAttestationFetcher.cs

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.IO;
using System.IO.Abstractions;
using System.IO.Compression;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using System.Net.Http;
using Microsoft.Extensions.Logging;
using StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest.Authentication;
using StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest.Configuration;
using StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest.Discovery;
using System.Formats.Tar;

namespace StellaOps.Excititor.Connectors.OCI.OpenVEX.Attest.Fetch;

public sealed class OciAttestationFetcher
{
    private readonly IHttpClientFactory _httpClientFactory;
    private readonly IFileSystem _fileSystem;
    private readonly ILogger<OciAttestationFetcher> _logger;

    public OciAttestationFetcher(
        IHttpClientFactory httpClientFactory,
        IFileSystem fileSystem,
        ILogger<OciAttestationFetcher> logger)
    {
        _httpClientFactory = httpClientFactory ?? throw new ArgumentNullException(nameof(httpClientFactory));
        _fileSystem = fileSystem ?? throw new ArgumentNullException(nameof(fileSystem));
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
    }

    public async IAsyncEnumerable<OciAttestationDocument> FetchAsync(
        OciAttestationDiscoveryResult discovery,
        OciOpenVexAttestationConnectorOptions options,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(discovery);
        ArgumentNullException.ThrowIfNull(options);

        foreach (var target in discovery.Targets)
        {
            cancellationToken.ThrowIfCancellationRequested();

            bool yieldedOffline = false;
            if (target.OfflineBundle is not null && target.OfflineBundle.Exists)
            {
                await foreach (var offlineDocument in ReadOfflineAsync(target, cancellationToken))
                {
                    yieldedOffline = true;
                    yield return offlineDocument;
                }

                if (!discovery.AllowNetworkFallback)
                {
                    continue;
                }
            }

            if (discovery.PreferOffline && yieldedOffline && !discovery.AllowNetworkFallback)
            {
                continue;
            }

            if (!discovery.PreferOffline || discovery.AllowNetworkFallback || !yieldedOffline)
            {
                await foreach (var registryDocument in FetchFromRegistryAsync(discovery, options, target, cancellationToken))
                {
                    yield return registryDocument;
                }
            }
        }
    }

    private async IAsyncEnumerable<OciAttestationDocument> ReadOfflineAsync(
        OciAttestationTarget target,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var offline = target.OfflineBundle!;
        var path = _fileSystem.Path.GetFullPath(offline.Path);

        if (!_fileSystem.File.Exists(path))
        {
            if (offline.Exists)
            {
                _logger.LogWarning("Offline bundle {Path} disappeared before processing.", path);
            }
            yield break;
        }

        var extension = _fileSystem.Path.GetExtension(path).ToLowerInvariant();
        var subjectDigest = target.Image.Digest ?? target.ExpectedSubjectDigest;

        if (string.Equals(extension, ".json", StringComparison.OrdinalIgnoreCase) ||
            string.Equals(extension, ".dsse", StringComparison.OrdinalIgnoreCase))
        {
            var bytes = await _fileSystem.File.ReadAllBytesAsync(path, cancellationToken).ConfigureAwait(false);
            var metadata = BuildOfflineMetadata(target, path, entryName: null, subjectDigest);
            yield return new OciAttestationDocument(
                new Uri(path, UriKind.Absolute),
                bytes,
                metadata,
                subjectDigest,
                null,
                null,
                "offline");
            yield break;
        }

        if (string.Equals(extension, ".tgz", StringComparison.OrdinalIgnoreCase) ||
            string.Equals(extension, ".gz", StringComparison.OrdinalIgnoreCase) ||
            string.Equals(extension, ".tar", StringComparison.OrdinalIgnoreCase))
        {
            await foreach (var document in ReadTarArchiveAsync(target, path, subjectDigest, cancellationToken))
            {
                yield return document;
            }
            yield break;
        }

        // Default: treat as binary blob.
        var fallbackBytes = await _fileSystem.File.ReadAllBytesAsync(path, cancellationToken).ConfigureAwait(false);
        var fallbackMetadata = BuildOfflineMetadata(target, path, entryName: null, subjectDigest);
        yield return new OciAttestationDocument(
            new Uri(path, UriKind.Absolute),
            fallbackBytes,
            fallbackMetadata,
            subjectDigest,
            null,
            null,
            "offline");
    }

    private async IAsyncEnumerable<OciAttestationDocument> ReadTarArchiveAsync(
        OciAttestationTarget target,
        string path,
        string? subjectDigest,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        await using var fileStream = _fileSystem.File.OpenRead(path);
        Stream archiveStream = fileStream;

        if (path.EndsWith(".gz", StringComparison.OrdinalIgnoreCase) ||
            path.EndsWith(".tgz", StringComparison.OrdinalIgnoreCase))
        {
            archiveStream = new GZipStream(fileStream, CompressionMode.Decompress, leaveOpen: false);
        }

        using var tarReader = new TarReader(archiveStream, leaveOpen: false);
        TarEntry? entry;

        while ((entry = await tarReader.GetNextEntryAsync(copyData: false, cancellationToken).ConfigureAwait(false)) is not null)
        {
            if (entry.EntryType is not TarEntryType.RegularFile || entry.DataStream is null)
            {
                continue;
            }

            await using var entryStream = entry.DataStream;
            using var buffer = new MemoryStream();
            await entryStream.CopyToAsync(buffer, cancellationToken).ConfigureAwait(false);

            var metadata = BuildOfflineMetadata(target, path, entry.Name, subjectDigest);
            var sourceUri = new Uri($"{_fileSystem.Path.GetFullPath(path)}#{entry.Name}", UriKind.Absolute);
            yield return new OciAttestationDocument(
                sourceUri,
                buffer.ToArray(),
                metadata,
                subjectDigest,
                null,
                null,
                "offline");
        }
    }

    private async IAsyncEnumerable<OciAttestationDocument> FetchFromRegistryAsync(
        OciAttestationDiscoveryResult discovery,
        OciOpenVexAttestationConnectorOptions options,
        OciAttestationTarget target,
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var registryClient = new OciRegistryClient(
            _httpClientFactory,
            _logger,
            discovery.RegistryAuthorization,
            options);

        var subjectDigest = target.Image.Digest ?? target.ExpectedSubjectDigest;
        if (string.IsNullOrWhiteSpace(subjectDigest))
        {
            subjectDigest = await registryClient.ResolveDigestAsync(target.Image, cancellationToken).ConfigureAwait(false);
        }

        if (string.IsNullOrWhiteSpace(subjectDigest))
        {
            _logger.LogWarning("Unable to resolve subject digest for {Reference}; skipping registry fetch.", target.Image.Canonical);
            yield break;
        }

        if (!string.IsNullOrWhiteSpace(target.ExpectedSubjectDigest) &&
            !string.Equals(target.ExpectedSubjectDigest, subjectDigest, StringComparison.OrdinalIgnoreCase))
        {
            _logger.LogWarning(
                "Resolved digest {Resolved} does not match expected digest {Expected} for {Reference}.",
                subjectDigest,
                target.ExpectedSubjectDigest,
                target.Image.Canonical);
        }

        var descriptors = await registryClient.ListReferrersAsync(target.Image, subjectDigest, cancellationToken).ConfigureAwait(false);
        if (descriptors.Count == 0)
        {
            yield break;
        }

        foreach (var descriptor in descriptors)
        {
            cancellationToken.ThrowIfCancellationRequested();
            var document = await registryClient.DownloadAttestationAsync(target.Image, descriptor, subjectDigest, cancellationToken).ConfigureAwait(false);
            if (document is not null)
            {
                yield return document;
            }
        }
    }

    private static ImmutableDictionary<string, string> BuildOfflineMetadata(
        OciAttestationTarget target,
        string bundlePath,
        string? entryName,
        string? subjectDigest)
    {
        var builder = ImmutableDictionary.CreateBuilder<string, string>(StringComparer.Ordinal);
        builder["oci.image.registry"] = target.Image.Registry;
        builder["oci.image.repository"] = target.Image.Repository;
        builder["oci.image.reference"] = target.Image.Canonical;
        if (!string.IsNullOrWhiteSpace(subjectDigest))
        {
            builder["oci.image.subjectDigest"] = subjectDigest;
        }
        if (!string.IsNullOrWhiteSpace(target.ExpectedSubjectDigest))
        {
            builder["oci.image.expectedSubjectDigest"] = target.ExpectedSubjectDigest!;
        }

        builder["oci.attestation.sourceKind"] = "offline";
        builder["oci.attestation.source"] = bundlePath;
        if (!string.IsNullOrWhiteSpace(entryName))
        {
            builder["oci.attestation.bundleEntry"] = entryName!;
        }

        return builder.ToImmutable();
    }
}