palaryn 0.3.6 → 0.4.2

package/src/server/gateway.ts

@@ -9,6 +9,8 @@ import { DLPScanner } from '../dlp/scanner';
 import { CompositeDLPScanner } from '../dlp/composite-scanner';
 import { DLPBackend } from '../dlp/interfaces';
 import { PromptInjectionBackend } from '../dlp/prompt-injection-backend';
+import { HeuristicScorerBackend } from '../dlp/heuristic-scorer';
+import { scorePromptInjection } from '../dlp/heuristic-scorer';
 import { TruffleHogBackend } from '../dlp/trufflehog-backend';
 import { BudgetManager, CostRecord } from '../budget/manager';
 import { UsageExtractor } from '../budget/usage-extractor';
@@ -29,6 +31,7 @@ import { UsageData } from '../types/tool-result';
 import { GatewayMetrics } from '../metrics';
 import { GatewayTracer } from '../tracing';
 import { AnomalyDetector } from '../anomaly';
+import { LlmPromptInjectionClassifier } from '../dlp/llm-classifier';
 import { log as devLog, logger } from './logger';
 
 export interface PreExecuteResult {
@@ -118,6 +121,7 @@ export class Gateway {
   private rateLimitConfigStore?: RateLimitConfigStore;
   private budgetConfigStore?: BudgetConfigStore;
   private usageExtractor: UsageExtractor;
+  private llmClassifier?: LlmPromptInjectionClassifier;
   private inFlightCleanupInterval?: ReturnType<typeof setInterval>;
   /**
    * Tracks tool_call_ids currently being processed to prevent TOCTOU races.
@@ -137,6 +141,7 @@ export class Gateway {
       dlpBackends.push(new PromptInjectionBackend({
         scan_output: config.dlp.scan_output,
       }));
+      dlpBackends.push(new HeuristicScorerBackend());
     }
     if (config.dlp.trufflehog?.enabled) {
       dlpBackends.push(new TruffleHogBackend({
@@ -165,6 +170,11 @@ export class Gateway {
       this.opaEngine = new OPAEngine(config.policy.opa);
     }
 
+    // Set up LLM-based prompt injection classifier if enabled
+    if (config.dlp.llm_classifier?.enabled) {
+      this.llmClassifier = new LlmPromptInjectionClassifier(config.dlp.llm_classifier);
+    }
+
     this.usageExtractor = new UsageExtractor(config.budget.token_pricing);
 
     // Set up executor registry with HTTP as default + catch-all fallback
@@ -355,6 +365,68 @@ export class Gateway {
       }
     }
 
+    // Heuristic scoring — force LLM classifier when structural signals are elevated
+    let forceLlmClassification = false;
+    const inputText = JSON.stringify(toolCall);
+    const heuristicResult = scorePromptInjection(inputText);
+    if (heuristicResult.score >= 0.4) {
+      forceLlmClassification = true;
+      devLog.pipelineStep('🔍', 'HEURISTIC_SCORER',
+        `score=${heuristicResult.score.toFixed(2)} signals=[${heuristicResult.signals.join(',')}]`);
+    }
+
+    // LLM-based prompt injection classification on INPUT (async, runs after sync DLP scan)
+    if ((this.llmClassifier && this.config.dlp.llm_classifier?.scan_input !== false) || (forceLlmClassification && this.llmClassifier)) {
+      const llmInputStart = Date.now();
+      const llmInputClassifications = await asyncChildSpan(otel, 'gateway.llm_classifier_input', () =>
+        this.llmClassifier!.classify(inputText)
+      );
+      stepTimings.llm_classifier_input = Date.now() - llmInputStart;
+
+      if (llmInputClassifications.length > 0) {
+        const llmDetections = LlmPromptInjectionClassifier.toDLPDetections(llmInputClassifications, inputText);
+        const llmDetectionNames = llmDetections.map(d => d.pattern_name);
+        argsDlp.detected.push(...llmDetectionNames);
+
+        for (const d of llmDetections) {
+          argsDlp.severity = this.maxSeverity(argsDlp.severity, d.severity);
+        }
+
+        this.auditLogger.logDLPScanned(toolCall, llmDetectionNames, argsDlp.severity, 0);
+        for (const name of llmDetectionNames) {
+          this.metrics?.recordDLPDetection(name, argsDlp.severity);
+        }
+
+        // Re-check prompt injection blocking with LLM classifier detections
+        if (piAction === 'block') {
+          const llmPiDetections = argsDlp.detected.filter((d: string) => d.startsWith('llm_classifier_'));
+          if (llmPiDetections.length > 0) {
+            const threshold = this.config.dlp.prompt_injection_block_threshold || 'high';
+            const severityRank: Record<string, number> = { low: 0, medium: 1, high: 2 };
+            const thresholdRank = severityRank[threshold] ?? 2;
+            const maxSeverityRank = severityRank[argsDlp.severity] ?? 0;
+
+            if (maxSeverityRank >= thresholdRank) {
+              devLog.pipelineStep('🛡️', 'LLM_CLASSIFIER_INPUT_BLOCK',
+                `Blocked: ${llmPiDetections.join(', ')} (severity: ${argsDlp.severity}, threshold: ${threshold})`);
+              const durationSec = (Date.now() - startTime) / 1000;
+              this.metrics?.recordRequest('blocked', toolCall.tool.name, toolCall.tool.capability, durationSec);
+              const result = this.buildResult(toolCall, 'blocked', {
+                decision: 'deny',
+                rule_id: 'llm_classifier_input_block',
+                rule_name: 'LLM classifier detected prompt injection in input',
+                reasons: [`LLM classifier detected: ${llmPiDetections.join(', ')}`],
+              }, startTime, undefined,
+                `Blocked by LLM classifier: ${llmPiDetections.join(', ')} (severity: ${argsDlp.severity})`,
+                undefined, argsDlp);
+              devLog.pipelineEnd('blocked', Date.now() - startTime);
+              return { allowed: false, result, stepTimings, startTime };
+            }
+          }
+        }
+      }
+    }
+
     // Policy evaluation — DLP context is passed so DLP-conditioned rules
     // compete with all other rules in a single priority-ordered pass.
     stepStart = Date.now();
@@ -533,6 +605,76 @@ export class Gateway {
     }
     stepTimings.dlp_out = Date.now() - stepStart;
 
+    // LLM-based prompt injection classification (async, runs after sync DLP scan)
+    if (this.llmClassifier && output.body && this.config.dlp.llm_classifier?.scan_output !== false) {
+      const llmStart = Date.now();
+      const text = typeof output.body === 'string' ? output.body : JSON.stringify(output.body);
+      const llmClassifications = await asyncChildSpan(otel, 'gateway.llm_classifier', () =>
+        this.llmClassifier!.classify(text)
+      );
+      stepTimings.llm_classifier = Date.now() - llmStart;
+
+      if (llmClassifications.length > 0) {
+        const llmDetections = LlmPromptInjectionClassifier.toDLPDetections(llmClassifications, text);
+        const llmDetectionNames = llmDetections.map(d => d.pattern_name);
+        outputDlp.detected.push(...llmDetectionNames);
+
+        // Recalculate severity: take the max of existing and LLM detections
+        for (const d of llmDetections) {
+          outputDlp.severity = this.maxSeverity(outputDlp.severity, d.severity);
+        }
+
+        this.auditLogger.logDLPScanned(toolCall, llmDetectionNames, outputDlp.severity, 0);
+        for (const name of llmDetectionNames) {
+          this.metrics?.recordDLPDetection(name, outputDlp.severity);
+        }
+      }
+    }
+
+    // Prompt injection blocking check on output (regex + LLM classifier detections)
+    const piOutputAction = this.config.dlp.prompt_injection_action || 'log';
+    if (piOutputAction === 'block' && outputDlp.detected.length > 0) {
+      const piDetections = outputDlp.detected.filter((d: string) =>
+        d.startsWith('prompt_injection_') || d.startsWith('llm_classifier_')
+      );
+      if (piDetections.length > 0) {
+        const threshold = this.config.dlp.prompt_injection_block_threshold || 'medium';
+        const severityRank: Record<string, number> = { low: 0, medium: 1, high: 2 };
+        const thresholdRank = severityRank[threshold] ?? 1;
+        const maxSeverityRank = severityRank[outputDlp.severity] ?? 0;
+
+        if (maxSeverityRank >= thresholdRank) {
+          // Release budget reservation since we're blocking
+          if (reservationKey) {
+            this.budgetManager.commitReservation(reservationKey, 0);
+          }
+
+          devLog.pipelineStep('🛡️', 'PROMPT_INJECTION_OUTPUT_BLOCK',
+            `Blocked output: ${piDetections.join(', ')} (severity: ${outputDlp.severity}, threshold: ${threshold})`);
+          const durationSec = (Date.now() - startTime) / 1000;
+          this.metrics?.recordRequest('blocked', toolCall.tool.name, toolCall.tool.capability, durationSec);
+          const defaultPolicy = policyResult || { decision: 'allow' as const, rule_id: 'passthrough', rule_name: 'Passthrough', reasons: [] };
+          const argsDlpSafe = argsDlp || { detected: [], redactions: [], severity: 'low' as DLPSeverity };
+          const mergedDlp = {
+            detected: [...new Set([...argsDlpSafe.detected, ...outputDlp.detected])],
+            redactions: [...argsDlpSafe.redactions, ...outputDlp.redactions],
+            severity: this.maxSeverity(argsDlpSafe.severity, outputDlp.severity),
+          };
+          const result = this.buildResult(toolCall, 'blocked', {
+            decision: 'deny',
+            rule_id: 'prompt_injection_output_block',
+            rule_name: 'Prompt injection detected in output',
+            reasons: [`Prompt injection detected in output: ${piDetections.join(', ')}`],
+          }, startTime, undefined,
+            `Blocked by prompt injection detection in output: ${piDetections.join(', ')} (severity: ${outputDlp.severity})`,
+            undefined, mergedDlp);
+          this.auditLogger.logToolResultReturned(toolCall, 'blocked', Date.now() - startTime, { prompt_injection_blocked: true, detections: piDetections });
+          devLog.pipelineEnd('blocked', Date.now() - startTime);
+          return result;
+        }
+      }
+    }
+
     // Extract usage data from response
     const headerUsage = this.usageExtractor.extractFromHeaders(output.headers);
     const bodyUsage = this.usageExtractor.extractFromBody(output.body);

package/src/types/config.ts

@@ -140,6 +140,16 @@ export interface OPAConfig {
   package_name?: string;
 }
 
+export interface LlmClassifierConfig {
+  enabled: boolean;
+  model?: string;
+  confidence_threshold?: number;
+  /** Run LLM classifier on input/args (default: true) */
+  scan_input?: boolean;
+  /** Run LLM classifier on output (default: true) */
+  scan_output?: boolean;
+}
+
 export interface DLPConfig {
   enabled: boolean;
   scan_args: boolean;
@@ -149,7 +159,7 @@ export interface DLPConfig {
   prompt_injection_detection?: boolean;
   /** Action when prompt injection is detected: 'log' (default), 'flag', or 'block' */
   prompt_injection_action?: 'log' | 'flag' | 'block';
-  /** Minimum severity to trigger blocking (when action is 'block'). Default: 'high' */
+  /** Minimum severity to trigger blocking (when action is 'block'). Default: 'medium' */
   prompt_injection_block_threshold?: 'medium' | 'high';
   /** Response mode when injection is blocked: 'deny' (default), 'sanitize', or 'require_approval' */
   prompt_injection_response?: 'deny' | 'sanitize' | 'require_approval';
@@ -159,6 +169,10 @@ export interface DLPConfig {
     binary_path?: string;
     timeout_ms?: number;
   };
+  /** Maximum scan depth for nested structures (default: 64) */
+  max_scan_depth?: number;
+  /** LLM-based prompt injection classifier (async, semantic analysis) */
+  llm_classifier?: LlmClassifierConfig;
 }
 
 export interface AuditConfig {