nomyo-js/tests/unit/secure_client.test.ts

/**
 * Unit tests for SecureCompletionClient (with mocked HTTP)
 */

import { SecureCompletionClient } from '../../src/core/SecureCompletionClient';
import {
    SecurityError,
    APIConnectionError,
    AuthenticationError,
    ForbiddenError,
    ServiceUnavailableError,
    RateLimitError,
} from '../../src/errors';
import { stringToArrayBuffer } from '../../src/core/crypto/utils';

// ---- helpers ---------------------------------------------------------------

function makeJsonResponse(statusCode: number, body: object): { statusCode: number; headers: Record<string, string>; body: ArrayBuffer } {
    return {
        statusCode,
        headers: { 'content-type': 'application/json' },
        body: stringToArrayBuffer(JSON.stringify(body)),
    };
}

function mockHttpClient(handler: (url: string, opts: unknown) => Promise<{ statusCode: number; headers: Record<string, string>; body: ArrayBuffer }>) {
    return {
        post: jest.fn((url: string, opts: unknown) => handler(url, opts)),
        get: jest.fn((url: string, _opts?: unknown) => handler(url, _opts)),
    };
}

// ---------------------------------------------------------------------------

describe('SecureCompletionClient constructor', () => {
    test('warns about HTTP if allowHttp is false', () => {
        const spy = jest.spyOn(console, 'warn').mockImplementation(() => undefined);
        new SecureCompletionClient({ routerUrl: 'http://localhost:1234' });
        expect(spy).toHaveBeenCalledWith(expect.stringContaining('INSECURE'));
        spy.mockRestore();
    });

    test('does not warn when allowHttp is true', () => {
        const spy = jest.spyOn(console, 'warn').mockImplementation(() => undefined);
        new SecureCompletionClient({ routerUrl: 'http://localhost:1234', allowHttp: true });
        // Only the HTTP-mode log should appear, not a warning
        expect(spy).not.toHaveBeenCalledWith(expect.stringContaining('Man-in-the-middle'));
        spy.mockRestore();
    });

    test('removes trailing slash from routerUrl', () => {
        const client = new SecureCompletionClient({ routerUrl: 'https://api.example.com:12434/' });
        // We can verify indirectly via fetchServerPublicKey URL construction
        expect((client as unknown as { routerUrl: string }).routerUrl).toBe('https://api.example.com:12434');
    });
});

describe('SecureCompletionClient.fetchServerPublicKey', () => {
    test('throws SecurityError over HTTP without allowHttp', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'http://localhost:1234',
            allowHttp: false,
        });
        // Suppress console.warn from constructor
        jest.spyOn(console, 'warn').mockImplementation(() => undefined);
        await expect(client.fetchServerPublicKey()).rejects.toBeInstanceOf(SecurityError);
    });
});

describe('SecureCompletionClient.sendSecureRequest — security tier validation', () => {
    test('throws for invalid security tier', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com:12434',
        });

        await expect(
            client.sendSecureRequest({}, 'test-id', undefined, 'ultra')
        ).rejects.toThrow("Invalid securityTier: 'ultra'");
    });

    test('accepts valid security tiers', async () => {
        // We just need to verify no validation error is thrown at the tier check stage
        // (subsequent network call will fail, which is expected in unit tests)
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com:12434',
        });

        for (const tier of ['standard', 'high', 'maximum']) {
            // Should not throw a tier validation error (will throw something else)
            await expect(
                client.sendSecureRequest({}, 'test-id', undefined, tier)
            ).rejects.not.toThrow("Invalid securityTier");
        }
    });
});

describe('SecureCompletionClient.buildErrorFromResponse (via sendSecureRequest)', () => {
    // We can test error mapping by making the HTTP mock return specific status codes
    // and verifying the correct typed error is thrown.

    async function clientWithMockedHttp(statusCode: number, body: object) {
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com:12434',
        });

        // Inject mocked HTTP client
        const http = mockHttpClient(async (url: string) => {
            if (url.includes('/pki/public_key')) {
                // Should not be reached in error tests
                throw new Error('unexpected pki call');
            }
            return makeJsonResponse(statusCode, body);
        });
        (client as unknown as { httpClient: typeof http }).httpClient = http;

        return client;
    }

    test('401 → AuthenticationError', async () => {
        const client = await clientWithMockedHttp(401, { detail: 'bad key' });
        // Keys must be generated first, so inject a pre-generated key set
        await (client as unknown as { generateKeys: () => Promise<void> }).generateKeys();

        // Mock encryptPayload to skip actual encryption
        jest.spyOn(client as unknown as { encryptPayload: () => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({ model: 'test', messages: [] }, 'id-1')
        ).rejects.toBeInstanceOf(AuthenticationError);
    }, 30000);

    test('403 → ForbiddenError', async () => {
        const client = await clientWithMockedHttp(403, { detail: 'not allowed' });
        await (client as unknown as { generateKeys: () => Promise<void> }).generateKeys();
        jest.spyOn(client as unknown as { encryptPayload: () => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({ model: 'test', messages: [] }, 'id-1')
        ).rejects.toBeInstanceOf(ForbiddenError);
    }, 30000);

    test('429 → RateLimitError', async () => {
        const client = await clientWithMockedHttp(429, { detail: 'too many' });
        await (client as unknown as { generateKeys: () => Promise<void> }).generateKeys();
        jest.spyOn(client as unknown as { encryptPayload: () => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({ model: 'test', messages: [] }, 'id-1')
        ).rejects.toBeInstanceOf(RateLimitError);
    }, 30000);

    test('503 → ServiceUnavailableError', async () => {
        const client = await clientWithMockedHttp(503, { detail: 'down' });
        await (client as unknown as { generateKeys: () => Promise<void> }).generateKeys();
        jest.spyOn(client as unknown as { encryptPayload: () => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({ model: 'test', messages: [] }, 'id-1')
        ).rejects.toBeInstanceOf(ServiceUnavailableError);
    }, 30000);

    test('network error → APIConnectionError (not wrapping typed errors)', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com:12434',
        });

        const http = mockHttpClient(async () => { throw new Error('ECONNREFUSED'); });
        (client as unknown as { httpClient: typeof http }).httpClient = http;

        await (client as unknown as { generateKeys: () => Promise<void> }).generateKeys();
        jest.spyOn(client as unknown as { encryptPayload: () => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({ model: 'test', messages: [] }, 'id-1')
        ).rejects.toBeInstanceOf(APIConnectionError);
    }, 30000);
});

describe('SecureCompletionClient encrypt/decrypt roundtrip', () => {
    test('encryptPayload + decryptResponse roundtrip', async () => {
        // Use two clients: one for "client", one to simulate "server"
        const clientA = new SecureCompletionClient({ routerUrl: 'https://x', allowHttp: true });
        const clientB = new SecureCompletionClient({ routerUrl: 'https://x', allowHttp: true });

        await (clientA as unknown as { generateKeys: () => Promise<void> }).generateKeys();
        await (clientB as unknown as { generateKeys: () => Promise<void> }).generateKeys();

        const payload = { model: 'test', messages: [{ role: 'user', content: 'hi' }] };

        // clientA encrypts, clientB decrypts (simulating server responding)
        // We can only test the client-side encrypt → client-side decrypt roundtrip
        // because the server uses its own key pair to encrypt the response.

        // Directly test encryptPayload → decryptResponse using the SAME client's keys
        // (as the server would decrypt with its private key and re-encrypt with client's public key)
        // For a full roundtrip test we encrypt with clientB's public key and decrypt with clientB's private key.
        const serverPublicKeyPem = await (clientB as unknown as { keyManager: { getPublicKeyPEM: () => Promise<string> } }).keyManager.getPublicKeyPEM();

        // Mock fetchServerPublicKey on clientA to return clientB's public key
        jest.spyOn(clientA as unknown as { fetchServerPublicKey: () => Promise<string> }, 'fetchServerPublicKey')
            .mockResolvedValue(serverPublicKeyPem);

        const encrypted = await clientA.encryptPayload(payload);
        expect(encrypted.byteLength).toBeGreaterThan(0);

        // Now simulate clientB decrypting (server decrypts the payload — we can only test
        // structure here since decryptResponse expects server-format encrypted response)
        const pkg = JSON.parse(new TextDecoder().decode(encrypted));
        expect(pkg.version).toBe('1.0');
        expect(pkg.algorithm).toBe('hybrid-aes256-rsa4096');
        expect(pkg.encrypted_payload.ciphertext).toBeTruthy();
        expect(pkg.encrypted_payload.nonce).toBeTruthy();
        expect(pkg.encrypted_payload.tag).toBeTruthy(); // tag must be present
        expect(pkg.encrypted_aes_key).toBeTruthy();
    }, 60000);
});