use axum::{ body::Body, extract::{Request, State}, http::StatusCode, response::Response, }; use bytes::Bytes; use serde_json::Value; use std::{ collections::HashMap, sync::{Arc, Mutex}, time::Instant, }; use tokio::sync::broadcast; use crate::{ capture::{save_capture, CapturedRequest}, config::Config, error::BridgeResult, mapper::{map_to_core, placeholder_response}, shaper::shape_response, }; const CAPTURE_DEDUP_SECS: u64 = 1; #[derive(Clone)] pub struct ProxyState { pub config: Arc, pub http_client: reqwest::Client, /// Broadcast channel for streaming new captures to SSE subscribers. pub capture_tx: Arc>, /// Deduplication window: (method+path) → last saved instant. pub dedup: Arc>>, /// PEM-encoded TLS certificate for download. None when TLS is disabled. pub cert_pem: Option>>, } impl ProxyState { pub fn new(config: Config) -> Self { let (capture_tx, _) = broadcast::channel(256); Self { config: Arc::new(config), http_client: reqwest::Client::builder() .danger_accept_invalid_certs(true) .timeout(std::time::Duration::from_secs(10)) .build() .expect("failed to build HTTP client"), capture_tx: Arc::new(capture_tx), dedup: Arc::new(Mutex::new(HashMap::new())), cert_pem: None, } } pub fn with_cert(mut self, cert_pem: Vec) -> Self { self.cert_pem = Some(Arc::new(cert_pem)); self } } /// Returns true if this (method, path) pair was already saved within the dedup window. fn is_duplicate(dedup: &Mutex>, method: &str, path: &str) -> bool { let key = format!("{method} {path}"); let mut map = dedup.lock().unwrap(); let threshold = std::time::Duration::from_secs(CAPTURE_DEDUP_SECS); if let Some(last) = map.get(&key) { if last.elapsed() < threshold { return true; } } map.insert(key, Instant::now()); false } pub async fn catch_all_handler( State(state): State, req: Request, ) -> BridgeResult> { let method = req.method().to_string(); let uri = req.uri().clone(); let path = uri.path().to_string(); let query = uri.query().map(String::from); let headers: Vec<(String, String)> = req .headers() .iter() .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string())) .collect(); let (_parts, body) = req.into_parts(); let body_bytes: Bytes = axum::body::to_bytes(body, 1024 * 1024) .await .unwrap_or_default(); let body_str = if body_bytes.is_empty() { None } else { Some(String::from_utf8_lossy(&body_bytes).to_string()) }; tracing::info!( "→ {} {}{}", method, path, query .as_deref() .map(|q| format!("?{q}")) .unwrap_or_default() ); let mut capture = CapturedRequest::new( &method, &path, query.as_deref(), headers.clone(), body_str.clone(), ); let (response_body, status_code): (Value, u16) = if let Some(mapping) = map_to_core(&method, &path) { tracing::info!( " ↳ Mapped to Core: {} {}", mapping.method, mapping.core_path ); capture.mapped_to_core = Some(mapping.core_path.to_string()); match forward_to_core( &state, mapping.method, &mapping.core_path, body_str.as_deref(), &headers, ) .await { Ok((body, status)) => (shape_response(&mapping.core_path, body), status), Err(e) => { tracing::error!("Core request failed: {e}"); (serde_json::json!({ "error": e.to_string() }), 502) } } } else if state.config.placeholder_mode { (placeholder_response(&method, &path), 200) } else { (serde_json::json!({ "error": "endpoint not mapped" }), 404) }; capture = capture.with_response(status_code, Some(response_body.to_string())); // Deduplicate: skip saving if same method+path was saved within 1 second let should_save = !is_duplicate(&state.dedup, &method, &path); if should_save { let captures_dir = state.config.captures_dir.clone(); let capture_clone = capture.clone(); let capture_tx = state.capture_tx.clone(); tokio::spawn(async move { if let Err(e) = save_capture(&captures_dir, &capture_clone) { tracing::warn!("Failed to save capture: {e}"); } // Broadcast to SSE subscribers (ignore send errors — no subscribers is OK) let _ = capture_tx.send(capture_clone); }); } let status = StatusCode::from_u16(status_code).unwrap_or(StatusCode::OK); let json_bytes = serde_json::to_vec(&response_body).map_err(crate::error::BridgeError::Serialization)?; let response = Response::builder() .status(status) .header("content-type", "application/json") .body(Body::from(json_bytes)) .map_err(|e| anyhow::anyhow!("response build error: {e}"))?; Ok(response) } /// Public alias for use by the replay route in admin.rs. pub async fn forward_to_core_pub( state: &ProxyState, method: &str, path: &str, body: Option<&str>, headers: &[(String, String)], ) -> anyhow::Result<(Value, u16)> { forward_to_core(state, method, path, body, headers).await } async fn forward_to_core( state: &ProxyState, method: &str, path: &str, body: Option<&str>, incoming_headers: &[(String, String)], ) -> anyhow::Result<(Value, u16)> { let url = format!("{}{}", state.config.core_url, path); let mut builder = match method { "POST" => state.http_client.post(&url), "PUT" => state.http_client.put(&url), "DELETE" => state.http_client.delete(&url), _ => state.http_client.get(&url), }; // Pass through FUT session and anti-phishing headers so Core can log/correlate them. // These are used in #19 (X-UT-SID) and #20 (X-UT-PHISHING-TOKEN) passthrough. for (k, v) in incoming_headers { let lower = k.to_lowercase(); if lower == "x-ut-sid" || lower == "x-ut-phishing-token" || lower == "x-request-id" { builder = builder.header(k.as_str(), v.as_str()); } } builder = if let Some(b) = body { builder .header("content-type", "application/json") .body(b.to_string()) } else { builder }; let resp = builder.send().await?; let status = resp.status().as_u16(); let body: Value = resp.json().await.unwrap_or(Value::Null); Ok((body, status)) }