mvroz/ai-firewall-gateway

AI Firewall & Data Privacy Gateway for LLM Inference

Upstream inference proxy that inspects raw incoming prompts and outgoing completions across four sequential stages:

raw prompt -> [S1 Pre-hook] -> [S2 RU-PII Engine] -> [S3 OWASP Classifier] -> [S4 Jailbreak Guard]
              deterministic     NER + checksums      trained multi-label      intent scoring
              ~1-3 ms           < 25 ms (CPU)        < 35 ms p95 (INT8 CPU)   < 45 ms p95 (CPU)
                          |             |                    |                     |
                          +-----------> POLICY ENGINE (allow / redact / block) <---+

Stages

1. S1 Deterministic pre-hook (src/privacy/prehook.py) — NFKC canonicalization, zero-width/format char stripping, Cyrillic/Latin homoglyph fold, plausibility-gated Base64/hex/ROT13 decode, context-gated leetspeak, chat-template delimiter escaping (never deletion — audit fidelity).
2. S2 Russian morphology-aware PII engine (src/privacy/) — checksum-validated identifiers (SNILS weighted-mod-101, INN-10/12 dual-coefficient mod-11, passport region plausibility, Luhn bank cards, phones, emails) + two interchangeable NER backends: Natasha 1.6.0 hybrid (fast path) and ONNX INT8 transformer (Gherman/bert-base-NER-Russian, name-part + address labels). Grammatical-agreement-preserving pseudonymizer via pymorphy3 (all 6 cases: nomn/gent/datv/accs/ablt/loct).
3. S3 Trained OWASP multi-label classifier (src/owasp_model/) — microsoft/deberta-v3-base fine-tuned with BCEWithLogitsLoss + per-class pos_weight on labels OWASP-LLM01 / OWASP-LLM02 / OWASP-LLM06 / BENIGN. No regex, no zero-shot. Served as INT8 dynamic-quantized ONNX.
4. S4 Adversarial jailbreak guard (src/guardrails/guard.py) — binary intent scorer warm-started from protectai/deberta-v3-base-prompt-injection-v2, scores the pre-hook decoded view, with multi-turn payload-splitting defense via a per-conversation rolling window (max of per-turn and window scores).

SLA targets (wired into harnesses as assertions)

Metric	Target	Harness
PER recall, oblique Russian cases	>= 0.94 per case	tests/test_ru_cases.py
S2 latency / sentence (CPU)	< 25 ms (Natasha)	scripts/benchmark_ner.py
S3 micro/macro F1	>= 0.90	src/owasp_model/evaluate.py
S3 FPR on benign engineering prompts	<= 2%	evaluate.py (threshold tuning constraint)
S3 FNR on active injections	<= 5%	evaluate.py
S3 INT8 p95 latency (CPU, seq 256)	< 35 ms	scripts/benchmark_onnx.py
S4 p95 latency	< 45 ms	scripts/benchmark_onnx.py

Data curation (OWASP corpus)

Verified public sources: deepset/prompt-injections (LLM01), xTRam1/safe-guard-prompt-injection (LLM01), jayavibhav/prompt-injection-safety (LLM01), benign: HuggingFaceH4/CodeAlpaca_20K, iamtarun/python_code_instructions_18k_alpaca. LLM02 from an embedded SecLists-derived payload seed set + mutation engine; LLM06 from a synthetic extraction-template family. Opt-in gated/non-commercial sources (allenai/wildjailbreak, allenai/wildguardmix, hackaprompt) behind --use_gated; lmsys/toxic-chat excluded (CC-BY-NC).

Discipline: MinHash near-dup removal -> iterative-stratified 70/15/15 split (scikit-multilearn) -> train-split-only augmentation (typo injection, case mutation, homoglyphs, base64/hex/rot13 wrapping; <= 4 variants/seed, synthetic capped at 35% of any class).

Runbook

pip install -r requirements.txt
python -m src.owasp_model.dataset --out_dir data/owasp --report
python -m src.owasp_model.train --model_name microsoft/deberta-v3-base \
    --data_dir data/owasp --output_dir runs/owasp-v1
python -m src.owasp_model.evaluate --model_dir runs/owasp-v1 --data_dir data/owasp
python scripts/export_onnx.py --model_dir runs/owasp-v1 --out artifacts/owasp-int8
python scripts/benchmark_onnx.py --model_dir runs/owasp-v1 --onnx_dir artifacts/owasp-int8
OWASP_MODEL_DIR=artifacts/owasp-int8/int8 GUARD_MODEL_DIR=artifacts/guard-int8/int8 \
    uvicorn src.api.gateway:app --port 8080
pytest tests/ -m "not slow" -v

API

• POST /v1/sanitize — full pipeline, blocking enforced, fail-closed on model errors. Body: {text, conversation_id?, direction: input|output}.
• POST /v1/audit — same signals, no mutation/blocking, fail-open with alarm (observability must not become a DoS lever).
• GET /healthz.

Security notes

• Audit log is append-only and hash-chained, records redacted text only; the per-request entity vault is memory-only with TTL.
• Model artifacts ship as (ONNX INT8 + thresholds.json + label_map.json + eval report) pinned by revision SHA.
• Known limits: multi-turn defense bounded by window size; homoglyph fold covers Cyrillic/Latin; scheduled re-training recommended (data flywheel via /v1/audit).