ScadaBridge/tests/ZB.MOM.WW.ScadaBridge.Host.Tests/ControlPlaneAuthEndToEndTests.cs

using Grpc.Core;
using Grpc.Net.Client;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.TestHost;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using NSubstitute;
using ZB.MOM.WW.ScadaBridge.Communication;
using ZB.MOM.WW.ScadaBridge.Communication.Grpc;
using ZB.MOM.WW.Secrets.Abstractions;

namespace ZB.MOM.WW.ScadaBridge.Host.Tests;

/// <summary>
/// End-to-end proof that the two halves of the control-plane PSK actually interoperate:
/// <see cref="ControlPlaneCredentials"/> on the client and
/// <see cref="ControlPlaneAuthInterceptor"/> on the server, over a real gRPC stack.
/// </summary>
/// <remarks>
/// <para>
/// The unit tests either side of this file each test one half against a hand-built input, and
/// would both stay green if the halves disagreed — if the credentials never attached to a
/// streaming call, if the metadata key case differed, or if attaching call credentials to a
/// plaintext channel were rejected outright (gRPC refuses that by default; the code opts in with
/// <c>UnsafeUseInsecureChannelCallCredentials</c>, and nothing but a real call proves the opt-in
/// works). Getting that wrong takes down every site's streaming and audit-pull path at once,
/// which is a bad thing to discover on the rig.
/// </para>
/// <para>
/// Runs entirely in-process over <see cref="TestServer"/>: no ports, no containers. The service
/// is a stub rather than the real <c>SiteStreamGrpcServer</c> — this is a test of the auth
/// pipeline, and the real server would drag in an actor system for no added coverage. The method
/// paths and message types are the real generated ones.
/// </para>
/// </remarks>
public class ControlPlaneAuthEndToEndTests : IAsyncLifetime
{
    private IHost _host = null!;
    private TestServer _server = null!;

    /// <summary>Boots the in-process gRPC host with the real interceptor.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async Task InitializeAsync()
    {
        _host = await new HostBuilder()
            .ConfigureWebHost(web => web
                .UseTestServer()
                .ConfigureServices(services =>
                {
                    // Registered exactly as Program.cs does: by TYPE on AddGrpc, with the
                    // interceptor itself NOT in DI. That is load-bearing. An earlier version of
                    // this test added it as a singleton, which let DI hand back the instance and
                    // bypassed Grpc.AspNetCore's own activation — hiding a defect where the
                    // interceptor had two public constructors and
                    // InterceptorRegistration.GetFactory() threw on every single call. The rig
                    // caught it; this test did not. Do not pre-register it.
                    services.AddGrpc(o => o.Interceptors.Add<ControlPlaneAuthInterceptor>());
                    services.AddSingleton(Options.Create(
                        new CommunicationOptions { GrpcPsk = SiteKey }));
                    services.AddSingleton<EchoSiteStreamService>();
                })
                .Configure(app =>
                {
                    app.UseRouting();
                    app.UseEndpoints(e => e.MapGrpcService<EchoSiteStreamService>());
                }))
            .StartAsync();

        _server = _host.GetTestServer();
    }

    /// <summary>Stops and disposes the in-process gRPC host.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    public async Task DisposeAsync()
    {
        await _host.StopAsync();
        _host.Dispose();
    }

    private const string SiteKey = "the-site-a-preshared-key";

    /// <summary>
    /// Builds a channel through the test server, credentialed exactly as production does.
    /// </summary>
    private GrpcChannel Channel(string? key, string siteId = "site-a")
    {
        var options = new GrpcChannelOptions { HttpHandler = _server.CreateHandler() }.WithMessageLimits();
        if (key is not null)
        {
            options.WithSiteCredentials(new FixedPskProvider(key), siteId);
        }
        return GrpcChannel.ForAddress(_server.BaseAddress, options);
    }

    private sealed class FixedPskProvider(string key) : ISitePskProvider
    {
        /// <inheritdoc />
        public ValueTask<string> GetAsync(string siteId, CancellationToken ct) => new(key);
        /// <inheritdoc />
        public void Invalidate(string siteId) { }
    }

    /// <summary>Stub service: echoes back what the auth pipeline let through.</summary>
    private sealed class EchoSiteStreamService : SiteStreamService.SiteStreamServiceBase
    {
        /// <summary>The site header the last accepted call carried.</summary>
        public string? LastSiteHeader { get; private set; }

        /// <inheritdoc />
        public override Task<PullAuditEventsResponse> PullAuditEvents(
            PullAuditEventsRequest request, ServerCallContext context)
        {
            LastSiteHeader = context.RequestHeaders
                .FirstOrDefault(h => h.Key == ControlPlaneCredentials.SiteHeader)?.Value;
            return Task.FromResult(new PullAuditEventsResponse { MoreAvailable = false });
        }

        /// <inheritdoc />
        public override async Task SubscribeInstance(
            InstanceStreamRequest request,
            IServerStreamWriter<SiteStreamEvent> responseStream,
            ServerCallContext context)
        {
            await responseStream.WriteAsync(new SiteStreamEvent
            {
                CorrelationId = request.CorrelationId == "large-audit-window"
                    ? new string('x', 17 * 1024 * 1024)
                    : request.CorrelationId,
            });
        }
    }

    /// <summary>Verifies that the control-plane channel can receive a worst-case audit-sized message.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CM1_ControlPlaneChannelReceivesWorstCaseAuditSizedMessage()
    {
        using var channel = Channel(SiteKey);
        var client = new SiteStreamService.SiteStreamServiceClient(channel);
        using var call = client.SubscribeInstance(new InstanceStreamRequest { CorrelationId = "large-audit-window" });
        Assert.True(await call.ResponseStream.MoveNext(CancellationToken.None));
        Assert.Equal(17 * 1024 * 1024, call.ResponseStream.Current.CorrelationId.Length);
    }

    /// <summary>Verifies that a call carrying the correct key is accepted on a unary call.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CorrectKey_IsAccepted_OnAUnaryCall()
    {
        using var channel = Channel(SiteKey);
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        var reply = await client.PullAuditEventsAsync(new PullAuditEventsRequest { BatchSize = 1 });

        Assert.False(reply.MoreAvailable);
    }

    /// <summary>Verifies that a call carrying the wrong key is rejected with a permission-denied status.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WrongKey_IsRejected_WithPermissionDenied()
    {
        using var channel = Channel("some-other-sites-key");
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        var ex = await Assert.ThrowsAsync<RpcException>(
            async () => await client.PullAuditEventsAsync(new PullAuditEventsRequest()));

        Assert.Equal(StatusCode.PermissionDenied, ex.StatusCode);
    }

    /// <summary>Verifies that a call with no credentials is rejected.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task NoCredentialsAtAll_IsRejected()
    {
        // The pre-T0.3 client shape. This is the case that proves the gap is actually closed.
        using var channel = Channel(key: null);
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        var ex = await Assert.ThrowsAsync<RpcException>(
            async () => await client.PullAuditEventsAsync(new PullAuditEventsRequest()));

        Assert.Equal(StatusCode.PermissionDenied, ex.StatusCode);
    }

    /// <summary>Verifies that the credentials are applied to streaming calls as well as unary calls.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task CredentialsApplyToStreamingCalls_NotJustUnaryOnes()
    {
        // CallCredentials cover every call on the channel; a client interceptor that only
        // handled the unary path would pass the test above and still break every subscription.
        using var channel = Channel(SiteKey);
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        using var call = client.SubscribeInstance(
            new InstanceStreamRequest { CorrelationId = "c1", InstanceUniqueName = "i1" });

        Assert.True(await call.ResponseStream.MoveNext(CancellationToken.None));
        Assert.Equal("c1", call.ResponseStream.Current.CorrelationId);
    }

    /// <summary>Verifies that the wrong key is also rejected on streaming calls.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task WrongKey_IsRejected_OnStreamingCallsToo()
    {
        using var channel = Channel("wrong");
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        using var call = client.SubscribeInstance(new InstanceStreamRequest { CorrelationId = "c1" });

        var ex = await Assert.ThrowsAsync<RpcException>(
            async () => await call.ResponseStream.MoveNext(CancellationToken.None));

        Assert.Equal(StatusCode.PermissionDenied, ex.StatusCode);
    }

    /// <summary>Verifies that a key that cannot be resolved fails the call instead of sending it without credentials.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task AnUnresolvableKey_FailsTheCall_RatherThanDialingWithoutOne()
    {
        // SitePskProvider throws when a site has no key anywhere. What matters here is that the
        // throw stops the call: the alternative — swallowing it and sending the request
        // unauthenticated — is the exact failure this design exists to prevent. The status code
        // is gRPC's choice, so assert the RpcException and record what it actually is rather
        // than pinning a guess: callers already treat every non-OK status as a failed call, and
        // the diagnosable signal is SitePskProvider's own LogError, not this code.
        var options = new GrpcChannelOptions { HttpHandler = _server.CreateHandler() }
            .WithSiteCredentials(new ThrowingPskProvider(), "site-a");
        using var channel = GrpcChannel.ForAddress(_server.BaseAddress, options);
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        var ex = await Assert.ThrowsAsync<RpcException>(
            async () => await client.PullAuditEventsAsync(new PullAuditEventsRequest()));

        Assert.NotEqual(StatusCode.OK, ex.StatusCode);

        // And nothing reached the service.
        Assert.Null(_host.Services.GetRequiredService<EchoSiteStreamService>().LastSiteHeader);
    }

    private sealed class ThrowingPskProvider : ISitePskProvider
    {
        /// <inheritdoc />
        public ValueTask<string> GetAsync(string siteId, CancellationToken ct)
            => throw new InvalidOperationException($"no key for '{siteId}'");
        /// <inheritdoc />
        public void Invalidate(string siteId) { }
    }

    /// <summary>Verifies that the site header is sent with the call so central can choose a per-site key.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task TheSiteHeaderTravels_SoCentralCanPickAPerSiteKeyInPhase1A()
    {
        // Central's own interceptor (T1A.2) verifies against the key for the site named in this
        // header. Shipping it now means Phase 1A adds a lookup, not a wire change.
        using var channel = Channel(SiteKey, siteId: "site-a");
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        await client.PullAuditEventsAsync(new PullAuditEventsRequest());

        var service = _host.Services.GetRequiredService<EchoSiteStreamService>();
        Assert.Equal("site-a", service.LastSiteHeader);
    }
    // --- Gitea #96: a site's rejection drops central's cached key ---------------------------

    /// <summary>
    /// Central's real <see cref="SitePskProvider"/> over a store that served a stale key first and
    /// the rotated (site-accepted) key after: the in-place rotation scenario.
    /// </summary>
    private static (SitePskProvider Provider, ISecretResolver Resolver) RotatedProvider()
    {
        var resolver = Substitute.For<ISecretResolver>();
        resolver.GetAsync(Arg.Any<SecretName>(), Arg.Any<CancellationToken>())
            .Returns("the-pre-rotation-key", SiteKey);
        var options = Substitute.For<IOptionsMonitor<CommunicationOptions>>();
        options.CurrentValue.Returns(new CommunicationOptions());
        return (new SitePskProvider(resolver, options, NullLogger<SitePskProvider>.Instance), resolver);
    }

    private GrpcChannel Channel(ISitePskProvider provider) =>
        GrpcChannel.ForAddress(_server.BaseAddress, new GrpcChannelOptions { HttpHandler = _server.CreateHandler() }
            .WithMessageLimits().WithSiteCredentials(provider, "site-a"));

    /// <summary>
    /// Verifies that after a site rejects the cached (pre-rotation) key on a unary call, the next
    /// call re-reads the key and is accepted — no restart needed.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task RejectedUnaryCall_DropsTheCachedKey_SoTheNextCallUsesTheRotatedOne()
    {
        var (provider, _) = RotatedProvider();
        using var channel = Channel(provider);
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        var ex = await Assert.ThrowsAsync<RpcException>(
            async () => await client.PullAuditEventsAsync(new PullAuditEventsRequest()));
        Assert.Equal(StatusCode.PermissionDenied, ex.StatusCode);

        var reply = await client.PullAuditEventsAsync(new PullAuditEventsRequest { BatchSize = 1 });
        Assert.False(reply.MoreAvailable);
    }

    /// <summary>Verifies the same refresh for a rejected streaming call.</summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task RejectedStreamingCall_DropsTheCachedKey_SoTheNextCallUsesTheRotatedOne()
    {
        var (provider, _) = RotatedProvider();
        using var channel = Channel(provider);
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        using (var rejected = client.SubscribeInstance(new InstanceStreamRequest { CorrelationId = "c1" }))
        {
            var ex = await Assert.ThrowsAsync<RpcException>(
                async () => await rejected.ResponseStream.MoveNext(CancellationToken.None));
            Assert.Equal(StatusCode.PermissionDenied, ex.StatusCode);
        }

        using var accepted = client.SubscribeInstance(new InstanceStreamRequest { CorrelationId = "c2" });
        Assert.True(await accepted.ResponseStream.MoveNext(CancellationToken.None));
        Assert.Equal("c2", accepted.ResponseStream.Current.CorrelationId);
    }

    /// <summary>
    /// Negative control: accepted calls never drop the cache, so the store is read once no
    /// matter how many calls succeed.
    /// </summary>
    /// <returns>A task that represents the asynchronous operation.</returns>
    [Fact]
    public async Task AcceptedCalls_KeepTheCachedKey()
    {
        var resolver = Substitute.For<ISecretResolver>();
        resolver.GetAsync(Arg.Any<SecretName>(), Arg.Any<CancellationToken>()).Returns(SiteKey);
        var options = Substitute.For<IOptionsMonitor<CommunicationOptions>>();
        options.CurrentValue.Returns(new CommunicationOptions());
        var provider = new SitePskProvider(resolver, options, NullLogger<SitePskProvider>.Instance);
        using var channel = Channel(provider);
        var client = new SiteStreamService.SiteStreamServiceClient(channel);

        await client.PullAuditEventsAsync(new PullAuditEventsRequest());
        await client.PullAuditEventsAsync(new PullAuditEventsRequest());
        await client.PullAuditEventsAsync(new PullAuditEventsRequest());

        await resolver.Received(1).GetAsync(Arg.Any<SecretName>(), Arg.Any<CancellationToken>());
    }
}