diff --git a/Cargo.lock b/Cargo.lock index 304da4fb3..d6a1f96d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11209,6 +11209,7 @@ dependencies = [ "webbrowser", "yaak", "yaak-api", + "yaak-common", "yaak-core", "yaak-crypto", "yaak-http", @@ -11240,6 +11241,7 @@ dependencies = [ name = "yaak-common" version = "0.1.0" dependencies = [ + "dirs 7.0.0", "serde_json", "tokio", ] @@ -11516,17 +11518,27 @@ dependencies = [ name = "yaak-proxy-lib" version = "0.0.0" dependencies = [ + "bytes", "chrono", + "flate2", + "http-body-util", + "hyper", + "hyper-util", "include_dir", "log 0.4.34", "r2d2", "r2d2_sqlite", + "reqwest", "rusqlite", "sea-query", "serde", "serde_json", + "tempfile", + "tokio", "ts-rs", "yaak-database", + "yaak-http", + "yaak-models", "yaak-proxy", "yaak-rpc", ] diff --git a/apps/yaak-proxy/components/ExchangesTable.tsx b/apps/yaak-proxy/components/ExchangesTable.tsx index d87d5d551..532d948d3 100644 --- a/apps/yaak-proxy/components/ExchangesTable.tsx +++ b/apps/yaak-proxy/components/ExchangesTable.tsx @@ -1,4 +1,4 @@ -import type { HttpExchange, ProxyHeader } from "@yaakapp-internal/proxy-lib"; +import type { Flow } from "@yaakapp-internal/proxy-lib"; import { Table, TableBody, @@ -11,7 +11,7 @@ import { import classNames from "classnames"; interface Props { - exchanges: HttpExchange[]; + exchanges: Flow[]; style?: React.CSSProperties; className?: string; } @@ -40,10 +40,10 @@ export function ExchangesTable({ exchanges, style, className }: Props) { {ex.url} - + - {getContentType(ex.resHeaders)} + {getContentType(ex.contentType)} ))} @@ -69,8 +69,7 @@ function StatusBadge({ status, error }: { status: number | null; error: string | return {status}; } -function getContentType(headers: ProxyHeader[]): string { - const ct = headers.find((h) => h.name.toLowerCase() === "content-type")?.value; +function getContentType(ct: string | null): string { if (ct == null) return "—"; // Strip parameters (e.g. "; charset=utf-8") return ct.split(";")[0]?.trim() ?? ct; diff --git a/apps/yaak-proxy/components/Sidebar.tsx b/apps/yaak-proxy/components/Sidebar.tsx index 74c19cfc3..4d1631d82 100644 --- a/apps/yaak-proxy/components/Sidebar.tsx +++ b/apps/yaak-proxy/components/Sidebar.tsx @@ -1,10 +1,10 @@ -import type { HttpExchange } from "@yaakapp-internal/proxy-lib"; +import type { Flow } from "@yaakapp-internal/proxy-lib"; import type { TreeNode } from "@yaakapp-internal/ui"; import { selectedIdsFamily, Tree } from "@yaakapp-internal/ui"; import { atom, useAtomValue } from "jotai"; import { atomFamily } from "jotai-family"; import { useCallback } from "react"; -import { httpExchangesAtom } from "../lib/store"; +import { flowsAtom } from "../lib/store"; /** A node in the sidebar tree — either a domain or a path segment. */ export type SidebarItem = { @@ -18,13 +18,13 @@ const collapsedAtom = atomFamily((_treeId: string) => atom>((get) => { - const exchanges = get(httpExchangesAtom); + const exchanges = get(flowsAtom); return buildTree(exchanges); }); /** Exchanges filtered by the currently selected sidebar node(s). */ export const filteredExchangesAtom = atom((get) => { - const exchanges = get(httpExchangesAtom); + const exchanges = get(flowsAtom); const tree = get(sidebarTreeAtom); const selectedIds = get(selectedIdsFamily(SIDEBAR_TREE_ID)); @@ -72,7 +72,7 @@ function collectNodes(node: TreeNode, map: Map * /123 * /orders */ -function buildTree(exchanges: HttpExchange[]): TreeNode { +function buildTree(exchanges: Flow[]): TreeNode { const root: SidebarItem = { id: "root", label: "All Traffic", exchangeIds: [] }; const rootNode: TreeNode = { item: root, diff --git a/apps/yaak-proxy/lib/store.ts b/apps/yaak-proxy/lib/store.ts index fea413cf7..5fbea515d 100644 --- a/apps/yaak-proxy/lib/store.ts +++ b/apps/yaak-proxy/lib/store.ts @@ -1,11 +1,11 @@ import { createModelStore } from "@yaakapp-internal/model-store"; -import type { HttpExchange } from "@yaakapp-internal/proxy-lib"; +import type { Flow } from "@yaakapp-internal/proxy-lib"; type ProxyModels = { - http_exchange: HttpExchange; + flow: Flow; }; export const { dataAtom, applyChange, replaceAll, listAtom, orderedListAtom } = - createModelStore(["http_exchange"]); + createModelStore(["flow"]); -export const httpExchangesAtom = orderedListAtom("http_exchange", "createdAt", "desc"); +export const flowsAtom = orderedListAtom("flow", "startedAt", "desc"); diff --git a/apps/yaak-proxy/main.tsx b/apps/yaak-proxy/main.tsx index 1e197fb0e..323c86350 100644 --- a/apps/yaak-proxy/main.tsx +++ b/apps/yaak-proxy/main.tsx @@ -14,7 +14,7 @@ const jotaiStore = createStore(); // Load initial models from the database void rpc("list_models", {}).then((res) => { - jotaiStore.set(dataAtom, (prev) => replaceAll(prev, "http_exchange", res.httpExchanges)); + jotaiStore.set(dataAtom, (prev) => replaceAll(prev, "flow", res.flows)); }); // Register hotkeys from action metadata @@ -22,9 +22,7 @@ void initHotkeys(); // Subscribe to model change events from the backend void listen("model_write", (payload) => { - jotaiStore.set(dataAtom, (prev) => - applyChange(prev, "http_exchange", payload.model, payload.change), - ); + jotaiStore.set(dataAtom, (prev) => applyChange(prev, "flow", payload.model, payload.change)); }); const motionFeatures = () => import("framer-motion").then((mod) => mod.domAnimation); diff --git a/crates-cli/yaak-cli/Cargo.toml b/crates-cli/yaak-cli/Cargo.toml index fb275424b..50db83428 100644 --- a/crates-cli/yaak-cli/Cargo.toml +++ b/crates-cli/yaak-cli/Cargo.toml @@ -45,6 +45,7 @@ webbrowser = "1" zip = "7" yaak = { workspace = true } yaak-api = { workspace = true } +yaak-common = { workspace = true } yaak-core = { workspace = true } yaak-crypto = { workspace = true } yaak-http = { workspace = true } diff --git a/crates-cli/yaak-cli/src/main.rs b/crates-cli/yaak-cli/src/main.rs index b6020378a..de9dfdb1a 100644 --- a/crates-cli/yaak-cli/src/main.rs +++ b/crates-cli/yaak-cli/src/main.rs @@ -10,7 +10,7 @@ mod version_check; use clap::{CommandFactory, FromArgMatches}; use cli::{AGENT_HINTS, Cli, Commands, PluginCommands, RequestCommands}; use context::{CliContext, CliExecutionContext}; -use std::path::PathBuf; +use yaak_common::data_dir::resolve_data_dir; use yaak_models::queries::any_request::AnyRequest; /// Built at runtime so root `--help` can report whether the installed agent skill is @@ -286,46 +286,3 @@ fn resolve_cookie_jar_id( .map(|jar| jar.id); Ok(default_cookie_jar) } - -fn resolve_data_dir(app_id: &str) -> PathBuf { - if let Some(dir) = wsl_data_dir(app_id) { - return dir; - } - dirs::data_dir().expect("Could not determine data directory").join(app_id) -} - -/// Detect WSL and resolve the Windows AppData\Roaming path for the Yaak data directory. -fn wsl_data_dir(app_id: &str) -> Option { - if !cfg!(target_os = "linux") { - return None; - } - - let proc_version = std::fs::read_to_string("/proc/version").ok()?; - let is_wsl = proc_version.to_lowercase().contains("microsoft"); - if !is_wsl { - return None; - } - - // We're in WSL, so try to resolve the Yaak app's data directory in Windows - - // Get the Windows %APPDATA% path via cmd.exe - let appdata_output = - std::process::Command::new("cmd.exe").args(["/C", "echo", "%APPDATA%"]).output().ok()?; - - let win_path = String::from_utf8(appdata_output.stdout).ok()?.trim().to_string(); - if win_path.is_empty() || win_path == "%APPDATA%" { - return None; - } - - // Convert Windows path to WSL path using wslpath (handles custom mount points) - let wslpath_output = std::process::Command::new("wslpath").arg(&win_path).output().ok()?; - - let wsl_appdata = String::from_utf8(wslpath_output.stdout).ok()?.trim().to_string(); - if wsl_appdata.is_empty() { - return None; - } - - let wsl_path = PathBuf::from(wsl_appdata).join(app_id); - - if wsl_path.exists() { Some(wsl_path) } else { None } -} diff --git a/crates-proxy/yaak-proxy-lib/Cargo.toml b/crates-proxy/yaak-proxy-lib/Cargo.toml index 0eb2dd8ca..3753adedf 100644 --- a/crates-proxy/yaak-proxy-lib/Cargo.toml +++ b/crates-proxy/yaak-proxy-lib/Cargo.toml @@ -15,7 +15,20 @@ rusqlite = { version = "0.38", features = ["bundled", "chrono"] } sea-query = { version = "1.0", features = ["with-chrono", "attr"] } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +tokio = { workspace = true, features = ["io-util"] } ts-rs = { workspace = true, features = ["chrono-impl"] } yaak-database = { workspace = true } +yaak-http = { workspace = true } +yaak-models = { workspace = true } yaak-proxy = { workspace = true } yaak-rpc = { workspace = true } + +[dev-dependencies] +bytes = "1" +flate2 = "1" +http-body-util = "0.1" +hyper = { version = "1", features = ["http1", "server"] } +hyper-util = { version = "0.1", features = ["tokio"] } +reqwest = { workspace = true } +tempfile = "3" +tokio = { workspace = true, features = ["macros", "net", "rt-multi-thread", "time"] } diff --git a/crates-proxy/yaak-proxy-lib/bindings/gen_models.ts b/crates-proxy/yaak-proxy-lib/bindings/gen_models.ts index ba452b1cf..917c14541 100644 --- a/crates-proxy/yaak-proxy-lib/bindings/gen_models.ts +++ b/crates-proxy/yaak-proxy-lib/bindings/gen_models.ts @@ -1,8 +1,47 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { ModelChangeEvent } from "./ModelChangeEvent"; -export type HttpExchange = { id: string, createdAt: string, updatedAt: string, url: string, method: string, reqHeaders: Array, reqBody: Array | null, resStatus: number | null, resHeaders: Array, resBody: Array | null, error: string | null, }; +/** + * One proxied request and its response. The bodies are in `blobs.sqlite`, not the row. + */ +export type Flow = { id: string, createdAt: string, updatedAt: string, startedAt: string, method: string, url: string, +/** + * Version spoken by the client. + */ +clientHttpVersion: string, +/** + * Version spoken with the upstream server. + */ +upstreamHttpVersion: string | null, remoteAddr: string | null, status: number | null, +/** + * The response's Content-Type. + */ +contentType: string | null, requestHeaders: Array, responseHeaders: Array, +/** + * Bytes of the request body on the wire, captured or not. + */ +requestSize: number, +/** + * Bytes of the response body on the wire, captured or not. + */ +responseSize: number, +/** + * Only the first `capture_body_limit` bytes of the body are stored. + */ +requestTruncated: boolean, +/** + * Only the first `capture_body_limit` bytes of the body are stored. + */ +responseTruncated: boolean, +/** + * Milliseconds from `started_at` to the response headers. + */ +elapsedHeadersMs: number | null, +/** + * Milliseconds from `started_at` to the end of the response. + */ +elapsedMs: number | null, error: string | null, }; -export type ModelPayload = { model: HttpExchange, change: ModelChangeEvent, }; +export type ModelPayload = { model: Flow, change: ModelChangeEvent, }; export type ProxyHeader = { name: string, value: string, }; diff --git a/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts b/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts index 850717c35..6ce078dc4 100644 --- a/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts +++ b/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts @@ -1,5 +1,5 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -import type { HttpExchange, ModelPayload } from "./gen_models"; +import type { Flow, ModelPayload } from "./gen_models"; export type ActionInvocation = { "scope": "global", action: GlobalAction, }; @@ -17,7 +17,7 @@ export type ListActionsResponse = { actions: Array<[ActionInvocation, ActionMeta export type ListModelsRequest = Record; -export type ListModelsResponse = { httpExchanges: Array, }; +export type ListModelsResponse = { flows: Array, }; export type ProxyState = "running" | "stopped"; diff --git a/crates-proxy/yaak-proxy-lib/migrations/00000002_flows.sql b/crates-proxy/yaak-proxy-lib/migrations/00000002_flows.sql new file mode 100644 index 000000000..ba31ed55e --- /dev/null +++ b/crates-proxy/yaak-proxy-lib/migrations/00000002_flows.sql @@ -0,0 +1,27 @@ +DROP TABLE http_exchanges; + +CREATE TABLE flows +( + id TEXT NOT NULL PRIMARY KEY, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL, + updated_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL, + started_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL, + method TEXT NOT NULL DEFAULT '', + url TEXT NOT NULL DEFAULT '', + client_http_version TEXT NOT NULL DEFAULT '', + upstream_http_version TEXT, + remote_addr TEXT, + status INTEGER, + content_type TEXT, + request_headers TEXT NOT NULL DEFAULT '[]', + response_headers TEXT NOT NULL DEFAULT '[]', + request_size INTEGER NOT NULL DEFAULT 0, + response_size INTEGER NOT NULL DEFAULT 0, + request_truncated BOOLEAN NOT NULL DEFAULT FALSE, + response_truncated BOOLEAN NOT NULL DEFAULT FALSE, + elapsed_headers_ms INTEGER, + elapsed_ms INTEGER, + error TEXT +); + +CREATE INDEX idx_flows_started_at ON flows (started_at); diff --git a/crates-proxy/yaak-proxy-lib/src/bodies.rs b/crates-proxy/yaak-proxy-lib/src/bodies.rs new file mode 100644 index 000000000..9d7add1bf --- /dev/null +++ b/crates-proxy/yaak-proxy-lib/src/bodies.rs @@ -0,0 +1,37 @@ +use crate::models::{Flow, FlowBody, find_header}; +use std::io::Cursor; +use tokio::io::AsyncReadExt; +use yaak_http::decompress::{ContentEncoding, streaming_decoder}; +use yaak_models::blob_manager::BlobManager; + +/// A flow's body as stored: the bytes on the wire, up to the capture limit. +pub fn read_raw_body( + blobs: &BlobManager, + flow_id: &str, + body: FlowBody, +) -> Result, String> { + let chunks = + blobs.connect().get_chunks(&Flow::body_id(flow_id, body)).map_err(|e| e.to_string())?; + Ok(chunks.into_iter().flat_map(|c| c.data).collect()) +} + +/// A flow's body decoded with its Content-Encoding. A truncated body decodes as far as its +/// stored bytes go. +pub async fn read_body( + blobs: &BlobManager, + flow: &Flow, + body: FlowBody, +) -> Result, String> { + let raw = read_raw_body(blobs, &flow.id, body)?; + let encoding = ContentEncoding::from_header(find_header(flow.headers(body), "content-encoding")); + if encoding == ContentEncoding::Identity { + return Ok(raw); + } + + let mut decoded = Vec::new(); + match streaming_decoder(Cursor::new(raw), encoding).read_to_end(&mut decoded).await { + Ok(_) => Ok(decoded), + Err(_) if flow.truncated(body) => Ok(decoded), + Err(e) => Err(format!("Failed to decode {encoding:?} body: {e}")), + } +} diff --git a/crates-proxy/yaak-proxy-lib/src/db.rs b/crates-proxy/yaak-proxy-lib/src/db.rs index e5bba29d2..8e71c5e46 100644 --- a/crates-proxy/yaak-proxy-lib/src/db.rs +++ b/crates-proxy/yaak-proxy-lib/src/db.rs @@ -1,31 +1,131 @@ use include_dir::{Dir, include_dir}; use r2d2::Pool; use r2d2_sqlite::SqliteConnectionManager; +use rusqlite::{Transaction, TransactionBehavior}; use std::path::Path; -use yaak_database::{ConnectionOrTx, DbContext, run_migrations}; +use std::time::Duration; +use yaak_database::{ConnectionOrTx, DbContext, Error, Result, SqlitePool, run_migrations}; static MIGRATIONS: Dir<'static> = include_dir!("$CARGO_MANIFEST_DIR/migrations"); +/// Reads come from a pool; writes go through one connection, as in yaak-models' +/// `QueryManager`. WAL lets the CLI read `proxy.db` while the app writes it. #[derive(Clone)] pub struct ProxyQueryManager { - pool: Pool, + readers: SqlitePool, + writer: SqlitePool, } impl ProxyQueryManager { - pub fn new(db_path: &Path) -> Self { - let manager = SqliteConnectionManager::file(db_path); - let pool = - Pool::builder().max_size(5).build(manager).expect("Failed to create proxy DB pool"); - run_migrations(&pool, &MIGRATIONS).expect("Failed to run proxy DB migrations"); - Self { pool } + pub fn new(db_path: &Path) -> Result { + if let Some(parent) = db_path.parent() { + std::fs::create_dir_all(parent)?; + } + let writer = file_pool(db_path, 1)?; + run_migrations(&writer, &MIGRATIONS)?; + Ok(Self { readers: file_pool(db_path, 5)?, writer }) } pub fn with_conn(&self, func: F) -> T where F: FnOnce(&DbContext) -> T, { - let conn = self.pool.get().expect("Failed to get proxy DB connection"); + let conn = self.readers.get().expect("Failed to get proxy DB connection"); let ctx = DbContext::new(ConnectionOrTx::Connection(conn)); func(&ctx) } + + /// Run `func` in a transaction on the writer connection. + pub fn with_tx( + &self, + func: impl FnOnce(&DbContext) -> std::result::Result, + ) -> std::result::Result + where + E: From, + { + let conn = self.writer.get().map_err(Error::from)?; + let tx = Transaction::new_unchecked(&conn, TransactionBehavior::Immediate) + .map_err(Error::from)?; + let ctx = DbContext::new(ConnectionOrTx::Transaction(&tx)); + let result = func(&ctx); + drop(ctx); + match result { + Ok(val) => { + tx.commit().map_err(Error::from)?; + Ok(val) + } + Err(e) => { + tx.rollback().map_err(Error::from)?; + Err(e) + } + } + } +} + +fn file_pool(path: &Path, max_size: u32) -> Result { + let manager = SqliteConnectionManager::file(path).with_init(|conn| { + conn.pragma_update(None, "journal_mode", "WAL")?; + conn.pragma_update(None, "synchronous", "NORMAL")?; + conn.busy_timeout(Duration::from_millis(5000)) + }); + Pool::builder() + .max_size(max_size) + .connection_timeout(Duration::from_secs(10)) + .build(manager) + .map_err(|e| Error::Database(e.to_string())) +} + +#[cfg(test)] +mod tests { + use super::*; + use rusqlite::Connection; + + fn columns(conn: &Connection, table: &str) -> Vec { + let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})")).unwrap(); + stmt.query_map([], |r| r.get::<_, String>("name")).unwrap().map(|c| c.unwrap()).collect() + } + + #[test] + fn migrates_an_old_schema_db() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("proxy.db"); + + // The database as the first migration left it, with a row in it. + let conn = Connection::open(&path).unwrap(); + conn.execute_batch(include_str!("../migrations/00000001_init.sql")).unwrap(); + conn.execute_batch( + "CREATE TABLE _sqlx_migrations ( + version TEXT PRIMARY KEY, + description TEXT NOT NULL, + applied_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL + ); + INSERT INTO _sqlx_migrations (version, description) + VALUES ('00000001', '00000001_init.sql'); + INSERT INTO http_exchanges (id, url, method, req_body, res_body) + VALUES ('he_old', 'https://example.com', 'GET', x'00', x'01');", + ) + .unwrap(); + drop(conn); + + let db = ProxyQueryManager::new(&path).unwrap(); + drop(db); + + let conn = Connection::open(&path).unwrap(); + let tables: Vec = conn + .prepare("SELECT name FROM sqlite_master WHERE type = 'table' ORDER BY name") + .unwrap() + .query_map([], |r| r.get(0)) + .unwrap() + .map(|t| t.unwrap()) + .collect(); + assert_eq!(tables, vec!["_sqlx_migrations", "flows"]); + + let flow_columns = columns(&conn, "flows"); + assert!(flow_columns.contains(&"response_truncated".to_string())); + assert!(!flow_columns.iter().any(|c| c.contains("body")), "{flow_columns:?}"); + + let journal_mode: String = + conn.query_row("PRAGMA journal_mode", [], |r| r.get(0)).unwrap(); + assert_eq!(journal_mode, "wal"); + } } diff --git a/crates-proxy/yaak-proxy-lib/src/lib.rs b/crates-proxy/yaak-proxy-lib/src/lib.rs index de186e97d..0f2bfbf2f 100644 --- a/crates-proxy/yaak-proxy-lib/src/lib.rs +++ b/crates-proxy/yaak-proxy-lib/src/lib.rs @@ -1,18 +1,23 @@ pub mod actions; +pub mod bodies; pub mod db; pub mod models; +pub mod settings; use crate::actions::{ActionInvocation, ActionMetadata, GlobalAction}; use crate::db::ProxyQueryManager; -use crate::models::{HttpExchange, ModelPayload, ProxyHeader}; +use crate::models::{Flow, FlowBody, ModelPayload, ProxyHeader, find_header}; +use crate::settings::ProxySettings; +use chrono::{DateTime, Utc}; use log::warn; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; use ts_rs::TS; -use yaak_database::{ModelChangeEvent, UpdateSource}; -use yaak_proxy::{CaStore, CapturedRequest, ProxyConfig, ProxyEvent, ProxyHandle, RequestState}; +use yaak_database::{ModelChangeEvent, UpdateSource, UpsertModelInfo}; +use yaak_models::blob_manager::{BlobManager, BodyChunkWriter}; +use yaak_proxy::{CaStore, ProxyConfig, ProxyEvent, ProxyHandle}; use yaak_rpc::{RpcError, RpcEventEmitter, define_rpc}; // -- Context -- @@ -21,20 +26,68 @@ use yaak_rpc::{RpcError, RpcEventEmitter, define_rpc}; #[derive(Clone)] pub struct ProxyCtx { handle: Arc>>, - ca_dir: PathBuf, + data_dir: PathBuf, pub db: ProxyQueryManager, + pub blobs: BlobManager, pub events: RpcEventEmitter, } impl ProxyCtx { - pub fn new(db_path: &Path, events: RpcEventEmitter) -> Self { - let ca_dir = db_path.parent().unwrap_or(Path::new(".")).join("ca"); - Self { + /// Opens `proxy.db` and `blobs.sqlite` in `data_dir`, the proxy app's data directory. + pub fn new(data_dir: &Path, events: RpcEventEmitter) -> Result { + let db = ProxyQueryManager::new(&data_dir.join("proxy.db")) + .map_err(|e| format!("Failed to open proxy database: {e}"))?; + let blobs = yaak_models::init_blob_store(data_dir.join("blobs.sqlite")) + .map_err(|e| format!("Failed to open proxy blob store: {e}"))?; + Ok(Self { handle: Arc::new(Mutex::new(None)), - ca_dir, - db: ProxyQueryManager::new(db_path), + data_dir: data_dir.to_path_buf(), + db, + blobs, events, + }) + } + + pub fn settings_path(&self) -> PathBuf { + self.data_dir.join("settings.json") + } + + /// Starts the engine with the current `settings.json`. Does nothing if it is running. + pub fn start(&self) -> Result<(), String> { + let mut handle = self.handle.lock().map_err(|_| "lock poisoned".to_string())?; + if handle.is_some() { + return Ok(()); } + + let settings = ProxySettings::load(&self.settings_path())?; + let ca = CaStore::load_or_create(&self.data_dir.join("ca")) + .map_err(|e| format!("Failed to load proxy CA: {e}"))?; + let mut config = ProxyConfig::new(settings.port, Arc::new(ca)); + config.capture_body_limit = settings.capture_body_limit; + let mut proxy_handle = yaak_proxy::start_proxy(config)?; + + if let Some(event_rx) = proxy_handle.take_event_rx() { + let db = self.db.clone(); + let blobs = self.blobs.clone(); + let events = self.events.clone(); + std::thread::spawn(move || run_event_loop(event_rx, db, blobs, events)); + } + + *handle = Some(proxy_handle); + self.events.emit("proxy_state_changed", &ProxyStatePayload { state: ProxyState::Running }); + Ok(()) + } + + pub fn stop(&self) -> Result<(), String> { + let mut handle = self.handle.lock().map_err(|_| "lock poisoned".to_string())?; + handle.take(); + self.events.emit("proxy_state_changed", &ProxyStatePayload { state: ProxyState::Stopped }); + Ok(()) + } + + /// The port the engine is bound to, while it runs. + pub fn port(&self) -> Option { + self.handle.lock().ok()?.as_ref().map(|h| h.port) } } @@ -74,7 +127,7 @@ pub struct ListModelsRequest {} #[ts(export, export_to = "gen_rpc.ts")] #[serde(rename_all = "camelCase")] pub struct ListModelsResponse { - pub http_exchanges: Vec, + pub flows: Vec, } #[derive(Deserialize, TS)] @@ -92,41 +145,11 @@ pub struct GetProxyStateResponse { fn execute_action(ctx: &ProxyCtx, invocation: ActionInvocation) -> Result { match invocation { ActionInvocation::Global { action } => match action { - GlobalAction::ProxyStart => { - let mut handle = - ctx.handle.lock().map_err(|_| RpcError { message: "lock poisoned".into() })?; - - if handle.is_some() { - return Ok(true); // already running - } - - let ca = CaStore::load_or_create(&ctx.ca_dir) - .map_err(|e| RpcError { message: format!("Failed to load proxy CA: {e}") })?; - let mut proxy_handle = - yaak_proxy::start_proxy(ProxyConfig::new(9090, Arc::new(ca))) - .map_err(|e| RpcError { message: e })?; - - if let Some(event_rx) = proxy_handle.take_event_rx() { - let db = ctx.db.clone(); - let events = ctx.events.clone(); - std::thread::spawn(move || run_event_loop(event_rx, db, events)); - } - - *handle = Some(proxy_handle); - ctx.events - .emit("proxy_state_changed", &ProxyStatePayload { state: ProxyState::Running }); - Ok(true) - } - GlobalAction::ProxyStop => { - let mut handle = - ctx.handle.lock().map_err(|_| RpcError { message: "lock poisoned".into() })?; - handle.take(); - ctx.events - .emit("proxy_state_changed", &ProxyStatePayload { state: ProxyState::Stopped }); - Ok(true) - } + GlobalAction::ProxyStart => ctx.start().map_err(|message| RpcError { message })?, + GlobalAction::ProxyStop => ctx.stop().map_err(|message| RpcError { message })?, }, } + Ok(true) } fn get_proxy_state( @@ -148,128 +171,151 @@ fn list_actions( fn list_models(ctx: &ProxyCtx, _req: ListModelsRequest) -> Result { ctx.db.with_conn(|db| { Ok(ListModelsResponse { - http_exchanges: db - .find_all::() - .map_err(|e| RpcError { message: e.to_string() })?, + flows: db.find_all::().map_err(|e| RpcError { message: e.to_string() })?, }) }) } // -- Event loop -- +/// A flow between its `RequestStart` and its end, with its bodies being written as they +/// stream. +struct InFlight { + flow: Flow, + request_body: Option, + response_body: Option, +} + fn run_event_loop( rx: std::sync::mpsc::Receiver, db: ProxyQueryManager, + blobs: BlobManager, events: RpcEventEmitter, ) { - let mut in_flight: HashMap = HashMap::new(); + let mut in_flight: HashMap = HashMap::new(); while let Ok(event) = rx.recv() { match event { - ProxyEvent::RequestStart { id, method, url, http_version, .. } => { + ProxyEvent::RequestStart { id, started_at, method, url, http_version } => { + let flow_id = Flow::generate_id(); + let request_body = Flow::body_id(&flow_id, FlowBody::Request); + let response_body = Flow::body_id(&flow_id, FlowBody::Response); + let flow = Flow { + id: flow_id, + started_at: DateTime::::from(started_at).naive_utc(), + method, + url, + client_http_version: http_version, + ..Default::default() + }; in_flight.insert( id, - CapturedRequest { - id, - method, - url, - http_version, - status: None, - elapsed_ms: None, - remote_http_version: None, - request_headers: vec![], - request_body: None, - response_headers: vec![], - response_body: None, - response_body_size: 0, - state: RequestState::Sending, - error: None, + InFlight { + flow, + request_body: Some(BodyChunkWriter::new(blobs.clone(), request_body)), + response_body: Some(BodyChunkWriter::new(blobs.clone(), response_body)), }, ); } ProxyEvent::RequestHeader { id, name, value } => { - if let Some(r) = in_flight.get_mut(&id) { - r.request_headers.push((name, value)); + if let Some(f) = in_flight.get_mut(&id) { + f.flow.request_headers.push(ProxyHeader { name, value }); } } - ProxyEvent::RequestBodyChunk { id, data, .. } => { - if let Some(r) = in_flight.get_mut(&id) - && !data.is_empty() - { - r.request_body.get_or_insert_default().extend_from_slice(&data); + // Counted per chunk too: the response can end, and the flow with it, before the + // upload does. + ProxyEvent::RequestBodyChunk { id, len, data } => { + if let Some(f) = in_flight.get_mut(&id) { + f.flow.request_size += len as i64; + f.flow.request_truncated |= data.len() < len; + write_body_chunk(&mut f.request_body, &data); } } - ProxyEvent::RequestBodyComplete { .. } => {} - ProxyEvent::ResponseStart { id, status, http_version, elapsed_ms, .. } => { - if let Some(r) = in_flight.get_mut(&id) { - r.status = Some(status); - r.remote_http_version = Some(http_version); - r.elapsed_ms = Some(elapsed_ms); - r.state = RequestState::Receiving; + ProxyEvent::RequestBodyComplete { id, size, truncated } => { + if let Some(f) = in_flight.get_mut(&id) { + f.flow.request_size = size as i64; + f.flow.request_truncated = truncated; + finish_body(f.request_body.take()); + } + } + ProxyEvent::ResponseStart { id, status, http_version, remote_addr, elapsed_ms } => { + if let Some(f) = in_flight.get_mut(&id) { + f.flow.status = Some(status as i32); + f.flow.upstream_http_version = Some(http_version); + f.flow.remote_addr = remote_addr.map(|a| a.to_string()); + f.flow.elapsed_headers_ms = Some(clamp_ms(elapsed_ms)); } } ProxyEvent::ResponseHeader { id, name, value } => { - if let Some(r) = in_flight.get_mut(&id) { - r.response_headers.push((name, value)); + if let Some(f) = in_flight.get_mut(&id) { + f.flow.response_headers.push(ProxyHeader { name, value }); } } ProxyEvent::ResponseBodyChunk { id, data, .. } => { - if let Some(r) = in_flight.get_mut(&id) - && !data.is_empty() - { - r.response_body.get_or_insert_default().extend_from_slice(&data); + if let Some(f) = in_flight.get_mut(&id) { + write_body_chunk(&mut f.response_body, &data); } } - ProxyEvent::ResponseBodyComplete { id, size, elapsed_ms, error, .. } => { - if let Some(mut r) = in_flight.remove(&id) { - r.response_body_size = size; - r.elapsed_ms = r.elapsed_ms.or(Some(elapsed_ms)); - r.state = - if error.is_some() { RequestState::Error } else { RequestState::Complete }; - r.error = error; - write_entry(&db, &events, &r); + ProxyEvent::ResponseBodyComplete { id, size, truncated, elapsed_ms, error } => { + if let Some(mut f) = in_flight.remove(&id) { + f.flow.response_size = size as i64; + f.flow.response_truncated = truncated; + f.flow.elapsed_ms = Some(clamp_ms(elapsed_ms)); + f.flow.error = error; + finish_flow(&db, &events, f); } } - ProxyEvent::Error { id, error, .. } => { - if let Some(mut r) = in_flight.remove(&id) { - r.error = Some(error); - r.state = RequestState::Error; - write_entry(&db, &events, &r); + ProxyEvent::Error { id, error, elapsed_ms } => { + if let Some(mut f) = in_flight.remove(&id) { + f.flow.elapsed_ms = Some(clamp_ms(elapsed_ms)); + f.flow.error = Some(error); + finish_flow(&db, &events, f); } } } } } -fn write_entry(db: &ProxyQueryManager, events: &RpcEventEmitter, r: &CapturedRequest) { - let entry = HttpExchange { - url: r.url.clone(), - method: r.method.clone(), - req_headers: r - .request_headers - .iter() - .map(|(n, v)| ProxyHeader { name: n.clone(), value: v.clone() }) - .collect(), - req_body: r.request_body.clone(), - res_status: r.status.map(|s| s as i32), - res_headers: r - .response_headers - .iter() - .map(|(n, v)| ProxyHeader { name: n.clone(), value: v.clone() }) - .collect(), - res_body: r.response_body.clone(), - error: r.error.clone(), - ..Default::default() - }; - db.with_conn(|ctx| match ctx.upsert(&entry, &UpdateSource::Background) { +fn clamp_ms(ms: u64) -> i32 { + i32::try_from(ms).unwrap_or(i32::MAX) +} + +/// Stops writing a body after its first failed write, keeping the chunks before it. +fn write_body_chunk(writer: &mut Option, data: &[u8]) { + if data.is_empty() { + return; + } + if let Some(w) = writer.as_mut() + && let Err(e) = w.write(data) + { + warn!("Failed to store proxy body chunk: {e}"); + *writer = None; + } +} + +fn finish_body(writer: Option) { + if let Some(w) = writer + && let Err(e) = w.finish() + { + warn!("Failed to store proxy body chunk: {e}"); + } +} + +fn finish_flow(db: &ProxyQueryManager, events: &RpcEventEmitter, f: InFlight) { + finish_body(f.request_body); + finish_body(f.response_body); + + let mut flow = f.flow; + flow.content_type = find_header(&flow.response_headers, "content-type").map(String::from); + match db.with_tx(|ctx| ctx.upsert(&flow, &UpdateSource::Background)) { Ok((saved, created)) => { events.emit( "model_write", &ModelPayload { model: saved, change: ModelChangeEvent::Upsert { created } }, ); } - Err(e) => warn!("Failed to write proxy entry: {e}"), - }); + Err(e) => warn!("Failed to write proxy flow: {e}"), + } } // -- Router + Schema -- diff --git a/crates-proxy/yaak-proxy-lib/src/models.rs b/crates-proxy/yaak-proxy-lib/src/models.rs index 0b6249ae0..5fa05c507 100644 --- a/crates-proxy/yaak-proxy-lib/src/models.rs +++ b/crates-proxy/yaak-proxy-lib/src/models.rs @@ -16,47 +16,102 @@ pub struct ProxyHeader { pub value: String, } +/// One proxied request and its response. The bodies are in `blobs.sqlite`, not the row. #[derive(Debug, Clone, Default, Serialize, Deserialize, TS)] #[serde(default, rename_all = "camelCase")] #[ts(export, export_to = "gen_models.ts")] -#[enum_def(table_name = "http_exchanges")] -pub struct HttpExchange { +#[enum_def(table_name = "flows")] +pub struct Flow { pub id: String, pub created_at: NaiveDateTime, pub updated_at: NaiveDateTime, - pub url: String, + pub started_at: NaiveDateTime, pub method: String, - pub req_headers: Vec, - pub req_body: Option>, - pub res_status: Option, - pub res_headers: Vec, - pub res_body: Option>, + pub url: String, + /// Version spoken by the client. + pub client_http_version: String, + /// Version spoken with the upstream server. + pub upstream_http_version: Option, + pub remote_addr: Option, + pub status: Option, + /// The response's Content-Type. + pub content_type: Option, + pub request_headers: Vec, + pub response_headers: Vec, + /// Bytes of the request body on the wire, captured or not. + #[ts(type = "number")] + pub request_size: i64, + /// Bytes of the response body on the wire, captured or not. + #[ts(type = "number")] + pub response_size: i64, + /// Only the first `capture_body_limit` bytes of the body are stored. + pub request_truncated: bool, + /// Only the first `capture_body_limit` bytes of the body are stored. + pub response_truncated: bool, + /// Milliseconds from `started_at` to the response headers. + pub elapsed_headers_ms: Option, + /// Milliseconds from `started_at` to the end of the response. + pub elapsed_ms: Option, pub error: Option, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FlowBody { + Request, + Response, +} + +impl Flow { + /// The blob store id of one of a flow's bodies. + pub fn body_id(flow_id: &str, body: FlowBody) -> String { + match body { + FlowBody::Request => format!("{flow_id}.request"), + FlowBody::Response => format!("{flow_id}.response"), + } + } + + pub fn headers(&self, body: FlowBody) -> &[ProxyHeader] { + match body { + FlowBody::Request => &self.request_headers, + FlowBody::Response => &self.response_headers, + } + } + + pub fn truncated(&self, body: FlowBody) -> bool { + match body { + FlowBody::Request => self.request_truncated, + FlowBody::Response => self.response_truncated, + } + } +} + +pub fn find_header<'a>(headers: &'a [ProxyHeader], name: &str) -> Option<&'a str> { + headers.iter().find(|h| h.name.eq_ignore_ascii_case(name)).map(|h| h.value.as_str()) +} + #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[serde(rename_all = "camelCase")] #[ts(export, export_to = "gen_models.ts")] pub struct ModelPayload { - pub model: HttpExchange, + pub model: Flow, pub change: ModelChangeEvent, } -impl UpsertModelInfo for HttpExchange { +impl UpsertModelInfo for Flow { fn table_name() -> impl IntoTableRef + IntoIden { - HttpExchangeIden::Table + FlowIden::Table } fn id_column() -> impl IntoIden + Eq + Clone { - HttpExchangeIden::Id + FlowIden::Id } fn generate_id() -> String { - generate_prefixed_id("he") + generate_prefixed_id("fl") } fn order_by() -> (impl IntoColumnRef, Order) { - (HttpExchangeIden::CreatedAt, Order::Desc) + (FlowIden::StartedAt, Order::Desc) } fn get_id(&self) -> String { @@ -67,32 +122,50 @@ impl UpsertModelInfo for HttpExchange { self, source: &UpdateSource, ) -> DbResult)>> { - use HttpExchangeIden::*; + use FlowIden::*; Ok(vec![ (CreatedAt, upsert_date(source, self.created_at)), (UpdatedAt, upsert_date(source, self.updated_at)), - (Url, self.url.into()), + (StartedAt, self.started_at.into()), (Method, self.method.into()), - (ReqHeaders, serde_json::to_string(&self.req_headers)?.into()), - (ReqBody, self.req_body.into()), - (ResStatus, self.res_status.into()), - (ResHeaders, serde_json::to_string(&self.res_headers)?.into()), - (ResBody, self.res_body.into()), + (Url, self.url.into()), + (ClientHttpVersion, self.client_http_version.into()), + (UpstreamHttpVersion, self.upstream_http_version.into()), + (RemoteAddr, self.remote_addr.into()), + (Status, self.status.into()), + (ContentType, self.content_type.into()), + (RequestHeaders, serde_json::to_string(&self.request_headers)?.into()), + (ResponseHeaders, serde_json::to_string(&self.response_headers)?.into()), + (RequestSize, self.request_size.into()), + (ResponseSize, self.response_size.into()), + (RequestTruncated, self.request_truncated.into()), + (ResponseTruncated, self.response_truncated.into()), + (ElapsedHeadersMs, self.elapsed_headers_ms.into()), + (ElapsedMs, self.elapsed_ms.into()), (Error, self.error.into()), ]) } fn update_columns() -> Vec { vec![ - HttpExchangeIden::UpdatedAt, - HttpExchangeIden::Url, - HttpExchangeIden::Method, - HttpExchangeIden::ReqHeaders, - HttpExchangeIden::ReqBody, - HttpExchangeIden::ResStatus, - HttpExchangeIden::ResHeaders, - HttpExchangeIden::ResBody, - HttpExchangeIden::Error, + FlowIden::UpdatedAt, + FlowIden::StartedAt, + FlowIden::Method, + FlowIden::Url, + FlowIden::ClientHttpVersion, + FlowIden::UpstreamHttpVersion, + FlowIden::RemoteAddr, + FlowIden::Status, + FlowIden::ContentType, + FlowIden::RequestHeaders, + FlowIden::ResponseHeaders, + FlowIden::RequestSize, + FlowIden::ResponseSize, + FlowIden::RequestTruncated, + FlowIden::ResponseTruncated, + FlowIden::ElapsedHeadersMs, + FlowIden::ElapsedMs, + FlowIden::Error, ] } @@ -100,19 +173,28 @@ impl UpsertModelInfo for HttpExchange { where Self: Sized, { - let req_headers: String = r.get("req_headers")?; - let res_headers: String = r.get("res_headers")?; + let request_headers: String = r.get("request_headers")?; + let response_headers: String = r.get("response_headers")?; Ok(Self { id: r.get("id")?, created_at: r.get("created_at")?, updated_at: r.get("updated_at")?, - url: r.get("url")?, + started_at: r.get("started_at")?, method: r.get("method")?, - req_headers: serde_json::from_str(&req_headers).unwrap_or_default(), - req_body: r.get("req_body")?, - res_status: r.get("res_status")?, - res_headers: serde_json::from_str(&res_headers).unwrap_or_default(), - res_body: r.get("res_body")?, + url: r.get("url")?, + client_http_version: r.get("client_http_version")?, + upstream_http_version: r.get("upstream_http_version")?, + remote_addr: r.get("remote_addr")?, + status: r.get("status")?, + content_type: r.get("content_type")?, + request_headers: serde_json::from_str(&request_headers).unwrap_or_default(), + response_headers: serde_json::from_str(&response_headers).unwrap_or_default(), + request_size: r.get("request_size")?, + response_size: r.get("response_size")?, + request_truncated: r.get("request_truncated")?, + response_truncated: r.get("response_truncated")?, + elapsed_headers_ms: r.get("elapsed_headers_ms")?, + elapsed_ms: r.get("elapsed_ms")?, error: r.get("error")?, }) } diff --git a/crates-proxy/yaak-proxy-lib/src/settings.rs b/crates-proxy/yaak-proxy-lib/src/settings.rs new file mode 100644 index 000000000..3cd4ffb21 --- /dev/null +++ b/crates-proxy/yaak-proxy-lib/src/settings.rs @@ -0,0 +1,85 @@ +use serde::{Deserialize, Serialize}; +use std::path::Path; +use yaak_proxy::DEFAULT_CAPTURE_BODY_LIMIT; + +/// `settings.json` in the proxy data directory. Missing keys take their defaults. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default, rename_all = "camelCase")] +pub struct ProxySettings { + pub port: u16, + /// Whether the system proxy points at the engine while it runs. + pub system_proxy: bool, + /// Bytes of each request and response body stored per flow. + pub capture_body_limit: usize, +} + +impl Default for ProxySettings { + fn default() -> Self { + Self { port: 9090, system_proxy: false, capture_body_limit: DEFAULT_CAPTURE_BODY_LIMIT } + } +} + +impl ProxySettings { + /// The settings at `path`, or the defaults when there is no file. + pub fn load(path: &Path) -> Result { + match std::fs::read_to_string(path) { + Ok(contents) => serde_json::from_str(&contents) + .map_err(|e| format!("Invalid settings in {}: {e}", path.display())), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()), + Err(e) => Err(format!("Failed to read {}: {e}", path.display())), + } + } + + pub fn save(&self, path: &Path) -> Result<(), String> { + let contents = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?; + std::fs::write(path, contents) + .map_err(|e| format!("Failed to write {}: {e}", path.display())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn defaults_without_a_file() { + let dir = tempfile::tempdir().unwrap(); + let settings = ProxySettings::load(&dir.path().join("settings.json")).unwrap(); + assert_eq!(settings.port, 9090); + assert!(!settings.system_proxy); + assert_eq!(settings.capture_body_limit, 10 * 1024 * 1024); + } + + #[test] + fn defaults_for_missing_keys() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("settings.json"); + std::fs::write(&path, r#"{"port": 8888}"#).unwrap(); + let settings = ProxySettings::load(&path).unwrap(); + assert_eq!(settings, ProxySettings { port: 8888, ..Default::default() }); + } + + #[test] + fn round_trip() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("settings.json"); + let settings = ProxySettings { port: 8080, system_proxy: true, capture_body_limit: 1024 }; + settings.save(&path).unwrap(); + assert_eq!(ProxySettings::load(&path).unwrap(), settings); + + let json: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({"port": 8080, "systemProxy": true, "captureBodyLimit": 1024}) + ); + } + + #[test] + fn invalid_json_is_an_error() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("settings.json"); + std::fs::write(&path, "{").unwrap(); + assert!(ProxySettings::load(&path).is_err()); + } +} diff --git a/crates-proxy/yaak-proxy-lib/tests/flow.rs b/crates-proxy/yaak-proxy-lib/tests/flow.rs new file mode 100644 index 000000000..8afc34824 --- /dev/null +++ b/crates-proxy/yaak-proxy-lib/tests/flow.rs @@ -0,0 +1,168 @@ +use std::convert::Infallible; +use std::io::Write; +use std::net::SocketAddr; +use std::time::{Duration, Instant}; + +use bytes::Bytes; +use flate2::Compression; +use flate2::write::GzEncoder; +use http_body_util::{BodyExt, Full}; +use hyper::body::Incoming; +use hyper::server::conn::http1; +use hyper::service::service_fn; +use hyper::{Request, Response}; +use hyper_util::rt::TokioIo; +use tokio::net::TcpListener; +use yaak_proxy_lib::ProxyCtx; +use yaak_proxy_lib::bodies::{read_body, read_raw_body}; +use yaak_proxy_lib::models::{Flow, FlowBody}; +use yaak_proxy_lib::settings::ProxySettings; +use yaak_rpc::RpcEventEmitter; + +const CAPTURE_LIMIT: usize = 1024; +const BIG_BODY: usize = 4096; + +fn gzip(data: &[u8]) -> Vec { + let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(data).unwrap(); + encoder.finish().unwrap() +} + +async fn route(req: Request) -> Result>, Infallible> { + Ok(match req.uri().path() { + "/echo" => { + let body = req.into_body().collect().await.unwrap().to_bytes(); + Response::builder() + .header("content-type", "text/plain") + .header("content-encoding", "gzip") + .body(Full::from(gzip(&body))) + .unwrap() + } + // Answers at once and drains the upload afterwards, so the response ends first. + "/early" => { + let mut body = req.into_body(); + tokio::spawn(async move { + while let Some(Ok(_)) = body.frame().await { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }); + Response::new(Full::from("ok")) + } + _ => Response::new(Full::from(vec![b'x'; BIG_BODY])), + }) +} + +async fn spawn_upstream() -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + loop { + let (tcp, _) = listener.accept().await.unwrap(); + tokio::spawn(http1::Builder::new().serve_connection(TokioIo::new(tcp), service_fn(route))); + } + }); + addr +} + +/// Waits for the event loop to write the flow for `url`. +async fn wait_for_flow(ctx: &ProxyCtx, url: &str) -> Flow { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let flows = ctx.db.with_conn(|db| db.find_all::().unwrap()); + if let Some(flow) = flows.into_iter().find(|f| f.url == url) { + return flow; + } + assert!(Instant::now() < deadline, "no flow for {url}"); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn flow_is_stored_with_its_bodies() { + let dir = tempfile::tempdir().unwrap(); + let settings = + ProxySettings { port: 0, capture_body_limit: CAPTURE_LIMIT, ..Default::default() }; + settings.save(&dir.path().join("settings.json")).unwrap(); + + let (events, _event_rx) = RpcEventEmitter::new(); + let ctx = ProxyCtx::new(dir.path(), events).unwrap(); + ctx.start().unwrap(); + let port = ctx.port().unwrap(); + assert!(dir.path().join("ca/ca.pem").exists()); + assert!(dir.path().join("blobs.sqlite").exists()); + + let upstream = spawn_upstream().await; + let client = reqwest::Client::builder() + .proxy(reqwest::Proxy::all(format!("http://127.0.0.1:{port}")).unwrap()) + .build() + .unwrap(); + + let echo_url = format!("http://{upstream}/echo"); + let resp = client.post(&echo_url).body("hello proxy").send().await.unwrap(); + assert_eq!(resp.status(), 200); + resp.bytes().await.unwrap(); + + let flow = wait_for_flow(&ctx, &echo_url).await; + let compressed = gzip(b"hello proxy"); + assert!(flow.id.starts_with("fl_"), "{}", flow.id); + assert_eq!(flow.method, "POST"); + assert_eq!(flow.status, Some(200)); + assert_eq!(flow.client_http_version, "HTTP/1.1"); + assert_eq!(flow.upstream_http_version.as_deref(), Some("HTTP/1.1")); + assert_eq!(flow.remote_addr, Some(upstream.to_string())); + assert_eq!(flow.content_type.as_deref(), Some("text/plain")); + assert_eq!(flow.request_size, 11); + assert_eq!(flow.response_size, compressed.len() as i64); + assert!(!flow.request_truncated && !flow.response_truncated); + assert!(flow.elapsed_headers_ms.is_some() && flow.elapsed_ms >= flow.elapsed_headers_ms); + assert!(flow.request_headers.iter().any(|h| h.name == "content-length")); + assert_eq!(flow.error, None); + + assert_eq!(read_raw_body(&ctx.blobs, &flow.id, FlowBody::Request).unwrap(), b"hello proxy"); + assert_eq!(read_raw_body(&ctx.blobs, &flow.id, FlowBody::Response).unwrap(), compressed); + assert_eq!(read_body(&ctx.blobs, &flow, FlowBody::Response).await.unwrap(), b"hello proxy"); + + let big_url = format!("http://{upstream}/big"); + let resp = client.get(&big_url).send().await.unwrap(); + assert_eq!(resp.bytes().await.unwrap().len(), BIG_BODY); + + let flow = wait_for_flow(&ctx, &big_url).await; + assert_eq!(flow.request_size, 0); + assert_eq!(flow.response_size, BIG_BODY as i64); + assert!(flow.response_truncated); + let stored = read_raw_body(&ctx.blobs, &flow.id, FlowBody::Response).unwrap(); + assert_eq!(stored, vec![b'x'; CAPTURE_LIMIT]); + assert!(!ctx.blobs.connect().body_exists(&Flow::body_id(&flow.id, FlowBody::Request)).unwrap()); + + ctx.stop().unwrap(); + assert_eq!(ctx.port(), None); +} + +#[tokio::test(flavor = "multi_thread")] +async fn early_response_keeps_request_totals() { + let dir = tempfile::tempdir().unwrap(); + let settings = + ProxySettings { port: 0, capture_body_limit: CAPTURE_LIMIT, ..Default::default() }; + settings.save(&dir.path().join("settings.json")).unwrap(); + + let (events, _event_rx) = RpcEventEmitter::new(); + let ctx = ProxyCtx::new(dir.path(), events).unwrap(); + ctx.start().unwrap(); + let port = ctx.port().unwrap(); + + let upstream = spawn_upstream().await; + let client = reqwest::Client::builder() + .proxy(reqwest::Proxy::all(format!("http://127.0.0.1:{port}")).unwrap()) + .build() + .unwrap(); + + let url = format!("http://{upstream}/early"); + let _ = client.post(&url).body(vec![b'u'; 2 * 1024 * 1024]).send().await; + + let flow = wait_for_flow(&ctx, &url).await; + assert_eq!(flow.status, Some(200)); + let stored = read_raw_body(&ctx.blobs, &flow.id, FlowBody::Request).unwrap().len() as i64; + assert!(stored > 0); + assert!(flow.request_size >= stored, "{} < {stored}", flow.request_size); + assert_eq!(flow.request_truncated, flow.request_size > stored); +} diff --git a/crates-tauri/yaak-app-proxy/src/lib.rs b/crates-tauri/yaak-app-proxy/src/lib.rs index d5d3b4acb..f893ae553 100644 --- a/crates-tauri/yaak-app-proxy/src/lib.rs +++ b/crates-tauri/yaak-app-proxy/src/lib.rs @@ -67,7 +67,7 @@ pub fn run() { std::fs::create_dir_all(&data_dir).expect("failed to create app data dir"); let (emitter, event_rx) = RpcEventEmitter::new(); - app.manage(ProxyCtx::new(&data_dir.join("proxy.db"), emitter)); + app.manage(ProxyCtx::new(&data_dir, emitter)?); app.manage(yaak_proxy_lib::build_router()); // Drain RPC events and forward as Tauri events diff --git a/crates/yaak-common/Cargo.toml b/crates/yaak-common/Cargo.toml index 42bdf949e..168b9f2d3 100644 --- a/crates/yaak-common/Cargo.toml +++ b/crates/yaak-common/Cargo.toml @@ -5,5 +5,6 @@ edition = "2024" publish = false [dependencies] +dirs = "7" serde_json = { workspace = true } tokio = { workspace = true, features = ["process"] } diff --git a/crates/yaak-common/src/data_dir.rs b/crates/yaak-common/src/data_dir.rs new file mode 100644 index 000000000..b77521e87 --- /dev/null +++ b/crates/yaak-common/src/data_dir.rs @@ -0,0 +1,46 @@ +use std::path::PathBuf; + +/// The data directory of the app with identifier `app_id`, as Tauri resolves it. Under WSL, +/// the Windows app's directory when it exists. +pub fn resolve_data_dir(app_id: &str) -> PathBuf { + if let Some(dir) = wsl_data_dir(app_id) { + return dir; + } + dirs::data_dir().expect("Could not determine data directory").join(app_id) +} + +/// Detect WSL and resolve the Windows AppData\Roaming path for the app's data directory. +fn wsl_data_dir(app_id: &str) -> Option { + if !cfg!(target_os = "linux") { + return None; + } + + let proc_version = std::fs::read_to_string("/proc/version").ok()?; + let is_wsl = proc_version.to_lowercase().contains("microsoft"); + if !is_wsl { + return None; + } + + // We're in WSL, so try to resolve the app's data directory in Windows + + // Get the Windows %APPDATA% path via cmd.exe + let appdata_output = + std::process::Command::new("cmd.exe").args(["/C", "echo", "%APPDATA%"]).output().ok()?; + + let win_path = String::from_utf8(appdata_output.stdout).ok()?.trim().to_string(); + if win_path.is_empty() || win_path == "%APPDATA%" { + return None; + } + + // Convert Windows path to WSL path using wslpath (handles custom mount points) + let wslpath_output = std::process::Command::new("wslpath").arg(&win_path).output().ok()?; + + let wsl_appdata = String::from_utf8(wslpath_output.stdout).ok()?.trim().to_string(); + if wsl_appdata.is_empty() { + return None; + } + + let wsl_path = PathBuf::from(wsl_appdata).join(app_id); + + if wsl_path.exists() { Some(wsl_path) } else { None } +} diff --git a/crates/yaak-common/src/lib.rs b/crates/yaak-common/src/lib.rs index 16a3846f6..f755fd4d0 100644 --- a/crates/yaak-common/src/lib.rs +++ b/crates/yaak-common/src/lib.rs @@ -1,3 +1,4 @@ pub mod command; +pub mod data_dir; pub mod platform; pub mod serde; diff --git a/crates/yaak-models/src/blob_manager.rs b/crates/yaak-models/src/blob_manager.rs index 41ba8bc5d..4f79df7b1 100644 --- a/crates/yaak-models/src/blob_manager.rs +++ b/crates/yaak-models/src/blob_manager.rs @@ -179,6 +179,58 @@ impl<'a> BlobWriter<'a> { } } +/// Size of each stored chunk. Bodies arrive in small reads (8 to 16 KiB is typical), so they +/// are grouped into chunks of this size to avoid thousands of tiny inserts for large bodies. +pub const BODY_CHUNK_SIZE: usize = 1024 * 1024; + +/// Writes one body as it streams, grouping its pieces into [`BODY_CHUNK_SIZE`] chunks. Each +/// full chunk is inserted as soon as it fills; [`BodyChunkWriter::finish`] inserts the rest. +pub struct BodyChunkWriter { + blob_manager: BlobManager, + body_id: String, + buf: Vec, + chunk_index: i32, + total_bytes: usize, +} + +impl BodyChunkWriter { + pub fn new(blob_manager: BlobManager, body_id: impl Into) -> Self { + Self { + blob_manager, + body_id: body_id.into(), + buf: Vec::new(), + chunk_index: 0, + total_bytes: 0, + } + } + + pub fn write(&mut self, data: &[u8]) -> Result<()> { + self.total_bytes += data.len(); + self.buf.extend_from_slice(data); + while self.buf.len() >= BODY_CHUNK_SIZE { + let data = self.buf.drain(..BODY_CHUNK_SIZE).collect(); + self.insert(data)?; + } + Ok(()) + } + + /// Inserts the partial last chunk and returns the bytes written in all. + pub fn finish(mut self) -> Result { + if !self.buf.is_empty() { + let data = std::mem::take(&mut self.buf); + self.insert(data)?; + } + Ok(self.total_bytes) + } + + fn insert(&mut self, data: Vec) -> Result<()> { + let chunk = BodyChunk::new(&self.body_id, self.chunk_index, data); + self.blob_manager.with_tx(|b| b.insert_chunk(&chunk))?; + self.chunk_index += 1; + Ok(()) + } +} + /// Run migrations for the blob database. pub fn migrate_blob_db(pool: &SqlitePool) -> Result<()> { info!("Running blob database migrations"); @@ -397,6 +449,34 @@ mod tests { assert_eq!(chunks2[0].data, b"data2"); } + #[test] + fn test_chunk_writer_groups_pieces() { + let manager = create_test_manager(); + let body_id = "fl_test.response"; + let piece = vec![7u8; 16 * 1024]; + let pieces = BODY_CHUNK_SIZE * 2 / piece.len() + 3; + + let mut writer = BodyChunkWriter::new(manager.clone(), body_id); + for _ in 0..pieces { + writer.write(&piece).unwrap(); + } + assert_eq!(manager.connect().get_chunks(body_id).unwrap().len(), 2); + assert_eq!(writer.finish().unwrap(), pieces * piece.len()); + + let chunks = manager.connect().get_chunks(body_id).unwrap(); + let sizes: Vec<_> = chunks.iter().map(|c| c.data.len()).collect(); + assert_eq!(sizes, vec![BODY_CHUNK_SIZE, BODY_CHUNK_SIZE, 3 * piece.len()]); + assert_eq!(chunks.iter().map(|c| c.chunk_index).collect::>(), vec![0, 1, 2]); + } + + #[test] + fn test_chunk_writer_empty_body() { + let manager = create_test_manager(); + let writer = BodyChunkWriter::new(manager.clone(), "fl_empty.request"); + assert_eq!(writer.finish().unwrap(), 0); + assert!(!manager.connect().body_exists("fl_empty.request").unwrap()); + } + #[test] fn test_large_chunk() { let manager = create_test_manager(); diff --git a/crates/yaak-models/src/lib.rs b/crates/yaak-models/src/lib.rs index 60eab2e3c..fdcc8f1ee 100644 --- a/crates/yaak-models/src/lib.rs +++ b/crates/yaak-models/src/lib.rs @@ -141,18 +141,25 @@ pub fn init_standalone( let (readers, writer) = open::file_pools(db_path, 20, 2)?; migrate_db(&writer)?; - info!("Initializing blobs database {blob_path:?}"); - let (blob_readers, blob_writer) = open::file_pools(blob_path, 10, 1)?; - migrate_blob_db(&blob_writer)?; + let blob_manager = init_blob_store(blob_path)?; let (tx, rx) = mpsc::channel(); let query_manager = QueryManager::new(readers, writer, tx); - let blob_manager = BlobManager::new(blob_readers, blob_writer); bootstrap(&query_manager)?; Ok((query_manager, blob_manager, rx)) } +/// Open (creating and migrating as needed) a blobs database on its own, for a process that +/// stores bodies without the app database. +pub fn init_blob_store(blob_path: impl AsRef) -> Result { + let blob_path = blob_path.as_ref(); + info!("Initializing blobs database {blob_path:?}"); + let (blob_readers, blob_writer) = open::file_pools(blob_path, 10, 1)?; + migrate_blob_db(&blob_writer)?; + Ok(BlobManager::new(blob_readers, blob_writer)) +} + /// Initialize the database managers with in-memory SQLite databases. /// Useful for testing and CI environments. /// diff --git a/crates/yaak-proxy/src/lib.rs b/crates/yaak-proxy/src/lib.rs index 43a8a5733..5a1d35678 100644 --- a/crates/yaak-proxy/src/lib.rs +++ b/crates/yaak-proxy/src/lib.rs @@ -128,33 +128,6 @@ pub enum ProxyEvent { }, } -/// Accumulated view of a proxied request, built from `ProxyEvent`s. -#[derive(Debug, Clone)] -pub struct CapturedRequest { - pub id: u64, - pub method: String, - pub url: String, - pub status: Option, - pub elapsed_ms: Option, - pub http_version: String, - pub remote_http_version: Option, - pub request_headers: Vec<(String, String)>, - pub request_body: Option>, - pub response_headers: Vec<(String, String)>, - pub response_body: Option>, - pub response_body_size: u64, - pub state: RequestState, - pub error: Option, -} - -#[derive(Debug, Clone, PartialEq)] -pub enum RequestState { - Sending, - Receiving, - Complete, - Error, -} - /// A running engine. Dropping it stops the listener and joins its thread. pub struct ProxyHandle { shutdown_tx: Option>, diff --git a/crates/yaak/src/send.rs b/crates/yaak/src/send.rs index 026c60d4e..4ac8d31ec 100644 --- a/crates/yaak/src/send.rs +++ b/crates/yaak/src/send.rs @@ -21,7 +21,7 @@ use yaak_http::transaction::HttpTransaction; use yaak_http::types::{ SendableBody, SendableHttpRequest, SendableHttpRequestOptions, append_query_params, }; -use yaak_models::blob_manager::{BlobManager, BodyChunk}; +use yaak_models::blob_manager::{BODY_CHUNK_SIZE, BlobManager, BodyChunk, BodyChunkWriter}; use yaak_models::models::{ ClientCertificate, Cookie, CookieJar, DnsOverride, Environment, HttpRequest, HttpResponse, HttpResponseEvent, HttpResponseEventData, HttpResponseHeader, HttpResponseState, ProxySetting, @@ -39,7 +39,6 @@ use yaak_templates::{RenderOptions, TemplateCallback}; use yaak_tls::find_client_certificate; const HTTP_EVENT_CHANNEL_CAPACITY: usize = 100; -const REQUEST_BODY_CHUNK_SIZE: usize = 1024 * 1024; const RESPONSE_PROGRESS_UPDATE_INTERVAL_MS: u128 = 100; const MAX_AUTH_BODY_BYTES: usize = 10 * 1024 * 1024; @@ -1038,7 +1037,7 @@ fn persist_request_body_bytes( blob_manager .with_tx(|b| { - for (chunk_index, data) in bytes.chunks(REQUEST_BODY_CHUNK_SIZE).enumerate() { + for (chunk_index, data) in bytes.chunks(BODY_CHUNK_SIZE).enumerate() { b.insert_chunk(&BodyChunk::new(body_id, chunk_index as i32, data.to_vec()))?; } Ok::<_, yaak_models::error::Error>(()) @@ -1051,29 +1050,11 @@ async fn persist_request_body_stream( body_id: String, mut rx: tokio::sync::mpsc::UnboundedReceiver>, ) -> std::result::Result { - let mut chunk_index: i32 = 0; - let mut total_bytes = 0usize; - - // Stream reads arrive in small (eg. 8-16 KiB) pieces, so accumulate them into - // full-size chunks to avoid thousands of tiny inserts for large bodies - let mut buf: Vec = Vec::with_capacity(REQUEST_BODY_CHUNK_SIZE); + let mut writer = BodyChunkWriter::new(blob_manager, body_id); while let Some(data) = rx.recv().await { - total_bytes += data.len(); - buf.extend_from_slice(&data); - while buf.len() >= REQUEST_BODY_CHUNK_SIZE { - let data = buf.drain(..REQUEST_BODY_CHUNK_SIZE).collect(); - let chunk = BodyChunk::new(&body_id, chunk_index, data); - blob_manager.with_tx(|b| b.insert_chunk(&chunk)).map_err(|e| e.to_string())?; - chunk_index += 1; - } + writer.write(&data).map_err(|e| e.to_string())?; } - - if !buf.is_empty() { - let chunk = BodyChunk::new(&body_id, chunk_index, buf); - blob_manager.with_tx(|b| b.insert_chunk(&chunk)).map_err(|e| e.to_string())?; - } - - Ok(total_bytes) + writer.finish().map_err(|e| e.to_string()) } fn append_error_message(existing_error: Option, message: String) -> String {