mirror of
https://github.com/mountain-loop/yaak.git
synced 2026-11-13 12:13:42 +01:00
Proxy storage: flows table, bodies in blobs.sqlite, settings (#751)
This commit is contained in:
Gregory Schier authored and 
GitHub committed
1 parent
694c0272cd
commit
b5508e6bb1
25 files changed
+942
-289
No files matched your search
Generated
+12
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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}
|
||||
</TruncatedWideTableCell>
|
||||
<TableCell>
|
||||
<StatusBadge status={ex.resStatus} error={ex.error} />
|
||||
<StatusBadge status={ex.status} error={ex.error} />
|
||||
</TableCell>
|
||||
<TableCell className="text-text-subtle text-xs">
|
||||
{getContentType(ex.resHeaders)}
|
||||
{getContentType(ex.contentType)}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
))}
|
||||
@@ -69,8 +69,7 @@ function StatusBadge({ status, error }: { status: number | null; error: string |
|
||||
return <span className={classNames("text-xs font-mono", color)}>{status}</span>;
|
||||
}
|
||||
|
||||
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;
|
||||
|
||||
@@ -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<Record<string, boolea
|
||||
export const SIDEBAR_TREE_ID = "proxy-sidebar";
|
||||
|
||||
const sidebarTreeAtom = atom<TreeNode<SidebarItem>>((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<SidebarItem>, map: Map<string, SidebarItem>
|
||||
* /123
|
||||
* /orders
|
||||
*/
|
||||
function buildTree(exchanges: HttpExchange[]): TreeNode<SidebarItem> {
|
||||
function buildTree(exchanges: Flow[]): TreeNode<SidebarItem> {
|
||||
const root: SidebarItem = { id: "root", label: "All Traffic", exchangeIds: [] };
|
||||
const rootNode: TreeNode<SidebarItem> = {
|
||||
item: root,
|
||||
|
||||
@@ -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<ProxyModels>(["http_exchange"]);
|
||||
createModelStore<ProxyModels>(["flow"]);
|
||||
|
||||
export const httpExchangesAtom = orderedListAtom("http_exchange", "createdAt", "desc");
|
||||
export const flowsAtom = orderedListAtom("flow", "startedAt", "desc");
|
||||
@@ -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);
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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<PathBuf> {
|
||||
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 }
|
||||
}
|
||||
@@ -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"] }
|
||||
+41
-2
@@ -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<ProxyHeader>, reqBody: Array<number> | null, resStatus: number | null, resHeaders: Array<ProxyHeader>, resBody: Array<number> | 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<ProxyHeader>, responseHeaders: Array<ProxyHeader>,
|
||||
/**
|
||||
* 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, };
|
||||
+2
-2
@@ -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<string, never>;
|
||||
|
||||
export type ListModelsResponse = { httpExchanges: Array<HttpExchange>, };
|
||||
export type ListModelsResponse = { flows: Array<Flow>, };
|
||||
|
||||
export type ProxyState = "running" | "stopped";
|
||||
|
||||
|
||||
@@ -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);
|
||||
@@ -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<Vec<u8>, 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<Vec<u8>, 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}")),
|
||||
}
|
||||
}
|
||||
@@ -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<SqliteConnectionManager>,
|
||||
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<Self> {
|
||||
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<F, T>(&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<T, E>(
|
||||
&self,
|
||||
func: impl FnOnce(&DbContext) -> std::result::Result<T, E>,
|
||||
) -> std::result::Result<T, E>
|
||||
where
|
||||
E: From<Error>,
|
||||
{
|
||||
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<SqlitePool> {
|
||||
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<String> {
|
||||
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<String> = 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");
|
||||
}
|
||||
}
|
||||
@@ -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<Mutex<Option<ProxyHandle>>>,
|
||||
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<Self, String> {
|
||||
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<u16> {
|
||||
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<HttpExchange>,
|
||||
pub flows: Vec<Flow>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, TS)]
|
||||
@@ -92,41 +145,11 @@ pub struct GetProxyStateResponse {
|
||||
fn execute_action(ctx: &ProxyCtx, invocation: ActionInvocation) -> Result<bool, RpcError> {
|
||||
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<ListModelsResponse, RpcError> {
|
||||
ctx.db.with_conn(|db| {
|
||||
Ok(ListModelsResponse {
|
||||
http_exchanges: db
|
||||
.find_all::<HttpExchange>()
|
||||
.map_err(|e| RpcError { message: e.to_string() })?,
|
||||
flows: db.find_all::<Flow>().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<BodyChunkWriter>,
|
||||
response_body: Option<BodyChunkWriter>,
|
||||
}
|
||||
|
||||
fn run_event_loop(
|
||||
rx: std::sync::mpsc::Receiver<ProxyEvent>,
|
||||
db: ProxyQueryManager,
|
||||
blobs: BlobManager,
|
||||
events: RpcEventEmitter,
|
||||
) {
|
||||
let mut in_flight: HashMap<u64, CapturedRequest> = HashMap::new();
|
||||
let mut in_flight: HashMap<u64, InFlight> = 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::<Utc>::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<BodyChunkWriter>, 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<BodyChunkWriter>) {
|
||||
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 --
|
||||
|
||||
@@ -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<ProxyHeader>,
|
||||
pub req_body: Option<Vec<u8>>,
|
||||
pub res_status: Option<i32>,
|
||||
pub res_headers: Vec<ProxyHeader>,
|
||||
pub res_body: Option<Vec<u8>>,
|
||||
pub url: String,
|
||||
/// Version spoken by the client.
|
||||
pub client_http_version: String,
|
||||
/// Version spoken with the upstream server.
|
||||
pub upstream_http_version: Option<String>,
|
||||
pub remote_addr: Option<String>,
|
||||
pub status: Option<i32>,
|
||||
/// The response's Content-Type.
|
||||
pub content_type: Option<String>,
|
||||
pub request_headers: Vec<ProxyHeader>,
|
||||
pub response_headers: Vec<ProxyHeader>,
|
||||
/// 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<i32>,
|
||||
/// Milliseconds from `started_at` to the end of the response.
|
||||
pub elapsed_ms: Option<i32>,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
#[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<Vec<(impl IntoIden + Eq, impl Into<SimpleExpr>)>> {
|
||||
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<impl IntoIden> {
|
||||
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")?,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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<Self, String> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -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<u8> {
|
||||
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
|
||||
encoder.write_all(data).unwrap();
|
||||
encoder.finish().unwrap()
|
||||
}
|
||||
|
||||
async fn route(req: Request<Incoming>) -> Result<Response<Full<Bytes>>, 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::<Flow>().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);
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -5,5 +5,6 @@ edition = "2024"
|
||||
publish = false
|
||||
|
||||
[dependencies]
|
||||
dirs = "7"
|
||||
serde_json = { workspace = true }
|
||||
tokio = { workspace = true, features = ["process"] }
|
||||
@@ -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<PathBuf> {
|
||||
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 }
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod command;
|
||||
pub mod data_dir;
|
||||
pub mod platform;
|
||||
pub mod serde;
|
||||
@@ -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<u8>,
|
||||
chunk_index: i32,
|
||||
total_bytes: usize,
|
||||
}
|
||||
|
||||
impl BodyChunkWriter {
|
||||
pub fn new(blob_manager: BlobManager, body_id: impl Into<String>) -> 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<usize> {
|
||||
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<u8>) -> 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<_>>(), 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();
|
||||
|
||||
@@ -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<Path>) -> Result<BlobManager> {
|
||||
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.
|
||||
///
|
||||
|
||||
@@ -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<u16>,
|
||||
pub elapsed_ms: Option<u64>,
|
||||
pub http_version: String,
|
||||
pub remote_http_version: Option<String>,
|
||||
pub request_headers: Vec<(String, String)>,
|
||||
pub request_body: Option<Vec<u8>>,
|
||||
pub response_headers: Vec<(String, String)>,
|
||||
pub response_body: Option<Vec<u8>>,
|
||||
pub response_body_size: u64,
|
||||
pub state: RequestState,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
#[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<tokio::sync::oneshot::Sender<()>>,
|
||||
|
||||
+5
-24
@@ -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<Vec<u8>>,
|
||||
) -> std::result::Result<usize, String> {
|
||||
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<u8> = 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<String>, message: String) -> String {
|
||||
|
||||
Reference in new issue
Block a user