vikramdgx

injection-vulnerability-detector

README

License: apache-2.0

What it does

Given a code snippet, the model identifies whether it contains an injection vulnerability and, if so, classifies the specific CWE (Common Weakness Enumeration) type. It outputs structured JSON with a verdict, CWE ID, vulnerability type, tainted data-flow analysis, an explanation, and a fix suggestion.

Supported vulnerability types (9 CWEs)

Table with columns: CWE, Vulnerability Type
CWE	Vulnerability Type
CWE-89	SQL Injection
CWE-78	OS Command Injection
CWE-94	Code Injection
CWE-611	XXE (XML External Entity)
CWE-90	LDAP Injection
CWE-643	XPath Injection
CWE-917	Expression Language Injection
CWE-1336	Server-Side Template Injection (SSTI)
CWE-113	CRLF Injection

Evaluation results

Our primary metrics are computed on novel code only (template clones removed; the full-set figure is shown for comparison) — test samples verified to be structurally distinct from all training data. Of 1572 total test samples, 1007 (64.1%) were identified as structural clones of training data and excluded (see Evaluation Methodology below).

Evaluation: three levels of rigor

Table with columns: Evaluation set, Precision, Recall, F1
Evaluation set	Precision	Recall	F1
Full test set (includes template clones)	96.2%	83.6%	89.5%
Novel code, all injection-related CWEs (n=565)	82.8%	46.6%	59.7%
Novel code, in-scope 9 target CWEs (n=496)	81.7%	58.3%	68.0%

Our primary metric is the in-scope novel row — performance on the 9 CWEs the model targets, on code it has never seen in any structural form. The full-set figure is what most detectors report; we consider it inflated by template-clone memorisation (see Evaluation Methodology). The all-CWE novel row additionally counts out-of-scope vulnerability types the model does not target.

Primary metrics (novel code, in-scope 9 CWEs, n=496)

Table with columns: Metric, Score
Metric	Score
Precision	81.7%
Recall	58.3%
F1 Score	68.0%
Accuracy	78.0%
JSON Parse Rate	94.9%

Per-CWE performance (novel code)

Table with columns: CWE, Vulnerability, Precision, Recall, F1, Samples
CWE	Vulnerability	Precision	Recall	F1	Samples
CWE-89	SQL Injection	76.5%	67.2%	71.6%	58
CWE-78	OS Command Injection	75.9%	42.3%	54.3%	52
CWE-94	Code Injection	77.4%

High-confidence mode (zero false positives)

The tainted_flow field doubles as a confidence signal. Accepting a VULNERABLE verdict only when the model names a concrete source and sink filters out its weakest calls:

Table with columns: Mode, Precision, Recall, F1
Mode	Precision	Recall	F1
Standard (all VULNERABLE verdicts)	81.7%	58.3%	68.0%
High-confidence (concrete `tainted_flow` required)	100.0%	22.1%	36.2%

On the in-scope novel test set this yielded 44 true positives and 0 false positives (n=496). Use high-confidence mode where alert fatigue matters more than coverage (CI gating, auto-filing issues); use standard mode for triage sweeps where a human reviews each finding.

python
import json

r = json.loads(response)
flow = r.get("tainted_flow") or {}
high_confidence = r.get("verdict") == "VULNERABLE" and flow.get("source") and flow.get("sink")

Evaluation methodology

Standard group-aware splitting (by CVE ID / project) prevents the same code from appearing in both train and test, but it does not prevent template clones — structurally identical code differing only in variable names, string literals, and numeric constants — from leaking across the split. A skeleton-hashing analysis revealed that 64.1% of the initial test set were structural clones of training samples, scoring near-perfect F1 from memorisation rather than generalisation.

Skeleton hashing method: strip comments, string literals, identifiers, and numeric literals from each code sample, then SHA-256 hash the normalised skeleton. Any test sample whose skeleton matches a training sample is classified as a clone and excluded from the metrics reported above.

The scores above therefore reflect performance on genuinely novel code that the model has never seen in any structural form during training. We report these numbers — not the inflated full-test-set figures — because they are what matters for real-world deployment.

Usage

Quick start with transformers + PEFT

python
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load base model + LoRA adapter
base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-Coder-7B-Instruct",
    torch_dtype="auto",
    device_map="auto",
)
model = PeftModel.from_pretrained(base_model, "vikramdgx/injection-vulnerability-detector")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")

# Analyze code — use the EXACT prompt format the model was trained on
code_snippet = '''import sqlite3
def get_user(username):
    conn = sqlite3.connect("app.db")
    query = f"SELECT * FROM users WHERE name = '{username}'"
    return conn.execute(query).fetchone()
'''

SYSTEM_PROMPT = (
    "You are a code security analyzer specialized in detecting injection "
    "vulnerabilities. Analyze the provided code and respond with a JSON "
    "object containing: verdict, cwe, vulnerability_type, tainted_flow, "
    "explanation, and fix_suggestion."
)

user_message = (
    "Analyze the following code for injection vulnerabilities. "
    "Respond with JSON.\n\n"
    + code_snippet
)

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": user_message},
]

text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=768, do_sample=False)
response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

Expected output format

The model returns structured JSON matching its training schema:

json
{
  "verdict": "VULNERABLE",
  "cwe": "CWE-89",
  "vulnerability_type": "SQL Injection",
  "tainted_flow": {
    "source": "username parameter",
    "sink": "conn.execute(query)",
    "sanitizer": null
  },
  "explanation": "User input is directly interpolated into SQL query string via f-string without parameterization, enabling SQL injection.",
  "fix_suggestion": "Use parameterized queries: conn.execute('SELECT * FROM users WHERE name = ?', (username,))"
}

For safe code, the model returns:

json
{
  "verdict": "SAFE",
  "cwe": null,
  "vulnerability_type": null,
  "tainted_flow": null,
  "explanation": "No injection vulnerability is present; untrusted input is not passed unsanitized to a sensitive sink.",
  "fix_suggestion": null
}

Training details

Hardware

NVIDIA DGX Spark (Grace Blackwell GB10)
128 GB unified LPDDR5X memory
CUDA 13.0, compute capability sm_121

Configuration

Table with columns: Parameter, Value
Parameter	Value
Base model	Qwen/Qwen2.5-Coder-7B-Instruct
Base model checkpoint used	`unsloth/Qwen2.5-Coder-7B-Instruct` (Unsloth mirror of the same Apache-2.0 weights; the adapter loads against either)
Method	LoRA (PEFT)
LoRA rank (r)	16
LoRA alpha	32
Target modules	q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
Trainable parameters	40.3M (0.53% of total)
Epochs

Why we built a custom dataset

Existing open-source vulnerability datasets (CVEfixes, DiverseVul, BigVul, etc.) provide valuable real-world code samples, but they share a common gap: they lack fine-grained, CWE-specific labels for injection subtypes. Most label code broadly as "vulnerable" or group all injections under generic categories like CWE-74 or CWE-20 — without distinguishing SQL injection (CWE-89) from command injection (CWE-78) from SSTI (CWE-1336). For niche injection types like LDAP injection (CWE-90), XPath injection (CWE-643), or EL injection (CWE-917), labeled samples in public datasets are extremely scarce (often single digits).

This makes it impossible to train a detector that both identifies vulnerabilities AND classifies the specific CWE type — which is what security engineers actually need for triage.

Our approach: template-based synthetic generation

To solve this, we built a procedural code generator that produces CWE-specific injection samples at scale — no LLM in the loop, pure template expansion:

151 handcrafted code templates covering all 9 target CWEs, each with @@PLACEHOLDER@@ tokens for variable names, function names, table names, database fields, and code patterns
Randomized variable pools (200+ variable names, 100+ function names, 50+ table names per CWE) ensure each generated sample is syntactically unique
Paired generation: every vulnerable template has a corresponding safe version that uses parameterized queries, input validation, or proper escaping — teaching the model the difference, not just the pattern
~5,000 synthetic samples generated in seconds, balanced across all 9 CWEs including the rare ones that have near-zero representation in public datasets

This approach is deterministic, reproducible, and produces exactly the CWE distribution the model needs — no data collection bottleneck, no labeling errors, no class imbalance.

Dataset composition

The final training set combines three sources:

Real-world vulnerabilities — filtered from open-source CVE/vulnerability datasets, keeping only injection-related CWEs. Provides realistic code patterns from production software.
Template-based synthetic data (~5,000 samples) — our original procedural generator. Fills the CWE-specific gap that public datasets leave open, especially for rare injection types (LDAP, XPath, EL, SSTI, Header Injection).
Hard negatives — safe code samples including post-patch fixes and non-vulnerable functions from vulnerability-adjacent codebases. Teaches the model what secure code looks like.

Group-aware splitting ensures no data leakage between train and test sets (samples sharing a CVE ID or project stay together). Additionally, a skeleton-hashing pass excludes structural template clones from the evaluation set (see Evaluation Methodology above).

Dataset sources and attribution

This model was trained on data from the following open-source datasets, combined with our original synthetic generation. We gratefully acknowledge the dataset creators:

Table with columns: Dataset, Source, License, Role
Dataset	Source	License	Role
CVEfixes	Bhandari et al.	Apache 2.0 (data: CC BY 4.0)	Real-world vuln + patch pairs
DiverseVul	Chen & Bhatt (RAID 2023)	Not specified on HF card	Safe code (non-vuln functions)
Code Vulnerability Security DPO	CyberNative AI

License notes:

The DiverseVul dataset does not declare an explicit license on its HuggingFace card as of this writing. This model's own adapter weights are original work released under Apache 2.0, but users should check the current licensing status of these upstream datasets before commercial deployment.

Limitations

Injection-only scope: This model detects 9 injection-related CWEs. It does not cover other vulnerability classes (buffer overflow, authentication, crypto, etc.).
Synthetic training bias: The model is trained partly on template-generated code. While this solves the CWE-distribution problem, performance on novel real-world patterns is lower than on template-similar code, reflecting the generalisation gap that template-based training introduces. See Evaluation Results for exact numbers.
Rare CWEs are synthetic-validated only: The five rare injection types (LDAP, XPath, EL, SSTI, Header) have near-zero representation in public vulnerability datasets. Our evaluation on novel code therefore cannot validate these CWEs — their coverage relies entirely on synthetic templates. Real-world performance on these types is unknown.
CWE classification accuracy: While detection (vulnerable vs. safe) is reliable, the specific CWE label assigned to a detected vulnerability may be incorrect in some cases — particularly between similar injection types (e.g. CWE-94 code injection vs CWE-1336 template injection, or CWE-78 vs CWE-94).
Code context: The model analyzes individual functions/snippets. It cannot trace data flow across files or understand application-level sanitization.
Language coverage: Primarily trained on Python, Java, PHP, C/C++, and JavaScript. Performance on other languages may vary.
Not a replacement for manual review: Use as a triage/prioritization tool alongside established SAST tooling and expert code review.

License

This model adapter is released under the Apache 2.0 license. The base model (Qwen2.5-Coder-7B-Instruct) is also Apache 2.0.

Citation

If you use this model in your research, please cite:

bibtex
@misc{injection-detector-v10,
  title={Injection Vulnerability Detector v1.0},
  author={Thrivikram Gujarathi},
  year={2026},
  publisher={HuggingFace},
  url={https://huggingface.co/vikramdgx/injection-vulnerability-detector}
}

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.