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

/**
 * Unit tests for secure memory handling.
 *
 * Mirrors Python's SecureBuffer/secure_bytearray lifecycle: lock on entry,
 * then zero *before* unlocking on exit, with every step best-effort so that a
 * platform without mlock still gets zeroing.
 */

import {
    SecureByteContext,
    createSecureMemory,
    getMemoryProtectionInfo,
    disableSecureMemory,
    enableSecureMemory,
    SecureMemory,
} from '../../src/core/memory/secure';
import { NodeSecureMemory } from '../../src/core/memory/node';
import { BrowserSecureMemory } from '../../src/core/memory/browser';

/** A SecureMemory stub that records the order of operations. */
function recordingMemory(overrides: Partial<SecureMemory> = {}) {
    const calls: string[] = [];
    const impl: SecureMemory = {
        zeroMemory: jest.fn((data: ArrayBuffer) => {
            calls.push('zero');
            new Uint8Array(data).fill(0);
        }),
        lockMemory: jest.fn(() => {
            calls.push('lock');
            return true;
        }),
        unlockMemory: jest.fn(() => {
            calls.push('unlock');
            return true;
        }),
        getProtectionInfo: jest.fn(() => ({
            canLock: true,
            isPlatformSecure: true,
            method: 'mlock' as const,
        })),
        ...overrides,
    };
    return { impl, calls };
}

function contextWith(data: ArrayBuffer, memory: SecureMemory, useSecure = true, lock = true) {
    const ctx = new SecureByteContext(data, useSecure, lock);
    (ctx as unknown as { secureMemory: SecureMemory }).secureMemory = memory;
    return ctx;
}

function filledBuffer(byte = 0xab, size = 32): ArrayBuffer {
    const buf = new ArrayBuffer(size);
    new Uint8Array(buf).fill(byte);
    return buf;
}

describe('SecureByteContext lifecycle', () => {
    test('locks before use and zeroes before unlocking', async () => {
        const { impl, calls } = recordingMemory();
        const buf = filledBuffer();

        await contextWith(buf, impl).use(async (data) => {
            calls.push('use');
            // Still intact while in scope
            expect(new Uint8Array(data).every(b => b === 0xab)).toBe(true);
        });

        // Zeroing must happen while the pages are still locked
        expect(calls).toEqual(['lock', 'use', 'zero', 'unlock']);
        expect(new Uint8Array(buf).every(b => b === 0)).toBe(true);
    });

    test('zeroes and unlocks even when the callback throws', async () => {
        const { impl, calls } = recordingMemory();
        const buf = filledBuffer();

        await expect(
            contextWith(buf, impl).use(async () => {
                throw new Error('boom');
            })
        ).rejects.toThrow('boom');

        expect(calls).toEqual(['lock', 'zero', 'unlock']);
        expect(new Uint8Array(buf).every(b => b === 0)).toBe(true);
    });

    test('a failed lock still zeroes, and does not unlock', async () => {
        const { impl, calls } = recordingMemory({
            lockMemory: jest.fn(() => false),
        });
        const buf = filledBuffer();
        const ctx = contextWith(buf, impl);

        await ctx.use(async () => undefined);

        expect(ctx.locked).toBe(false);
        expect(calls).toEqual(['zero']);
        expect(impl.unlockMemory).not.toHaveBeenCalled();
        expect(new Uint8Array(buf).every(b => b === 0)).toBe(true);
    });

    test('a throwing lockMemory degrades to zeroing', async () => {
        const { impl } = recordingMemory({
            lockMemory: jest.fn(() => { throw new Error('mlock exploded'); }),
        });
        const buf = filledBuffer();

        await expect(contextWith(buf, impl).use(async () => 'ok')).resolves.toBe('ok');
        expect(impl.zeroMemory).toHaveBeenCalled();
        expect(new Uint8Array(buf).every(b => b === 0)).toBe(true);
    });

    test('a throwing zeroMemory does not mask the callback error', async () => {
        const { impl } = recordingMemory({
            zeroMemory: jest.fn(() => { throw new Error('zero exploded'); }),
        });

        await expect(
            contextWith(filledBuffer(), impl).use(async () => {
                throw new Error('original');
            })
        ).rejects.toThrow('original');
    });

    test('lock: false skips locking but still zeroes', async () => {
        const { impl, calls } = recordingMemory();
        const buf = filledBuffer();

        await contextWith(buf, impl, true, false).use(async () => undefined);

        expect(calls).toEqual(['zero']);
        expect(impl.lockMemory).not.toHaveBeenCalled();
        expect(new Uint8Array(buf).every(b => b === 0)).toBe(true);
    });

    test('useSecure: false leaves the buffer untouched', async () => {
        const { impl, calls } = recordingMemory();
        const buf = filledBuffer();

        await contextWith(buf, impl, false).use(async () => undefined);

        expect(calls).toEqual([]);
        expect(new Uint8Array(buf).every(b => b === 0xab)).toBe(true);
    });

    test('locked reflects state across the lifecycle', async () => {
        const { impl } = recordingMemory();
        const ctx = contextWith(filledBuffer(), impl);

        expect(ctx.locked).toBe(false);
        await ctx.use(async () => {
            expect(ctx.locked).toBe(true);
        });
        expect(ctx.locked).toBe(false);
    });

    test('honours the global disable/enable flag', async () => {
        disableSecureMemory();
        try {
            const buf = filledBuffer();
            // No explicit useSecure: falls back to the global flag
            const ctx = new SecureByteContext(buf);
            const { impl } = recordingMemory();
            (ctx as unknown as { secureMemory: SecureMemory }).secureMemory = impl;

            await ctx.use(async () => undefined);
            expect(new Uint8Array(buf).every(b => b === 0xab)).toBe(true);
        } finally {
            enableSecureMemory();
        }
    });
});

describe('NodeSecureMemory', () => {
    const memory = new NodeSecureMemory();

    test('lock/unlock round-trips or degrades cleanly', () => {
        const buf = filledBuffer(0x11, 4096);
        const locked = memory.lockMemory(buf);

        // Whatever the platform allows, the result must agree with what we report
        expect(locked).toBe(getMemoryProtectionInfo().canLock);

        if (locked) {
            expect(memory.unlockMemory(buf)).toBe(true);
        }
    });

    test('zeroMemory clears the buffer', () => {
        const buf = filledBuffer(0xff, 64);
        memory.zeroMemory(buf);
        expect(new Uint8Array(buf).every(b => b === 0)).toBe(true);
    });

    test('reported protection info is self-consistent', () => {
        const info = memory.getProtectionInfo();

        // canLock is probed, so it must never disagree with the method claimed
        expect(info.method).toBe(info.canLock ? 'mlock' : 'zero-only');
        expect(info.isPlatformSecure).toBe(info.canLock);
        expect(info.platform).toBe(process.platform);

        if (info.canLock) {
            expect(info.hasSecureZeroing).toBe(true);
            expect(info.pageSize).toBeGreaterThan(0);
        }
    });

    test('claims mlock only when locking genuinely works', () => {
        const info = memory.getProtectionInfo();
        const buf = filledBuffer(0x22, 4096);

        const locked = memory.lockMemory(buf);
        if (locked) memory.unlockMemory(buf);

        expect(info.method === 'mlock').toBe(locked);
    });
});

describe('BrowserSecureMemory', () => {
    const memory = new BrowserSecureMemory();

    test('never claims to lock', () => {
        expect(memory.lockMemory(filledBuffer())).toBe(false);
        expect(memory.unlockMemory(filledBuffer())).toBe(false);

        const info = memory.getProtectionInfo();
        expect(info.canLock).toBe(false);
        expect(info.method).toBe('zero-only');
        expect(info.platform).toBe('browser');
    });

    test('still zeroes', () => {
        const buf = filledBuffer(0x7f);
        memory.zeroMemory(buf);
        expect(new Uint8Array(buf).every(b => b === 0)).toBe(true);
    });
});

describe('createSecureMemory', () => {
    test('returns an implementation satisfying the full interface', () => {
        const memory = createSecureMemory();
        expect(typeof memory.zeroMemory).toBe('function');
        expect(typeof memory.lockMemory).toBe('function');
        expect(typeof memory.unlockMemory).toBe('function');
        expect(typeof memory.getProtectionInfo).toBe('function');
    });
});