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:
2026-03-08 06:25:17 +00:00
co-authored by Claude Opus 4.6
parent e6115fbfd7
commit ec464e3d9e
12 changed files with 6789 additions and 0 deletions
File diff suppressed because it is too large Load Diff
+67
View File
@@ -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
+336
View File
@@ -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(&params.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(&params.domain)
.bind(&pii_json)
.bind(&params.action)
.bind(&doc_hash)
.bind(params.file_count)
.bind(&params.client_ip)
.bind(&params.request_url)
.bind(&params.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
}
}
}
+192
View File
@@ -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
}
}
+176
View File
@@ -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());
}
}
+266
View File
@@ -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(),
))
}
+349
View File
@@ -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()
}
+310
View File
@@ -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()
}
}
}
+320
View File
@@ -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());
}
}
+172
View File
@@ -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);
}
}