git.stella-ops.org/src/StellaOps.Feedser.Exporter.TrivyDb.Tests/TrivyDbOciWriterTests.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Reflection;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using StellaOps.Feedser.Storage.Mongo.Exporting;
using Xunit;

namespace StellaOps.Feedser.Exporter.TrivyDb.Tests;

public sealed class TrivyDbOciWriterTests : IDisposable
{
    private readonly string _root;

    public TrivyDbOciWriterTests()
    {
        _root = Directory.CreateTempSubdirectory("trivy-writer-tests").FullName;
    }

    [Fact]
    public async Task WriteAsync_ReusesBlobsFromBaseLayout_WhenDigestMatches()
    {
        var baseLayout = Path.Combine(_root, "base");
        Directory.CreateDirectory(Path.Combine(baseLayout, "blobs", "sha256"));

        var configBytes = Encoding.UTF8.GetBytes("base-config");
        var configDigest = ComputeDigest(configBytes);
        WriteBlob(baseLayout, configDigest, configBytes);

        var layerBytes = Encoding.UTF8.GetBytes("base-layer");
        var layerDigest = ComputeDigest(layerBytes);
        WriteBlob(baseLayout, layerDigest, layerBytes);

        var manifest = CreateManifest(configDigest, layerDigest);
        var manifestBytes = SerializeManifest(manifest);
        var manifestDigest = ComputeDigest(manifestBytes);
        WriteBlob(baseLayout, manifestDigest, manifestBytes);

        var plan = new TrivyDbExportPlan(
            TrivyDbExportMode.Delta,
            TreeDigest: "sha256:tree",
            BaseExportId: "20241101T000000Z",
            BaseManifestDigest: manifestDigest,
            ResetBaseline: false,
            Manifest: Array.Empty<ExportFileRecord>(),
            ChangedFiles: new[] { new ExportFileRecord("data.json", 1, "sha256:data") },
            RemovedPaths: Array.Empty<string>());

        var configDescriptor = new OciDescriptor(TrivyDbMediaTypes.TrivyConfig, configDigest, configBytes.Length);
        var layerDescriptor = new OciDescriptor(TrivyDbMediaTypes.TrivyLayer, layerDigest, layerBytes.Length);
        var package = new TrivyDbPackage(
            manifest,
            new TrivyConfigDocument(
                TrivyDbMediaTypes.TrivyConfig,
                DateTimeOffset.Parse("2024-11-01T00:00:00Z"),
                "20241101T000000Z",
                layerDigest,
                layerBytes.Length),
            new Dictionary<string, TrivyDbBlob>(StringComparer.Ordinal)
            {
                [configDigest] = CreateThrowingBlob(),
                [layerDigest] = CreateThrowingBlob(),
            },
            JsonSerializer.SerializeToUtf8Bytes(new { mode = "delta" }));

        var writer = new TrivyDbOciWriter();
        var destination = Path.Combine(_root, "delta");
        await writer.WriteAsync(package, destination, reference: "example/trivy:delta", plan, baseLayout, CancellationToken.None);

        var reusedConfig = File.ReadAllBytes(GetBlobPath(destination, configDigest));
        Assert.Equal(configBytes, reusedConfig);

        var reusedLayer = File.ReadAllBytes(GetBlobPath(destination, layerDigest));
        Assert.Equal(layerBytes, reusedLayer);
    }

    private static TrivyDbBlob CreateThrowingBlob()
    {
        var ctor = typeof(TrivyDbBlob).GetConstructor(
            BindingFlags.NonPublic | BindingFlags.Instance,
            binder: null,
            new[] { typeof(Func<CancellationToken, ValueTask<Stream>>), typeof(long) },
            modifiers: null)
            ?? throw new InvalidOperationException("Unable to access TrivyDbBlob constructor.");

        Func<CancellationToken, ValueTask<Stream>> factory = _ => throw new InvalidOperationException("Blob should have been reused from base layout.");
        return (TrivyDbBlob)ctor.Invoke(new object[] { factory, 0L });
    }

    private static OciManifest CreateManifest(string configDigest, string layerDigest)
    {
        var configDescriptor = new OciDescriptor(TrivyDbMediaTypes.TrivyConfig, configDigest, 0);
        var layerDescriptor = new OciDescriptor(TrivyDbMediaTypes.TrivyLayer, layerDigest, 0);
        return new OciManifest(
            SchemaVersion: 2,
            MediaType: TrivyDbMediaTypes.OciManifest,
            Config: configDescriptor,
            Layers: new[] { layerDescriptor });
    }

    private static byte[] SerializeManifest(OciManifest manifest)
    {
        var options = new JsonSerializerOptions
        {
            PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
            DefaultIgnoreCondition = System.Text.Json.Serialization.JsonIgnoreCondition.WhenWritingNull,
            WriteIndented = false,
        };
        return JsonSerializer.SerializeToUtf8Bytes(manifest, options);
    }

    private static void WriteBlob(string layoutRoot, string digest, byte[] payload)
    {
        var path = GetBlobPath(layoutRoot, digest);
        Directory.CreateDirectory(Path.GetDirectoryName(path)!);
        File.WriteAllBytes(path, payload);
    }

    private static string GetBlobPath(string layoutRoot, string digest)
    {
        var fileName = digest[7..];
        return Path.Combine(layoutRoot, "blobs", "sha256", fileName);
    }

    private static string ComputeDigest(byte[] payload)
    {
        var hash = SHA256.HashData(payload);
        return "sha256:" + Convert.ToHexString(hash).ToLowerInvariant();
    }

    public void Dispose()
    {
        try
        {
            if (Directory.Exists(_root))
            {
                Directory.Delete(_root, recursive: true);
            }
        }
        catch
        {
            // best effort cleanup
        }
    }
}