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

/**
 * Unit tests for key persistence vs. ephemeral key handling.
 *
 * Mirrors the Python SDK's `key_dir` semantics: omitting it (the default)
 * generates a key pair in memory for this session only and never writes to
 * disk; supplying it reuses or creates a pair in that directory.
 */

import * as fs from 'fs';
import * as os from 'os';
import * as path from 'path';
import { SecureCompletionClient } from '../../src/core/SecureCompletionClient';

/** Drive key initialisation without any network traffic. */
async function initKeys(client: SecureCompletionClient): Promise<void> {
    await (client as unknown as { ensureKeys: () => Promise<void> }).ensureKeys();
}

function publicKeyOf(client: SecureCompletionClient): Promise<string> {
    return (client as unknown as {
        keyManager: { getPublicKeyPEM: () => Promise<string> };
    }).keyManager.getPublicKeyPEM();
}

/**
 * Watch for any attempt to write keys to disk.
 *
 * Asserting on the filesystem alone is unreliable here — a stale `client_keys/`
 * from before ephemeral keys became the default would make such a check pass or
 * fail for the wrong reason. This asserts the intent directly.
 */
function watchSaves(client: SecureCompletionClient): jest.SpyInstance {
    const keyManager = (client as unknown as { keyManager: object }).keyManager;
    return jest.spyOn(keyManager as { saveKeys: (...args: never[]) => Promise<void> }, 'saveKeys');
}

let tmpDir: string;

beforeEach(() => {
    tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'nomyo-keys-'));
});

afterEach(() => {
    fs.rmSync(tmpDir, { recursive: true, force: true });
});

describe('ephemeral keys (default)', () => {
    test('writes nothing to disk when no keyDir is configured', async () => {
        const cwdBefore = fs.readdirSync(process.cwd());

        const client = new SecureCompletionClient({
            routerUrl: 'https://router.test',
            keySize: 2048,
            keyRotationInterval: 0,
        });
        const saveSpy = watchSaves(client);
        await initKeys(client);

        expect(saveSpy).not.toHaveBeenCalled();
        // And no new directory materialised in the working directory
        expect(fs.readdirSync(process.cwd())).toEqual(cwdBefore);
        client.dispose();
    });

    test('each client gets its own key pair', async () => {
        const a = new SecureCompletionClient({
            routerUrl: 'https://router.test', keySize: 2048, keyRotationInterval: 0,
        });
        const b = new SecureCompletionClient({
            routerUrl: 'https://router.test', keySize: 2048, keyRotationInterval: 0,
        });
        await initKeys(a);
        await initKeys(b);

        expect(await publicKeyOf(a)).not.toBe(await publicKeyOf(b));
        a.dispose();
        b.dispose();
    });

    test('keyDir: null is treated as ephemeral', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://router.test',
            keySize: 2048,
            keyRotationInterval: 0,
            keyDir: null,
        });
        await initKeys(client);

        expect(fs.readdirSync(tmpDir)).toEqual([]);
        client.dispose();
    });

    test('key rotation stays in memory', async () => {
        const client = new SecureCompletionClient({
            routerUrl: 'https://router.test',
            keySize: 2048,
            keyRotationInterval: 0,
        });
        await initKeys(client);
        const before = await publicKeyOf(client);
        const saveSpy = watchSaves(client);

        await (client as unknown as { rotateKeys: () => Promise<void> }).rotateKeys();

        // Rotated to a genuinely new pair, still with nothing written to disk
        expect(await publicKeyOf(client)).not.toBe(before);
        expect(saveSpy).not.toHaveBeenCalled();
        client.dispose();
    });
});

describe('persistent keys (explicit keyDir)', () => {
    test('generates and saves a key pair on first use', async () => {
        const keyDir = path.join(tmpDir, 'keys');
        const client = new SecureCompletionClient({
            routerUrl: 'https://router.test',
            keySize: 2048,
            keyRotationInterval: 0,
            keyDir,
        });
        await initKeys(client);

        expect(fs.existsSync(path.join(keyDir, 'private_key.pem'))).toBe(true);
        expect(fs.existsSync(path.join(keyDir, 'public_key.pem'))).toBe(true);
        client.dispose();
    });

    test('private key is written owner-only (0600)', async () => {
        const keyDir = path.join(tmpDir, 'keys');
        const client = new SecureCompletionClient({
            routerUrl: 'https://router.test',
            keySize: 2048,
            keyRotationInterval: 0,
            keyDir,
        });
        await initKeys(client);

        const mode = fs.statSync(path.join(keyDir, 'private_key.pem')).mode & 0o777;
        expect(mode).toBe(0o600);
        client.dispose();
    });

    test('reuses the existing key pair across clients', async () => {
        const keyDir = path.join(tmpDir, 'keys');
        const config = {
            routerUrl: 'https://router.test',
            keySize: 2048 as const,
            keyRotationInterval: 0,
            keyDir,
        };

        const first = new SecureCompletionClient(config);
        await initKeys(first);
        const firstKey = await publicKeyOf(first);
        first.dispose();

        const second = new SecureCompletionClient(config);
        await initKeys(second);

        expect(await publicKeyOf(second)).toBe(firstKey);
        second.dispose();
    });

    test('rotation persists into the configured directory', async () => {
        const keyDir = path.join(tmpDir, 'keys');
        const client = new SecureCompletionClient({
            routerUrl: 'https://router.test',
            keySize: 2048,
            keyRotationInterval: 0,
            keyDir,
        });
        await initKeys(client);
        const before = fs.readFileSync(path.join(keyDir, 'public_key.pem'), 'utf-8');

        await (client as unknown as { rotateKeys: () => Promise<void> }).rotateKeys();

        const after = fs.readFileSync(path.join(keyDir, 'public_key.pem'), 'utf-8');
        expect(after).not.toBe(before);
        expect(after.trim()).toBe((await publicKeyOf(client)).trim());
        client.dispose();
    });
});