git.stella-ops.org/src/StellaOps.Scheduler.Worker/ImpactTargetingService.cs

using System.Collections.Generic;
using System.Collections.Immutable;
using System.Text.RegularExpressions;
using System.Linq;
using StellaOps.Scheduler.ImpactIndex;
using StellaOps.Scheduler.Models;

namespace StellaOps.Scheduler.Worker;

public interface IImpactTargetingService
{
    ValueTask<ImpactSet> ResolveByPurlsAsync(
        IEnumerable<string> productKeys,
        bool usageOnly,
        Selector selector,
        CancellationToken cancellationToken = default);

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

    ValueTask<ImpactSet> ResolveAllAsync(
        Selector selector,
        bool usageOnly,
        CancellationToken cancellationToken = default);
}

public sealed class ImpactTargetingService : IImpactTargetingService
{
    private readonly IImpactIndex _impactIndex;
    private readonly TimeProvider _timeProvider;

    public ImpactTargetingService(IImpactIndex impactIndex, TimeProvider? timeProvider = null)
    {
        _impactIndex = impactIndex ?? throw new ArgumentNullException(nameof(impactIndex));
        _timeProvider = timeProvider ?? TimeProvider.System;
    }

    public async ValueTask<ImpactSet> ResolveByPurlsAsync(
        IEnumerable<string> productKeys,
        bool usageOnly,
        Selector selector,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(productKeys);
        ArgumentNullException.ThrowIfNull(selector);

        var distinct = productKeys
            .Where(static key => !string.IsNullOrWhiteSpace(key))
            .Select(static key => key.Trim())
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToArray();

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

        var impactSet = await _impactIndex.ResolveByPurlsAsync(distinct, usageOnly, selector, cancellationToken).ConfigureAwait(false);
        return SanitizeImpactSet(impactSet, selector);
    }

    public async ValueTask<ImpactSet> ResolveByVulnerabilitiesAsync(
        IEnumerable<string> vulnerabilityIds,
        bool usageOnly,
        Selector selector,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(vulnerabilityIds);
        ArgumentNullException.ThrowIfNull(selector);

        var distinct = vulnerabilityIds
            .Where(static id => !string.IsNullOrWhiteSpace(id))
            .Select(static id => id.Trim())
            .Distinct(StringComparer.OrdinalIgnoreCase)
            .ToArray();

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

        var impactSet = await _impactIndex.ResolveByVulnerabilitiesAsync(distinct, usageOnly, selector, cancellationToken).ConfigureAwait(false);
        return SanitizeImpactSet(impactSet, selector);
    }

    public async ValueTask<ImpactSet> ResolveAllAsync(
        Selector selector,
        bool usageOnly,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(selector);

        var impactSet = await _impactIndex.ResolveAllAsync(selector, usageOnly, cancellationToken).ConfigureAwait(false);
        return SanitizeImpactSet(impactSet, selector);
    }

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

    private static ImpactSet SanitizeImpactSet(ImpactSet impactSet, Selector selector)
    {
        ArgumentNullException.ThrowIfNull(impactSet);
        ArgumentNullException.ThrowIfNull(selector);

        if (impactSet.Images.Length == 0)
        {
            return impactSet;
        }

        var filteredImages = FilterAndDeduplicate(impactSet.Images, selector);
        if (filteredImages.Length == impactSet.Images.Length && filteredImages.SequenceEqual(impactSet.Images))
        {
            return impactSet;
        }

        return new ImpactSet(
            impactSet.Selector,
            filteredImages,
            impactSet.UsageOnly,
            impactSet.GeneratedAt,
            impactSet.Total,
            impactSet.SnapshotId,
            impactSet.SchemaVersion);
    }

    private static ImmutableArray<ImpactImage> FilterAndDeduplicate(
        IReadOnlyList<ImpactImage> images,
        Selector selector)
    {
        var digestFilter = selector.Digests.Length == 0
            ? null
            : new HashSet<string>(selector.Digests, StringComparer.OrdinalIgnoreCase);
        var namespaceFilter = selector.Namespaces.Length == 0
            ? null
            : new HashSet<string>(selector.Namespaces, StringComparer.Ordinal);
        var repositoryFilter = selector.Repositories.Length == 0
            ? null
            : new HashSet<string>(selector.Repositories, StringComparer.Ordinal);
        var tagMatchers = BuildTagMatchers(selector.IncludeTags);
        var labelFilters = BuildLabelFilters(selector.Labels);

        var filtered = new List<ImpactImage>(images.Count);
        foreach (var image in images)
        {
            if (image is null)
            {
                continue;
            }

            if (!MatchesSelector(image, digestFilter, namespaceFilter, repositoryFilter, tagMatchers, labelFilters))
            {
                continue;
            }

            filtered.Add(image);
        }

        if (filtered.Count == 0)
        {
            return ImmutableArray<ImpactImage>.Empty;
        }

        return DeduplicateByDigest(filtered);
    }

    private static bool MatchesSelector(
        ImpactImage image,
        HashSet<string>? digestFilter,
        HashSet<string>? namespaceFilter,
        HashSet<string>? repositoryFilter,
        IReadOnlyList<Func<string, bool>> tagMatchers,
        IReadOnlyList<LabelFilter> labelFilters)
    {
        if (digestFilter is not null && !digestFilter.Contains(image.ImageDigest))
        {
            return false;
        }

        if (namespaceFilter is not null)
        {
            var matchesNamespace = image.Namespaces.Any(namespaceFilter.Contains);
            if (!matchesNamespace)
            {
                return false;
            }
        }

        if (repositoryFilter is not null && !repositoryFilter.Contains(image.Repository))
        {
            return false;
        }

        if (tagMatchers.Count > 0)
        {
            var tagMatches = image.Tags.Any(tag => tagMatchers.Any(matcher => matcher(tag)));
            if (!tagMatches)
            {
                return false;
            }
        }

        if (labelFilters.Count > 0)
        {
            foreach (var labelFilter in labelFilters)
            {
                if (!image.Labels.TryGetValue(labelFilter.Key, out var value))
                {
                    return false;
                }

                if (labelFilter.AcceptedValues is not null && !labelFilter.AcceptedValues.Contains(value))
                {
                    return false;
                }
            }
        }

        return true;
    }

    private static IReadOnlyList<Func<string, bool>> BuildTagMatchers(ImmutableArray<string> includeTags)
    {
        if (includeTags.Length == 0)
        {
            return Array.Empty<Func<string, bool>>();
        }

        var matchers = new List<Func<string, bool>>(includeTags.Length);
        foreach (var pattern in includeTags)
        {
            if (string.IsNullOrWhiteSpace(pattern))
            {
                continue;
            }

            matchers.Add(CreateTagMatcher(pattern));
        }

        return matchers;
    }

    private static Func<string, bool> CreateTagMatcher(string pattern)
    {
        if (pattern == "*")
        {
            return static _ => true;
        }

        if (!pattern.Contains('*', StringComparison.Ordinal))
        {
            return tag => string.Equals(tag, pattern, StringComparison.OrdinalIgnoreCase);
        }

        var regexPattern = "^" + Regex.Escape(pattern).Replace("\\*", ".*", StringComparison.Ordinal) + "$";
        var regex = new Regex(regexPattern, RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);
        return tag => regex.IsMatch(tag);
    }

    private static IReadOnlyList<LabelFilter> BuildLabelFilters(ImmutableArray<LabelSelector> labelSelectors)
    {
        if (labelSelectors.Length == 0)
        {
            return Array.Empty<LabelFilter>();
        }

        var filters = new List<LabelFilter>(labelSelectors.Length);
        foreach (var selector in labelSelectors)
        {
            var key = selector.Key.ToLowerInvariant();
            HashSet<string>? values = null;
            if (selector.Values.Length > 0)
            {
                values = new HashSet<string>(selector.Values, StringComparer.OrdinalIgnoreCase);
            }

            filters.Add(new LabelFilter(key, values));
        }

        return filters;
    }

    private static ImmutableArray<ImpactImage> DeduplicateByDigest(IEnumerable<ImpactImage> images)
    {
        var aggregators = new Dictionary<string, ImpactImageAggregator>(StringComparer.OrdinalIgnoreCase);

        foreach (var image in images)
        {
            if (!aggregators.TryGetValue(image.ImageDigest, out var aggregator))
            {
                aggregator = new ImpactImageAggregator(image.ImageDigest);
                aggregators.Add(image.ImageDigest, aggregator);
            }

            aggregator.Add(image);
        }

        return aggregators.Values
            .Select(static aggregator => aggregator.Build())
            .OrderBy(static image => image.ImageDigest, StringComparer.OrdinalIgnoreCase)
            .ToImmutableArray();
    }

    private sealed record LabelFilter(string Key, HashSet<string>? AcceptedValues);

    private sealed class ImpactImageAggregator
    {
        private readonly string _digest;
        private readonly SortedSet<string> _registries = new(StringComparer.Ordinal);
        private readonly SortedSet<string> _repositories = new(StringComparer.Ordinal);
        private readonly SortedSet<string> _namespaces = new(StringComparer.Ordinal);
        private readonly SortedSet<string> _tags = new(StringComparer.OrdinalIgnoreCase);
        private readonly SortedDictionary<string, string> _labels = new(StringComparer.Ordinal);
        private bool _usedByEntrypoint;

        public ImpactImageAggregator(string digest)
        {
            _digest = digest;
        }

        public void Add(ImpactImage image)
        {
            _registries.Add(image.Registry);
            _repositories.Add(image.Repository);

            foreach (var ns in image.Namespaces)
            {
                _namespaces.Add(ns);
            }

            foreach (var tag in image.Tags)
            {
                _tags.Add(tag);
            }

            foreach (var label in image.Labels)
            {
                _labels[label.Key] = label.Value;
            }

            _usedByEntrypoint |= image.UsedByEntrypoint;
        }

        public ImpactImage Build()
        {
            var registry = _registries.Count > 0 ? _registries.Min! : string.Empty;
            var repository = _repositories.Count > 0 ? _repositories.Min! : string.Empty;

            var namespaces = _namespaces.Count == 0 ? Enumerable.Empty<string>() : _namespaces;
            var tags = _tags.Count == 0 ? Enumerable.Empty<string>() : _tags;

            return new ImpactImage(
                _digest,
                registry,
                repository,
                namespaces,
                tags,
                _usedByEntrypoint,
                _labels);
        }
    }
}