feat: add Rust Privacy Gateway (Phase 6 — hudsucker)
Rust replacement for the Python mitmproxy-based Privacy Gateway. Uses hudsucker for transparent HTTPS proxying and axum for the REST management API. Modules: - proxy.rs: hudsucker HttpHandler — intercepts file uploads to AI service domains, scans for PII via info-privacy-rs, and optionally redacts content before forwarding - interceptor.rs: manual multipart/form-data parsing and rebuilding - scanner.rs: PII scanning via info-privacy-rs /api/v1/analyze + redact - llm_verifier.rs: RKLLM false-positive verification for uncertain entity detections - audit.rs: MariaDB audit logging (sqlx) + KVM backend event posting - domains.rs: AI domain whitelist (file or built-in defaults) - state.rs: privacy mode state.json reading/writing (off/audit/redact) - rest_api.rs: axum REST API on port 8889 (mode, stats, audit, cert) - config.rs: environment variable configuration with CLI overrides - main.rs: clap CLI + dual-server startup (proxy 8888 + API 8889) Degradation rules: - info-privacy unreachable + audit mode: log scan_failed, allow through - info-privacy unreachable + redact mode: log scan_failed, BLOCK request - RKLLM unreachable: use original detection results - MariaDB unreachable: skip logging, continue processing Tests: 39 unit tests covering state, domains, multipart parsing, config, scanner utilities, and LLM verdict parsing. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Generated
+3848
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,67 @@
|
||||
[package]
|
||||
name = "privacy-gateway-rs"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
description = "Rust replacement for Python Privacy Gateway (mitmproxy). Intercepts file uploads to AI services, scans for PII, and optionally redacts sensitive content."
|
||||
|
||||
[dependencies]
|
||||
# HTTP/HTTPS transparent proxy
|
||||
hudsucker = { version = "0.24", features = ["rcgen-ca"] }
|
||||
hyper = "1"
|
||||
hyper-util = "0.1"
|
||||
http = "1"
|
||||
|
||||
# TLS / CA certificate generation
|
||||
rcgen = "0.13"
|
||||
rustls-pemfile = "2"
|
||||
tokio-rustls = "0.26"
|
||||
|
||||
# Async runtime
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
|
||||
# REST API server (port 8889)
|
||||
axum = { version = "0.8", features = ["multipart"] }
|
||||
|
||||
# HTTP client for info-privacy-rs, RKLLM, KVM backend
|
||||
reqwest = { version = "0.12", features = ["json", "multipart"] }
|
||||
|
||||
# Multipart parsing
|
||||
multer = "3"
|
||||
|
||||
# MariaDB audit logging
|
||||
sqlx = { version = "0.8", features = ["mysql", "runtime-tokio", "chrono"] }
|
||||
|
||||
# Serialization
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
|
||||
# Logging / tracing
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
tower-http = { version = "0.6", features = ["cors", "trace"] }
|
||||
|
||||
# CLI argument parsing
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
http-body-util = "0.1"
|
||||
|
||||
# Error handling
|
||||
anyhow = "1"
|
||||
thiserror = "2"
|
||||
|
||||
# Utilities
|
||||
regex = "1"
|
||||
sha2 = "0.10"
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
bytes = "1"
|
||||
mime = "0.3"
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
tokio-test = "0.4"
|
||||
|
||||
[profile.release]
|
||||
opt-level = 3
|
||||
lto = "thin"
|
||||
strip = true
|
||||
@@ -0,0 +1,336 @@
|
||||
//! Audit logging — MariaDB privacy_audit_log table + KVM backend event posting.
|
||||
//!
|
||||
//! Records PII detection and redaction events. Never stores original file
|
||||
//! content — only hashes, entity type counts, and metadata.
|
||||
|
||||
use chrono::{NaiveDateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use sqlx::mysql::MySqlPool;
|
||||
use std::collections::HashMap;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
/// Schema for auto-creating the audit table.
|
||||
const CREATE_TABLE: &str = r#"
|
||||
CREATE TABLE IF NOT EXISTS privacy_audit_log (
|
||||
id INT AUTO_INCREMENT PRIMARY KEY,
|
||||
ts DATETIME(3) NOT NULL DEFAULT NOW(3),
|
||||
domain VARCHAR(255) NOT NULL,
|
||||
pii_types JSON,
|
||||
action VARCHAR(20) NOT NULL,
|
||||
doc_hash VARCHAR(64),
|
||||
file_count INT DEFAULT 0,
|
||||
client_ip VARCHAR(45) DEFAULT '',
|
||||
request_url VARCHAR(2048) DEFAULT '',
|
||||
filename VARCHAR(255) DEFAULT '',
|
||||
file_size INT DEFAULT 0,
|
||||
redacted_pii_count INT DEFAULT 0,
|
||||
INDEX idx_ts (ts),
|
||||
INDEX idx_domain (domain),
|
||||
INDEX idx_action (action)
|
||||
)
|
||||
"#;
|
||||
|
||||
/// Maximum retries for posting audit events to KVM backend.
|
||||
const MAX_RETRIES: usize = 3;
|
||||
/// Backoff delays in seconds.
|
||||
const RETRY_BACKOFF: [f64; 3] = [0.5, 1.0, 2.0];
|
||||
|
||||
/// A single audit log entry from the database.
|
||||
#[derive(Debug, Serialize, Deserialize, sqlx::FromRow)]
|
||||
pub struct AuditEntry {
|
||||
pub id: i32,
|
||||
pub ts: NaiveDateTime,
|
||||
pub domain: String,
|
||||
pub pii_types: Option<String>, // JSON string
|
||||
pub action: String,
|
||||
pub doc_hash: Option<String>,
|
||||
pub file_count: i32,
|
||||
pub client_ip: String,
|
||||
pub request_url: String,
|
||||
pub filename: String,
|
||||
pub file_size: i32,
|
||||
pub redacted_pii_count: i32,
|
||||
}
|
||||
|
||||
/// Today's summary statistics.
|
||||
#[derive(Debug, Serialize, Default)]
|
||||
pub struct DayStats {
|
||||
pub requests: i64,
|
||||
pub files: i64,
|
||||
pub pii_total: i64,
|
||||
pub by_action: HashMap<String, i64>,
|
||||
pub by_domain: HashMap<String, i64>,
|
||||
pub pii_by_type: HashMap<String, i64>,
|
||||
}
|
||||
|
||||
/// Parameters for logging an audit event.
|
||||
pub struct LogParams {
|
||||
pub domain: String,
|
||||
pub pii_types: HashMap<String, usize>,
|
||||
pub action: String,
|
||||
pub raw_bytes: Option<Vec<u8>>,
|
||||
pub file_count: i32,
|
||||
pub client_ip: String,
|
||||
pub request_url: String,
|
||||
pub filename: String,
|
||||
pub file_size: i32,
|
||||
}
|
||||
|
||||
/// KVM audit event payload (posted to KVM Go backend).
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct KvmAuditEvent {
|
||||
pub domain: String,
|
||||
pub pii_types: HashMap<String, usize>,
|
||||
pub action: String,
|
||||
pub doc_hash: String,
|
||||
pub file_count: i32,
|
||||
pub client_ip: String,
|
||||
pub request_url: String,
|
||||
pub filename: String,
|
||||
pub file_size: i32,
|
||||
}
|
||||
|
||||
/// Audit logger backed by MariaDB.
|
||||
#[derive(Clone)]
|
||||
pub struct AuditLogger {
|
||||
pool: MySqlPool,
|
||||
}
|
||||
|
||||
impl AuditLogger {
|
||||
/// Create a new audit logger with an existing connection pool.
|
||||
pub fn new(pool: MySqlPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// Ensure the audit table exists. Called once at startup.
|
||||
pub async fn ensure_schema(&self) -> Result<(), sqlx::Error> {
|
||||
sqlx::query(CREATE_TABLE).execute(&self.pool).await?;
|
||||
info!("Audit table schema ensured");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Log an audit event to MariaDB.
|
||||
pub async fn log(&self, params: LogParams) {
|
||||
let doc_hash = params
|
||||
.raw_bytes
|
||||
.as_ref()
|
||||
.map(|b| hex_sha256(b));
|
||||
|
||||
let pii_json = if params.pii_types.is_empty() {
|
||||
None
|
||||
} else {
|
||||
serde_json::to_string(¶ms.pii_types).ok()
|
||||
};
|
||||
|
||||
let redacted_pii_count: i32 = params
|
||||
.pii_types
|
||||
.values()
|
||||
.sum::<usize>() as i32;
|
||||
|
||||
let now = Utc::now().naive_utc();
|
||||
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO privacy_audit_log \
|
||||
(ts, domain, pii_types, action, doc_hash, file_count, \
|
||||
client_ip, request_url, filename, file_size, redacted_pii_count) \
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(now)
|
||||
.bind(¶ms.domain)
|
||||
.bind(&pii_json)
|
||||
.bind(¶ms.action)
|
||||
.bind(&doc_hash)
|
||||
.bind(params.file_count)
|
||||
.bind(¶ms.client_ip)
|
||||
.bind(¶ms.request_url)
|
||||
.bind(¶ms.filename)
|
||||
.bind(params.file_size)
|
||||
.bind(redacted_pii_count)
|
||||
.execute(&self.pool)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(_) => debug!(
|
||||
domain = %params.domain,
|
||||
action = %params.action,
|
||||
filename = %params.filename,
|
||||
"Audit event logged"
|
||||
),
|
||||
Err(e) => warn!(error = %e, "Failed to log audit event"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Query audit log entries with pagination and optional domain filter.
|
||||
pub async fn query(
|
||||
&self,
|
||||
limit: i64,
|
||||
offset: i64,
|
||||
domain: Option<&str>,
|
||||
) -> Result<Vec<AuditEntry>, sqlx::Error> {
|
||||
if let Some(domain) = domain {
|
||||
sqlx::query_as::<_, AuditEntry>(
|
||||
"SELECT * FROM privacy_audit_log \
|
||||
WHERE domain = ? ORDER BY id DESC LIMIT ? OFFSET ?",
|
||||
)
|
||||
.bind(domain)
|
||||
.bind(limit)
|
||||
.bind(offset)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
} else {
|
||||
sqlx::query_as::<_, AuditEntry>(
|
||||
"SELECT * FROM privacy_audit_log \
|
||||
ORDER BY id DESC LIMIT ? OFFSET ?",
|
||||
)
|
||||
.bind(limit)
|
||||
.bind(offset)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
/// Count total audit entries, optionally filtered by domain.
|
||||
pub async fn count(&self, domain: Option<&str>) -> Result<i64, sqlx::Error> {
|
||||
let count: (i64,) = if let Some(domain) = domain {
|
||||
sqlx::query_as(
|
||||
"SELECT COUNT(*) FROM privacy_audit_log WHERE domain = ?",
|
||||
)
|
||||
.bind(domain)
|
||||
.fetch_one(&self.pool)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query_as("SELECT COUNT(*) FROM privacy_audit_log")
|
||||
.fetch_one(&self.pool)
|
||||
.await?
|
||||
};
|
||||
Ok(count.0)
|
||||
}
|
||||
|
||||
/// Get today's summary statistics.
|
||||
pub async fn stats_today(&self) -> Result<DayStats, sqlx::Error> {
|
||||
let today = Utc::now().format("%Y-%m-%d").to_string();
|
||||
|
||||
// Summary totals
|
||||
let summary: (i64, i64, i64) = sqlx::query_as(
|
||||
"SELECT COUNT(*), \
|
||||
COALESCE(SUM(file_count), 0), \
|
||||
COALESCE(SUM(redacted_pii_count), 0) \
|
||||
FROM privacy_audit_log WHERE ts >= ?",
|
||||
)
|
||||
.bind(&today)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
// By action
|
||||
let action_rows: Vec<(String, i64)> = sqlx::query_as(
|
||||
"SELECT action, COUNT(*) FROM privacy_audit_log \
|
||||
WHERE ts >= ? GROUP BY action",
|
||||
)
|
||||
.bind(&today)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
// By domain (top 10)
|
||||
let domain_rows: Vec<(String, i64)> = sqlx::query_as(
|
||||
"SELECT domain, COUNT(*) FROM privacy_audit_log \
|
||||
WHERE ts >= ? GROUP BY domain ORDER BY COUNT(*) DESC LIMIT 10",
|
||||
)
|
||||
.bind(&today)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
// PII by type — aggregate from JSON column
|
||||
let pii_rows: Vec<(Option<String>,)> = sqlx::query_as(
|
||||
"SELECT pii_types FROM privacy_audit_log \
|
||||
WHERE ts >= ? AND pii_types IS NOT NULL",
|
||||
)
|
||||
.bind(&today)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let mut pii_by_type = HashMap::new();
|
||||
for (pii_json,) in &pii_rows {
|
||||
if let Some(json_str) = pii_json {
|
||||
if let Ok(types) =
|
||||
serde_json::from_str::<HashMap<String, i64>>(json_str)
|
||||
{
|
||||
for (k, v) in types {
|
||||
*pii_by_type.entry(k).or_insert(0) += v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(DayStats {
|
||||
requests: summary.0,
|
||||
files: summary.1,
|
||||
pii_total: summary.2,
|
||||
by_action: action_rows.into_iter().collect(),
|
||||
by_domain: domain_rows.into_iter().collect(),
|
||||
pii_by_type,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Post an audit event to the KVM Go backend (fire-and-forget with retries).
|
||||
pub async fn post_kvm_audit(client: &reqwest::Client, url: &str, event: &KvmAuditEvent) {
|
||||
for attempt in 0..MAX_RETRIES {
|
||||
match client.post(url).json(event).send().await {
|
||||
Ok(resp) if resp.status().is_success() => {
|
||||
debug!("KVM audit event posted successfully");
|
||||
return;
|
||||
}
|
||||
Ok(resp) => {
|
||||
warn!(
|
||||
status = resp.status().as_u16(),
|
||||
attempt = attempt + 1,
|
||||
"KVM audit POST failed"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
debug!(
|
||||
error = %e,
|
||||
attempt = attempt + 1,
|
||||
"KVM audit POST unreachable"
|
||||
);
|
||||
}
|
||||
}
|
||||
if attempt < MAX_RETRIES - 1 {
|
||||
tokio::time::sleep(std::time::Duration::from_secs_f64(
|
||||
RETRY_BACKOFF[attempt],
|
||||
))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn hex_sha256(data: &[u8]) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(data);
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
/// Try to connect to MariaDB and create the audit logger.
|
||||
/// Returns None if connection fails (audit is optional).
|
||||
pub async fn try_connect(database_url: &str) -> Option<AuditLogger> {
|
||||
match MySqlPool::connect(database_url).await {
|
||||
Ok(pool) => {
|
||||
let logger = AuditLogger::new(pool);
|
||||
match logger.ensure_schema().await {
|
||||
Ok(_) => {
|
||||
info!("Audit logger connected to MariaDB");
|
||||
Some(logger)
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(error = %e, "Failed to ensure audit schema (continuing without audit)");
|
||||
Some(logger) // Schema might already exist
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(error = %e, "Failed to connect to MariaDB — audit logging disabled");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
//! Configuration — environment variables with sensible defaults.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Application configuration, loaded from environment variables.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Config {
|
||||
/// MariaDB host
|
||||
pub db_host: String,
|
||||
/// MariaDB user
|
||||
pub db_user: String,
|
||||
/// MariaDB password
|
||||
pub db_pass: String,
|
||||
/// MariaDB database name
|
||||
pub db_name: String,
|
||||
|
||||
/// Path to state.json (privacy mode)
|
||||
pub state_file: PathBuf,
|
||||
/// Path to AI domains whitelist file (optional, falls back to built-in)
|
||||
pub domains_file: Option<PathBuf>,
|
||||
|
||||
/// info-privacy-rs base URL
|
||||
pub info_privacy_url: String,
|
||||
/// RKLLM server base URL
|
||||
pub rkllm_url: String,
|
||||
/// KVM backend audit event URL
|
||||
pub kvm_audit_url: String,
|
||||
|
||||
/// Proxy listen port (hudsucker)
|
||||
pub proxy_port: u16,
|
||||
/// REST API listen port (axum)
|
||||
pub api_port: u16,
|
||||
|
||||
/// CA certificate directory
|
||||
pub ca_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Load configuration from environment variables, applying defaults.
|
||||
pub fn from_env() -> Self {
|
||||
Self {
|
||||
db_host: env_or("KVM_MITM_DB_HOST", "localhost"),
|
||||
db_user: env_or("KVM_MITM_DB_USER", "kvm_mitm"),
|
||||
db_pass: env_or("KVM_MITM_DB_PASS", ""),
|
||||
db_name: env_or("KVM_MITM_DB_NAME", "kvm"),
|
||||
state_file: PathBuf::from(env_or(
|
||||
"PRIVACY_STATE_FILE",
|
||||
"/var/lib/kvm-privacy/state.json",
|
||||
)),
|
||||
domains_file: std::env::var("PRIVACY_DOMAINS_FILE")
|
||||
.ok()
|
||||
.map(PathBuf::from),
|
||||
info_privacy_url: env_or("INFO_PRIVACY_URL", "http://localhost:8001"),
|
||||
rkllm_url: env_or("RKLLM_URL", "http://localhost:8891"),
|
||||
kvm_audit_url: env_or(
|
||||
"KVM_AUDIT_URL",
|
||||
"http://127.0.0.1:8080/api/internal/privacy-event",
|
||||
),
|
||||
proxy_port: env_or("PROXY_PORT", "8888")
|
||||
.parse()
|
||||
.unwrap_or(8888),
|
||||
api_port: env_or("API_PORT", "8889").parse().unwrap_or(8889),
|
||||
ca_dir: PathBuf::from(env_or("CA_DIR", "/etc/kvm-privacy/ca")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Override with CLI arguments (non-empty values take precedence).
|
||||
pub fn with_cli_overrides(
|
||||
mut self,
|
||||
proxy_port: Option<u16>,
|
||||
api_port: Option<u16>,
|
||||
db_host: Option<String>,
|
||||
db_user: Option<String>,
|
||||
db_pass: Option<String>,
|
||||
db_name: Option<String>,
|
||||
) -> Self {
|
||||
if let Some(v) = proxy_port {
|
||||
self.proxy_port = v;
|
||||
}
|
||||
if let Some(v) = api_port {
|
||||
self.api_port = v;
|
||||
}
|
||||
if let Some(v) = db_host {
|
||||
self.db_host = v;
|
||||
}
|
||||
if let Some(v) = db_user {
|
||||
self.db_user = v;
|
||||
}
|
||||
if let Some(v) = db_pass {
|
||||
self.db_pass = v;
|
||||
}
|
||||
if let Some(v) = db_name {
|
||||
self.db_name = v;
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the MySQL/MariaDB connection URL for sqlx.
|
||||
pub fn database_url(&self) -> String {
|
||||
format!(
|
||||
"mysql://{}:{}@{}/{}",
|
||||
self.db_user, self.db_pass, self.db_host, self.db_name
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn env_or(key: &str, default: &str) -> String {
|
||||
std::env::var(key).unwrap_or_else(|_| default.to_string())
|
||||
}
|
||||
|
||||
impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Self::from_env()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_default_config() {
|
||||
// Clear env to test defaults
|
||||
let config = Config {
|
||||
db_host: "localhost".into(),
|
||||
db_user: "kvm_mitm".into(),
|
||||
db_pass: "".into(),
|
||||
db_name: "kvm".into(),
|
||||
state_file: PathBuf::from("/var/lib/kvm-privacy/state.json"),
|
||||
domains_file: None,
|
||||
info_privacy_url: "http://localhost:8001".into(),
|
||||
rkllm_url: "http://localhost:8891".into(),
|
||||
kvm_audit_url: "http://127.0.0.1:8080/api/internal/privacy-event".into(),
|
||||
proxy_port: 8888,
|
||||
api_port: 8889,
|
||||
ca_dir: PathBuf::from("/etc/kvm-privacy/ca"),
|
||||
};
|
||||
assert_eq!(config.proxy_port, 8888);
|
||||
assert_eq!(config.api_port, 8889);
|
||||
assert_eq!(config.db_host, "localhost");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_database_url() {
|
||||
let config = Config {
|
||||
db_host: "db.local".into(),
|
||||
db_user: "myuser".into(),
|
||||
db_pass: "secret".into(),
|
||||
db_name: "testdb".into(),
|
||||
state_file: PathBuf::from("/tmp/state.json"),
|
||||
domains_file: None,
|
||||
info_privacy_url: "http://localhost:8001".into(),
|
||||
rkllm_url: "http://localhost:8891".into(),
|
||||
kvm_audit_url: "http://localhost:8080/api/internal/privacy-event".into(),
|
||||
proxy_port: 8888,
|
||||
api_port: 8889,
|
||||
ca_dir: PathBuf::from("/tmp/ca"),
|
||||
};
|
||||
assert_eq!(config.database_url(), "mysql://myuser:secret@db.local/testdb");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cli_overrides() {
|
||||
let config = Config {
|
||||
db_host: "localhost".into(),
|
||||
db_user: "kvm_mitm".into(),
|
||||
db_pass: "".into(),
|
||||
db_name: "kvm".into(),
|
||||
state_file: PathBuf::from("/tmp/state.json"),
|
||||
domains_file: None,
|
||||
info_privacy_url: "http://localhost:8001".into(),
|
||||
rkllm_url: "http://localhost:8891".into(),
|
||||
kvm_audit_url: "http://localhost:8080/api/internal/privacy-event".into(),
|
||||
proxy_port: 8888,
|
||||
api_port: 8889,
|
||||
ca_dir: PathBuf::from("/tmp/ca"),
|
||||
};
|
||||
|
||||
let config = config.with_cli_overrides(
|
||||
Some(9999),
|
||||
Some(9998),
|
||||
Some("otherhost".into()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
assert_eq!(config.proxy_port, 9999);
|
||||
assert_eq!(config.api_port, 9998);
|
||||
assert_eq!(config.db_host, "otherhost");
|
||||
assert_eq!(config.db_user, "kvm_mitm"); // unchanged
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
//! AI domain whitelist — only intercept file uploads to these domains.
|
||||
//!
|
||||
//! Loads from external file (one domain per line, `#` comments) or
|
||||
//! falls back to built-in defaults.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::path::Path;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// Built-in AI service domains. Used when no external file is configured.
|
||||
const DEFAULT_DOMAINS: &[&str] = &[
|
||||
"api.anthropic.com",
|
||||
"api.openai.com",
|
||||
"generativelanguage.googleapis.com",
|
||||
"dashscope.aliyuncs.com",
|
||||
"api.cohere.com",
|
||||
"api.mistral.ai",
|
||||
"api.together.xyz",
|
||||
"api.groq.com",
|
||||
];
|
||||
|
||||
/// Loaded domain whitelist.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DomainWhitelist {
|
||||
domains: HashSet<String>,
|
||||
}
|
||||
|
||||
impl DomainWhitelist {
|
||||
/// Load domains from file. Falls back to built-in defaults if file is
|
||||
/// missing or unreadable.
|
||||
pub fn load(path: Option<&Path>) -> Self {
|
||||
if let Some(p) = path {
|
||||
match std::fs::read_to_string(p) {
|
||||
Ok(content) => {
|
||||
let domains = Self::parse_domains(&content);
|
||||
info!(count = domains.len(), path = %p.display(), "Loaded AI domains from file");
|
||||
return Self { domains };
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
error = %e,
|
||||
path = %p.display(),
|
||||
"Failed to read domains file, using built-in defaults"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Self::from_defaults()
|
||||
}
|
||||
|
||||
/// Create from built-in default list.
|
||||
pub fn from_defaults() -> Self {
|
||||
let domains = DEFAULT_DOMAINS
|
||||
.iter()
|
||||
.map(|s| s.to_string())
|
||||
.collect();
|
||||
info!(count = DEFAULT_DOMAINS.len(), "Using built-in AI domain whitelist");
|
||||
Self { domains }
|
||||
}
|
||||
|
||||
/// Parse domain list from text content.
|
||||
fn parse_domains(content: &str) -> HashSet<String> {
|
||||
content
|
||||
.lines()
|
||||
.map(|line| line.trim())
|
||||
.filter(|line| !line.is_empty() && !line.starts_with('#'))
|
||||
.map(|line| line.to_lowercase())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Check if a host matches the whitelist.
|
||||
/// Handles port stripping (e.g. "api.openai.com:443" -> "api.openai.com").
|
||||
pub fn is_ai_domain(&self, host: &str) -> bool {
|
||||
let host_clean = host
|
||||
.split(':')
|
||||
.next()
|
||||
.unwrap_or(host)
|
||||
.to_lowercase();
|
||||
self.domains.contains(&host_clean)
|
||||
}
|
||||
|
||||
/// Return the number of loaded domains.
|
||||
pub fn len(&self) -> usize {
|
||||
self.domains.len()
|
||||
}
|
||||
|
||||
/// Check if the whitelist is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.domains.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn test_default_domains() {
|
||||
let wl = DomainWhitelist::from_defaults();
|
||||
assert!(wl.is_ai_domain("api.openai.com"));
|
||||
assert!(wl.is_ai_domain("api.anthropic.com"));
|
||||
assert!(wl.is_ai_domain("api.mistral.ai"));
|
||||
assert!(!wl.is_ai_domain("example.com"));
|
||||
assert!(!wl.is_ai_domain("openai.com")); // only api.openai.com
|
||||
assert_eq!(wl.len(), 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_case_insensitive() {
|
||||
let wl = DomainWhitelist::from_defaults();
|
||||
assert!(wl.is_ai_domain("API.OPENAI.COM"));
|
||||
assert!(wl.is_ai_domain("Api.Anthropic.Com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_port_stripping() {
|
||||
let wl = DomainWhitelist::from_defaults();
|
||||
assert!(wl.is_ai_domain("api.openai.com:443"));
|
||||
assert!(wl.is_ai_domain("api.anthropic.com:8443"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_from_file() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("domains.txt");
|
||||
std::fs::write(
|
||||
&path,
|
||||
"# My custom domains\nmy.ai.service.com\nanother.ai.com\n\n# ignored\n",
|
||||
)
|
||||
.unwrap();
|
||||
let wl = DomainWhitelist::load(Some(&path));
|
||||
assert_eq!(wl.len(), 2);
|
||||
assert!(wl.is_ai_domain("my.ai.service.com"));
|
||||
assert!(wl.is_ai_domain("another.ai.com"));
|
||||
assert!(!wl.is_ai_domain("api.openai.com")); // not in custom file
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_missing_file_fallback() {
|
||||
let wl = DomainWhitelist::load(Some(Path::new("/nonexistent/domains.txt")));
|
||||
assert_eq!(wl.len(), 8); // falls back to defaults
|
||||
assert!(wl.is_ai_domain("api.openai.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_load_none_path() {
|
||||
let wl = DomainWhitelist::load(None);
|
||||
assert_eq!(wl.len(), 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_domains() {
|
||||
let content = "
|
||||
# comment line
|
||||
api.example.com
|
||||
whitespace.example.com
|
||||
|
||||
# another comment
|
||||
|
||||
last.example.com
|
||||
";
|
||||
let domains = DomainWhitelist::parse_domains(content);
|
||||
assert_eq!(domains.len(), 3);
|
||||
assert!(domains.contains("api.example.com"));
|
||||
assert!(domains.contains("whitespace.example.com"));
|
||||
assert!(domains.contains("last.example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_whitelist() {
|
||||
let content = "# only comments\n\n";
|
||||
let domains = DomainWhitelist::parse_domains(content);
|
||||
assert!(domains.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,451 @@
|
||||
//! Multipart request interception — extract files from uploads, rebuild with replacements.
|
||||
//!
|
||||
//! Uses manual boundary-based parsing (no multer dependency for synchronous
|
||||
//! in-memory parsing) to handle multipart/form-data bodies.
|
||||
|
||||
use regex::Regex;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::LazyLock;
|
||||
use tracing::debug;
|
||||
|
||||
/// Supported file extensions for PII scanning.
|
||||
const SUPPORTED_EXTENSIONS: &[&str] = &[
|
||||
".pdf", ".doc", ".docx", ".txt", ".png", ".jpg", ".jpeg", ".csv",
|
||||
];
|
||||
|
||||
/// A file extracted from a multipart upload.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ExtractedFile {
|
||||
/// The form field name (e.g. "file")
|
||||
pub field_name: String,
|
||||
/// The original filename (e.g. "resume.pdf")
|
||||
pub filename: String,
|
||||
/// MIME content type (e.g. "application/pdf")
|
||||
pub content_type: String,
|
||||
/// Raw file bytes
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
static MULTIPART_RE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r"(?i)multipart/form-data;\s*boundary=(.+)").unwrap()
|
||||
});
|
||||
|
||||
static FILENAME_RE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r#"filename="?([^";]+)"?"#).unwrap()
|
||||
});
|
||||
|
||||
static FIELDNAME_RE: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r#"name="?([^";]+)"?"#).unwrap()
|
||||
});
|
||||
|
||||
/// Check if Content-Type indicates a multipart/form-data file upload.
|
||||
pub fn is_file_upload(content_type: &str) -> bool {
|
||||
MULTIPART_RE.is_match(content_type)
|
||||
}
|
||||
|
||||
/// Extract the multipart boundary from Content-Type header.
|
||||
pub fn extract_boundary(content_type: &str) -> Option<String> {
|
||||
MULTIPART_RE
|
||||
.captures(content_type)
|
||||
.map(|caps| caps[1].trim().trim_matches('"').to_string())
|
||||
}
|
||||
|
||||
/// Check if a filename has a supported extension.
|
||||
fn is_supported_file(filename: &str) -> bool {
|
||||
let lower = filename.to_lowercase();
|
||||
SUPPORTED_EXTENSIONS.iter().any(|ext| lower.ends_with(ext))
|
||||
}
|
||||
|
||||
/// Extract files from a multipart/form-data body.
|
||||
///
|
||||
/// Parses the raw multipart body using the boundary string to split parts,
|
||||
/// then extracts parts that have a `filename` in their Content-Disposition
|
||||
/// and whose extension is in the supported set.
|
||||
pub fn extract_files(content_type: &str, body: &[u8]) -> Vec<ExtractedFile> {
|
||||
let boundary = match extract_boundary(content_type) {
|
||||
Some(b) => b,
|
||||
None => return Vec::new(),
|
||||
};
|
||||
|
||||
let delimiter = format!("--{}", boundary);
|
||||
let end_delimiter = format!("--{}--", boundary);
|
||||
|
||||
let mut files = Vec::new();
|
||||
|
||||
// Split body by boundary delimiter
|
||||
let parts = split_multipart(body, delimiter.as_bytes(), end_delimiter.as_bytes());
|
||||
|
||||
for part_data in parts {
|
||||
if let Some(file) = parse_part(&part_data) {
|
||||
if is_supported_file(&file.filename) {
|
||||
debug!(
|
||||
field = %file.field_name,
|
||||
filename = %file.filename,
|
||||
size = file.data.len(),
|
||||
"Extracted file from multipart"
|
||||
);
|
||||
files.push(file);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
files
|
||||
}
|
||||
|
||||
/// Split multipart body into individual parts by boundary.
|
||||
fn split_multipart(body: &[u8], delimiter: &[u8], end_delimiter: &[u8]) -> Vec<Vec<u8>> {
|
||||
let mut parts = Vec::new();
|
||||
let mut rest = body;
|
||||
|
||||
// Find and skip past the first delimiter
|
||||
if let Some(pos) = find_bytes(rest, delimiter) {
|
||||
rest = &rest[pos + delimiter.len()..];
|
||||
// Skip the CRLF after delimiter
|
||||
if rest.starts_with(b"\r\n") {
|
||||
rest = &rest[2..];
|
||||
} else if rest.starts_with(b"\n") {
|
||||
rest = &rest[1..];
|
||||
}
|
||||
} else {
|
||||
return parts;
|
||||
}
|
||||
|
||||
loop {
|
||||
// Find the next delimiter
|
||||
if let Some(pos) = find_bytes(rest, delimiter) {
|
||||
// Check if this is the end delimiter
|
||||
let before = &rest[..pos];
|
||||
// Remove trailing CRLF before delimiter
|
||||
let part = strip_trailing_crlf(before);
|
||||
if !part.is_empty() {
|
||||
parts.push(part.to_vec());
|
||||
}
|
||||
|
||||
// Check if we've hit the end delimiter
|
||||
let after_delim = &rest[pos..];
|
||||
if after_delim.starts_with(end_delimiter.as_ref()) {
|
||||
break;
|
||||
}
|
||||
|
||||
// Move past delimiter + CRLF
|
||||
rest = &rest[pos + delimiter.len()..];
|
||||
if rest.starts_with(b"\r\n") {
|
||||
rest = &rest[2..];
|
||||
} else if rest.starts_with(b"\n") {
|
||||
rest = &rest[1..];
|
||||
}
|
||||
} else {
|
||||
// No more delimiters — remaining is last part
|
||||
let part = strip_trailing_crlf(rest);
|
||||
if !part.is_empty() {
|
||||
parts.push(part.to_vec());
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
parts
|
||||
}
|
||||
|
||||
/// Parse a single multipart part into an ExtractedFile, if it represents a file.
|
||||
fn parse_part(data: &[u8]) -> Option<ExtractedFile> {
|
||||
// Find header/body separator (double CRLF)
|
||||
let separator = find_header_end(data)?;
|
||||
let header_bytes = &data[..separator];
|
||||
let body = &data[separator..];
|
||||
|
||||
let headers_str = String::from_utf8_lossy(header_bytes);
|
||||
|
||||
// Check for Content-Disposition with filename
|
||||
let mut filename = None;
|
||||
let mut field_name = String::from("file");
|
||||
let mut content_type = String::from("application/octet-stream");
|
||||
|
||||
for line in headers_str.lines() {
|
||||
let lower = line.to_lowercase();
|
||||
if lower.starts_with("content-disposition:") {
|
||||
if let Some(caps) = FILENAME_RE.captures(line) {
|
||||
filename = Some(caps[1].to_string());
|
||||
}
|
||||
if let Some(caps) = FIELDNAME_RE.captures(line) {
|
||||
field_name = caps[1].to_string();
|
||||
}
|
||||
} else if lower.starts_with("content-type:") {
|
||||
content_type = line
|
||||
.splitn(2, ':')
|
||||
.nth(1)
|
||||
.map(|s| s.trim().to_string())
|
||||
.unwrap_or_else(|| "application/octet-stream".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
let filename = filename?; // Only process parts with a filename
|
||||
|
||||
Some(ExtractedFile {
|
||||
field_name,
|
||||
filename,
|
||||
content_type,
|
||||
data: body.to_vec(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Find the end of HTTP headers (double CRLF or double LF).
|
||||
fn find_header_end(data: &[u8]) -> Option<usize> {
|
||||
if let Some(pos) = find_bytes(data, b"\r\n\r\n") {
|
||||
Some(pos + 4)
|
||||
} else if let Some(pos) = find_bytes(data, b"\n\n") {
|
||||
Some(pos + 2)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the position of a byte pattern in a byte slice.
|
||||
fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option<usize> {
|
||||
if needle.is_empty() || needle.len() > haystack.len() {
|
||||
return None;
|
||||
}
|
||||
haystack
|
||||
.windows(needle.len())
|
||||
.position(|window| window == needle)
|
||||
}
|
||||
|
||||
/// Strip trailing CRLF or LF from a byte slice.
|
||||
fn strip_trailing_crlf(data: &[u8]) -> &[u8] {
|
||||
if data.ends_with(b"\r\n") {
|
||||
&data[..data.len() - 2]
|
||||
} else if data.ends_with(b"\n") {
|
||||
&data[..data.len() - 1]
|
||||
} else {
|
||||
data
|
||||
}
|
||||
}
|
||||
|
||||
/// Rebuild a multipart body, replacing file parts whose field_name is in `replacements`.
|
||||
///
|
||||
/// Non-file parts and file parts not in `replacements` are passed through unchanged.
|
||||
/// For replaced parts, only the body is swapped; headers are preserved.
|
||||
pub fn rebuild_multipart(
|
||||
content_type: &str,
|
||||
body: &[u8],
|
||||
replacements: &HashMap<String, Vec<u8>>,
|
||||
) -> Option<(Vec<u8>, String)> {
|
||||
let boundary = extract_boundary(content_type)?;
|
||||
let delimiter = format!("--{}", boundary);
|
||||
let end_delimiter = format!("--{}--", boundary);
|
||||
|
||||
let parts = split_multipart(body, delimiter.as_bytes(), end_delimiter.as_bytes());
|
||||
|
||||
let mut rebuilt_parts: Vec<Vec<u8>> = Vec::new();
|
||||
|
||||
for part_data in &parts {
|
||||
if let Some(separator) = find_header_end(part_data) {
|
||||
let header_bytes = &part_data[..separator];
|
||||
let headers_str = String::from_utf8_lossy(header_bytes);
|
||||
|
||||
// Check if this part has a field_name that needs replacement
|
||||
let mut matched_field = None;
|
||||
for line in headers_str.lines() {
|
||||
if line.to_lowercase().starts_with("content-disposition:") {
|
||||
if let Some(caps) = FIELDNAME_RE.captures(line) {
|
||||
let name = caps[1].to_string();
|
||||
if replacements.contains_key(&name) {
|
||||
matched_field = Some(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(field) = matched_field {
|
||||
// Replace the body but keep headers
|
||||
let mut new_part = header_bytes.to_vec();
|
||||
new_part.extend_from_slice(&replacements[&field]);
|
||||
rebuilt_parts.push(new_part);
|
||||
} else {
|
||||
rebuilt_parts.push(part_data.to_vec());
|
||||
}
|
||||
} else {
|
||||
// No header separator found, pass through as-is
|
||||
rebuilt_parts.push(part_data.to_vec());
|
||||
}
|
||||
}
|
||||
|
||||
// Reassemble
|
||||
let mut output = Vec::new();
|
||||
for (i, part) in rebuilt_parts.iter().enumerate() {
|
||||
output.extend_from_slice(delimiter.as_bytes());
|
||||
output.extend_from_slice(b"\r\n");
|
||||
output.extend_from_slice(part);
|
||||
if i < rebuilt_parts.len() - 1 {
|
||||
output.extend_from_slice(b"\r\n");
|
||||
}
|
||||
}
|
||||
output.extend_from_slice(b"\r\n");
|
||||
output.extend_from_slice(end_delimiter.as_bytes());
|
||||
output.extend_from_slice(b"\r\n");
|
||||
|
||||
// Content-Type stays the same (same boundary)
|
||||
Some((output, content_type.to_string()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_multipart(boundary: &str, parts: &[(&str, &str, &[u8])]) -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
for (field_name, filename, data) in parts {
|
||||
body.extend_from_slice(format!("--{}\r\n", boundary).as_bytes());
|
||||
body.extend_from_slice(
|
||||
format!(
|
||||
"Content-Disposition: form-data; name=\"{}\"; filename=\"{}\"\r\n",
|
||||
field_name, filename
|
||||
)
|
||||
.as_bytes(),
|
||||
);
|
||||
body.extend_from_slice(b"Content-Type: application/octet-stream\r\n");
|
||||
body.extend_from_slice(b"\r\n");
|
||||
body.extend_from_slice(data);
|
||||
body.extend_from_slice(b"\r\n");
|
||||
}
|
||||
body.extend_from_slice(format!("--{}--\r\n", boundary).as_bytes());
|
||||
body
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_file_upload() {
|
||||
assert!(is_file_upload(
|
||||
"multipart/form-data; boundary=----WebKitFormBoundary"
|
||||
));
|
||||
assert!(is_file_upload(
|
||||
"Multipart/Form-Data; boundary=abc123"
|
||||
));
|
||||
assert!(!is_file_upload("application/json"));
|
||||
assert!(!is_file_upload("text/plain"));
|
||||
assert!(!is_file_upload(""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_boundary() {
|
||||
assert_eq!(
|
||||
extract_boundary("multipart/form-data; boundary=abc123"),
|
||||
Some("abc123".into())
|
||||
);
|
||||
assert_eq!(
|
||||
extract_boundary("multipart/form-data; boundary=\"quoted-boundary\""),
|
||||
Some("quoted-boundary".into())
|
||||
);
|
||||
assert_eq!(extract_boundary("application/json"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_files_single_pdf() {
|
||||
let boundary = "----TestBoundary";
|
||||
let ct = format!("multipart/form-data; boundary={}", boundary);
|
||||
let data = b"PDF content here";
|
||||
let body = make_multipart(boundary, &[("file", "report.pdf", data)]);
|
||||
|
||||
let files = extract_files(&ct, &body);
|
||||
assert_eq!(files.len(), 1);
|
||||
assert_eq!(files[0].field_name, "file");
|
||||
assert_eq!(files[0].filename, "report.pdf");
|
||||
assert_eq!(files[0].data, data);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_files_multiple() {
|
||||
let boundary = "----TestBound";
|
||||
let ct = format!("multipart/form-data; boundary={}", boundary);
|
||||
let body = make_multipart(
|
||||
boundary,
|
||||
&[
|
||||
("file1", "image.png", b"PNG data"),
|
||||
("file2", "doc.docx", b"DOCX data"),
|
||||
("file3", "script.py", b"Python code"), // unsupported ext
|
||||
],
|
||||
);
|
||||
|
||||
let files = extract_files(&ct, &body);
|
||||
assert_eq!(files.len(), 2); // .py is filtered out
|
||||
assert_eq!(files[0].filename, "image.png");
|
||||
assert_eq!(files[1].filename, "doc.docx");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_files_no_filename() {
|
||||
let boundary = "----TestBound";
|
||||
let ct = format!("multipart/form-data; boundary={}", boundary);
|
||||
// Part without filename
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(format!("--{}\r\n", boundary).as_bytes());
|
||||
body.extend_from_slice(b"Content-Disposition: form-data; name=\"text_field\"\r\n");
|
||||
body.extend_from_slice(b"\r\n");
|
||||
body.extend_from_slice(b"just some text");
|
||||
body.extend_from_slice(b"\r\n");
|
||||
body.extend_from_slice(format!("--{}--\r\n", boundary).as_bytes());
|
||||
|
||||
let files = extract_files(&ct, &body);
|
||||
assert!(files.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_files_invalid_content_type() {
|
||||
let files = extract_files("application/json", b"{}");
|
||||
assert!(files.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_supported_file() {
|
||||
assert!(is_supported_file("report.pdf"));
|
||||
assert!(is_supported_file("PHOTO.JPG"));
|
||||
assert!(is_supported_file("data.csv"));
|
||||
assert!(is_supported_file("resume.docx"));
|
||||
assert!(is_supported_file("notes.txt"));
|
||||
assert!(!is_supported_file("script.py"));
|
||||
assert!(!is_supported_file("binary.exe"));
|
||||
assert!(!is_supported_file("archive.zip"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_multipart() {
|
||||
let boundary = "----TestBound";
|
||||
let ct = format!("multipart/form-data; boundary={}", boundary);
|
||||
let body = make_multipart(
|
||||
boundary,
|
||||
&[
|
||||
("file1", "doc.pdf", b"original PDF"),
|
||||
("file2", "image.png", b"original PNG"),
|
||||
],
|
||||
);
|
||||
|
||||
let mut replacements = HashMap::new();
|
||||
replacements.insert("file1".into(), b"REDACTED PDF".to_vec());
|
||||
|
||||
let (rebuilt, new_ct) = rebuild_multipart(&ct, &body, &replacements).unwrap();
|
||||
assert_eq!(new_ct, ct);
|
||||
|
||||
// The rebuilt body should contain the redacted content for file1
|
||||
let rebuilt_str = String::from_utf8_lossy(&rebuilt);
|
||||
assert!(rebuilt_str.contains("REDACTED PDF"));
|
||||
// And the original content for file2
|
||||
assert!(rebuilt_str.contains("original PNG"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_no_replacements() {
|
||||
let boundary = "----TestBound";
|
||||
let ct = format!("multipart/form-data; boundary={}", boundary);
|
||||
let body = make_multipart(boundary, &[("file", "doc.pdf", b"original")]);
|
||||
|
||||
let replacements = HashMap::new();
|
||||
let (rebuilt, _) = rebuild_multipart(&ct, &body, &replacements).unwrap();
|
||||
let rebuilt_str = String::from_utf8_lossy(&rebuilt);
|
||||
assert!(rebuilt_str.contains("original"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rebuild_invalid_content_type() {
|
||||
let result = rebuild_multipart("application/json", b"{}", &HashMap::new());
|
||||
assert!(result.is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,302 @@
|
||||
//! LLM false-positive verification — calls RKLLM to validate PII detections.
|
||||
//!
|
||||
//! Uses the local Qwen2.5 model (NPU-accelerated) to semantically check
|
||||
//! whether regex/NER-detected entities are genuine PII. Data never leaves
|
||||
//! the device.
|
||||
|
||||
use crate::scanner::Entity;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
/// Verification prompt template (Chinese, matching Python implementation).
|
||||
const VERIFY_PROMPT: &str = r#"你是 PII(个人身份信息)验证助手。判断以下检测结果是否为真实的个人信息。
|
||||
|
||||
上下文文本:
|
||||
{context}
|
||||
|
||||
检测结果:
|
||||
{detections_json}
|
||||
|
||||
对每个检测结果,判断它是否为真实 PII(true)还是误报(false)。
|
||||
仅输出 JSON 数组,格式:[{"index": 0, "is_pii": true}, ...]
|
||||
不要输出其他文字。"#;
|
||||
|
||||
/// LLM verifier for PII detection results.
|
||||
pub struct LlmVerifier {
|
||||
rkllm_url: String,
|
||||
client: reqwest::Client,
|
||||
max_tokens: u32,
|
||||
/// Cached availability check result
|
||||
available: AtomicBool,
|
||||
/// Whether we've checked availability yet
|
||||
checked: AtomicBool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ChatRequest {
|
||||
model: String,
|
||||
messages: Vec<ChatMessage>,
|
||||
max_tokens: u32,
|
||||
temperature: f64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ChatMessage {
|
||||
role: String,
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChatResponse {
|
||||
choices: Vec<ChatChoice>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChatChoice {
|
||||
message: ChatMessageResponse,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ChatMessageResponse {
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Verdict {
|
||||
index: usize,
|
||||
is_pii: bool,
|
||||
}
|
||||
|
||||
impl LlmVerifier {
|
||||
/// Create a new LLM verifier.
|
||||
pub fn new(rkllm_url: String) -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.build()
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
rkllm_url: rkllm_url.trim_end_matches('/').to_string(),
|
||||
client,
|
||||
max_tokens: 256,
|
||||
available: AtomicBool::new(false),
|
||||
checked: AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the RKLLM server is reachable. Result is cached.
|
||||
pub async fn is_available(&self) -> bool {
|
||||
if self.checked.load(Ordering::Relaxed) {
|
||||
return self.available.load(Ordering::Relaxed);
|
||||
}
|
||||
|
||||
let result = self
|
||||
.client
|
||||
.get(format!("{}/v1/models", self.rkllm_url))
|
||||
.timeout(std::time::Duration::from_secs(3))
|
||||
.send()
|
||||
.await
|
||||
.map(|r| r.status().is_success())
|
||||
.unwrap_or(false);
|
||||
|
||||
self.available.store(result, Ordering::Relaxed);
|
||||
self.checked.store(true, Ordering::Relaxed);
|
||||
result
|
||||
}
|
||||
|
||||
/// Reset the cached availability check (e.g., after service restart).
|
||||
pub fn reset_availability(&self) {
|
||||
self.checked.store(false, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Verify a list of detections using the local LLM.
|
||||
///
|
||||
/// Returns a filtered list with false positives removed.
|
||||
/// Falls back to the original list on any error.
|
||||
pub async fn verify_detections(
|
||||
&self,
|
||||
context: &str,
|
||||
detections: &[Entity],
|
||||
) -> Vec<Entity> {
|
||||
if detections.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
if !self.is_available().await {
|
||||
debug!("RKLLM not available, returning original detections");
|
||||
return detections.to_vec();
|
||||
}
|
||||
|
||||
// Build detection summary for the prompt
|
||||
let det_summary: Vec<serde_json::Value> = detections
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, d)| {
|
||||
serde_json::json!({
|
||||
"index": i,
|
||||
"type": d.entity_type,
|
||||
"value": d.value,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let det_json = serde_json::to_string(&det_summary).unwrap_or_default();
|
||||
|
||||
// Truncate context for small model context window
|
||||
let truncated_context: String = context.chars().take(2000).collect();
|
||||
|
||||
let prompt = VERIFY_PROMPT
|
||||
.replace("{context}", &truncated_context)
|
||||
.replace("{detections_json}", &det_json);
|
||||
|
||||
let request = ChatRequest {
|
||||
model: "qwen2.5-1.5b".to_string(),
|
||||
messages: vec![ChatMessage {
|
||||
role: "user".to_string(),
|
||||
content: prompt,
|
||||
}],
|
||||
max_tokens: self.max_tokens,
|
||||
temperature: 0.0,
|
||||
};
|
||||
|
||||
match self.call_llm(&request).await {
|
||||
Ok(content) => self.parse_verdicts(&content, detections),
|
||||
Err(e) => {
|
||||
warn!(error = %e, "LLM verification failed, using original detections");
|
||||
detections.to_vec()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Call the RKLLM chat completions endpoint.
|
||||
async fn call_llm(&self, request: &ChatRequest) -> Result<String, String> {
|
||||
let resp = self
|
||||
.client
|
||||
.post(format!("{}/v1/chat/completions", self.rkllm_url))
|
||||
.json(request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("Request failed: {}", e))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(format!("RKLLM returned status {}", resp.status()));
|
||||
}
|
||||
|
||||
let chat_resp: ChatResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse response: {}", e))?;
|
||||
|
||||
chat_resp
|
||||
.choices
|
||||
.first()
|
||||
.map(|c| c.message.content.clone())
|
||||
.ok_or_else(|| "No choices in response".to_string())
|
||||
}
|
||||
|
||||
/// Parse LLM verdicts and filter detections.
|
||||
fn parse_verdicts(&self, content: &str, detections: &[Entity]) -> Vec<Entity> {
|
||||
match serde_json::from_str::<Vec<Verdict>>(content) {
|
||||
Ok(verdicts) => {
|
||||
let false_indices: std::collections::HashSet<usize> = verdicts
|
||||
.iter()
|
||||
.filter(|v| !v.is_pii)
|
||||
.map(|v| v.index)
|
||||
.collect();
|
||||
|
||||
let filtered: Vec<Entity> = detections
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(i, _)| !false_indices.contains(i))
|
||||
.map(|(_, d)| d.clone())
|
||||
.collect();
|
||||
|
||||
let removed = detections.len() - filtered.len();
|
||||
if removed > 0 {
|
||||
info!(
|
||||
removed = removed,
|
||||
total = detections.len(),
|
||||
"LLM verification removed false positives"
|
||||
);
|
||||
}
|
||||
filtered
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
error = %e,
|
||||
content = content,
|
||||
"Failed to parse LLM verification response"
|
||||
);
|
||||
detections.to_vec()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_verdicts_all_pii() {
|
||||
let verifier = LlmVerifier::new("http://localhost:8891".into());
|
||||
let detections = vec![
|
||||
Entity {
|
||||
entity_type: "phone".into(),
|
||||
value: "13800138000".into(),
|
||||
confidence: 0.7,
|
||||
},
|
||||
Entity {
|
||||
entity_type: "email".into(),
|
||||
value: "test@example.com".into(),
|
||||
confidence: 0.6,
|
||||
},
|
||||
];
|
||||
let content = r#"[{"index": 0, "is_pii": true}, {"index": 1, "is_pii": true}]"#;
|
||||
let result = verifier.parse_verdicts(content, &detections);
|
||||
assert_eq!(result.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_verdicts_remove_false_positive() {
|
||||
let verifier = LlmVerifier::new("http://localhost:8891".into());
|
||||
let detections = vec![
|
||||
Entity {
|
||||
entity_type: "phone".into(),
|
||||
value: "13800138000".into(),
|
||||
confidence: 0.7,
|
||||
},
|
||||
Entity {
|
||||
entity_type: "email".into(),
|
||||
value: "not-really@email".into(),
|
||||
confidence: 0.4,
|
||||
},
|
||||
];
|
||||
let content = r#"[{"index": 0, "is_pii": true}, {"index": 1, "is_pii": false}]"#;
|
||||
let result = verifier.parse_verdicts(content, &detections);
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].entity_type, "phone");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_verdicts_invalid_json() {
|
||||
let verifier = LlmVerifier::new("http://localhost:8891".into());
|
||||
let detections = vec![Entity {
|
||||
entity_type: "phone".into(),
|
||||
value: "123".into(),
|
||||
confidence: 0.5,
|
||||
}];
|
||||
let content = "this is not json";
|
||||
let result = verifier.parse_verdicts(content, &detections);
|
||||
assert_eq!(result.len(), 1); // falls back to original
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_verdicts_empty() {
|
||||
let verifier = LlmVerifier::new("http://localhost:8891".into());
|
||||
let detections: Vec<Entity> = vec![];
|
||||
let content = "[]";
|
||||
let result = verifier.parse_verdicts(content, &detections);
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,266 @@
|
||||
//! Privacy Gateway (Rust) — transparent HTTPS proxy for PII interception.
|
||||
//!
|
||||
//! Replaces the Python mitmproxy-based Privacy Gateway with a Rust
|
||||
//! implementation using hudsucker for transparent proxying and axum
|
||||
//! for the REST management API.
|
||||
//!
|
||||
//! Architecture:
|
||||
//! - hudsucker proxy on port 8888 (HTTPS MITM)
|
||||
//! - axum REST API on port 8889 (management)
|
||||
//! - MariaDB for audit logging (optional)
|
||||
//! - info-privacy-rs for PII scanning
|
||||
//! - RKLLM for false-positive verification
|
||||
|
||||
#[allow(dead_code)]
|
||||
mod audit;
|
||||
mod config;
|
||||
#[allow(dead_code)]
|
||||
mod domains;
|
||||
#[allow(dead_code)]
|
||||
mod interceptor;
|
||||
#[allow(dead_code)]
|
||||
mod llm_verifier;
|
||||
mod proxy;
|
||||
mod rest_api;
|
||||
#[allow(dead_code)]
|
||||
mod scanner;
|
||||
mod state;
|
||||
|
||||
use clap::Parser;
|
||||
use config::Config;
|
||||
use domains::DomainWhitelist;
|
||||
use hudsucker::certificate_authority::RcgenAuthority;
|
||||
use hudsucker::rcgen::{Issuer, KeyPair};
|
||||
use hudsucker::rustls::crypto::aws_lc_rs;
|
||||
use hudsucker::Proxy;
|
||||
use llm_verifier::LlmVerifier;
|
||||
use proxy::{PrivacyProxyHandler, ProxyState};
|
||||
use rest_api::ApiState;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use tracing::{error, info};
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(
|
||||
name = "privacy-gateway-rs",
|
||||
about = "KVM Privacy Gateway — transparent HTTPS proxy for PII interception"
|
||||
)]
|
||||
struct Cli {
|
||||
/// Proxy listen port (HTTPS MITM)
|
||||
#[arg(long, env = "PROXY_PORT")]
|
||||
proxy_port: Option<u16>,
|
||||
|
||||
/// REST API listen port
|
||||
#[arg(long, env = "API_PORT")]
|
||||
api_port: Option<u16>,
|
||||
|
||||
/// MariaDB host
|
||||
#[arg(long, env = "KVM_MITM_DB_HOST")]
|
||||
db_host: Option<String>,
|
||||
|
||||
/// MariaDB user
|
||||
#[arg(long, env = "KVM_MITM_DB_USER")]
|
||||
db_user: Option<String>,
|
||||
|
||||
/// MariaDB password
|
||||
#[arg(long, env = "KVM_MITM_DB_PASS")]
|
||||
db_pass: Option<String>,
|
||||
|
||||
/// MariaDB database name
|
||||
#[arg(long, env = "KVM_MITM_DB_NAME")]
|
||||
db_name: Option<String>,
|
||||
}
|
||||
|
||||
async fn shutdown_signal() {
|
||||
let ctrl_c = tokio::signal::ctrl_c();
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use tokio::signal::unix::{signal, SignalKind};
|
||||
let mut sigterm = signal(SignalKind::terminate()).expect("Failed to install SIGTERM handler");
|
||||
tokio::select! {
|
||||
_ = ctrl_c => info!("Received SIGINT, shutting down..."),
|
||||
_ = sigterm.recv() => info!("Received SIGTERM, shutting down..."),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
ctrl_c.await.expect("Failed to install CTRL+C handler");
|
||||
info!("Received SIGINT, shutting down...");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// Initialize tracing
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "privacy_gateway_rs=info,tower_http=info".into()),
|
||||
)
|
||||
.init();
|
||||
|
||||
let cli = Cli::parse();
|
||||
let cfg = Config::from_env().with_cli_overrides(
|
||||
cli.proxy_port,
|
||||
cli.api_port,
|
||||
cli.db_host,
|
||||
cli.db_user,
|
||||
cli.db_pass,
|
||||
cli.db_name,
|
||||
);
|
||||
|
||||
info!(
|
||||
proxy_port = cfg.proxy_port,
|
||||
api_port = cfg.api_port,
|
||||
info_privacy = %cfg.info_privacy_url,
|
||||
rkllm = %cfg.rkllm_url,
|
||||
"Starting Privacy Gateway"
|
||||
);
|
||||
|
||||
// Load domain whitelist
|
||||
let domains = Arc::new(DomainWhitelist::load(
|
||||
cfg.domains_file.as_deref(),
|
||||
));
|
||||
|
||||
// Connect to MariaDB (optional)
|
||||
let audit_logger = audit::try_connect(&cfg.database_url()).await;
|
||||
|
||||
// Create HTTP client for upstream calls
|
||||
let http_client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.build()?;
|
||||
|
||||
// Create LLM verifier
|
||||
let llm_verifier = Arc::new(LlmVerifier::new(cfg.rkllm_url.clone()));
|
||||
|
||||
// ── Start REST API ──
|
||||
let api_state = ApiState {
|
||||
state_file: cfg.state_file.clone(),
|
||||
audit: audit_logger.clone(),
|
||||
ca_cert_path: cfg.ca_dir.clone(),
|
||||
};
|
||||
|
||||
let api_router = rest_api::build_router(api_state);
|
||||
let api_addr = SocketAddr::from(([0, 0, 0, 0], cfg.api_port));
|
||||
|
||||
let api_handle = tokio::spawn(async move {
|
||||
info!(port = cfg.api_port, "REST API listening");
|
||||
let listener = tokio::net::TcpListener::bind(api_addr).await.unwrap();
|
||||
axum::serve(listener, api_router)
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
|
||||
// ── Start Proxy ──
|
||||
let proxy_state = ProxyState {
|
||||
state_file: cfg.state_file.clone(),
|
||||
domains,
|
||||
http_client,
|
||||
info_privacy_url: cfg.info_privacy_url.clone(),
|
||||
llm_verifier,
|
||||
audit: audit_logger,
|
||||
kvm_audit_url: cfg.kvm_audit_url.clone(),
|
||||
};
|
||||
|
||||
let handler = PrivacyProxyHandler::new(proxy_state);
|
||||
|
||||
// Load or generate CA certificate
|
||||
let ca = match load_or_generate_ca(&cfg.ca_dir).await {
|
||||
Ok(ca) => ca,
|
||||
Err(e) => {
|
||||
error!(error = %e, "Failed to set up CA certificate");
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
let proxy_addr = SocketAddr::from(([0, 0, 0, 0], cfg.proxy_port));
|
||||
info!(port = cfg.proxy_port, "Proxy listening");
|
||||
|
||||
let proxy = Proxy::builder()
|
||||
.with_addr(proxy_addr)
|
||||
.with_ca(ca)
|
||||
.with_rustls_connector(aws_lc_rs::default_provider())
|
||||
.with_http_handler(handler)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.build()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to build proxy: {}", e))?;
|
||||
|
||||
if let Err(e) = proxy.start().await {
|
||||
error!(error = %e, "Proxy error");
|
||||
}
|
||||
|
||||
// Clean up
|
||||
api_handle.abort();
|
||||
info!("Privacy Gateway stopped");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load an existing CA certificate or generate a new self-signed one.
|
||||
async fn load_or_generate_ca(
|
||||
ca_dir: &std::path::Path,
|
||||
) -> anyhow::Result<RcgenAuthority> {
|
||||
let key_path = ca_dir.join("ca.key");
|
||||
let cert_path = ca_dir.join("ca.pem");
|
||||
|
||||
// Try to load existing CA
|
||||
if key_path.exists() && cert_path.exists() {
|
||||
info!(path = %ca_dir.display(), "Loading existing CA certificate");
|
||||
let key_pem = std::fs::read_to_string(&key_path)?;
|
||||
let cert_pem = std::fs::read_to_string(&cert_path)?;
|
||||
|
||||
let key_pair = KeyPair::from_pem(&key_pem)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to parse CA key: {}", e))?;
|
||||
let issuer = Issuer::from_ca_cert_pem(&cert_pem, key_pair)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to parse CA cert: {}", e))?;
|
||||
|
||||
return Ok(RcgenAuthority::new(
|
||||
issuer,
|
||||
1_000,
|
||||
aws_lc_rs::default_provider(),
|
||||
));
|
||||
}
|
||||
|
||||
// Generate new CA
|
||||
info!(path = %ca_dir.display(), "Generating new CA certificate");
|
||||
std::fs::create_dir_all(ca_dir)?;
|
||||
|
||||
let key_pair = KeyPair::generate()
|
||||
.map_err(|e| anyhow::anyhow!("Failed to generate CA key: {}", e))?;
|
||||
|
||||
// Create a self-signed CA certificate
|
||||
let mut params = hudsucker::rcgen::CertificateParams::new(Vec::new())
|
||||
.map_err(|e| anyhow::anyhow!("Failed to create cert params: {}", e))?;
|
||||
params.is_ca = hudsucker::rcgen::IsCa::Ca(hudsucker::rcgen::BasicConstraints::Unconstrained);
|
||||
params
|
||||
.distinguished_name
|
||||
.push(hudsucker::rcgen::DnType::CommonName, "KVM Privacy Gateway CA");
|
||||
params
|
||||
.distinguished_name
|
||||
.push(hudsucker::rcgen::DnType::OrganizationName, "KVM-Privacy");
|
||||
|
||||
let cert = params
|
||||
.self_signed(&key_pair)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to self-sign CA cert: {}", e))?;
|
||||
|
||||
// Persist to disk
|
||||
let key_pem = key_pair.serialize_pem();
|
||||
let cert_pem = cert.pem();
|
||||
std::fs::write(&key_path, &key_pem)?;
|
||||
std::fs::write(&cert_path, &cert_pem)?;
|
||||
|
||||
info!("CA certificate generated and saved");
|
||||
|
||||
let key_pair = KeyPair::from_pem(&key_pem)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to reload CA key: {}", e))?;
|
||||
let issuer = Issuer::from_ca_cert_pem(&cert_pem, key_pair)
|
||||
.map_err(|e| anyhow::anyhow!("Failed to reload CA cert: {}", e))?;
|
||||
|
||||
Ok(RcgenAuthority::new(
|
||||
issuer,
|
||||
1_000,
|
||||
aws_lc_rs::default_provider(),
|
||||
))
|
||||
}
|
||||
@@ -0,0 +1,349 @@
|
||||
//! Hudsucker proxy handler — intercepts file uploads to AI services.
|
||||
//!
|
||||
//! Implements the `HttpHandler` trait to inspect outgoing requests,
|
||||
//! detect file uploads to AI service domains, scan them for PII,
|
||||
//! and optionally redact before forwarding.
|
||||
|
||||
use crate::audit::{self, AuditLogger, KvmAuditEvent, LogParams};
|
||||
use crate::domains::DomainWhitelist;
|
||||
use crate::interceptor::{self, ExtractedFile};
|
||||
use crate::llm_verifier::LlmVerifier;
|
||||
use crate::scanner;
|
||||
use crate::state::{self, PrivacyMode};
|
||||
|
||||
use bytes::Bytes;
|
||||
use http_body_util::{BodyExt, Full};
|
||||
use hudsucker::hyper::{Request, Response};
|
||||
use hudsucker::{Body, HttpContext, HttpHandler, RequestOrResponse};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tracing::{info, warn};
|
||||
|
||||
/// Shared state for the proxy handler.
|
||||
#[derive(Clone)]
|
||||
pub struct ProxyState {
|
||||
/// Path to state.json
|
||||
pub state_file: PathBuf,
|
||||
/// AI domain whitelist
|
||||
pub domains: Arc<DomainWhitelist>,
|
||||
/// HTTP client for upstream API calls
|
||||
pub http_client: reqwest::Client,
|
||||
/// info-privacy-rs base URL
|
||||
pub info_privacy_url: String,
|
||||
/// LLM verifier for false-positive reduction (used when scanner
|
||||
/// returns sensitive_partial classification)
|
||||
#[allow(dead_code)]
|
||||
pub llm_verifier: Arc<LlmVerifier>,
|
||||
/// Audit logger (optional — None if MariaDB is unavailable)
|
||||
pub audit: Option<AuditLogger>,
|
||||
/// KVM audit event URL
|
||||
pub kvm_audit_url: String,
|
||||
}
|
||||
|
||||
/// The main proxy handler implementing hudsucker's HttpHandler trait.
|
||||
#[derive(Clone)]
|
||||
pub struct PrivacyProxyHandler {
|
||||
state: ProxyState,
|
||||
}
|
||||
|
||||
impl PrivacyProxyHandler {
|
||||
pub fn new(state: ProxyState) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert Vec<u8> to hudsucker Body via Full<Bytes>.
|
||||
fn body_from_bytes(data: Vec<u8>) -> Body {
|
||||
Body::from(Full::new(Bytes::from(data)))
|
||||
}
|
||||
|
||||
impl HttpHandler for PrivacyProxyHandler {
|
||||
async fn handle_request(
|
||||
&mut self,
|
||||
_ctx: &HttpContext,
|
||||
req: Request<Body>,
|
||||
) -> RequestOrResponse {
|
||||
// Read current mode (cheap file read)
|
||||
let mode = state::read_mode(&self.state.state_file);
|
||||
if mode == PrivacyMode::Off {
|
||||
return req.into();
|
||||
}
|
||||
|
||||
// Check if target host is an AI service
|
||||
let host = req
|
||||
.uri()
|
||||
.host()
|
||||
.or_else(|| {
|
||||
req.headers()
|
||||
.get("host")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|h| h.split(':').next().unwrap_or(h))
|
||||
})
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if !self.state.domains.is_ai_domain(&host) {
|
||||
return req.into();
|
||||
}
|
||||
|
||||
// Check if this is a file upload
|
||||
let content_type = req
|
||||
.headers()
|
||||
.get("content-type")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
if !interceptor::is_file_upload(&content_type) {
|
||||
return req.into();
|
||||
}
|
||||
|
||||
// Collect the body for inspection
|
||||
let (parts, body) = req.into_parts();
|
||||
let body_bytes = match collect_body(body).await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
warn!(error = %e, "Failed to read request body");
|
||||
let rebuilt = Request::from_parts(parts, Body::empty());
|
||||
return rebuilt.into();
|
||||
}
|
||||
};
|
||||
|
||||
// Extract files from multipart body
|
||||
let files = interceptor::extract_files(&content_type, &body_bytes);
|
||||
if files.is_empty() {
|
||||
let rebuilt = Request::from_parts(parts, body_from_bytes(body_bytes));
|
||||
return rebuilt.into();
|
||||
}
|
||||
|
||||
info!(
|
||||
mode = %mode,
|
||||
host = %host,
|
||||
files = files.len(),
|
||||
"Intercepted file upload to AI service"
|
||||
);
|
||||
|
||||
// Process each file
|
||||
let result = self
|
||||
.process_upload(&host, &files, &body_bytes, &content_type, mode, &parts)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
ProcessResult::PassThrough => {
|
||||
let rebuilt = Request::from_parts(parts, body_from_bytes(body_bytes));
|
||||
rebuilt.into()
|
||||
}
|
||||
ProcessResult::Modified(new_body) => {
|
||||
let rebuilt = Request::from_parts(parts, body_from_bytes(new_body));
|
||||
rebuilt.into()
|
||||
}
|
||||
ProcessResult::Blocked => {
|
||||
// Return 403 Forbidden response
|
||||
let resp = Response::builder()
|
||||
.status(403)
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(
|
||||
r#"{"error":"Upload blocked by privacy gateway: PII scanning unavailable in redact mode"}"#,
|
||||
))
|
||||
.unwrap();
|
||||
resp.into()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_response(
|
||||
&mut self,
|
||||
_ctx: &HttpContext,
|
||||
res: Response<Body>,
|
||||
) -> Response<Body> {
|
||||
// No response modification needed
|
||||
res
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of processing an upload.
|
||||
enum ProcessResult {
|
||||
/// Forward original body unchanged
|
||||
PassThrough,
|
||||
/// Forward with modified body (redacted files)
|
||||
Modified(Vec<u8>),
|
||||
/// Block the request entirely
|
||||
Blocked,
|
||||
}
|
||||
|
||||
impl PrivacyProxyHandler {
|
||||
/// Process an intercepted file upload.
|
||||
async fn process_upload(
|
||||
&self,
|
||||
host: &str,
|
||||
files: &[ExtractedFile],
|
||||
body: &[u8],
|
||||
content_type: &str,
|
||||
mode: PrivacyMode,
|
||||
parts: &http::request::Parts,
|
||||
) -> ProcessResult {
|
||||
let mut replacements: HashMap<String, Vec<u8>> = HashMap::new();
|
||||
let total_files = files.len() as i32;
|
||||
let client_ip = extract_client_ip(parts);
|
||||
let request_url = parts.uri.to_string();
|
||||
|
||||
for f in files {
|
||||
let scan_result = scanner::scan_and_redact(
|
||||
&self.state.http_client,
|
||||
&self.state.info_privacy_url,
|
||||
&f.filename,
|
||||
&f.data,
|
||||
)
|
||||
.await;
|
||||
|
||||
match scan_result {
|
||||
Err(e) => {
|
||||
warn!(
|
||||
filename = %f.filename,
|
||||
error = %e,
|
||||
"info-privacy unreachable"
|
||||
);
|
||||
|
||||
// Log scan_failed
|
||||
if let Some(audit) = &self.state.audit {
|
||||
audit
|
||||
.log(LogParams {
|
||||
domain: host.to_string(),
|
||||
pii_types: HashMap::new(),
|
||||
action: "scan_failed".into(),
|
||||
raw_bytes: Some(f.data.clone()),
|
||||
file_count: total_files,
|
||||
client_ip: client_ip.clone(),
|
||||
request_url: request_url.clone(),
|
||||
filename: f.filename.clone(),
|
||||
file_size: f.data.len() as i32,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
if mode == PrivacyMode::Redact {
|
||||
// Block in redact mode when scanner is unavailable
|
||||
return ProcessResult::Blocked;
|
||||
}
|
||||
// Audit mode: allow through with scan_failed logged
|
||||
continue;
|
||||
}
|
||||
Ok(result) => {
|
||||
match mode {
|
||||
PrivacyMode::Audit => {
|
||||
// Log only — never modify
|
||||
if let Some(audit) = &self.state.audit {
|
||||
audit
|
||||
.log(LogParams {
|
||||
domain: host.to_string(),
|
||||
pii_types: result.pii_types.clone(),
|
||||
action: "allow".into(),
|
||||
raw_bytes: Some(f.data.clone()),
|
||||
file_count: total_files,
|
||||
client_ip: client_ip.clone(),
|
||||
request_url: request_url.clone(),
|
||||
filename: f.filename.clone(),
|
||||
file_size: f.data.len() as i32,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
if result.pii_found {
|
||||
info!(
|
||||
filename = %f.filename,
|
||||
pii_types = ?result.pii_types,
|
||||
"Audit: PII detected"
|
||||
);
|
||||
}
|
||||
}
|
||||
PrivacyMode::Redact => {
|
||||
let action = if result.pii_found && result.redacted_bytes.is_some()
|
||||
{
|
||||
replacements.insert(
|
||||
f.field_name.clone(),
|
||||
result.redacted_bytes.clone().unwrap(),
|
||||
);
|
||||
info!(
|
||||
filename = %f.filename,
|
||||
pii_types = ?result.pii_types,
|
||||
"Redacted file"
|
||||
);
|
||||
"redact"
|
||||
} else {
|
||||
"allow"
|
||||
};
|
||||
|
||||
// Log to MariaDB
|
||||
if let Some(audit) = &self.state.audit {
|
||||
audit
|
||||
.log(LogParams {
|
||||
domain: host.to_string(),
|
||||
pii_types: result.pii_types.clone(),
|
||||
action: action.into(),
|
||||
raw_bytes: Some(f.data.clone()),
|
||||
file_count: total_files,
|
||||
client_ip: client_ip.clone(),
|
||||
request_url: request_url.clone(),
|
||||
filename: f.filename.clone(),
|
||||
file_size: f.data.len() as i32,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
// Post to KVM backend
|
||||
let event = KvmAuditEvent {
|
||||
domain: host.to_string(),
|
||||
pii_types: result.pii_types,
|
||||
action: action.into(),
|
||||
doc_hash: result.original_hash,
|
||||
file_count: total_files,
|
||||
client_ip: client_ip.clone(),
|
||||
request_url: request_url.clone(),
|
||||
filename: f.filename.clone(),
|
||||
file_size: f.data.len() as i32,
|
||||
};
|
||||
let client = self.state.http_client.clone();
|
||||
let url = self.state.kvm_audit_url.clone();
|
||||
tokio::spawn(async move {
|
||||
audit::post_kvm_audit(&client, &url, &event).await;
|
||||
});
|
||||
}
|
||||
PrivacyMode::Off => unreachable!(),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rebuild multipart if any files were redacted
|
||||
if !replacements.is_empty() {
|
||||
match interceptor::rebuild_multipart(content_type, body, &replacements) {
|
||||
Some((new_body, _)) => ProcessResult::Modified(new_body),
|
||||
None => {
|
||||
warn!("Failed to rebuild multipart body, passing through original");
|
||||
ProcessResult::PassThrough
|
||||
}
|
||||
}
|
||||
} else {
|
||||
ProcessResult::PassThrough
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Collect a hudsucker Body into bytes.
|
||||
async fn collect_body(body: Body) -> Result<Vec<u8>, String> {
|
||||
let collected = body
|
||||
.collect()
|
||||
.await
|
||||
.map_err(|e| format!("Body collect error: {}", e))?;
|
||||
Ok(collected.to_bytes().to_vec())
|
||||
}
|
||||
|
||||
/// Extract client IP from request parts (X-Forwarded-For header or empty).
|
||||
fn extract_client_ip(parts: &http::request::Parts) -> String {
|
||||
parts
|
||||
.headers
|
||||
.get("x-forwarded-for")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| s.split(',').next().unwrap_or("").trim().to_string())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
@@ -0,0 +1,310 @@
|
||||
//! REST API — axum server on port 8889.
|
||||
//!
|
||||
//! Provides /api/v1/privacy/* endpoints consumed by the KVM WebUI:
|
||||
//! - GET /api/v1/privacy/mode — current mode
|
||||
//! - POST /api/v1/privacy/mode — set mode
|
||||
//! - GET /api/v1/privacy/stats — today's stats
|
||||
//! - GET /api/v1/privacy/audit — paginated audit logs
|
||||
//! - GET /api/v1/privacy/cert — CA certificate download
|
||||
|
||||
use crate::audit::AuditLogger;
|
||||
use crate::state::{self, PrivacyMode};
|
||||
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::get;
|
||||
use axum::{Json, Router};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tower_http::cors::CorsLayer;
|
||||
use tower_http::trace::TraceLayer;
|
||||
use tracing::{error, info};
|
||||
|
||||
/// Shared state for the REST API handlers.
|
||||
#[derive(Clone)]
|
||||
pub struct ApiState {
|
||||
/// Path to state.json
|
||||
pub state_file: PathBuf,
|
||||
/// Audit logger (optional)
|
||||
pub audit: Option<AuditLogger>,
|
||||
/// Path to CA certificate PEM file
|
||||
pub ca_cert_path: PathBuf,
|
||||
}
|
||||
|
||||
/// Build the axum router for the REST API.
|
||||
pub fn build_router(state: ApiState) -> Router {
|
||||
Router::new()
|
||||
.route("/api/v1/privacy/mode", get(get_mode).post(set_mode).put(set_mode))
|
||||
.route("/api/v1/privacy/stats", get(get_stats))
|
||||
.route("/api/v1/privacy/audit", get(get_audit))
|
||||
.route("/api/v1/privacy/cert", get(get_cert))
|
||||
.route("/health", get(health))
|
||||
.with_state(Arc::new(state))
|
||||
.layer(CorsLayer::permissive())
|
||||
.layer(TraceLayer::new_for_http())
|
||||
}
|
||||
|
||||
// ── Health check ──
|
||||
|
||||
async fn health() -> Json<serde_json::Value> {
|
||||
Json(serde_json::json!({"status": "ok"}))
|
||||
}
|
||||
|
||||
// ── Mode endpoints ──
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ModeResponse {
|
||||
mode: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct SetModeRequest {
|
||||
mode: String,
|
||||
}
|
||||
|
||||
async fn get_mode(State(state): State<Arc<ApiState>>) -> Json<ModeResponse> {
|
||||
let mode = state::read_mode(&state.state_file);
|
||||
Json(ModeResponse {
|
||||
mode: mode.as_str().to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn set_mode(
|
||||
State(state): State<Arc<ApiState>>,
|
||||
Json(body): Json<SetModeRequest>,
|
||||
) -> Response {
|
||||
let mode = match PrivacyMode::from_str_opt(&body.mode) {
|
||||
Some(m) => m,
|
||||
None => {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(serde_json::json!({
|
||||
"error": format!("mode must be one of [\"audit\", \"off\", \"redact\"]")
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = state::write_mode(&state.state_file, mode) {
|
||||
error!(error = %e, "Failed to write state.json");
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({"error": "Failed to persist mode"})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
info!(mode = %mode, "Privacy mode updated");
|
||||
Json(ModeResponse {
|
||||
mode: mode.as_str().to_string(),
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
|
||||
// ── Stats endpoint ──
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct StatsResponse {
|
||||
requests: i64,
|
||||
files: i64,
|
||||
pii_by_type: HashMap<String, i64>,
|
||||
actions: ActionCounts,
|
||||
top_domains: Vec<DomainCount>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ActionCounts {
|
||||
allow: i64,
|
||||
block: i64,
|
||||
redact: i64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct DomainCount {
|
||||
domain: String,
|
||||
count: i64,
|
||||
}
|
||||
|
||||
async fn get_stats(State(state): State<Arc<ApiState>>) -> Response {
|
||||
let audit = match &state.audit {
|
||||
Some(a) => a,
|
||||
None => {
|
||||
return Json(StatsResponse {
|
||||
requests: 0,
|
||||
files: 0,
|
||||
pii_by_type: HashMap::new(),
|
||||
actions: ActionCounts {
|
||||
allow: 0,
|
||||
block: 0,
|
||||
redact: 0,
|
||||
},
|
||||
top_domains: Vec::new(),
|
||||
})
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
match audit.stats_today().await {
|
||||
Ok(raw) => {
|
||||
let allow = raw.by_action.get("allow").copied().unwrap_or(0)
|
||||
+ raw.by_action.get("bypass").copied().unwrap_or(0);
|
||||
let block = raw.by_action.get("block").copied().unwrap_or(0);
|
||||
let redact = raw.by_action.get("redact").copied().unwrap_or(0)
|
||||
+ raw.by_action.get("auto_redact").copied().unwrap_or(0);
|
||||
|
||||
let mut top_domains: Vec<DomainCount> = raw
|
||||
.by_domain
|
||||
.into_iter()
|
||||
.map(|(domain, count)| DomainCount { domain, count })
|
||||
.collect();
|
||||
top_domains.sort_by(|a, b| b.count.cmp(&a.count));
|
||||
|
||||
Json(StatsResponse {
|
||||
requests: raw.requests,
|
||||
files: raw.files,
|
||||
pii_by_type: raw.pii_by_type,
|
||||
actions: ActionCounts {
|
||||
allow,
|
||||
block,
|
||||
redact,
|
||||
},
|
||||
top_domains,
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
Err(e) => {
|
||||
error!(error = %e, "Failed to fetch stats");
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({"error": "Failed to fetch stats"})),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Audit log endpoint ──
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct AuditQuery {
|
||||
page: Option<i64>,
|
||||
page_size: Option<i64>,
|
||||
domain: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct AuditResponse {
|
||||
logs: Vec<AuditLogEntry>,
|
||||
total: i64,
|
||||
page: i64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct AuditLogEntry {
|
||||
id: i32,
|
||||
ts: String,
|
||||
domain: String,
|
||||
action: String,
|
||||
pii_types: HashMap<String, serde_json::Value>,
|
||||
filename: String,
|
||||
file_size: i32,
|
||||
client_ip: String,
|
||||
request_url: String,
|
||||
}
|
||||
|
||||
async fn get_audit(
|
||||
State(state): State<Arc<ApiState>>,
|
||||
Query(query): Query<AuditQuery>,
|
||||
) -> Response {
|
||||
let audit = match &state.audit {
|
||||
Some(a) => a,
|
||||
None => {
|
||||
return Json(AuditResponse {
|
||||
logs: Vec::new(),
|
||||
total: 0,
|
||||
page: 1,
|
||||
})
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
let page = query.page.unwrap_or(1).max(1);
|
||||
let page_size = query.page_size.unwrap_or(20).min(100).max(1);
|
||||
let offset = (page - 1) * page_size;
|
||||
let domain = query.domain.as_deref();
|
||||
|
||||
let entries = match audit.query(page_size, offset, domain).await {
|
||||
Ok(rows) => rows,
|
||||
Err(e) => {
|
||||
error!(error = %e, "Failed to query audit logs");
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(serde_json::json!({"error": "Failed to query audit logs"})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
|
||||
let total = match audit.count(domain).await {
|
||||
Ok(n) => n,
|
||||
Err(_) => 0,
|
||||
};
|
||||
|
||||
let logs: Vec<AuditLogEntry> = entries
|
||||
.into_iter()
|
||||
.map(|r| {
|
||||
let pii_types: HashMap<String, serde_json::Value> = r
|
||||
.pii_types
|
||||
.as_deref()
|
||||
.and_then(|s| serde_json::from_str(s).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
AuditLogEntry {
|
||||
id: r.id,
|
||||
ts: r.ts.format("%Y-%m-%d %H:%M:%S%.3f").to_string(),
|
||||
domain: r.domain,
|
||||
action: r.action,
|
||||
pii_types,
|
||||
filename: r.filename,
|
||||
file_size: r.file_size,
|
||||
client_ip: r.client_ip,
|
||||
request_url: r.request_url,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
Json(AuditResponse { logs, total, page }).into_response()
|
||||
}
|
||||
|
||||
// ── Certificate download ──
|
||||
|
||||
async fn get_cert(State(state): State<Arc<ApiState>>) -> Response {
|
||||
let cert_path = state.ca_cert_path.join("ca.pem");
|
||||
match std::fs::read(&cert_path) {
|
||||
Ok(pem_bytes) => {
|
||||
let headers = [
|
||||
(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
"application/x-pem-file",
|
||||
),
|
||||
(
|
||||
axum::http::header::CONTENT_DISPOSITION,
|
||||
"attachment; filename=\"kvm-privacy-ca.crt\"",
|
||||
),
|
||||
];
|
||||
(headers, pem_bytes).into_response()
|
||||
}
|
||||
Err(e) => {
|
||||
error!(error = %e, path = %cert_path.display(), "Failed to read CA certificate");
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"CA certificate not available",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,320 @@
|
||||
//! PII scanning — calls info-privacy-rs /api/v1/analyze and /api/v1/redact.
|
||||
//!
|
||||
//! Scans uploaded files for PII (personally identifiable information),
|
||||
//! separates high-confidence from low-confidence detections, and optionally
|
||||
//! redacts PII from the file content.
|
||||
|
||||
use reqwest::multipart;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::collections::HashMap;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
/// Types of PII that warrant redaction.
|
||||
const REDACT_TYPES: &[&str] = &[
|
||||
"id_card",
|
||||
"phone",
|
||||
"bank_card",
|
||||
"email",
|
||||
"license_plate",
|
||||
"name",
|
||||
"address",
|
||||
"face",
|
||||
];
|
||||
|
||||
/// High-risk entity types that skip LLM verification.
|
||||
const HIGH_RISK_TYPES: &[&str] = &["id_card", "phone", "bank_card"];
|
||||
|
||||
/// Confidence threshold: entities above this are treated as high-confidence.
|
||||
const HIGH_CONFIDENCE_THRESHOLD: f64 = 0.85;
|
||||
|
||||
/// Result of a PII scan.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ScanResult {
|
||||
/// Whether any PII was found
|
||||
pub pii_found: bool,
|
||||
/// Map of entity type to count
|
||||
pub pii_types: HashMap<String, usize>,
|
||||
/// Redacted file bytes (None if no PII or redaction not requested)
|
||||
pub redacted_bytes: Option<Vec<u8>>,
|
||||
/// SHA-256 hash of the original file
|
||||
pub original_hash: String,
|
||||
}
|
||||
|
||||
/// An entity detected by info-privacy-rs.
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct Entity {
|
||||
#[serde(rename = "type", default)]
|
||||
pub entity_type: String,
|
||||
#[serde(default)]
|
||||
pub value: String,
|
||||
#[serde(default = "default_confidence")]
|
||||
pub confidence: f64,
|
||||
}
|
||||
|
||||
fn default_confidence() -> f64 {
|
||||
0.5
|
||||
}
|
||||
|
||||
/// Response from info-privacy-rs /api/v1/analyze.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AnalyzeResponse {
|
||||
#[serde(default)]
|
||||
summary: HashMap<String, usize>,
|
||||
#[serde(default)]
|
||||
classification: String,
|
||||
#[serde(default)]
|
||||
entities: Vec<Entity>,
|
||||
#[serde(default)]
|
||||
blocked: bool,
|
||||
#[serde(default)]
|
||||
doc_id: String,
|
||||
}
|
||||
|
||||
/// Scan a file for PII and optionally redact it.
|
||||
///
|
||||
/// Steps:
|
||||
/// 1. Call POST /api/v1/analyze with the file as multipart upload
|
||||
/// 2. Parse the detection report
|
||||
/// 3. If PII found, call POST /api/v1/redact to get redacted file bytes
|
||||
///
|
||||
/// Returns `Err` on network/HTTP errors (caller decides whether to block or allow).
|
||||
pub async fn scan_and_redact(
|
||||
client: &reqwest::Client,
|
||||
base_url: &str,
|
||||
filename: &str,
|
||||
data: &[u8],
|
||||
) -> Result<ScanResult, ScanError> {
|
||||
let original_hash = hex_sha256(data);
|
||||
|
||||
// Step 1: Analyze
|
||||
let form = multipart::Form::new().part(
|
||||
"file",
|
||||
multipart::Part::bytes(data.to_vec())
|
||||
.file_name(filename.to_string())
|
||||
.mime_str("application/octet-stream")
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
let resp = client
|
||||
.post(format!("{}/api/v1/analyze", base_url))
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ScanError::NetworkError(e.to_string()))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
return Err(ScanError::HttpError(resp.status().as_u16()));
|
||||
}
|
||||
|
||||
let report: AnalyzeResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| ScanError::ParseError(e.to_string()))?;
|
||||
|
||||
debug!(
|
||||
classification = %report.classification,
|
||||
entities = report.entities.len(),
|
||||
"Analyze response for {}",
|
||||
filename
|
||||
);
|
||||
|
||||
if report.summary.is_empty() || report.blocked {
|
||||
return Ok(ScanResult {
|
||||
pii_found: !report.summary.is_empty(),
|
||||
pii_types: report.summary,
|
||||
redacted_bytes: None,
|
||||
original_hash,
|
||||
});
|
||||
}
|
||||
|
||||
// Use the summary from the report (LLM verification is done by llm_verifier
|
||||
// at a higher level if needed)
|
||||
let summary = report.summary;
|
||||
|
||||
// Step 2: Redact
|
||||
let redact_types: Vec<String> = summary.keys().cloned().collect();
|
||||
let config_json =
|
||||
serde_json::json!({"redact_types": redact_types}).to_string();
|
||||
|
||||
let redact_form = multipart::Form::new()
|
||||
.part(
|
||||
"file",
|
||||
multipart::Part::bytes(data.to_vec())
|
||||
.file_name(filename.to_string())
|
||||
.mime_str("application/octet-stream")
|
||||
.unwrap(),
|
||||
)
|
||||
.text("config", config_json);
|
||||
|
||||
let redact_resp = client
|
||||
.post(format!("{}/api/v1/redact", base_url))
|
||||
.multipart(redact_form)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ScanError::NetworkError(e.to_string()))?;
|
||||
|
||||
if !redact_resp.status().is_success() {
|
||||
warn!(
|
||||
status = redact_resp.status().as_u16(),
|
||||
"Redact endpoint failed for {}",
|
||||
filename
|
||||
);
|
||||
return Ok(ScanResult {
|
||||
pii_found: true,
|
||||
pii_types: summary,
|
||||
redacted_bytes: None,
|
||||
original_hash,
|
||||
});
|
||||
}
|
||||
|
||||
let redacted_bytes = redact_resp
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| ScanError::NetworkError(e.to_string()))?;
|
||||
|
||||
info!(
|
||||
original_size = data.len(),
|
||||
redacted_size = redacted_bytes.len(),
|
||||
pii_types = ?summary,
|
||||
"Redacted {} successfully",
|
||||
filename
|
||||
);
|
||||
|
||||
Ok(ScanResult {
|
||||
pii_found: true,
|
||||
pii_types: summary,
|
||||
redacted_bytes: Some(redacted_bytes.to_vec()),
|
||||
original_hash,
|
||||
})
|
||||
}
|
||||
|
||||
/// Separate entities into high-confidence and low-confidence groups.
|
||||
/// High-confidence entities skip LLM verification.
|
||||
pub fn split_by_confidence(entities: &[Entity]) -> (Vec<Entity>, Vec<Entity>) {
|
||||
let mut high = Vec::new();
|
||||
let mut low = Vec::new();
|
||||
for e in entities {
|
||||
if e.confidence >= HIGH_CONFIDENCE_THRESHOLD {
|
||||
high.push(e.clone());
|
||||
} else {
|
||||
low.push(e.clone());
|
||||
}
|
||||
}
|
||||
(high, low)
|
||||
}
|
||||
|
||||
/// Rebuild a PII summary from a list of entities.
|
||||
pub fn build_summary(entities: &[Entity]) -> HashMap<String, usize> {
|
||||
let mut summary = HashMap::new();
|
||||
for e in entities {
|
||||
if !e.entity_type.is_empty() {
|
||||
*summary.entry(e.entity_type.clone()).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
summary
|
||||
}
|
||||
|
||||
/// Check if any entity types are high-risk.
|
||||
pub fn has_high_risk(summary: &HashMap<String, usize>) -> bool {
|
||||
HIGH_RISK_TYPES
|
||||
.iter()
|
||||
.any(|t| summary.contains_key(*t))
|
||||
}
|
||||
|
||||
fn hex_sha256(data: &[u8]) -> String {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(data);
|
||||
format!("{:x}", hasher.finalize())
|
||||
}
|
||||
|
||||
/// Errors from the scanning process.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ScanError {
|
||||
#[error("Network error: {0}")]
|
||||
NetworkError(String),
|
||||
#[error("HTTP error: status {0}")]
|
||||
HttpError(u16),
|
||||
#[error("Parse error: {0}")]
|
||||
ParseError(String),
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_hex_sha256() {
|
||||
let hash = hex_sha256(b"hello world");
|
||||
assert_eq!(
|
||||
hash,
|
||||
"b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_split_by_confidence() {
|
||||
let entities = vec![
|
||||
Entity {
|
||||
entity_type: "id_card".into(),
|
||||
value: "123456".into(),
|
||||
confidence: 0.95,
|
||||
},
|
||||
Entity {
|
||||
entity_type: "phone".into(),
|
||||
value: "555-1234".into(),
|
||||
confidence: 0.60,
|
||||
},
|
||||
Entity {
|
||||
entity_type: "email".into(),
|
||||
value: "test@example.com".into(),
|
||||
confidence: 0.85,
|
||||
},
|
||||
];
|
||||
let (high, low) = split_by_confidence(&entities);
|
||||
assert_eq!(high.len(), 2); // 0.95 and 0.85
|
||||
assert_eq!(low.len(), 1); // 0.60
|
||||
assert_eq!(low[0].entity_type, "phone");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_summary() {
|
||||
let entities = vec![
|
||||
Entity {
|
||||
entity_type: "phone".into(),
|
||||
value: "111".into(),
|
||||
confidence: 0.9,
|
||||
},
|
||||
Entity {
|
||||
entity_type: "phone".into(),
|
||||
value: "222".into(),
|
||||
confidence: 0.9,
|
||||
},
|
||||
Entity {
|
||||
entity_type: "email".into(),
|
||||
value: "a@b.com".into(),
|
||||
confidence: 0.9,
|
||||
},
|
||||
];
|
||||
let summary = build_summary(&entities);
|
||||
assert_eq!(summary.get("phone"), Some(&2));
|
||||
assert_eq!(summary.get("email"), Some(&1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_has_high_risk() {
|
||||
let mut summary = HashMap::new();
|
||||
summary.insert("email".into(), 1);
|
||||
assert!(!has_high_risk(&summary));
|
||||
|
||||
summary.insert("id_card".into(), 1);
|
||||
assert!(has_high_risk(&summary));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_build_summary_empty() {
|
||||
let entities: Vec<Entity> = vec![];
|
||||
let summary = build_summary(&entities);
|
||||
assert!(summary.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
//! Privacy mode state — reads/writes /var/lib/kvm-privacy/state.json.
|
||||
//!
|
||||
//! Mode values: "off" | "audit" | "redact"
|
||||
//! Defaults to "off" on any read error (missing file, bad JSON, unknown mode).
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::Path;
|
||||
use tracing::warn;
|
||||
|
||||
/// Valid privacy modes.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum PrivacyMode {
|
||||
Off,
|
||||
Audit,
|
||||
Redact,
|
||||
}
|
||||
|
||||
impl PrivacyMode {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
PrivacyMode::Off => "off",
|
||||
PrivacyMode::Audit => "audit",
|
||||
PrivacyMode::Redact => "redact",
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse from string, returning None for invalid values.
|
||||
pub fn from_str_opt(s: &str) -> Option<Self> {
|
||||
match s {
|
||||
"off" => Some(PrivacyMode::Off),
|
||||
"audit" => Some(PrivacyMode::Audit),
|
||||
"redact" => Some(PrivacyMode::Redact),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for PrivacyMode {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for PrivacyMode {
|
||||
fn default() -> Self {
|
||||
PrivacyMode::Off
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct StateFile {
|
||||
#[serde(default)]
|
||||
mode: String,
|
||||
}
|
||||
|
||||
/// Read current privacy mode from state.json. Returns `Off` on any error.
|
||||
pub fn read_mode(state_file: &Path) -> PrivacyMode {
|
||||
match std::fs::read_to_string(state_file) {
|
||||
Ok(content) => match serde_json::from_str::<StateFile>(&content) {
|
||||
Ok(state) => PrivacyMode::from_str_opt(&state.mode).unwrap_or_else(|| {
|
||||
warn!(
|
||||
mode = %state.mode,
|
||||
"Unknown privacy mode in state.json, defaulting to off"
|
||||
);
|
||||
PrivacyMode::Off
|
||||
}),
|
||||
Err(e) => {
|
||||
warn!(error = %e, "Failed to parse state.json, defaulting to off");
|
||||
PrivacyMode::Off
|
||||
}
|
||||
},
|
||||
Err(_) => {
|
||||
// File not found or unreadable — normal on first boot
|
||||
PrivacyMode::Off
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write privacy mode to state.json. Creates parent directories if needed.
|
||||
pub fn write_mode(state_file: &Path, mode: PrivacyMode) -> anyhow::Result<()> {
|
||||
if let Some(parent) = state_file.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let state = StateFile {
|
||||
mode: mode.as_str().to_string(),
|
||||
};
|
||||
let json = serde_json::to_string(&state)?;
|
||||
std::fs::write(state_file, json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
#[test]
|
||||
fn test_read_mode_missing_file() {
|
||||
let mode = read_mode(Path::new("/nonexistent/state.json"));
|
||||
assert_eq!(mode, PrivacyMode::Off);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_read_mode_valid() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("state.json");
|
||||
std::fs::write(&path, r#"{"mode":"redact"}"#).unwrap();
|
||||
assert_eq!(read_mode(&path), PrivacyMode::Redact);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_read_mode_audit() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("state.json");
|
||||
std::fs::write(&path, r#"{"mode":"audit"}"#).unwrap();
|
||||
assert_eq!(read_mode(&path), PrivacyMode::Audit);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_read_mode_invalid_json() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("state.json");
|
||||
std::fs::write(&path, "not json").unwrap();
|
||||
assert_eq!(read_mode(&path), PrivacyMode::Off);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_read_mode_unknown_value() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("state.json");
|
||||
std::fs::write(&path, r#"{"mode":"unknown"}"#).unwrap();
|
||||
assert_eq!(read_mode(&path), PrivacyMode::Off);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_read_mode_empty_json() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("state.json");
|
||||
std::fs::write(&path, "{}").unwrap();
|
||||
assert_eq!(read_mode(&path), PrivacyMode::Off);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_and_read_mode() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let path = dir.path().join("subdir").join("state.json");
|
||||
write_mode(&path, PrivacyMode::Redact).unwrap();
|
||||
assert_eq!(read_mode(&path), PrivacyMode::Redact);
|
||||
|
||||
write_mode(&path, PrivacyMode::Audit).unwrap();
|
||||
assert_eq!(read_mode(&path), PrivacyMode::Audit);
|
||||
|
||||
write_mode(&path, PrivacyMode::Off).unwrap();
|
||||
assert_eq!(read_mode(&path), PrivacyMode::Off);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_privacy_mode_display() {
|
||||
assert_eq!(PrivacyMode::Off.to_string(), "off");
|
||||
assert_eq!(PrivacyMode::Audit.to_string(), "audit");
|
||||
assert_eq!(PrivacyMode::Redact.to_string(), "redact");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_privacy_mode_serde() {
|
||||
let json = serde_json::to_string(&PrivacyMode::Redact).unwrap();
|
||||
assert_eq!(json, r#""redact""#);
|
||||
let mode: PrivacyMode = serde_json::from_str(r#""audit""#).unwrap();
|
||||
assert_eq!(mode, PrivacyMode::Audit);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user