Compare commits

...
Author SHA1 Message Date
Gergo Magyar 8229de1dad Harden transport security and improve type safety
Transport hardening:
- Add MAX_BUFFER_SIZE (10 MB) cap to prevent OOM from oversized
  Content-Length or unbounded newline-delimited input
- Replace recursive readNewlineMessage with iterative loop to prevent
  stack overflow from consecutive empty lines
- Tighten looksLikeContentLength to require 14+ bytes before matching
- Add closed-state guard and error handling to send()
- Simplify processReadBuffer loop to break on error
- Fix loose equality (==) to strict (===)
- Widen constructor param types to ReadableStream/WritableStream

Type safety:
- Constrain createLazyAction generics so export name is validated
  against the module's actual exports at compile time
- Use proper type guard instead of lint suppression
- Fix test tsconfig type errors

Regression tests for all hardening fixes (13 tests passing).
2026-03-07 07:44:21 +00:00
Shockang 1e2ed0c90c Harden MCP transport error handling 2026-03-07 12:41:57 +08:00
Shockang 22092b49ca Preserve CLI flags in MCP startup fix 2026-03-07 12:33:52 +08:00
Shockang 88b21d132b Fix MCP startup transport compatibility 2026-03-07 12:04:45 +08:00
7 changed files with 609 additions and 27 deletions
+16 -25
View File
@@ -4,18 +4,9 @@
// Removing it from here improves MCP server startup time significantly.
import { Command } from 'commander';
import { analyzeCommand } from './analyze.js';
import { serveCommand } from './serve.js';
import { listCommand } from './list.js';
import { statusCommand } from './status.js';
import { mcpCommand } from './mcp.js';
import { cleanCommand } from './clean.js';
import { setupCommand } from './setup.js';
import { augmentCommand } from './augment.js';
import { wikiCommand } from './wiki.js';
import { queryCommand, contextCommand, impactCommand, cypherCommand } from './tool.js';
import { evalServerCommand } from './eval-server.js';
import { createRequire } from 'node:module';
import { createLazyAction } from './lazy-action.js';
const _require = createRequire(import.meta.url);
const pkg = _require('../../package.json');
const program = new Command();
@@ -28,43 +19,43 @@ program
program
.command('setup')
.description('One-time setup: configure MCP for Cursor, Claude Code, OpenCode')
.action(setupCommand);
.action(createLazyAction(() => import('./setup.js'), 'setupCommand'));
program
.command('analyze [path]')
.description('Index a repository (full analysis)')
.option('-f, --force', 'Force full re-index even if up to date')
.option('--embeddings', 'Enable embedding generation for semantic search (off by default)')
.action(analyzeCommand);
.action(createLazyAction(() => import('./analyze.js'), 'analyzeCommand'));
program
.command('serve')
.description('Start local HTTP server for web UI connection')
.option('-p, --port <port>', 'Port number', '4747')
.option('--host <host>', 'Bind address (default: 127.0.0.1, use 0.0.0.0 for remote access)')
.action(serveCommand);
.action(createLazyAction(() => import('./serve.js'), 'serveCommand'));
program
.command('mcp')
.description('Start MCP server (stdio) — serves all indexed repos')
.action(mcpCommand);
.action(createLazyAction(() => import('./mcp.js'), 'mcpCommand'));
program
.command('list')
.description('List all indexed repositories')
.action(listCommand);
.action(createLazyAction(() => import('./list.js'), 'listCommand'));
program
.command('status')
.description('Show index status for current repo')
.action(statusCommand);
.action(createLazyAction(() => import('./status.js'), 'statusCommand'));
program
.command('clean')
.description('Delete GitNexus index for current repo')
.option('-f, --force', 'Skip confirmation prompt')
.option('--all', 'Clean all indexed repos')
.action(cleanCommand);
.action(createLazyAction(() => import('./clean.js'), 'cleanCommand'));
program
.command('wiki [path]')
@@ -75,12 +66,12 @@ program
.option('--api-key <key>', 'LLM API key (saved to ~/.gitnexus/config.json)')
.option('--concurrency <n>', 'Parallel LLM calls (default: 3)', '3')
.option('--gist', 'Publish wiki as a public GitHub Gist after generation')
.action(wikiCommand);
.action(createLazyAction(() => import('./wiki.js'), 'wikiCommand'));
program
.command('augment <pattern>')
.description('Augment a search pattern with knowledge graph context (used by hooks)')
.action(augmentCommand);
.action(createLazyAction(() => import('./augment.js'), 'augmentCommand'));
// ─── Direct Tool Commands (no MCP overhead) ────────────────────────
// These invoke LocalBackend directly for use in eval, scripts, and CI.
@@ -93,7 +84,7 @@ program
.option('-g, --goal <text>', 'What you want to find')
.option('-l, --limit <n>', 'Max processes to return (default: 5)')
.option('--content', 'Include full symbol source code')
.action(queryCommand);
.action(createLazyAction(() => import('./tool.js'), 'queryCommand'));
program
.command('context [name]')
@@ -102,7 +93,7 @@ program
.option('-u, --uid <uid>', 'Direct symbol UID (zero-ambiguity lookup)')
.option('-f, --file <path>', 'File path to disambiguate common names')
.option('--content', 'Include full symbol source code')
.action(contextCommand);
.action(createLazyAction(() => import('./tool.js'), 'contextCommand'));
program
.command('impact <target>')
@@ -111,13 +102,13 @@ program
.option('-r, --repo <name>', 'Target repository')
.option('--depth <n>', 'Max relationship depth (default: 3)')
.option('--include-tests', 'Include test files in results')
.action(impactCommand);
.action(createLazyAction(() => import('./tool.js'), 'impactCommand'));
program
.command('cypher <query>')
.description('Execute raw Cypher query against the knowledge graph')
.option('-r, --repo <name>', 'Target repository')
.action(cypherCommand);
.action(createLazyAction(() => import('./tool.js'), 'cypherCommand'));
// ─── Eval Server (persistent daemon for SWE-bench) ─────────────────
@@ -126,6 +117,6 @@ program
.description('Start lightweight HTTP server for fast tool calls during evaluation')
.option('-p, --port <port>', 'Port number', '4848')
.option('--idle-timeout <seconds>', 'Auto-shutdown after N seconds idle (0 = disabled)', '0')
.action(evalServerCommand);
.action(createLazyAction(() => import('./eval-server.js'), 'evalServerCommand'));
program.parse(process.argv);
+26
View File
@@ -0,0 +1,26 @@
/**
* Creates a lazy-loaded CLI action that defers module import until invocation.
* The generic constraints ensure the export name is a valid key of the module
* at compile time — catching typos when used with concrete module imports.
*/
function isCallable(value: unknown): value is (...args: unknown[]) => unknown {
return typeof value === 'function';
}
export function createLazyAction<
TModule extends Record<string, unknown>,
TKey extends string & keyof TModule,
>(
loader: () => Promise<TModule>,
exportName: TKey,
): (...args: unknown[]) => Promise<void> {
return async (...args: unknown[]): Promise<void> => {
const module = await loader();
const action = module[exportName];
if (!isCallable(action)) {
throw new Error(`Lazy action export not found: ${exportName}`);
}
await action(...args);
};
}
@@ -0,0 +1,240 @@
import process from 'node:process';
import type { Transport, TransportSendOptions } from '@modelcontextprotocol/sdk/shared/transport.js';
import { JSONRPCMessageSchema, type JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
export type StdioFraming = 'content-length' | 'newline';
function deserializeMessage(raw: string): JSONRPCMessage {
return JSONRPCMessageSchema.parse(JSON.parse(raw));
}
function serializeNewlineMessage(message: JSONRPCMessage): string {
return `${JSON.stringify(message)}\n`;
}
function serializeContentLengthMessage(message: JSONRPCMessage): string {
const body = JSON.stringify(message);
return `Content-Length: ${Buffer.byteLength(body, 'utf8')}\r\n\r\n${body}`;
}
function findHeaderEnd(buffer: Buffer): { index: number; separatorLength: number } | null {
const crlfEnd = buffer.indexOf('\r\n\r\n');
if (crlfEnd !== -1) {
return { index: crlfEnd, separatorLength: 4 };
}
const lfEnd = buffer.indexOf('\n\n');
if (lfEnd !== -1) {
return { index: lfEnd, separatorLength: 2 };
}
return null;
}
function looksLikeContentLength(buffer: Buffer): boolean {
if (buffer.length < 14) {
return false;
}
const probe = buffer.toString('utf8', 0, Math.min(buffer.length, 32));
return /^content-length\s*:/i.test(probe);
}
const MAX_BUFFER_SIZE = 10 * 1024 * 1024; // 10 MB — generous for JSON-RPC
export class CompatibleStdioServerTransport implements Transport {
private _readBuffer: Buffer | undefined;
private _started = false;
private _framing: StdioFraming | null = null;
onmessage?: (message: JSONRPCMessage) => void;
onerror?: (error: Error) => void;
onclose?: () => void;
constructor(
private readonly _stdin: NodeJS.ReadableStream = process.stdin,
private readonly _stdout: NodeJS.WritableStream = process.stdout,
) {}
private readonly _ondata = (chunk: Buffer) => {
this._readBuffer = this._readBuffer ? Buffer.concat([this._readBuffer, chunk]) : chunk;
if (this._readBuffer.length > MAX_BUFFER_SIZE) {
this.onerror?.(new Error(`Read buffer exceeded maximum size (${MAX_BUFFER_SIZE} bytes)`));
this.discardBufferedInput();
return;
}
this.processReadBuffer();
};
private readonly _onerror = (error: Error) => {
this.onerror?.(error);
};
async start() {
if (this._started) {
throw new Error('CompatibleStdioServerTransport already started!');
}
this._started = true;
this._stdin.on('data', this._ondata);
this._stdin.on('error', this._onerror);
}
private detectFraming(): StdioFraming | null {
if (!this._readBuffer || this._readBuffer.length === 0) {
return null;
}
const firstByte = this._readBuffer[0];
if (firstByte === 0x7b || firstByte === 0x5b) {
return 'newline';
}
if (looksLikeContentLength(this._readBuffer)) {
return 'content-length';
}
return null;
}
private discardBufferedInput() {
this._readBuffer = undefined;
this._framing = null;
}
private readContentLengthMessage(): JSONRPCMessage | null {
if (!this._readBuffer) {
return null;
}
const header = findHeaderEnd(this._readBuffer);
if (header === null) {
return null;
}
const headerText = this._readBuffer
.toString('utf8', 0, header.index)
.replace(/\r\n/g, '\n')
.replace(/\r/g, '\n');
const match = headerText.match(/(?:^|\n)content-length\s*:\s*(\d+)/i);
if (!match) {
this.discardBufferedInput();
throw new Error('Missing Content-Length header from MCP client');
}
const contentLength = Number.parseInt(match[1], 10);
if (!Number.isFinite(contentLength) || contentLength < 0) {
this.discardBufferedInput();
throw new Error('Invalid Content-Length header from MCP client');
}
if (contentLength > MAX_BUFFER_SIZE) {
this.discardBufferedInput();
throw new Error(`Content-Length ${contentLength} exceeds maximum allowed size (${MAX_BUFFER_SIZE} bytes)`);
}
const bodyStart = header.index + header.separatorLength;
const bodyEnd = bodyStart + contentLength;
if (this._readBuffer.length < bodyEnd) {
return null;
}
const body = this._readBuffer.toString('utf8', bodyStart, bodyEnd);
this._readBuffer = this._readBuffer.subarray(bodyEnd);
return deserializeMessage(body);
}
private readNewlineMessage(): JSONRPCMessage | null {
if (!this._readBuffer) {
return null;
}
while (true) {
const newlineIndex = this._readBuffer.indexOf('\n');
if (newlineIndex === -1) {
return null;
}
const line = this._readBuffer.toString('utf8', 0, newlineIndex).replace(/\r$/, '');
this._readBuffer = this._readBuffer.subarray(newlineIndex + 1);
if (line.trim().length === 0) {
continue;
}
return deserializeMessage(line);
}
}
private readMessage(): JSONRPCMessage | null {
if (!this._readBuffer || this._readBuffer.length === 0) {
return null;
}
if (this._framing === null) {
this._framing = this.detectFraming();
if (this._framing === null) {
return null;
}
}
return this._framing === 'content-length'
? this.readContentLengthMessage()
: this.readNewlineMessage();
}
private processReadBuffer() {
while (true) {
try {
const message = this.readMessage();
if (message === null) {
break;
}
this.onmessage?.(message);
} catch (error) {
this.onerror?.(error as Error);
break;
}
}
}
async close() {
this._stdin.off('data', this._ondata);
this._stdin.off('error', this._onerror);
const remainingDataListeners = this._stdin.listenerCount('data');
if (remainingDataListeners === 0) {
this._stdin.pause();
}
this._started = false;
this._readBuffer = undefined;
this.onclose?.();
}
send(message: JSONRPCMessage, _options?: TransportSendOptions) {
return new Promise<void>((resolve, reject) => {
if (!this._started) {
reject(new Error('Transport is closed'));
return;
}
const payload = this._framing === 'newline'
? serializeNewlineMessage(message)
: serializeContentLengthMessage(message);
const onError = (error: Error) => {
this._stdout.removeListener('error', onError);
reject(error);
};
this._stdout.on('error', onError);
if (this._stdout.write(payload)) {
this._stdout.removeListener('error', onError);
resolve();
} else {
this._stdout.once('drain', () => {
this._stdout.removeListener('error', onError);
resolve();
});
}
});
}
}
+2 -2
View File
@@ -13,7 +13,7 @@
import { createRequire } from 'module';
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { CompatibleStdioServerTransport } from './compatible-stdio-transport.js';
import {
CallToolRequestSchema,
ListToolsRequestSchema,
@@ -277,7 +277,7 @@ export async function startMCPServer(backend: LocalBackend): Promise<void> {
const server = createMCPServer(backend);
// Connect to stdio transport
const transport = new StdioServerTransport();
const transport = new CompatibleStdioServerTransport();
await server.connect(transport);
// Graceful shutdown helper
+45
View File
@@ -0,0 +1,45 @@
import { spawnSync } from 'node:child_process';
import path from 'node:path';
import { fileURLToPath } from 'node:url';
import { describe, expect, it } from 'vitest';
const testDir = path.dirname(fileURLToPath(import.meta.url));
const repoRoot = path.resolve(testDir, '../..');
const cliEntry = path.join(repoRoot, 'src/cli/index.ts');
function runHelp(command: string) {
return spawnSync(process.execPath, ['--import', 'tsx', cliEntry, command, '--help'], {
cwd: repoRoot,
encoding: 'utf8',
});
}
describe('CLI help surface', () => {
it('query help keeps advanced search options without importing analyze deps', () => {
const result = runHelp('query');
expect(result.status).toBe(0);
expect(result.stdout).toContain('--context <text>');
expect(result.stdout).toContain('--goal <text>');
expect(result.stdout).toContain('--content');
expect(result.stderr).not.toContain('tree-sitter-kotlin');
});
it('context help keeps optional name and disambiguation flags', () => {
const result = runHelp('context');
expect(result.status).toBe(0);
expect(result.stdout).toContain('context [options] [name]');
expect(result.stdout).toContain('--uid <uid>');
expect(result.stdout).toContain('--file <path>');
});
it('impact help keeps repo and include-tests flags', () => {
const result = runHelp('impact');
expect(result.status).toBe(0);
expect(result.stdout).toContain('--depth <n>');
expect(result.stdout).toContain('--include-tests');
expect(result.stdout).toContain('--repo <name>');
});
});
@@ -0,0 +1,259 @@
import { PassThrough } from 'node:stream';
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { CompatibleStdioServerTransport } from '../../src/mcp/compatible-stdio-transport.js';
function onceMessage(transport: CompatibleStdioServerTransport): Promise<any> {
return new Promise((resolve, reject) => {
transport.onmessage = (message) => resolve(message);
transport.onerror = (error) => reject(error);
});
}
describe('CompatibleStdioServerTransport', () => {
let stdin: PassThrough;
let stdout: PassThrough;
let transport: CompatibleStdioServerTransport;
beforeEach(() => {
stdin = new PassThrough();
stdout = new PassThrough();
transport = new CompatibleStdioServerTransport(stdin, stdout);
});
it('parses Content-Length framed initialize requests', async () => {
await transport.start();
const messagePromise = onceMessage(transport);
const body = JSON.stringify({
jsonrpc: '2.0',
id: 1,
method: 'initialize',
params: {
protocolVersion: '2024-11-05',
capabilities: {},
clientInfo: { name: 'codex', version: '0.1' },
},
});
stdin.write(`Content-Length: ${Buffer.byteLength(body, 'utf8')}\r\n\r\n${body}`);
await expect(messagePromise).resolves.toMatchObject({
method: 'initialize',
params: { clientInfo: { name: 'codex' } },
});
});
it('parses newline-delimited initialize requests', async () => {
await transport.start();
const messagePromise = onceMessage(transport);
stdin.write(`${JSON.stringify({
jsonrpc: '2.0',
id: 1,
method: 'initialize',
params: {
protocolVersion: '2024-11-05',
capabilities: {},
clientInfo: { name: 'cursor', version: '0.1' },
},
})}\n`);
await expect(messagePromise).resolves.toMatchObject({
method: 'initialize',
params: { clientInfo: { name: 'cursor' } },
});
});
it('responds with Content-Length framing after Content-Length input', async () => {
await transport.start();
const body = JSON.stringify({
jsonrpc: '2.0',
id: 1,
method: 'initialize',
params: {
protocolVersion: '2024-11-05',
capabilities: {},
clientInfo: { name: 'codex', version: '0.1' },
},
});
const messagePromise = onceMessage(transport);
stdin.write(`Content-Length: ${Buffer.byteLength(body, 'utf8')}\n\n${body}`);
await messagePromise;
const chunks: Buffer[] = [];
stdout.on('data', (chunk) => chunks.push(Buffer.from(chunk)));
await transport.send({ jsonrpc: '2.0', id: 1, result: { ok: true } });
const raw = Buffer.concat(chunks).toString('utf8');
expect(raw).toMatch(/^Content-Length: \d+\r\n\r\n/);
expect(raw).toContain('"ok":true');
});
it('reports malformed Content-Length headers once without looping forever', async () => {
await transport.start();
const onError = vi.fn();
transport.onerror = onError;
stdin.write('Content-Length:\r\n\r\n{}');
await new Promise((resolve) => setTimeout(resolve, 25));
expect(onError).toHaveBeenCalledTimes(1);
expect(onError.mock.calls[0]?.[0]).toBeInstanceOf(Error);
});
it('recovers after discarding a malformed Content-Length frame', async () => {
await transport.start();
const onError = vi.fn();
transport.onerror = onError;
stdin.write('Content-Length:\r\n\r\n{}');
await new Promise((resolve) => setTimeout(resolve, 25));
const body = JSON.stringify({
jsonrpc: '2.0',
id: 2,
method: 'initialize',
params: {
protocolVersion: '2024-11-05',
capabilities: {},
clientInfo: { name: 'recovery-client', version: '0.1' },
},
});
const messagePromise = onceMessage(transport);
stdin.write(`Content-Length: ${Buffer.byteLength(body, 'utf8')}\r\n\r\n${body}`);
await expect(messagePromise).resolves.toMatchObject({
method: 'initialize',
params: { clientInfo: { name: 'recovery-client' } },
});
expect(onError).toHaveBeenCalledTimes(1);
});
// ─── Security hardening regressions ──────────────────────────────
it('rejects Content-Length values exceeding the buffer cap', async () => {
await transport.start();
const onError = vi.fn();
transport.onerror = onError;
// 20 MB — exceeds the 10 MB MAX_BUFFER_SIZE
stdin.write('Content-Length: 20971520\r\n\r\n{}');
await new Promise((resolve) => setTimeout(resolve, 25));
expect(onError).toHaveBeenCalledTimes(1);
expect(onError.mock.calls[0]?.[0]?.message).toMatch(/exceeds maximum/i);
});
it('errors when read buffer exceeds maximum size in newline mode', async () => {
await transport.start();
const onError = vi.fn();
transport.onerror = onError;
// Send a JSON-starting chunk (triggers newline mode) with no newline,
// then keep appending until we exceed the 10 MB cap
const chunkSize = 1024 * 1024; // 1 MB
const chunk = Buffer.alloc(chunkSize, 0x61); // 'a' repeated
// First byte must be '{' to trigger newline framing detection
const first = Buffer.from('{' + 'a'.repeat(chunkSize - 1));
stdin.write(first);
for (let i = 0; i < 10; i++) {
stdin.write(chunk);
}
await new Promise((resolve) => setTimeout(resolve, 25));
expect(onError).toHaveBeenCalled();
const hasMaxSizeError = onError.mock.calls.some(
(call) => call[0] instanceof Error && /maximum size/i.test(call[0].message),
);
expect(hasMaxSizeError).toBe(true);
});
it('handles many consecutive empty lines without stack overflow', async () => {
await transport.start();
const onError = vi.fn();
transport.onerror = onError;
// First, seed the framing mode with a valid newline-delimited message
const seed = JSON.stringify({
jsonrpc: '2.0',
id: 1,
method: 'initialize',
params: {
protocolVersion: '2024-11-05',
capabilities: {},
clientInfo: { name: 'seed', version: '0.1' },
},
});
const seedPromise = onceMessage(transport);
stdin.write(seed + '\n');
await seedPromise;
// Now send 15K empty lines followed by a real message — this would
// stack-overflow with the old recursive readNewlineMessage
const followup = JSON.stringify({
jsonrpc: '2.0',
id: 2,
method: 'notifications/initialized',
params: {},
});
const messagePromise = onceMessage(transport);
stdin.write('\n'.repeat(15_000) + followup + '\n');
await expect(messagePromise).resolves.toMatchObject({
method: 'notifications/initialized',
});
expect(onError).not.toHaveBeenCalled();
});
it('rejects send() after transport is closed', async () => {
await transport.start();
await transport.close();
await expect(
transport.send({ jsonrpc: '2.0', id: 1, result: { ok: true } }),
).rejects.toThrow(/closed/i);
});
it('does not detect content-length framing from short ambiguous prefix', async () => {
await transport.start();
const onError = vi.fn();
transport.onerror = onError;
// Write only "cont" — fewer than 14 bytes, should NOT trigger
// content-length detection. Transport should wait for more data.
stdin.write(Buffer.from('cont'));
await new Promise((resolve) => setTimeout(resolve, 25));
// No message and no error — transport is waiting for more data
expect(onError).not.toHaveBeenCalled();
});
it('responds with newline framing after newline input', async () => {
await transport.start();
const messagePromise = onceMessage(transport);
stdin.write(`${JSON.stringify({
jsonrpc: '2.0',
id: 1,
method: 'initialize',
params: {
protocolVersion: '2024-11-05',
capabilities: {},
clientInfo: { name: 'cursor', version: '0.1' },
},
})}\n`);
await messagePromise;
const chunks: Buffer[] = [];
stdout.on('data', (chunk) => chunks.push(Buffer.from(chunk)));
await transport.send({ jsonrpc: '2.0', id: 1, result: { ok: true } });
const raw = Buffer.concat(chunks).toString('utf8');
expect(raw).toBe('{"jsonrpc":"2.0","id":1,"result":{"ok":true}}\n');
});
});
+21
View File
@@ -0,0 +1,21 @@
import { describe, expect, it, vi } from 'vitest';
import { createLazyAction } from '../../src/cli/lazy-action.js';
describe('createLazyAction', () => {
it('does not import target module until invoked', async () => {
const loader = vi.fn(async () => ({
run: vi.fn(async () => 'ok'),
}));
const action = createLazyAction(loader, 'run');
expect(loader).not.toHaveBeenCalled();
await expect(action('arg-1')).resolves.toBeUndefined();
expect(loader).toHaveBeenCalledTimes(1);
});
it('throws a clear error when export is not a function', async () => {
const action = createLazyAction(async () => ({ notAFunction: 'string-value' }), 'notAFunction');
await expect(action()).rejects.toThrow('notAFunction');
});
});