natsdotnet/tests/NATS.Server.Auth.Tests/ImportExportTests.cs

using Microsoft.Extensions.Logging.Abstractions;
using NATS.Server;
using NATS.Server.Auth;
using NATS.Server.Imports;
using NATS.Server.Subscriptions;

using NATS.Server.TestUtilities;

namespace NATS.Server.Auth.Tests;

public class ImportExportTests
{
    [Fact]
    public void ExportAuth_public_export_authorizes_any_account()
    {
        var auth = new ExportAuth();
        var account = new Account("test");
        auth.IsAuthorized(account).ShouldBeTrue();
    }

    [Fact]
    public void ExportAuth_approved_accounts_restricts_access()
    {
        var auth = new ExportAuth { ApprovedAccounts = ["allowed"] };
        var allowed = new Account("allowed");
        var denied = new Account("denied");
        auth.IsAuthorized(allowed).ShouldBeTrue();
        auth.IsAuthorized(denied).ShouldBeFalse();
    }

    [Fact]
    public void ExportAuth_revoked_account_denied()
    {
        var auth = new ExportAuth
        {
            ApprovedAccounts = ["test"],
            RevokedAccounts = new() { ["test"] = DateTimeOffset.UtcNow.ToUnixTimeSeconds() },
        };
        var account = new Account("test");
        auth.IsAuthorized(account).ShouldBeFalse();
    }

    [Fact]
    public void ServiceResponseType_defaults_to_singleton()
    {
        var import = new ServiceImport
        {
            DestinationAccount = new Account("dest"),
            From = "requests.>",
            To = "api.>",
        };
        import.ResponseType.ShouldBe(ServiceResponseType.Singleton);
    }

    [Fact]
    public void ExportMap_stores_and_retrieves_exports()
    {
        var map = new ExportMap();
        map.Services["api.>"] = new ServiceExport { Account = new Account("svc") };
        map.Streams["events.>"] = new StreamExport();

        map.Services.ShouldContainKey("api.>");
        map.Streams.ShouldContainKey("events.>");
    }

    [Fact]
    public void ImportMap_stores_service_imports()
    {
        var map = new ImportMap();
        var si = new ServiceImport
        {
            DestinationAccount = new Account("dest"),
            From = "requests.>",
            To = "api.>",
        };
        map.AddServiceImport(si);
        map.Services.ShouldContainKey("requests.>");
        map.Services["requests.>"].Count.ShouldBe(1);
    }

    [Fact]
    public void Account_add_service_export_and_import()
    {
        var exporter = new Account("exporter");
        var importer = new Account("importer");

        exporter.AddServiceExport("api.>", ServiceResponseType.Singleton, null);
        exporter.Exports.Services.ShouldContainKey("api.>");

        var si = importer.AddServiceImport(exporter, "requests.>", "api.>");
        si.ShouldNotBeNull();
        si.From.ShouldBe("requests.>");
        si.To.ShouldBe("api.>");
        si.DestinationAccount.ShouldBe(exporter);
        importer.Imports.Services.ShouldContainKey("requests.>");
    }

    [Fact]
    public void Account_add_stream_export_and_import()
    {
        var exporter = new Account("exporter");
        var importer = new Account("importer");

        exporter.AddStreamExport("events.>", null);
        exporter.Exports.Streams.ShouldContainKey("events.>");

        importer.AddStreamImport(exporter, "events.>", "imported.events.>");
        importer.Imports.Streams.Count.ShouldBe(1);
        importer.Imports.Streams[0].From.ShouldBe("events.>");
        importer.Imports.Streams[0].To.ShouldBe("imported.events.>");
    }

    [Fact]
    public void Account_service_import_auth_rejected()
    {
        var exporter = new Account("exporter");
        var importer = new Account("importer");

        exporter.AddServiceExport("api.>", ServiceResponseType.Singleton, [new Account("other")]);

        Should.Throw<UnauthorizedAccessException>(() =>
            importer.AddServiceImport(exporter, "requests.>", "api.>"));
    }

    [Fact]
    public void Account_lazy_creates_internal_client()
    {
        var account = new Account("test");
        var client = account.GetOrCreateInternalClient(99);
        client.ShouldNotBeNull();
        client.Kind.ShouldBe(ClientKind.Account);
        client.Account.ShouldBe(account);

        // Second call returns same instance
        var client2 = account.GetOrCreateInternalClient(100);
        client2.ShouldBeSameAs(client);
    }

    [Fact]
    public async Task Service_import_forwards_message_to_export_account()
    {
        using var server = CreateTestServer();
        _ = server.StartAsync(CancellationToken.None);
        await server.WaitForReadyAsync();

        // Set up exporter and importer accounts
        var exporter = server.GetOrCreateAccount("exporter");
        var importer = server.GetOrCreateAccount("importer");

        exporter.AddServiceExport("api.>", ServiceResponseType.Singleton, null);
        importer.AddServiceImport(exporter, "requests.>", "api.>");

        // Wire the import subscriptions into the importer account
        server.WireServiceImports(importer);

        // Subscribe in exporter account to receive forwarded message
        var exportSub = new Subscription { Subject = "api.test", Sid = "export-1", Client = null };
        exporter.SubList.Insert(exportSub);

        // Verify import infrastructure is wired: the importer should have service import entries
        importer.Imports.Services.ShouldContainKey("requests.>");
        importer.Imports.Services["requests.>"].Count.ShouldBe(1);
        importer.Imports.Services["requests.>"][0].DestinationAccount.ShouldBe(exporter);

        await server.ShutdownAsync();
    }

    [Fact]
    public void ProcessServiceImport_delivers_to_destination_account_subscribers()
    {
        using var server = CreateTestServer();

        var exporter = server.GetOrCreateAccount("exporter");
        var importer = server.GetOrCreateAccount("importer");

        exporter.AddServiceExport("api.>", ServiceResponseType.Singleton, null);
        importer.AddServiceImport(exporter, "requests.>", "api.>");

        // Add a subscriber in the exporter account's SubList
        var received = new List<(string Subject, string Sid)>();
        var mockClient = new TestNatsClient(1, exporter);
        mockClient.OnMessage = (subject, sid, _, _, _) =>
            received.Add((subject, sid));

        var exportSub = new Subscription { Subject = "api.test", Sid = "s1", Client = mockClient };
        exporter.SubList.Insert(exportSub);

        // Process a service import directly
        var si = importer.Imports.Services["requests.>"][0];
        server.ProcessServiceImport(si, "requests.test", null,
            ReadOnlyMemory<byte>.Empty, ReadOnlyMemory<byte>.Empty);

        received.Count.ShouldBe(1);
        received[0].Subject.ShouldBe("api.test");
        received[0].Sid.ShouldBe("s1");
    }

    [Fact]
    public void ProcessServiceImport_with_transform_applies_subject_mapping()
    {
        using var server = CreateTestServer();

        var exporter = server.GetOrCreateAccount("exporter");
        var importer = server.GetOrCreateAccount("importer");

        exporter.AddServiceExport("api.>", ServiceResponseType.Singleton, null);
        var si = importer.AddServiceImport(exporter, "requests.>", "api.>");

        // Create a transform from requests.> to api.>
        var transform = SubjectTransform.Create("requests.>", "api.>");
        transform.ShouldNotBeNull();

        // Create a new import with the transform set
        var siWithTransform = new ServiceImport
        {
            DestinationAccount = exporter,
            From = "requests.>",
            To = "api.>",
            Transform = transform,
        };

        var received = new List<string>();
        var mockClient = new TestNatsClient(1, exporter);
        mockClient.OnMessage = (subject, _, _, _, _) =>
            received.Add(subject);

        var exportSub = new Subscription { Subject = "api.hello", Sid = "s1", Client = mockClient };
        exporter.SubList.Insert(exportSub);

        server.ProcessServiceImport(siWithTransform, "requests.hello", null,
            ReadOnlyMemory<byte>.Empty, ReadOnlyMemory<byte>.Empty);

        received.Count.ShouldBe(1);
        received[0].ShouldBe("api.hello");
    }

    [Fact]
    public void ProcessServiceImport_skips_invalid_imports()
    {
        using var server = CreateTestServer();

        var exporter = server.GetOrCreateAccount("exporter");
        var importer = server.GetOrCreateAccount("importer");

        exporter.AddServiceExport("api.>", ServiceResponseType.Singleton, null);
        importer.AddServiceImport(exporter, "requests.>", "api.>");

        // Mark the import as invalid
        var si = importer.Imports.Services["requests.>"][0];
        si.Invalid = true;

        // Add a subscriber in the exporter account
        var received = new List<string>();
        var mockClient = new TestNatsClient(1, exporter);
        mockClient.OnMessage = (subject, _, _, _, _) =>
            received.Add(subject);

        var exportSub = new Subscription { Subject = "api.test", Sid = "s1", Client = mockClient };
        exporter.SubList.Insert(exportSub);

        // ProcessServiceImport should be a no-op for invalid imports
        server.ProcessServiceImport(si, "requests.test", null,
            ReadOnlyMemory<byte>.Empty, ReadOnlyMemory<byte>.Empty);

        received.Count.ShouldBe(0);
    }

    [Fact]
    public void ProcessServiceImport_delivers_to_queue_groups()
    {
        using var server = CreateTestServer();

        var exporter = server.GetOrCreateAccount("exporter");
        var importer = server.GetOrCreateAccount("importer");

        exporter.AddServiceExport("api.>", ServiceResponseType.Singleton, null);
        importer.AddServiceImport(exporter, "requests.>", "api.>");

        // Add queue group subscribers in the exporter account
        var received = new List<(string Subject, string Sid)>();
        var mockClient1 = new TestNatsClient(1, exporter);
        mockClient1.OnMessage = (subject, sid, _, _, _) =>
            received.Add((subject, sid));
        var mockClient2 = new TestNatsClient(2, exporter);
        mockClient2.OnMessage = (subject, sid, _, _, _) =>
            received.Add((subject, sid));

        var qSub1 = new Subscription { Subject = "api.test", Sid = "q1", Queue = "workers", Client = mockClient1 };
        var qSub2 = new Subscription { Subject = "api.test", Sid = "q2", Queue = "workers", Client = mockClient2 };
        exporter.SubList.Insert(qSub1);
        exporter.SubList.Insert(qSub2);

        var si = importer.Imports.Services["requests.>"][0];
        server.ProcessServiceImport(si, "requests.test", null,
            ReadOnlyMemory<byte>.Empty, ReadOnlyMemory<byte>.Empty);

        // One member of the queue group should receive the message
        received.Count.ShouldBe(1);
    }

    private static NatsServer CreateTestServer()
    {
        var port = TestPortAllocator.GetFreePort();
        return new NatsServer(new NatsOptions { Port = port }, NullLoggerFactory.Instance);
    }
    /// <summary>
    /// Minimal test double for INatsClient used in import/export tests.
    /// </summary>
    private sealed class TestNatsClient(ulong id, Account account) : INatsClient
    {
        public ulong Id => id;
        public ClientKind Kind => ClientKind.Client;
        public Account? Account => account;
        public Protocol.ClientOptions? ClientOpts => null;
        public ClientPermissions? Permissions => null;

        public Action<string, string, string?, ReadOnlyMemory<byte>, ReadOnlyMemory<byte>>? OnMessage { get; set; }

        public void SendMessage(string subject, string sid, string? replyTo,
            ReadOnlyMemory<byte> headers, ReadOnlyMemory<byte> payload)
        {
            OnMessage?.Invoke(subject, sid, replyTo, headers, payload);
        }

        public bool QueueOutbound(ReadOnlyMemory<byte> data) => true;

        public void RemoveSubscription(string sid) { }
    }
}