security-classifier.ts
258 lines · browse/src/security-classifier.ts at 3f79d81
1/**
2 * Security classifier — ML prompt injection detection (L4, TestSavantAI).
3 *
4 * This module is IMPORTED ONLY BY security-sidecar-entry.ts and runs inside
5 * the security sidecar subprocess (plain Node, spawned lazily by
6 * security-sidecar-client.ts). It CANNOT be imported by server.ts or any
7 * other module that ends up in the compiled browse binary, because
8 * @huggingface/transformers requires onnxruntime-node at runtime and that
9 * native module fails to dlopen from Bun's compiled-binary temp extraction
10 * dir.
11 *
12 * See: 2026-04-19-prompt-injection-guard.md Pre-Impl Gate 1 outcome.
13 *
14 * Layer:
15 * L4 (testsavant_content) — TestSavantAI BERT-small ONNX classifier on page
16 * snapshots and tool outputs. Detects indirect
17 * prompt injection + jailbreak attempts.
18 *
19 * The classifier degrades gracefully — if the model fails to load, the layer
20 * reports status 'degraded' and returns verdict 'safe' (fail-open). The
21 * caller (server.ts's /pty-inject-scan path) falls through to its
22 * L1-L3-only verdict; only the extra ML defense disappears.
23 */
25import * as fs from 'fs';
26import * as path from 'path';
27import { mkdirSecure } from './file-permissions';
28import { type LayerSignal } from './security';
29import { resolveStateRoot } from '../../lib/state-root';
31// ─── Model location + packaging ──────────────────────────────
33/**
34 * TestSavantAI prompt-injection-defender-small-v0-onnx.
35 *
36 * The HuggingFace repo stores model.onnx at the root, but @huggingface/transformers
37 * v4 expects it under an `onnx/` subdirectory. We stage the files into the expected
38 * layout at ~/.gstack/models/testsavant-small/ on first use.
39 *
40 * Files (fetched from HF on first use, cached for lifetime of install):
41 * config.json
42 * tokenizer.json
43 * tokenizer_config.json
44 * special_tokens_map.json
45 * vocab.txt
46 * onnx/model.onnx (~112MB)
47 */
48const MODELS_DIR = path.join(resolveStateRoot(), 'models');
49const TESTSAVANT_DIR = path.join(MODELS_DIR, 'testsavant-small');
50const TESTSAVANT_HF_URL = 'https://huggingface.co/testsavantai/prompt-injection-defender-small-v0-onnx/resolve/main';
51const TESTSAVANT_FILES = [
52 'config.json',
53 'tokenizer.json',
54 'tokenizer_config.json',
55 'special_tokens_map.json',
56 'vocab.txt',
57];
59// ─── Load state ──────────────────────────────────────────────
61type LoadState = 'uninitialized' | 'loading' | 'loaded' | 'failed';
63let testsavantState: LoadState = 'uninitialized';
64let testsavantClassifier: any = null;
65let testsavantLoadError: string | null = null;
67export interface ClassifierStatus {
68 testsavant: 'ok' | 'degraded' | 'off';
69}
71export function getClassifierStatus(): ClassifierStatus {
72 const testsavant =
73 testsavantState === 'loaded' ? 'ok' :
74 testsavantState === 'failed' ? 'degraded' :
75 'off';
76 return { testsavant };
77}
79// ─── Model download + staging ────────────────────────────────
81export async function downloadFile(url: string, dest: string): Promise<void> {
82 const res = await fetch(url);
83 if (!res.ok || !res.body) {
84 throw new Error(`Failed to fetch ${url}: ${res.status} ${res.statusText}`);
85 }
86 const tmp = `${dest}.tmp.${process.pid}`;
87 const writer = fs.createWriteStream(tmp);
88 // @ts-ignore — Node stream compat
89 const reader = res.body.getReader();
90 try {
91 let done = false;
92 while (!done) {
93 const chunk = await reader.read();
94 if (chunk.done) { done = true; break; }
95 writer.write(chunk.value);
96 }
97 await new Promise<void>((resolve, reject) => {
98 writer.end((err?: Error | null) => (err ? reject(err) : resolve()));
99 });
100 fs.renameSync(tmp, dest);
101 } catch (err) {
102 // Drop the half-written tmp so we don't ship a truncated model file to
103 // a retry's renameSync. Wait for the writer to close fully before
104 // unlinking: Node's createWriteStream lazily opens the FD and flushes
105 // buffered writes during destroy(), so a naive unlinkSync hits ENOENT
106 // first and the writer re-creates the file on the next tick.
107 await new Promise<void>((resolve) => {
108 writer.once('close', () => resolve());
109 writer.destroy();
110 });
111 try { fs.unlinkSync(tmp); } catch { /* nothing to clean */ }
112 throw err;
113 }
114}
116async function ensureTestsavantStaged(onProgress?: (msg: string) => void): Promise<void> {
117 mkdirSecure(path.join(TESTSAVANT_DIR, 'onnx'));
119 // Small config/tokenizer files
120 for (const f of TESTSAVANT_FILES) {
121 const dst = path.join(TESTSAVANT_DIR, f);
122 if (fs.existsSync(dst)) continue;
123 onProgress?.(`downloading ${f}`);
124 await downloadFile(`${TESTSAVANT_HF_URL}/${f}`, dst);
125 }
127 // Large model file — only download if missing. Put under onnx/ to match the
128 // layout @huggingface/transformers v4 expects.
129 const modelDst = path.join(TESTSAVANT_DIR, 'onnx', 'model.onnx');
130 if (!fs.existsSync(modelDst)) {
131 onProgress?.('downloading model.onnx (112MB) — first run only');
132 await downloadFile(`${TESTSAVANT_HF_URL}/model.onnx`, modelDst);
133 }
134}
136// ─── L4: TestSavantAI content classifier ─────────────────────
138/**
139 * Load the TestSavantAI classifier. Idempotent — concurrent calls share the
140 * same in-flight promise. Sets state to 'loaded' on success or 'failed' on error.
141 *
142 * Called by the sidecar on the first scan-page-content request to warm up.
143 * First call triggers the model download (~112MB from HuggingFace).
144 * Subsequent calls reuse the cached instance.
145 */
146let loadPromise: Promise<void> | null = null;
148export function loadTestsavant(onProgress?: (msg: string) => void): Promise<void> {
149 if (process.env.GSTACK_SECURITY_OFF === '1') {
150 testsavantState = 'failed';
151 testsavantLoadError = 'GSTACK_SECURITY_OFF=1 — ML classifier kill switch engaged';
152 return Promise.resolve();
153 }
154 if (testsavantState === 'loaded') return Promise.resolve();
155 if (loadPromise) return loadPromise;
156 testsavantState = 'loading';
157 loadPromise = (async () => {
158 try {
159 await ensureTestsavantStaged(onProgress);
160 // Dynamic import — keeps the module boundary clean so static analyzers
161 // don't pull @huggingface/transformers into compiled contexts.
162 onProgress?.('initializing classifier');
163 const { pipeline, env } = await import('@huggingface/transformers');
164 env.allowLocalModels = true;
165 env.allowRemoteModels = false;
166 env.localModelPath = MODELS_DIR;
167 testsavantClassifier = await pipeline(
168 'text-classification',
169 'testsavant-small',
170 { dtype: 'fp32' },
171 );
172 // TestSavantAI's tokenizer_config.json ships with model_max_length
173 // set to a huge placeholder (1e18) which disables automatic truncation
174 // in the TextClassificationPipeline. The underlying BERT-small has
175 // max_position_embeddings: 512 — passing anything longer throws a
176 // broadcast error. Override via _tokenizerConfig (the internal source
177 // the computed model_max_length getter reads from) so the pipeline's
178 // implicit truncation: true actually kicks in.
179 const tok = testsavantClassifier?.tokenizer as any;
180 if (tok?._tokenizerConfig) {
181 tok._tokenizerConfig.model_max_length = 512;
182 }
183 testsavantState = 'loaded';
184 } catch (err: any) {
185 testsavantState = 'failed';
186 testsavantLoadError = err?.message ?? String(err);
187 console.error('[security-classifier] Failed to load TestSavantAI:', testsavantLoadError);
188 }
189 })();
190 return loadPromise;
191}
193/**
194 * Strip HTML tags and collapse whitespace. TestSavantAI was trained on
195 * plain text, not markup — feeding it raw HTML massively reduces recall
196 * because all the tag noise dilutes the injection signal. Callers that
197 * already have plain text (page snapshot innerText, tool output strings)
198 * get no-op behavior; callers with HTML get the markup stripped.
199 */
200function htmlToPlainText(input: string): string {
201 // Fast path: if no angle brackets, it's already plain text.
202 if (!input.includes('<')) return input;
203 return input
204 .replace(/<(script|style)[^>]*>[\s\S]*?<\/\1>/gi, ' ') // drop script/style bodies entirely
205 .replace(/<[^>]+>/g, ' ') // drop tags
206 .replace(/ /g, ' ')
207 .replace(/&/g, '&')
208 .replace(/</g, '<')
209 .replace(/>/g, '>')
210 .replace(/"/g, '"')
211 .replace(/\s+/g, ' ')
212 .trim();
213}
215/**
216 * Scan text content for prompt injection. Intended for page snapshots, tool
217 * outputs, and other untrusted content blocks.
218 *
219 * Returns a LayerSignal. On load failure or classification error, returns
220 * confidence=0 with status flagged degraded — the verdict combiner in
221 * security.ts then falls through to 'safe' (fail-open by design).
222 *
223 * Note: TestSavantAI returns {label: 'INJECTION'|'SAFE', score: 0-1}. When
224 * label is 'SAFE', we return confidence=0 to the combiner. When label is
225 * 'INJECTION', we return the score directly.
226 */
227export async function scanPageContent(text: string): Promise<LayerSignal> {
228 if (!text || text.length === 0) {
229 return { layer: 'testsavant_content', confidence: 0 };
230 }
231 if (testsavantState !== 'loaded') {
232 return { layer: 'testsavant_content', confidence: 0, meta: { degraded: true } };
233 }
234 try {
235 // Normalize to plain text first — the classifier is trained on natural
236 // language, not HTML markup. A page with an injection buried in tag
237 // soup won't fire until we strip the noise.
238 const plain = htmlToPlainText(text);
239 // Character-level cap to avoid pathological memory use. The pipeline
240 // applies tokenizer truncation at 512 tokens (the BERT-small context
241 // limit — enforced via the model_max_length override in loadTestsavant)
242 // so the 4000-char cap is just a cheap upper bound. Real-world
243 // injection signals land in the first few hundred tokens anyway.
244 const input = plain.slice(0, 4000);
245 const raw = await testsavantClassifier(input);
246 const top = Array.isArray(raw) ? raw[0] : raw;
247 const label = top?.label ?? 'SAFE';
248 const score = Number(top?.score ?? 0);
249 if (label === 'INJECTION') {
250 return { layer: 'testsavant_content', confidence: score, meta: { label } };
251 }
252 return { layer: 'testsavant_content', confidence: 0, meta: { label, safeScore: score } };
253 } catch (err: any) {
254 testsavantState = 'failed';
255 testsavantLoadError = err?.message ?? String(err);
256 return { layer: 'testsavant_content', confidence: 0, meta: { degraded: true, error: testsavantLoadError } };
257 }
258}