/**
* postgres-mcp - Code Mode Security
*
* Input validation, rate limiting, and audit logging for code execution.
*/
import { logger } from '../utils/logger.js';
import {
DEFAULT_SECURITY_CONFIG,
type SecurityConfig,
type ValidationResult,
type ExecutionRecord,
type SandboxResult
} from './types.js';
/**
* Security manager for Code Mode executions
*/
export class CodeModeSecurityManager {
private readonly config: SecurityConfig;
private readonly rateLimitMap = new Map<string, { count: number; resetTime: number }>();
constructor(config?: Partial<SecurityConfig>) {
this.config = { ...DEFAULT_SECURITY_CONFIG, ...config };
}
/**
* Validate code before execution
*/
validateCode(code: string): ValidationResult {
const errors: string[] = [];
// Check code length
if (!code || typeof code !== 'string') {
errors.push('Code must be a non-empty string');
return { valid: false, errors };
}
if (code.length > this.config.maxCodeLength) {
errors.push(`Code exceeds maximum length of ${String(this.config.maxCodeLength)} bytes`);
return { valid: false, errors };
}
// Check for blocked patterns
for (const pattern of this.config.blockedPatterns) {
if (pattern.test(code)) {
errors.push(`Blocked pattern detected: ${pattern.source}`);
}
}
return {
valid: errors.length === 0,
errors
};
}
/**
* Check rate limit for a client
* @returns true if within limits, false if rate limited
*/
checkRateLimit(clientId: string): boolean {
const now = Date.now();
const windowMs = 60000; // 1 minute window
const existing = this.rateLimitMap.get(clientId);
if (!existing || now >= existing.resetTime) {
// Start new window
this.rateLimitMap.set(clientId, {
count: 1,
resetTime: now + windowMs
});
return true;
}
if (existing.count >= this.config.maxExecutionsPerMinute) {
return false;
}
existing.count++;
return true;
}
/**
* Get remaining rate limit for a client
*/
getRateLimitRemaining(clientId: string): number {
const existing = this.rateLimitMap.get(clientId);
if (!existing || Date.now() >= existing.resetTime) {
return this.config.maxExecutionsPerMinute;
}
return Math.max(0, this.config.maxExecutionsPerMinute - existing.count);
}
/**
* Sanitize and truncate result if too large
*/
sanitizeResult(result: unknown): unknown {
try {
const serialized = JSON.stringify(result);
if (serialized.length > this.config.maxResultSize) {
return {
_truncated: true,
_originalSize: serialized.length,
_maxSize: this.config.maxResultSize,
preview: serialized.substring(0, 1000) + '...'
};
}
return result;
} catch {
return {
_error: 'Result could not be serialized',
_type: typeof result
};
}
}
/**
* Log execution for audit purposes
*/
auditLog(execution: ExecutionRecord): void {
const { id, clientId, codePreview, result, readonly } = execution;
const logContext = {
module: 'CODEMODE' as const,
operation: 'execute',
entityId: id,
clientId: clientId ?? 'anonymous',
readonly,
success: result.success,
wallTimeMs: result.metrics.wallTimeMs,
memoryUsedMb: result.metrics.memoryUsedMb
};
if (result.success) {
logger.info(`Code execution completed: ${codePreview.substring(0, 50)}...`, logContext);
} else {
const errorContext = {
...logContext,
...(result.error !== undefined ? { error: result.error } : {}),
...(result.stack !== undefined ? { stack: result.stack } : {})
};
logger.warning(`Code execution failed: ${result.error ?? 'unknown error'}`, errorContext);
}
}
/**
* Create execution record for audit
*/
createExecutionRecord(
code: string,
result: SandboxResult,
readonly: boolean,
clientId?: string
): ExecutionRecord {
return {
id: crypto.randomUUID(),
clientId,
timestamp: new Date(),
codePreview: code.length > 200 ? code.substring(0, 200) + '...' : code,
result,
readonly
};
}
/**
* Clean up old rate limit entries
*/
cleanupRateLimits(): void {
const now = Date.now();
for (const [clientId, entry] of this.rateLimitMap) {
if (now >= entry.resetTime) {
this.rateLimitMap.delete(clientId);
}
}
}
}
