Lahiru-LK/codebert-vulnerability-api

CodeBERT Vulnerability Detection API

FastAPI-based code vulnerability detection using CodeBERT model trained for identifying SQL Injection and Certificate Validation vulnerabilities.

🚀 API Endpoints

• —GET / - API information
• —GET /health - Health check status
• —GET /docs - Interactive API documentation (Swagger UI)
• —POST /detect - Detect vulnerabilities in code

📝 Example Usage

Python

import requests

url = "https://your-username-codebert-vulnerability-api.hf.space/detect"

response = requests.post(url, json={
    "code": """
    String query = "SELECT * FROM users WHERE id = '" + userId + "'";
    Statement stmt = connection.createStatement();
    ResultSet rs = stmt.executeQuery(query);
    """,
    "max_length": 512
})

result = response.json()
print(f"Vulnerable: {result['is_vulnerable']}")
print(f"Type: {result['vulnerability_type']}")
print(f"Confidence: {result['confidence']:.2%}")

cURL

curl -X POST "https://your-username-codebert-vulnerability-api.hf.space/detect" \
  -H "Content-Type: application/json" \
  -d '{
    "code": "SELECT * FROM users WHERE id = " + user_input,
    "max_length": 512
  }'

🔍 Response Format

{
  "vulnerability_type": "SQL Injection",
  "confidence": 0.95,
  "is_vulnerable": true,
  "label": "s0"
}

🏷️ Vulnerability Labels

• —s0 / s1 - SQL Injection vulnerabilities
• —v0 / v1 - Certificate Validation vulnerabilities

🧠 Model Details

• —Base Model: microsoft/codebert-base
• —Architecture: RoBERTa with custom classification head
• —Model Size: 487 MB
• —Task: Binary classification for vulnerability detection
• —Categories: SQL Injection, Certificate Validation

⚡ Performance

The model uses CPU inference on Hugging Face Spaces free tier. For faster inference, consider upgrading to GPU hardware.

📄 License

Apache 2.0