git.stella-ops.org/src/Scheduler/__Libraries/StellaOps.Scheduler.ImpactIndex/RoaringImpactIndex.cs

using System;
using System.Buffers.Binary;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using System.Text.RegularExpressions;
using Collections.Special;
using Microsoft.Extensions.Logging;
using StellaOps.Scheduler.ImpactIndex.Ingestion;
using StellaOps.Scheduler.Models;

namespace StellaOps.Scheduler.ImpactIndex;

/// <summary>
/// Roaring bitmap-backed implementation of the scheduler impact index.
/// </summary>
public sealed class RoaringImpactIndex : IImpactIndex
{
    private readonly object _gate = new();

    private readonly Dictionary<string, int> _imageIds = new(StringComparer.OrdinalIgnoreCase);
    private readonly Dictionary<int, ImpactImageRecord> _images = new();
    private readonly Dictionary<string, RoaringBitmap> _containsByPurl = new(StringComparer.OrdinalIgnoreCase);
    private readonly Dictionary<string, RoaringBitmap> _usedByEntrypointByPurl = new(StringComparer.OrdinalIgnoreCase);

    private readonly ILogger<RoaringImpactIndex> _logger;
    private readonly TimeProvider _timeProvider;
    private string? _snapshotId;

    public RoaringImpactIndex(ILogger<RoaringImpactIndex> logger, TimeProvider? timeProvider = null)
    {
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));
        _timeProvider = timeProvider ?? TimeProvider.System;
    }

    public ValueTask RemoveAsync(string imageDigest, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(imageDigest);

        lock (_gate)
        {
            if (!_imageIds.TryGetValue(imageDigest, out var imageId))
            {
                return ValueTask.CompletedTask;
            }

            if (_images.TryGetValue(imageId, out var record))
            {
                RemoveImageComponents(record);
                _images.Remove(imageId);
            }

            _imageIds.Remove(imageDigest);
            _snapshotId = null;
        }

        return ValueTask.CompletedTask;
    }

    public async Task IngestAsync(ImpactIndexIngestionRequest request, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(request);
        ArgumentNullException.ThrowIfNull(request.BomIndexStream);

        using var buffer = new MemoryStream();
        await request.BomIndexStream.CopyToAsync(buffer, cancellationToken).ConfigureAwait(false);
        buffer.Position = 0;

        var document = BomIndexReader.Read(buffer);
        if (!string.Equals(document.ImageDigest, request.ImageDigest, StringComparison.OrdinalIgnoreCase))
        {
            throw new InvalidOperationException($"BOM-Index digest mismatch. Header '{document.ImageDigest}', request '{request.ImageDigest}'.");
        }

        var tenantId = request.TenantId ?? throw new ArgumentNullException(nameof(request.TenantId));
        var registry = request.Registry ?? throw new ArgumentNullException(nameof(request.Registry));
        var repository = request.Repository ?? throw new ArgumentNullException(nameof(request.Repository));

        var namespaces = request.Namespaces.IsDefault ? ImmutableArray<string>.Empty : request.Namespaces;
        var tags = request.Tags.IsDefault ? ImmutableArray<string>.Empty : request.Tags;
        var labels = request.Labels.Count == 0
            ? ImmutableSortedDictionary<string, string>.Empty.WithComparers(StringComparer.OrdinalIgnoreCase)
            : request.Labels;

        var componentKeys = document.Components
            .Select(component => component.Key)
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToImmutableArray();
        var entrypointComponents = document.Components
            .Where(component => component.UsedByEntrypoint)
            .Select(component => component.Key)
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToImmutableArray();

        lock (_gate)
        {
            var imageId = EnsureImageId(request.ImageDigest);

            if (_images.TryGetValue(imageId, out var existing))
            {
                RemoveImageComponents(existing);
            }

            var metadata = new ImpactImageRecord(
                imageId,
                tenantId,
                request.ImageDigest,
                registry,
                repository,
                namespaces,
                tags,
                labels,
                document.GeneratedAt,
                componentKeys,
                entrypointComponents);

            _images[imageId] = metadata;
            _imageIds[request.ImageDigest] = imageId;

            foreach (var key in componentKeys)
            {
                var bitmap = _containsByPurl.GetValueOrDefault(key);
                _containsByPurl[key] = AddImageToBitmap(bitmap, imageId);
            }

            foreach (var key in entrypointComponents)
            {
                var bitmap = _usedByEntrypointByPurl.GetValueOrDefault(key);
                _usedByEntrypointByPurl[key] = AddImageToBitmap(bitmap, imageId);
            }
        }

        _logger.LogInformation(
            "ImpactIndex ingested BOM-Index for {Digest} ({TenantId}/{Repository}). Components={ComponentCount} EntrypointComponents={EntrypointCount}",
            request.ImageDigest,
            tenantId,
            repository,
            componentKeys.Length,
            entrypointComponents.Length);
    }

    public ValueTask<ImpactSet> ResolveByPurlsAsync(
        IEnumerable<string> purls,
        bool usageOnly,
        Selector selector,
        CancellationToken cancellationToken = default)
        => ValueTask.FromResult(ResolveByPurlsCore(purls, usageOnly, selector));

    public ValueTask<ImpactSet> ResolveByVulnerabilitiesAsync(
        IEnumerable<string> vulnerabilityIds,
        bool usageOnly,
        Selector selector,
        CancellationToken cancellationToken = default)
        => ValueTask.FromResult(CreateEmptyImpactSet(selector, usageOnly));

    public ValueTask<ImpactSet> ResolveAllAsync(
        Selector selector,
        bool usageOnly,
        CancellationToken cancellationToken = default)
        => ValueTask.FromResult(ResolveAllCore(selector, usageOnly));

    public ValueTask<ImpactIndexSnapshot> CreateSnapshotAsync(CancellationToken cancellationToken = default)
    {
        cancellationToken.ThrowIfCancellationRequested();

        lock (_gate)
        {
            var orderedImages = _images
                .Values
                .OrderBy(img => img.Digest, StringComparer.OrdinalIgnoreCase)
                .ThenBy(img => img.Repository, StringComparer.OrdinalIgnoreCase)
                .ToArray();

            var idMap = orderedImages
                .Select((image, index) => (image.ImageId, NewId: index))
                .ToDictionary(tuple => tuple.ImageId, tuple => tuple.NewId);

            var compactedImages = orderedImages
                .Select(image => image with { ImageId = idMap[image.ImageId] })
                .ToImmutableArray();

            ImmutableDictionary<string, ImmutableArray<int>> CompactBitmaps(Dictionary<string, RoaringBitmap> source)
            {
                var builder = ImmutableDictionary.CreateBuilder<string, ImmutableArray<int>>(StringComparer.OrdinalIgnoreCase);
                foreach (var (key, bitmap) in source)
                {
                    var remapped = bitmap
                        .Select(id => idMap.TryGetValue(id, out var newId) ? newId : (int?)null)
                        .Where(id => id.HasValue)
                        .Select(id => id!.Value)
                        .Distinct()
                        .OrderBy(id => id)
                        .ToImmutableArray();

                    if (remapped.Length > 0)
                    {
                        builder[key] = remapped;
                    }
                }

                return builder.ToImmutable();
            }

            var contains = CompactBitmaps(_containsByPurl);
            var usedBy = CompactBitmaps(_usedByEntrypointByPurl);

            var generatedAt = orderedImages.Length == 0
                ? _timeProvider.GetUtcNow()
                : orderedImages.Max(img => img.GeneratedAt);

            var snapshotId = ComputeSnapshotId(compactedImages, contains, usedBy);
            _snapshotId = snapshotId;

            var snapshot = new ImpactIndexSnapshot(
                generatedAt,
                snapshotId,
                compactedImages,
                contains,
                usedBy);

            return ValueTask.FromResult(snapshot);
        }
    }

    public ValueTask RestoreSnapshotAsync(ImpactIndexSnapshot snapshot, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(snapshot);
        cancellationToken.ThrowIfCancellationRequested();

        lock (_gate)
        {
            _images.Clear();
            _imageIds.Clear();
            _containsByPurl.Clear();
            _usedByEntrypointByPurl.Clear();

            foreach (var image in snapshot.Images)
            {
                _images[image.ImageId] = image;
                _imageIds[image.Digest] = image.ImageId;
            }

            foreach (var kvp in snapshot.ContainsByPurl)
            {
                _containsByPurl[kvp.Key] = RoaringBitmap.Create(kvp.Value.ToArray());
            }

            foreach (var kvp in snapshot.UsedByEntrypointByPurl)
            {
                _usedByEntrypointByPurl[kvp.Key] = RoaringBitmap.Create(kvp.Value.ToArray());
            }

            _snapshotId = snapshot.SnapshotId;
        }

        return ValueTask.CompletedTask;
    }

    private ImpactSet ResolveByPurlsCore(IEnumerable<string> purls, bool usageOnly, Selector selector)
    {
        ArgumentNullException.ThrowIfNull(purls);
        ArgumentNullException.ThrowIfNull(selector);

        var normalized = purls
            .Where(static purl => !string.IsNullOrWhiteSpace(purl))
            .Select(static purl => purl.Trim())
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToArray();

        if (normalized.Length == 0)
        {
            return CreateEmptyImpactSet(selector, usageOnly);
        }

        RoaringBitmap imageIds;
        lock (_gate)
        {
            imageIds = RoaringBitmap.Create(Array.Empty<int>());
            foreach (var purl in normalized)
            {
                if (_containsByPurl.TryGetValue(purl, out var bitmap))
                {
                    imageIds = imageIds | bitmap;
                }
            }
        }

        return BuildImpactSet(imageIds, selector, usageOnly);
    }

    private ImpactSet ResolveAllCore(Selector selector, bool usageOnly)
    {
        ArgumentNullException.ThrowIfNull(selector);

        RoaringBitmap bitmap;
        lock (_gate)
        {
            var ids = _images.Keys.OrderBy(id => id).ToArray();
            bitmap = RoaringBitmap.Create(ids);
        }

        return BuildImpactSet(bitmap, selector, usageOnly);
    }

    private ImpactSet BuildImpactSet(RoaringBitmap imageIds, Selector selector, bool usageOnly)
    {
        var images = new List<ImpactImage>();
        var latestGeneratedAt = DateTimeOffset.MinValue;

        lock (_gate)
        {
            foreach (var imageId in imageIds)
            {
                if (!_images.TryGetValue(imageId, out var metadata))
                {
                    continue;
                }

                if (!ImageMatchesSelector(metadata, selector))
                {
                    continue;
                }

                if (usageOnly && metadata.EntrypointComponents.Length == 0)
                {
                    continue;
                }

                if (metadata.GeneratedAt > latestGeneratedAt)
                {
                    latestGeneratedAt = metadata.GeneratedAt;
                }

                images.Add(new ImpactImage(
                    metadata.Digest,
                    metadata.Registry,
                    metadata.Repository,
                    metadata.Namespaces,
                    metadata.Tags,
                    metadata.EntrypointComponents.Length > 0,
                    metadata.Labels));
            }
        }

        if (images.Count == 0)
        {
            return CreateEmptyImpactSet(selector, usageOnly);
        }

        images.Sort(static (left, right) => string.Compare(left.ImageDigest, right.ImageDigest, StringComparison.Ordinal));

        var generatedAt = latestGeneratedAt == DateTimeOffset.MinValue ? _timeProvider.GetUtcNow() : latestGeneratedAt;

        return new ImpactSet(
            selector,
            images.ToImmutableArray(),
            usageOnly,
            generatedAt,
            images.Count,
            snapshotId: _snapshotId,
            schemaVersion: SchedulerSchemaVersions.ImpactSet);
    }

    private ImpactSet CreateEmptyImpactSet(Selector selector, bool usageOnly)
    {
        return new ImpactSet(
            selector,
            ImmutableArray<ImpactImage>.Empty,
            usageOnly,
            _timeProvider.GetUtcNow(),
            0,
            snapshotId: _snapshotId,
            schemaVersion: SchedulerSchemaVersions.ImpactSet);
    }

    private static bool ImageMatchesSelector(ImpactImageRecord image, Selector selector)
    {
        if (selector.TenantId is not null && !string.Equals(selector.TenantId, image.TenantId, StringComparison.Ordinal))
        {
            return false;
        }

        if (!MatchesScope(image, selector))
        {
            return false;
        }

        if (selector.Digests.Length > 0 && !selector.Digests.Contains(image.Digest, StringComparer.OrdinalIgnoreCase))
        {
            return false;
        }

        if (selector.Repositories.Length > 0)
        {
            var repoMatch = selector.Repositories.Any(repo =>
                string.Equals(repo, image.Repository, StringComparison.OrdinalIgnoreCase) ||
                string.Equals(repo, $"{image.Registry}/{image.Repository}", StringComparison.OrdinalIgnoreCase));
            if (!repoMatch)
            {
                return false;
            }
        }

        if (selector.Namespaces.Length > 0)
        {
            if (image.Namespaces.IsDefaultOrEmpty)
            {
                return false;
            }

            var namespaceMatch = selector.Namespaces.Any(ns => image.Namespaces.Contains(ns, StringComparer.OrdinalIgnoreCase));
            if (!namespaceMatch)
            {
                return false;
            }
        }

        if (selector.IncludeTags.Length > 0)
        {
            if (image.Tags.IsDefaultOrEmpty)
            {
                return false;
            }

            var tagMatch = selector.IncludeTags.Any(pattern => image.Tags.Any(tag => MatchesTagPattern(tag, pattern)));
            if (!tagMatch)
            {
                return false;
            }
        }

        if (selector.Labels.Length > 0)
        {
            if (image.Labels.Count == 0)
            {
                return false;
            }

            foreach (var label in selector.Labels)
            {
                if (!image.Labels.TryGetValue(label.Key, out var value))
                {
                    return false;
                }

                if (label.Values.Length > 0 && !label.Values.Contains(value, StringComparer.OrdinalIgnoreCase))
                {
                    return false;
                }
            }
        }

        return true;
    }

    private void RemoveImageComponents(ImpactImageRecord record)
    {
        foreach (var key in record.Components)
        {
            if (_containsByPurl.TryGetValue(key, out var bitmap))
            {
                var updated = RemoveImageFromBitmap(bitmap, record.ImageId);
                if (updated is null)
                {
                    _containsByPurl.Remove(key);
                }
                else
                {
                    _containsByPurl[key] = updated;
                }
            }
        }

        foreach (var key in record.EntrypointComponents)
        {
            if (_usedByEntrypointByPurl.TryGetValue(key, out var bitmap))
            {
                var updated = RemoveImageFromBitmap(bitmap, record.ImageId);
                if (updated is null)
                {
                    _usedByEntrypointByPurl.Remove(key);
                }
                else
                {
                    _usedByEntrypointByPurl[key] = updated;
                }
            }
        }
    }

    private static RoaringBitmap AddImageToBitmap(RoaringBitmap? bitmap, int imageId)
    {
        if (bitmap is null)
        {
            return RoaringBitmap.Create(new[] { imageId });
        }

        if (bitmap.Any(id => id == imageId))
        {
            return bitmap;
        }

        var merged = bitmap
            .Concat(new[] { imageId })
            .Distinct()
            .OrderBy(id => id)
            .ToArray();

        return RoaringBitmap.Create(merged);
    }

    private static RoaringBitmap? RemoveImageFromBitmap(RoaringBitmap bitmap, int imageId)
    {
        var remaining = bitmap
            .Where(id => id != imageId)
            .OrderBy(id => id)
            .ToArray();
        if (remaining.Length == 0)
        {
            return null;
        }

        return RoaringBitmap.Create(remaining);
    }

    private static bool MatchesScope(ImpactImageRecord image, Selector selector)
    {
        return selector.Scope switch
        {
            SelectorScope.AllImages => true,
            SelectorScope.ByDigest => selector.Digests.Contains(image.Digest, StringComparer.OrdinalIgnoreCase),
            SelectorScope.ByRepository => selector.Repositories.Any(repo =>
                string.Equals(repo, image.Repository, StringComparison.OrdinalIgnoreCase) ||
                string.Equals(repo, $"{image.Registry}/{image.Repository}", StringComparison.OrdinalIgnoreCase)),
            SelectorScope.ByNamespace => !image.Namespaces.IsDefaultOrEmpty && selector.Namespaces.Any(ns => image.Namespaces.Contains(ns, StringComparer.OrdinalIgnoreCase)),
            SelectorScope.ByLabels => selector.Labels.All(label =>
                image.Labels.TryGetValue(label.Key, out var value) &&
                (label.Values.Length == 0 || label.Values.Contains(value, StringComparer.OrdinalIgnoreCase))),
            _ => true,
        };
    }

    private static string ComputeSnapshotId(
        ImmutableArray<ImpactImageRecord> images,
        ImmutableDictionary<string, ImmutableArray<int>> contains,
        ImmutableDictionary<string, ImmutableArray<int>> usedBy)
    {
        var builder = new StringBuilder();

        foreach (var image in images.OrderBy(img => img.Digest, StringComparer.OrdinalIgnoreCase))
        {
            builder.Append(image.Digest).Append('|').Append(image.GeneratedAt.ToUnixTimeSeconds()).Append(';');
        }

        void AppendMap(ImmutableDictionary<string, ImmutableArray<int>> map)
        {
            foreach (var kvp in map.OrderBy(pair => pair.Key, StringComparer.OrdinalIgnoreCase))
            {
                builder.Append(kvp.Key).Append('=');
                foreach (var id in kvp.Value)
                {
                    builder.Append(id).Append(',');
                }
                builder.Append('|');
            }
        }

        AppendMap(contains);
        AppendMap(usedBy);

        var hash = SHA256.HashData(Encoding.UTF8.GetBytes(builder.ToString()));
        return "snap-" + Convert.ToHexString(hash).ToLowerInvariant();
    }

    private static bool MatchesTagPattern(string tag, string pattern)
    {
        if (string.IsNullOrWhiteSpace(pattern))
        {
            return false;
        }

        if (pattern == "*")
        {
            return true;
        }

        if (!pattern.Contains('*') && !pattern.Contains('?'))
        {
            return string.Equals(tag, pattern, StringComparison.OrdinalIgnoreCase);
        }

        var escaped = Regex.Escape(pattern)
            .Replace("\\*", ".*")
            .Replace("\\?", ".");
        return Regex.IsMatch(tag, $"^{escaped}$", RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
    }

    private int EnsureImageId(string digest)
    {
        if (_imageIds.TryGetValue(digest, out var existing))
        {
            return existing;
        }

        var candidate = ComputeDeterministicId(digest);
        while (_images.ContainsKey(candidate))
        {
            candidate = (candidate + 1) & int.MaxValue;
            if (candidate == 0)
            {
                candidate = 1;
            }
        }

        _imageIds[digest] = candidate;
        return candidate;
    }

    private static int ComputeDeterministicId(string digest)
    {
        var bytes = SHA256.HashData(Encoding.UTF8.GetBytes(digest));
        for (var offset = 0; offset <= bytes.Length - sizeof(int); offset += sizeof(int))
        {
            var value = BinaryPrimitives.ReadInt32LittleEndian(bytes.AsSpan(offset, sizeof(int))) & int.MaxValue;
            if (value != 0)
            {
                return value;
            }
        }

        return digest.GetHashCode(StringComparison.OrdinalIgnoreCase) & int.MaxValue;
    }
}