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,
    DisposedError,
} 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/' });
        expect((client as unknown as { routerUrl: string }).routerUrl).toBe('https://api.example.com:12434');
        client.dispose();
    });

    test('throws on invalid URL', () => {
        expect(() => new SecureCompletionClient({ routerUrl: 'not-a-url' }))
            .toThrow('Invalid routerUrl');
    });

    test('http:// URL with allowHttp=true does not throw', () => {
        expect(() => new SecureCompletionClient({
            routerUrl: 'http://localhost:1234',
            allowHttp: true,
        })).not.toThrow();
    });
});

describe('SecureCompletionClient.dispose()', () => {
    test('calling dispose() twice does not throw', () => {
        const client = new SecureCompletionClient({ routerUrl: 'https://api.example.com', keyRotationInterval: 0 });
        client.dispose();
        expect(() => client.dispose()).not.toThrow();
    });

    test('methods throw DisposedError after dispose()', async () => {
        const client = new SecureCompletionClient({ routerUrl: 'https://api.example.com', keyRotationInterval: 0 });
        client.dispose();
        await expect(client.fetchServerPublicKey()).rejects.toBeInstanceOf(DisposedError);
        await expect(client.encryptPayload({})).rejects.toBeInstanceOf(DisposedError);
        await expect(client.sendSecureRequest({}, 'id')).rejects.toBeInstanceOf(DisposedError);
    });

    test('dispose() clears key rotation timer', () => {
        jest.useFakeTimers();
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com',
            keyRotationInterval: 1000,
        });
        const timerBefore = (client as unknown as { keyRotationTimer: unknown }).keyRotationTimer;
        expect(timerBefore).toBeDefined();
        client.dispose();
        const timerAfter = (client as unknown as { keyRotationTimer: unknown }).keyRotationTimer;
        expect(timerAfter).toBeUndefined();
        jest.useRealTimers();
    });

    test('keyRotationInterval=0 does not start timer', () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com',
            keyRotationInterval: 0,
        });
        const timer = (client as unknown as { keyRotationTimer: unknown }).keyRotationTimer;
        expect(timer).toBeUndefined();
        client.dispose();
    });
});

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

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

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

    test('accepts valid security tiers', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com:12434',
            keyRotationInterval: 0,
        });

        for (const tier of ['standard', 'high', 'maximum']) {
            await expect(
                client.sendSecureRequest({}, 'test-id', undefined, tier)
            ).rejects.not.toThrow("Invalid securityTier");
        }
        client.dispose();
    });
});

describe('SecureCompletionClient — header injection validation', () => {
    test('apiKey containing CR throws SecurityError', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com',
            keyRotationInterval: 0,
        });
        await (client as unknown as { generateKeys: () => Promise<void> }).generateKeys();
        jest.spyOn(client as unknown as { encryptPayload: (p: object) => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({}, 'id', 'key\rwith\rcr')
        ).rejects.toBeInstanceOf(SecurityError);
        client.dispose();
    }, 30000);

    test('apiKey containing LF throws SecurityError', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com',
            keyRotationInterval: 0,
        });
        await (client as unknown as { generateKeys: () => Promise<void> }).generateKeys();
        jest.spyOn(client as unknown as { encryptPayload: (p: object) => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({}, 'id', 'key\nwith\nlf')
        ).rejects.toBeInstanceOf(SecurityError);
        client.dispose();
    }, 30000);
});

describe('SecureCompletionClient — error detail sanitization', () => {
    test('long server detail is truncated to ≤100 chars in error message', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com:12434',
            keyRotationInterval: 0,
        });

        const http = mockHttpClient(async () => makeJsonResponse(400, { detail: 'x'.repeat(200) }));
        (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: (p: object) => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        try {
            await client.sendSecureRequest({}, 'id');
        } catch (err) {
            expect(err).toBeInstanceOf(Error);
            // "Bad request: " prefix + max 100 char detail
            expect((err as Error).message.length).toBeLessThanOrEqual(115);
        }
        client.dispose();
    }, 30000);
});

describe('SecureCompletionClient — debug flag', () => {
    test('console.log not called during construction when debug=false', () => {
        const spy = jest.spyOn(console, 'log').mockImplementation(() => undefined);
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com',
            debug: false,
            keyRotationInterval: 0,
        });
        expect(spy).not.toHaveBeenCalled();
        spy.mockRestore();
        client.dispose();
    });

    test('console.log called during construction when debug=true', () => {
        const spy = jest.spyOn(console, 'log').mockImplementation(() => undefined);
        const client = new SecureCompletionClient({
            routerUrl: 'https://api.example.com',
            debug: true,
            keyRotationInterval: 0,
        });
        expect(spy).toHaveBeenCalled();
        spy.mockRestore();
        client.dispose();
    });
});

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

        const http = mockHttpClient(async (url: string) => {
            if (url.includes('/pki/public_key')) {
                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' });
        await (client as unknown as { generateKeys: () => Promise<void> }).generateKeys();
        jest.spyOn(client as unknown as { encryptPayload: (p: object) => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({ model: 'test', messages: [] }, 'id-1')
        ).rejects.toBeInstanceOf(AuthenticationError);
        client.dispose();
    }, 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: (p: object) => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({ model: 'test', messages: [] }, 'id-1')
        ).rejects.toBeInstanceOf(ForbiddenError);
        client.dispose();
    }, 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: (p: object) => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

        await expect(
            client.sendSecureRequest({ model: 'test', messages: [] }, 'id-1')
        ).rejects.toBeInstanceOf(RateLimitError);
        client.dispose();
    }, 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: (p: object) => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

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

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

        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: (p: object) => Promise<ArrayBuffer> }, 'encryptPayload')
            .mockResolvedValue(new ArrayBuffer(8));

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

describe('SecureCompletionClient encrypt/decrypt roundtrip', () => {
    test('encryptPayload + decryptResponse roundtrip', async () => {
        const clientA = new SecureCompletionClient({ routerUrl: 'https://x', allowHttp: true, keyRotationInterval: 0 });
        const clientB = new SecureCompletionClient({ routerUrl: 'https://x', allowHttp: true, keyRotationInterval: 0 });

        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' }] };

        const serverPublicKeyPem = await (clientB as unknown as { keyManager: { getPublicKeyPEM: () => Promise<string> } }).keyManager.getPublicKeyPEM();

        jest.spyOn(clientA as unknown as { fetchServerPublicKey: () => Promise<string> }, 'fetchServerPublicKey')
            .mockResolvedValue(serverPublicKeyPem);

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

        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();
        expect(pkg.encrypted_aes_key).toBeTruthy();

        clientA.dispose();
        clientB.dispose();
    }, 60000);
});