fix: bound resolver caches, validate MCP input sizes, add integration tests (#213)
Replace the 7 unbounded ReferenceResolver Map caches with a bounded LRU (env-tunable via CODEGRAPH_RESOLVER_CACHE_SIZE) so memory stays flat on large codebases, and add length caps on MCP tool string inputs (query/task/symbol + projectPath/path/pattern) to prevent oversized-payload DoS. Includes LRU, MCP-input-limit, and full-pipeline integration tests. Closes #213
This commit is contained in:
+71
-2
@@ -22,6 +22,22 @@ import { join } from 'path';
|
||||
/** Maximum output length to prevent context bloat (characters) */
|
||||
const MAX_OUTPUT_LENGTH = 15000;
|
||||
|
||||
/**
|
||||
* Maximum length for free-form string inputs (query, task, symbol).
|
||||
* Bounds memory and CPU when a buggy or hostile MCP client sends a
|
||||
* huge payload — without this an attacker could ship a 100MB string
|
||||
* and force a full FTS5 scan / OOM the server. 10 000 characters is
|
||||
* far beyond any realistic legitimate query.
|
||||
*/
|
||||
const MAX_INPUT_LENGTH = 10_000;
|
||||
|
||||
/**
|
||||
* Maximum length for path-like string inputs (projectPath, path
|
||||
* filter, glob pattern). Paths beyond a few thousand chars are
|
||||
* never legitimate and signal abuse or a bug upstream.
|
||||
*/
|
||||
const MAX_PATH_LENGTH = 4_096;
|
||||
|
||||
/**
|
||||
* Rust path roots that have no file-system equivalent — `crate` is the
|
||||
* current crate, `super` is the parent module, `self` is the current
|
||||
@@ -609,12 +625,46 @@ export class ToolHandler {
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate that a value is a non-empty string
|
||||
* Validate that a value is a non-empty string within length bounds.
|
||||
*
|
||||
* The `maxLength` cap protects against MCP clients that ship huge
|
||||
* payloads (10MB+ query strings either by accident or maliciously).
|
||||
* Without this, a single oversized input can pin the FTS5 index or
|
||||
* exhaust memory before any real work runs.
|
||||
*/
|
||||
private validateString(value: unknown, name: string): string | ToolResult {
|
||||
private validateString(
|
||||
value: unknown,
|
||||
name: string,
|
||||
maxLength: number = MAX_INPUT_LENGTH
|
||||
): string | ToolResult {
|
||||
if (typeof value !== 'string' || value.length === 0) {
|
||||
return this.errorResult(`${name} must be a non-empty string`);
|
||||
}
|
||||
if (value.length > maxLength) {
|
||||
return this.errorResult(
|
||||
`${name} exceeds maximum length of ${maxLength} characters (got ${value.length})`
|
||||
);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate an optional path-like string input. Returns the value if
|
||||
* valid (or undefined), or a ToolResult with the error.
|
||||
*/
|
||||
private validateOptionalPath(
|
||||
value: unknown,
|
||||
name: string
|
||||
): string | undefined | ToolResult {
|
||||
if (value === undefined || value === null) return undefined;
|
||||
if (typeof value !== 'string') {
|
||||
return this.errorResult(`${name} must be a string`);
|
||||
}
|
||||
if (value.length > MAX_PATH_LENGTH) {
|
||||
return this.errorResult(
|
||||
`${name} exceeds maximum length of ${MAX_PATH_LENGTH} characters (got ${value.length})`
|
||||
);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
@@ -623,6 +673,25 @@ export class ToolHandler {
|
||||
*/
|
||||
async execute(toolName: string, args: Record<string, unknown>): Promise<ToolResult> {
|
||||
try {
|
||||
// Cross-cutting input validation. All tools accept an optional
|
||||
// `projectPath` and most accept either `query`, `task`, or
|
||||
// `symbol` — bound their lengths centrally so individual handlers
|
||||
// can stay focused on tool-specific logic.
|
||||
const pathCheck = this.validateOptionalPath(args.projectPath, 'projectPath');
|
||||
if (typeof pathCheck === 'object' && pathCheck !== undefined) {
|
||||
return pathCheck;
|
||||
}
|
||||
// The `path` and `pattern` properties used by codegraph_files are
|
||||
// also path-shaped — apply the same cap.
|
||||
if (args.path !== undefined) {
|
||||
const check = this.validateOptionalPath(args.path, 'path');
|
||||
if (typeof check === 'object' && check !== undefined) return check;
|
||||
}
|
||||
if (args.pattern !== undefined) {
|
||||
const check = this.validateOptionalPath(args.pattern, 'pattern');
|
||||
if (typeof check === 'object' && check !== undefined) return check;
|
||||
}
|
||||
|
||||
switch (toolName) {
|
||||
case 'codegraph_search':
|
||||
return await this.handleSearch(args);
|
||||
|
||||
+41
-7
@@ -22,6 +22,24 @@ import { detectFrameworks } from './frameworks';
|
||||
import { loadProjectAliases, type AliasMap } from './path-aliases';
|
||||
import { logDebug } from '../errors';
|
||||
import type { ReExport } from './types';
|
||||
import { LRUCache } from './lru-cache';
|
||||
|
||||
/**
|
||||
* Cache size limits. Each per-resolver cache is bounded so memory
|
||||
* stays flat on large codebases (20k+ files). Sizes were chosen to
|
||||
* cover the working set for typical resolution batches without
|
||||
* exceeding a few hundred MB worst-case. Override via the env var
|
||||
* `CODEGRAPH_RESOLVER_CACHE_SIZE` (single integer applied to all
|
||||
* caches) when tuning for very large or very small projects.
|
||||
*/
|
||||
const DEFAULT_CACHE_LIMIT = 5_000;
|
||||
function resolveCacheLimit(): number {
|
||||
const raw = process.env.CODEGRAPH_RESOLVER_CACHE_SIZE;
|
||||
if (!raw) return DEFAULT_CACHE_LIMIT;
|
||||
const parsed = Number.parseInt(raw, 10);
|
||||
if (Number.isFinite(parsed) && parsed > 0) return parsed;
|
||||
return DEFAULT_CACHE_LIMIT;
|
||||
}
|
||||
|
||||
// Re-export types
|
||||
export * from './types';
|
||||
@@ -121,13 +139,16 @@ export class ReferenceResolver {
|
||||
private queries: QueryBuilder;
|
||||
private context: ResolutionContext;
|
||||
private frameworks: FrameworkResolver[] = [];
|
||||
private nodeCache: Map<string, Node[]> = new Map(); // per-file node cache (bounded)
|
||||
private fileCache: Map<string, string | null> = new Map(); // per-file content cache (bounded)
|
||||
private importMappingCache: Map<string, ImportMapping[]> = new Map();
|
||||
private reExportCache: Map<string, ReExport[]> = new Map();
|
||||
private nameCache: Map<string, Node[]> = new Map(); // name → nodes cache
|
||||
private lowerNameCache: Map<string, Node[]> = new Map(); // lower(name) → nodes cache
|
||||
private qualifiedNameCache: Map<string, Node[]> = new Map(); // qualified_name → nodes cache
|
||||
// All per-resolver caches are LRU-bounded. Previously these were
|
||||
// unbounded Maps that grew with every distinct lookup and OOM'd on
|
||||
// codebases with 20k+ files (see issue: unbounded cache growth).
|
||||
private nodeCache: LRUCache<string, Node[]>; // per-file node cache
|
||||
private fileCache: LRUCache<string, string | null>; // per-file content cache
|
||||
private importMappingCache: LRUCache<string, ImportMapping[]>;
|
||||
private reExportCache: LRUCache<string, ReExport[]>;
|
||||
private nameCache: LRUCache<string, Node[]>; // name → nodes cache
|
||||
private lowerNameCache: LRUCache<string, Node[]>; // lower(name) → nodes cache
|
||||
private qualifiedNameCache: LRUCache<string, Node[]>; // qualified_name → nodes cache
|
||||
private knownNames: Set<string> | null = null; // all known symbol names for fast pre-filtering
|
||||
private knownFiles: Set<string> | null = null;
|
||||
private cachesWarmed = false;
|
||||
@@ -139,6 +160,19 @@ export class ReferenceResolver {
|
||||
constructor(projectRoot: string, queries: QueryBuilder) {
|
||||
this.projectRoot = projectRoot;
|
||||
this.queries = queries;
|
||||
|
||||
const limit = resolveCacheLimit();
|
||||
// The content cache is heavier (full file text), so we give it a
|
||||
// smaller budget than the metadata caches.
|
||||
const contentLimit = Math.max(64, Math.floor(limit / 5));
|
||||
this.nodeCache = new LRUCache(limit);
|
||||
this.fileCache = new LRUCache(contentLimit);
|
||||
this.importMappingCache = new LRUCache(limit);
|
||||
this.reExportCache = new LRUCache(limit);
|
||||
this.nameCache = new LRUCache(limit);
|
||||
this.lowerNameCache = new LRUCache(limit);
|
||||
this.qualifiedNameCache = new LRUCache(limit);
|
||||
|
||||
this.context = this.createContext();
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
/**
|
||||
* Simple LRU cache backed by JavaScript's insertion-ordered Map.
|
||||
*
|
||||
* Used by ReferenceResolver to bound the per-resolver caches that
|
||||
* previously grew without limit and OOM'd on large codebases (20k+
|
||||
* files). Each cache is sized independently — see `index.ts` for
|
||||
* the chosen limits per cache type.
|
||||
*
|
||||
* Eviction is plain LRU: on `set`, if the cache is full, the
|
||||
* least-recently-used entry (the first one in iteration order) is
|
||||
* evicted. Touching via `get` moves the entry to the most-recently-used
|
||||
* position so hot keys survive eviction passes.
|
||||
*/
|
||||
export class LRUCache<K, V> {
|
||||
private readonly max: number;
|
||||
private readonly store = new Map<K, V>();
|
||||
|
||||
constructor(max: number) {
|
||||
if (!Number.isFinite(max) || max <= 0) {
|
||||
throw new Error(`LRUCache max must be a positive finite number, got ${max}`);
|
||||
}
|
||||
this.max = Math.floor(max);
|
||||
}
|
||||
|
||||
get size(): number {
|
||||
return this.store.size;
|
||||
}
|
||||
|
||||
get(key: K): V | undefined {
|
||||
const value = this.store.get(key);
|
||||
if (value === undefined) {
|
||||
// Distinguish "missing" from "stored undefined" by checking has().
|
||||
// We don't store undefined in practice, but be defensive.
|
||||
return this.store.has(key) ? value : undefined;
|
||||
}
|
||||
// Refresh recency by re-inserting.
|
||||
this.store.delete(key);
|
||||
this.store.set(key, value);
|
||||
return value;
|
||||
}
|
||||
|
||||
has(key: K): boolean {
|
||||
return this.store.has(key);
|
||||
}
|
||||
|
||||
set(key: K, value: V): void {
|
||||
if (this.store.has(key)) {
|
||||
this.store.delete(key);
|
||||
} else if (this.store.size >= this.max) {
|
||||
// Evict the oldest entry — first key in iteration order.
|
||||
const oldest = this.store.keys().next().value;
|
||||
if (oldest !== undefined) {
|
||||
this.store.delete(oldest);
|
||||
}
|
||||
}
|
||||
this.store.set(key, value);
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
this.store.clear();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user