/**
* Unit tests for Code Mode Security Manager
*
* Tests code validation, rate limiting, result sanitization,
* and audit logging.
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { CodeModeSecurityManager } from '../security.js';
import type { SandboxResult } from '../types.js';
describe('CodeModeSecurityManager', () => {
let security: CodeModeSecurityManager;
beforeEach(() => {
security = new CodeModeSecurityManager();
});
describe('validateCode()', () => {
it('should accept valid code', () => {
const result = security.validateCode(`
const tables = await pg.core.listTables();
return tables;
`);
expect(result.valid).toBe(true);
expect(result.errors).toHaveLength(0);
});
it('should reject empty code', () => {
const result = security.validateCode('');
expect(result.valid).toBe(false);
expect(result.errors.some(e => e.includes('non-empty'))).toBe(true);
});
it('should reject code exceeding max length', () => {
const longCode = 'x'.repeat(100000);
const result = security.validateCode(longCode);
expect(result.valid).toBe(false);
expect(result.errors.some(e => e.includes('exceeds maximum'))).toBe(true);
});
it('should block require() calls', () => {
const result = security.validateCode('const fs = require("fs");');
expect(result.valid).toBe(false);
expect(result.errors.some(e => e.includes('require'))).toBe(true);
});
it('should block process access', () => {
const result = security.validateCode('console.log(process.env);');
expect(result.valid).toBe(false);
expect(result.errors.some(e => e.includes('process'))).toBe(true);
});
it('should block eval()', () => {
const result = security.validateCode('eval("malicious code");');
expect(result.valid).toBe(false);
expect(result.errors.some(e => e.includes('eval'))).toBe(true);
});
it('should block Function constructor', () => {
const result = security.validateCode('new Function("return 1")();');
expect(result.valid).toBe(false);
expect(result.errors.some(e => e.includes('Function'))).toBe(true);
});
it('should block __proto__ manipulation', () => {
const result = security.validateCode('obj.__proto__ = malicious;');
expect(result.valid).toBe(false);
expect(result.errors.some(e => e.includes('__proto__'))).toBe(true);
});
it('should block child_process', () => {
const result = security.validateCode('require("child_process")');
expect(result.valid).toBe(false);
});
it('should block global access', () => {
const result = security.validateCode('global.process');
expect(result.valid).toBe(false);
expect(result.errors.some(e => e.includes('global'))).toBe(true);
});
it('should allow legitimate code patterns', () => {
const validCode = `
const results = [];
for (const table of await pg.core.listTables()) {
const stats = await pg.performance.tableStats({ table: table.name });
results.push({ name: table.name, rows: stats.row_count });
}
return results;
`;
const result = security.validateCode(validCode);
expect(result.valid).toBe(true);
});
});
describe('checkRateLimit()', () => {
it('should allow requests within limit', () => {
for (let i = 0; i < 5; i++) {
expect(security.checkRateLimit('client1')).toBe(true);
}
});
it('should track clients separately', () => {
for (let i = 0; i < 30; i++) {
security.checkRateLimit('client1');
security.checkRateLimit('client2');
}
expect(security.checkRateLimit('client1')).toBe(true);
expect(security.checkRateLimit('client2')).toBe(true);
});
it('should block when rate limit exceeded', () => {
const clientId = 'test-client';
// Exhaust rate limit (default is 60/min)
for (let i = 0; i < 60; i++) {
security.checkRateLimit(clientId);
}
expect(security.checkRateLimit(clientId)).toBe(false);
});
});
describe('getRateLimitRemaining()', () => {
it('should return max for new clients', () => {
const remaining = security.getRateLimitRemaining('new-client');
expect(remaining).toBe(60); // Default max executions per minute
});
it('should return reduced count after requests', () => {
const clientId = 'rate-limit-test';
security.checkRateLimit(clientId);
security.checkRateLimit(clientId);
security.checkRateLimit(clientId);
const remaining = security.getRateLimitRemaining(clientId);
expect(remaining).toBe(57); // 60 - 3
});
it('should return 0 when exhausted', () => {
const clientId = 'exhausted-client';
for (let i = 0; i < 60; i++) {
security.checkRateLimit(clientId);
}
expect(security.getRateLimitRemaining(clientId)).toBe(0);
});
});
describe('cleanupRateLimits()', () => {
it('should remove expired entries', () => {
// This just verifies the method runs without error
// Actual cleanup depends on timing
security.checkRateLimit('cleanup-client');
security.cleanupRateLimits();
// Method should complete without error
expect(true).toBe(true);
});
});
describe('sanitizeResult()', () => {
it('should return primitive values unchanged', () => {
expect(security.sanitizeResult(42)).toBe(42);
expect(security.sanitizeResult('hello')).toBe('hello');
expect(security.sanitizeResult(true)).toBe(true);
expect(security.sanitizeResult(null)).toBe(null);
});
it('should return small objects unchanged', () => {
const obj = { a: 1, b: 'test' };
expect(security.sanitizeResult(obj)).toEqual(obj);
});
it('should return small arrays unchanged', () => {
const arr = [1, 2, 3, 4, 5];
expect(security.sanitizeResult(arr)).toEqual(arr);
});
it('should handle large results', () => {
// Test that sanitizeResult handles large strings by returning something
const largeString = 'x'.repeat(100); // Smaller to test behavior
const result = security.sanitizeResult(largeString);
expect(result).toBeDefined();
});
});
describe('createExecutionRecord()', () => {
it('should create record from successful result', () => {
const successResult: SandboxResult = {
success: true,
result: { data: 'test' },
metrics: { wallTimeMs: 100, cpuTimeMs: 80, memoryUsedMb: 5 }
};
const record = security.createExecutionRecord(
'return pg.core.listTables();',
successResult,
true,
'client-123'
);
expect(record).toBeDefined();
expect(record.id).toBeDefined();
expect(record.codePreview).toBeDefined();
expect(record.timestamp).toBeDefined();
expect(record.clientId).toBe('client-123');
expect(record.readonly).toBe(true);
});
it('should create record from failed result', () => {
const failResult: SandboxResult = {
success: false,
error: 'Test error',
stack: 'Error stack trace',
metrics: { wallTimeMs: 10, cpuTimeMs: 5, memoryUsedMb: 1 }
};
const record = security.createExecutionRecord(
'throw new Error("test");',
failResult,
false
);
expect(record).toBeDefined();
expect(record.result.success).toBe(false);
expect(record.result.error).toBe('Test error');
});
});
describe('auditLog()', () => {
let consoleErrorSpy: ReturnType<typeof vi.spyOn>;
beforeEach(() => {
consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => { });
});
afterEach(() => {
consoleErrorSpy.mockRestore();
});
it('should log successful execution', () => {
const successResult: SandboxResult = {
success: true,
metrics: { wallTimeMs: 50, cpuTimeMs: 40, memoryUsedMb: 2 }
};
const record = security.createExecutionRecord(
'return 1;',
successResult,
true,
'test-client'
);
security.auditLog(record);
expect(consoleErrorSpy).toHaveBeenCalled();
});
it('should log failed execution', () => {
const failResult: SandboxResult = {
success: false,
error: 'Execution failed',
metrics: { wallTimeMs: 10, cpuTimeMs: 5, memoryUsedMb: 1 }
};
const record = security.createExecutionRecord(
'throw new Error();',
failResult,
false,
'test-client'
);
security.auditLog(record);
expect(consoleErrorSpy).toHaveBeenCalled();
});
});
});
