lmxopcua/tests/Server/ZB.MOM.WW.OtOpcUa.ControlPlane.Tests/Communication/MeshPskProviderRefreshTests.cs

using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using Shouldly;
using Xunit;
using ZB.MOM.WW.OtOpcUa.Cluster;
using ZB.MOM.WW.OtOpcUa.ControlPlane.Communication;
using ZB.MOM.WW.Secrets.Abstractions;

namespace ZB.MOM.WW.OtOpcUa.ControlPlane.Tests.Communication;

/// <summary>
///     R03:3.4 — central invalidated every cached key on every contact refresh (default 60 s), so the
///     first inbound ack per cluster per minute blocked a gRPC thread on a synchronous, 3 s-bounded store
///     read and was denied past the bound. <see cref="MeshPskProvider.RefreshAsync"/> pre-resolves the
///     keys off the request path and swaps them in atomically.
/// </summary>
public sealed class MeshPskProviderRefreshTests
{
    private static MeshPskProvider Provider(ISecretResolver resolver) =>
        new(Options.Create(new MeshTransportOptions()), resolver, NullLogger<MeshPskProvider>.Instance);

    [Fact]
    public async Task After_a_refresh_the_lookup_is_served_without_touching_the_store()
    {
        var resolver = new GatedResolver { Value = "k-1" };
        var provider = Provider(resolver);

        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);
        resolver.Calls.ShouldBe(1);

        provider.TryGetPsk("SITE-A", out var psk).ShouldBeTrue();
        psk.ShouldBe("k-1");
        resolver.Calls.ShouldBe(1, "the request path did not consult the store");
        provider.Matches("SITE-A", "k-1").ShouldBeTrue();
    }

    /// <summary>
    ///     While a refresh is in flight the PREVIOUS key is still served (no invalidate-then-reread window);
    ///     once it completes the rotated key is served. Positive control on both sides of the swap.
    /// </summary>
    [Fact]
    public async Task A_rotated_key_is_swapped_in_atomically_with_the_old_key_served_until_then()
    {
        var resolver = new GatedResolver { Value = "old" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);
        provider.TryGetPsk("SITE-A", out var before).ShouldBeTrue();
        before.ShouldBe("old");

        resolver.Value = "new";
        resolver.Hold();
        var refresh = provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);
        resolver.Entered.Wait(TimeSpan.FromSeconds(5)).ShouldBeTrue();

        // Mid-refresh: the old key, from memory — the store is NOT re-read on this path.
        var callsMidRefresh = resolver.Calls;
        provider.TryGetPsk("SITE-A", out var during).ShouldBeTrue();
        during.ShouldBe("old");
        resolver.Calls.ShouldBe(callsMidRefresh);

        resolver.Release();
        await refresh;
        provider.TryGetPsk("SITE-A", out var after).ShouldBeTrue();
        after.ShouldBe("new");
    }

    [Fact]
    public async Task A_cluster_whose_refresh_fails_keeps_its_previous_key()
    {
        var resolver = new GatedResolver { Value = "k-1" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        resolver.Throw = new InvalidOperationException("store offline");
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        provider.TryGetPsk("SITE-A", out var psk).ShouldBeTrue();
        psk.ShouldBe("k-1");
    }

    /// <summary>
    ///     Fix round 1 / I8 — "the store says there is no such secret" is a REVOCATION, not a read failure:
    ///     the previous key must stop being served (from the refreshed snapshot AND the lazy cache). A thrown
    ///     read keeps the previous key (the test above); a definitive not-found drops it.
    /// </summary>
    [Fact]
    public async Task A_secret_deleted_from_the_store_revokes_the_key_on_the_next_refresh()
    {
        var resolver = new GatedResolver { Value = "k-1" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);
        provider.Matches("SITE-A", "k-1").ShouldBeTrue("positive control: the key was in service");

        resolver.Value = null;   // the secret was deleted
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        provider.TryGetPsk("SITE-A", out _).ShouldBeFalse();
        provider.Matches("SITE-A", "k-1").ShouldBeFalse();
    }

    /// <summary>The same revocation reaches a key that was first served from the LAZY cache, not a refresh.</summary>
    [Fact]
    public async Task A_deleted_secret_also_evicts_a_lazily_cached_key()
    {
        var resolver = new GatedResolver { Value = "k-lazy" };
        var provider = Provider(resolver);
        provider.TryGetPsk("SITE-A", out var lazy).ShouldBeTrue();
        lazy.ShouldBe("k-lazy");

        resolver.Value = null;
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        provider.TryGetPsk("SITE-A", out _).ShouldBeFalse();
    }

    /// <summary>Fix round 1 minor — the refresh fans out to the store with BOUNDED concurrency, not one read per cluster at once.</summary>
    [Fact]
    public async Task Refresh_store_reads_are_concurrency_bounded()
    {
        var resolver = new ConcurrencyTrackingResolver();
        var provider = Provider(resolver);
        var clusters = Enumerable.Range(0, 40).Select(i => $"SITE-{i}").ToList();

        await provider.RefreshAsync(clusters, CancellationToken.None);

        resolver.Calls.ShouldBe(40);
        resolver.MaxConcurrent.ShouldBeLessThanOrEqualTo(MeshPskProvider.MaxConcurrentRefreshReads);
        resolver.MaxConcurrent.ShouldBeGreaterThan(1, "still concurrent, just bounded");
        foreach (var c in clusters) provider.TryGetPsk(c, out _).ShouldBeTrue();
    }

    /// <summary>
    ///     lmxopcua#784 — the docker-dev centrals supply every cluster's key through
    ///     <c>MeshTransport:ClusterPsks</c> and seed nothing in the store, and logged "absent in the store; the
    ///     key is REVOKED — commands … refused" for all three clusters on every contact refresh while every
    ///     dispatch and ack authenticated. The override wins on every lookup, so the key was never revoked: the
    ///     warning was false. The refresh now leaves override-covered clusters alone. SITE-B (no override) is the
    ///     positive control — a cluster with no key anywhere still gets the REVOKED warning and is refused.
    /// </summary>
    [Fact]
    public async Task A_cluster_keyed_by_the_config_override_is_not_reported_revoked_and_stays_authenticated()
    {
        var resolver = new GatedResolver { Value = null };   // the store holds nothing, as on the rig
        var options = new MeshTransportOptions();
        options.ClusterPsks["SITE-A"] = "rig-psk-site-a";
        var logger = new RecordingLogger();
        var provider = new MeshPskProvider(Options.Create(options), resolver, logger);

        await provider.RefreshAsync(new[] { "site-a", "SITE-B" }, CancellationToken.None);

        resolver.Calls.ShouldBe(1, "only SITE-B, which has no override, is read from the store");
        var revoked = logger.Entries.Where(e => e.Level == LogLevel.Warning && e.Message.Contains("REVOKED")).ToList();
        revoked.Count.ShouldBe(1);
        revoked[0].Message.ShouldContain("SITE-B");

        // Enforcement is real for both: SITE-A authenticates against the override key and nothing else,
        // SITE-B (no key anywhere) is refused.
        provider.Matches("SITE-A", "rig-psk-site-a").ShouldBeTrue();
        provider.Matches("SITE-A", "wrong").ShouldBeFalse();
        provider.TryGetPsk("SITE-B", out _).ShouldBeFalse();
    }

    private sealed class RecordingLogger : ILogger<MeshPskProvider>
    {
        public List<(LogLevel Level, string Message)> Entries { get; } = new();

        public IDisposable? BeginScope<TState>(TState state) where TState : notnull => null;

        public bool IsEnabled(LogLevel logLevel) => true;

        public void Log<TState>(
            LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
        {
            lock (Entries) Entries.Add((logLevel, formatter(state, exception)));
        }
    }

    private sealed class ConcurrencyTrackingResolver : ISecretResolver
    {
        private int _inFlight;
        private int _max;
        private int _calls;
        public int MaxConcurrent => Volatile.Read(ref _max);
        public int Calls => Volatile.Read(ref _calls);

        public async Task<string?> GetAsync(SecretName name, CancellationToken ct)
        {
            Interlocked.Increment(ref _calls);
            var now = Interlocked.Increment(ref _inFlight);
            int seen;
            while (now > (seen = Volatile.Read(ref _max)) && Interlocked.CompareExchange(ref _max, now, seen) != seen) { }
            await Task.Delay(20, ct);
            Interlocked.Decrement(ref _inFlight);
            return "k-" + name.Value;
        }
    }

    /// <summary>A cluster that was never refreshed still resolves lazily on the request path (unchanged behaviour).</summary>
    [Fact]
    public async Task A_cluster_outside_the_refresh_set_still_resolves_lazily()
    {
        var resolver = new GatedResolver { Value = "k-b" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        provider.TryGetPsk("SITE-B", out var psk).ShouldBeTrue();
        psk.ShouldBe("k-b");
        resolver.Calls.ShouldBe(2);
    }

    [Fact]
    public async Task InvalidateAll_drops_refreshed_keys_too()
    {
        var resolver = new GatedResolver { Value = "k-1" };
        var provider = Provider(resolver);
        await provider.RefreshAsync(new[] { "SITE-A" }, CancellationToken.None);

        provider.InvalidateAll();
        resolver.Value = "k-2";
        provider.TryGetPsk("SITE-A", out var psk).ShouldBeTrue();
        psk.ShouldBe("k-2");
    }

    private sealed class GatedResolver : ISecretResolver
    {
        private int _calls;
        private volatile TaskCompletionSource? _gate;

        public string? Value { get; set; }
        public Exception? Throw { get; set; }
        public int Calls => Volatile.Read(ref _calls);
        public ManualResetEventSlim Entered { get; } = new();

        public void Hold() => _gate = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        public void Release() => _gate?.TrySetResult();

        public async Task<string?> GetAsync(SecretName name, CancellationToken ct)
        {
            Interlocked.Increment(ref _calls);
            Entered.Set();
            if (_gate is { } gate) await gate.Task.WaitAsync(ct);
            if (Throw is not null) throw Throw;
            return Value;
        }
    }
}