PR 2: proxy services (#787)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Gregory SchierandClaude Opus 5.5 authored and GitHub committed 2026-10-27 15:19:50 -07:00
1 parent 686294939b
commit 9bbd28ce4a
26 files changed
+3802 -50

No files matched your search

+1
View File
@@ -5,3 +5,4 @@ crates/yaak-templates/pkg/* linguist-generated=true
# Ensure consistent line endings for test files that check exact content
crates/yaak-http/tests/test.txt text eol=lf
crates-proxy/yaak-proxy-lib/tests/fixtures/export.har text eol=lf
Generated
+4
View File
@@ -11763,6 +11763,7 @@ dependencies = [
"hyper",
"hyper-util",
"include_dir",
"libc",
"log 0.4.34",
"r2d2",
"r2d2_sqlite",
@@ -11774,9 +11775,12 @@ dependencies = [
"tempfile",
"tokio",
"ts-rs",
"url",
"yaak-common",
"yaak-database",
"yaak-filter",
"yaak-http",
"yaak-lifecycle",
"yaak-models",
"yaak-proxy",
"yaak-rpc",
+5 -4
View File
@@ -73,7 +73,8 @@ export function FlowDetail({ style, sideBySide }: { style?: CSSProperties; sideB
function FlowPane({ flow }: { flow: Flow }) {
const [pretty, setPretty] = useState(true);
const responseBody = useFlowBody(flow, "response", flow.responseTruncated);
const failed = flow.error != null && flow.status == null;
// Failed before any response; an error after the headers is a body cut short mid-way.
const noResponse = flow.error != null && flow.status == null;
const mime = mimeEssence(findHeader(flow.responseHeaders, "content-type"));
const queryCount = useMemo(() => queryParams(flow.url).length, [flow.url]);
@@ -153,7 +154,7 @@ function FlowPane({ flow }: { flow: Flow }) {
</span>
</>
)}
{!failed && (
{!noResponse && (
<>
<span>&bull;</span>
<SizeTag contentLength={flow.responseSize} />
@@ -178,7 +179,7 @@ function FlowPane({ flow }: { flow: Flow }) {
</div>
<div className="overflow-hidden flex flex-col min-h-0">
{failed && (
{flow.error != null && (
<Banner color="danger" className="mx-3 mt-1 shrink-0 font-mono text-sm">
{flow.error}
</Banner>
@@ -190,7 +191,7 @@ function FlowPane({ flow }: { flow: Flow }) {
tabListClassName="mt-0.5 -mb-1.5"
>
<TabContent value={TAB_RESPONSE}>
{failed ? (
{noResponse ? (
<EmptyState>
<div className="text-text-subtle mb-1">No response</div>
<div className="text-xs">
+8 -2
View File
@@ -212,7 +212,9 @@ const FlowRow = memo(function FlowRow({
onSelect,
}: RowProps) {
const { host, path, query } = splitUrl(flow.url);
const failed = flow.error != null && flow.status == null;
const failed = flow.error != null;
// Failed before any response; an error after the headers is a body cut short mid-way.
const noResponse = failed && flow.status == null;
const type = shortContentType(flow.contentType);
return (
@@ -261,7 +263,11 @@ const FlowRow = memo(function FlowRow({
flow.responseTruncated ? "text-notice" : "text-text-subtle",
)}
>
{failed ? <span className="text-text-subtlest">--</span> : formatSize(flow.responseSize)}
{noResponse ? (
<span className="text-text-subtlest">--</span>
) : (
formatSize(flow.responseSize)
)}
</span>
)}
<span className="font-mono text-right text-text-subtle">
+12 -1
View File
@@ -7,25 +7,36 @@ publish = false
[dependencies]
base64 = "0.23.1"
bytes = "1"
chrono = { workspace = true, features = ["serde"] }
log = { workspace = true }
http-body-util = "0.1"
hyper = { version = "1", features = ["http1", "server"] }
hyper-util = { version = "0.1", features = ["server-graceful", "tokio"] }
include_dir = "0.7"
r2d2 = "0.8.10"
r2d2_sqlite = "0.32"
reqwest = { workspace = true, features = ["json"] }
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", "process"] }
tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt-multi-thread"] }
ts-rs = { workspace = true, features = ["chrono-impl"] }
url = "2"
yaak-common = { workspace = true }
yaak-database = { workspace = true }
yaak-filter = { workspace = true }
yaak-http = { workspace = true }
yaak-lifecycle = { workspace = true }
yaak-models = { workspace = true }
yaak-proxy = { workspace = true }
yaak-rpc = { workspace = true }
yaak-templates = { workspace = true }
[target.'cfg(unix)'.dependencies]
libc = "0.2"
[dev-dependencies]
base64 = "0.23.1"
bytes = "1"
+5 -1
View File
@@ -40,7 +40,11 @@ elapsedHeadersMs: number | null,
/**
* Milliseconds from `started_at` to the end of the response.
*/
elapsedMs: number | null, error: string | null, };
elapsedMs: number | null, error: string | null,
/**
* The id of the Yaak request this flow was last saved as.
*/
savedRequestId: string | null, };
export type FlowBody = "request" | "response";
+37 -1
View File
@@ -69,10 +69,46 @@ export type ListModelsRequest = Record<string, never>;
export type ListModelsResponse = { flows: Array<Flow>, };
/**
* The process that runs the engine.
*/
export type Owner = "app" | "cli";
export type ProxySetting = { enabled: boolean, server: string, port: number, };
export type ProxyState = "running" | "stopped";
export type ProxyStatePayload = { state: ProxyState, };
export type RpcEventSchema = { model_write: ModelPayload, proxy_state_changed: ProxyStatePayload, };
export type RpcSchema = { execute_action: [ActionInvocation, boolean], format_body: [FormatBodyRequest, string], get_proxy_state: [GetProxyStateRequest, GetProxyStateResponse], list_actions: [ListActionsRequest, ListActionsResponse], list_models: [ListModelsRequest, ListModelsResponse], ca_status: [CaStatusRequest, CaStatus], ca_trust: [CaTrustRequest, CaStatus], filter_body: [FilterBodyRequest, FilterBodyResponse], get_flow_body: [GetFlowBodyRequest, GetFlowBodyResponse], };
export type RpcSchema = { execute_action: [ActionInvocation, boolean], format_body: [FormatBodyRequest, string], get_proxy_state: [GetProxyStateRequest, GetProxyStateResponse], list_actions: [ListActionsRequest, ListActionsResponse], list_models: [ListModelsRequest, ListModelsResponse], ca_status: [CaStatusRequest, CaStatus], ca_trust: [CaTrustRequest, CaStatus], filter_body: [FilterBodyRequest, FilterBodyResponse], get_flow_body: [GetFlowBodyRequest, GetFlowBodyResponse], system_proxy_set: [SystemProxySetRequest, SystemProxyStatus], system_proxy_status: [SystemProxyStatusRequest, SystemProxyStatus], };
export type ServiceProxy = {
/**
* The network service, as System Settings names it: "Wi-Fi", "USB 10/100/1000 LAN".
*/
name: string, http: ProxySetting, https: ProxySetting, };
/**
* How a change gets administrator rights.
*/
export type SystemProxyBackend = "helper" | "admin_prompt" | "sudo";
export type SystemProxySetRequest = { enabled: boolean, };
export type SystemProxyStatus = {
/**
* Every enabled network service sends HTTP and HTTPS to `127.0.0.1:<port>`.
*/
enabled: boolean,
/**
* The port checked against.
*/
port: number, services: Array<ServiceProxy>, backend: SystemProxyBackend,
/**
* This platform has a system proxy the app can set.
*/
supported: boolean, };
export type SystemProxyStatusRequest = Record<string, never>;
@@ -0,0 +1,2 @@
-- The Yaak request a flow was last saved as.
ALTER TABLE flows ADD COLUMN saved_request_id TEXT;
+48 -12
View File
@@ -20,12 +20,15 @@ pub struct DecodedBody {
pub data: Vec<u8>,
/// Bytes stored for the body, as they were on the wire.
pub stored_size: usize,
/// The body is incomplete: cut on the wire at the capture limit, or decoded past `limit`.
/// The body is incomplete: cut on the wire at the capture limit, decoded past `limit`, or
/// decoded only as far as `decode_error`.
pub truncated: bool,
/// Why decoding stopped before the end of the stored bytes.
pub decode_error: Option<String>,
}
/// A flow's body decoded with its Content-Encoding, at most `limit` bytes of it. A body cut
/// at the capture limit decodes as far as its stored bytes go.
/// A flow's body decoded with its Content-Encoding, at most `limit` bytes of it. A body that
/// fails to decode, as one cut short does, keeps what decoded before the failure.
pub async fn read_body(
blobs: &BlobManager,
flow: &Flow,
@@ -34,7 +37,9 @@ pub async fn read_body(
) -> Result<DecodedBody, String> {
let raw = read_raw_body(blobs, &flow.id, body)?;
let stored_size = raw.len();
let wire_truncated = flow.truncated(body);
// A response that failed after its headers ended short on the wire.
let cut = body == FlowBody::Response && flow.status.is_some() && flow.error.is_some();
let wire_truncated = flow.truncated(body) || cut;
let encoding =
ContentEncoding::from_header(find_header(flow.headers(body), "content-encoding"));
// A HEAD or 304 response names its encoding but has no body to decode.
@@ -42,19 +47,50 @@ pub async fn read_body(
let truncated = wire_truncated || raw.len() > limit;
let mut data = raw;
data.truncate(limit);
return Ok(DecodedBody { data, stored_size, truncated });
return Ok(DecodedBody { data, stored_size, truncated, decode_error: None });
}
// One byte past the limit tells a body that fills it from one that overflows it.
let take = u64::try_from(limit).unwrap_or(u64::MAX).saturating_add(1);
let mut data = Vec::new();
let mut decoder = streaming_decoder(Cursor::new(raw), encoding).take(take);
match decoder.read_to_end(&mut data).await {
Ok(_) => {}
Err(_) if wire_truncated => {}
Err(e) => return Err(format!("Failed to decode {encoding:?} body: {e}")),
}
let truncated = wire_truncated || data.len() > limit;
let decode_error = match decoder.read_to_end(&mut data).await {
Ok(_) => None,
// Expected of a body cut on the wire; anything else is worth saying.
Err(_) if wire_truncated => None,
Err(e) => Some(format!("Failed to decode {encoding:?} body: {e}")),
};
let truncated = wire_truncated || decode_error.is_some() || data.len() > limit;
data.truncate(limit);
Ok(DecodedBody { data, stored_size, truncated })
Ok(DecodedBody { data, stored_size, truncated, decode_error })
}
#[cfg(test)]
mod tests {
use super::*;
use yaak_models::blob_manager::BodyChunkWriter;
#[tokio::test]
async fn response_failed_after_its_headers_is_truncated() {
let (_db, blobs, _rx) = yaak_models::init_in_memory().unwrap();
let mut writer =
BodyChunkWriter::new(blobs.clone(), Flow::body_id("fl_cut", FlowBody::Response));
writer.write(b"{\"rows\": [").unwrap();
writer.finish().unwrap();
let cut = Flow {
id: "fl_cut".into(),
status: Some(200),
error: Some("error reading a body from connection".into()),
..Default::default()
};
let body = read_body(&blobs, &cut, FlowBody::Response, 1024).await.unwrap();
assert_eq!((body.data.as_slice(), body.truncated), (&b"{\"rows\": ["[..], true));
// The request went out whole; the error is the response's.
let request = read_body(&blobs, &cut, FlowBody::Request, 1024).await.unwrap();
assert!(!request.truncated);
// With no headers there was never a response body to cut.
let refused = Flow { status: None, ..cut };
assert!(!read_body(&blobs, &refused, FlowBody::Response, 1024).await.unwrap().truncated);
}
}
+2 -2
View File
@@ -1,8 +1,8 @@
use serde::Serialize;
use serde::{Deserialize, Serialize};
use std::path::Path;
use ts_rs::TS;
#[derive(Debug, Clone, Serialize, TS)]
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export, export_to = "gen_rpc.ts")]
pub struct CaStatus {
+229
View File
@@ -0,0 +1,229 @@
//! The other side of the control API, for the CLI and the MCP tools: find the owner through
//! `control.json` and call it with its token, or, when nothing owns the engine, read `proxy.db`
//! and `blobs.sqlite` read-only so captured traffic stays reachable.
use crate::bodies::{self, DecodedBody};
use crate::ca::CaStatus;
use crate::control::{
self, ControlInfo, ControlStatus, ErrorBody, StartRequest, SystemProxyRequest, TRUNCATED_HEADER,
};
use crate::db::ProxyQueryManager;
use crate::flows::{self, FlowDetail, FlowFilter};
use crate::har::{self, Har};
use crate::macos::SystemProxyStatus;
use crate::models::{Flow, FlowBody};
use crate::save::{SaveFlowRequest, SaveFlowResponse, WorkspaceSummary};
use crate::settings::ProxySettings;
use serde::Serialize;
use serde::de::DeserializeOwned;
use std::path::Path;
use yaak_models::blob_manager::BlobManager;
/// Where captured traffic is read from.
pub enum ProxyConnection {
/// A process owns the engine; everything goes through its control API.
Live(ControlClient),
/// Nothing owns the engine; flows are read from disk.
Offline(OfflineStore),
}
impl ProxyConnection {
/// The live owner in `data_dir`, or its databases when there is none.
pub fn connect(data_dir: &Path) -> Result<Self, String> {
match control::live_owner(data_dir) {
Some(info) => Ok(Self::Live(ControlClient::new(info))),
None => Ok(Self::Offline(OfflineStore::open(data_dir)?)),
}
}
pub async fn flows(&self, filter: &FlowFilter) -> Result<Vec<Flow>, String> {
match self {
Self::Live(c) => c.flows(filter).await,
Self::Offline(s) => s.flows(filter),
}
}
pub async fn flow(&self, id: &str, body_limit: usize) -> Result<FlowDetail, String> {
match self {
Self::Live(c) => c.flow(id, body_limit).await,
Self::Offline(s) => s.flow(id, body_limit).await,
}
}
pub async fn body(&self, id: &str, which: FlowBody) -> Result<DecodedBody, String> {
match self {
Self::Live(c) => c.body(id, which).await,
Self::Offline(s) => s.body(id, which).await,
}
}
pub async fn har(&self, filter: &FlowFilter) -> Result<Har, String> {
match self {
Self::Live(c) => c.har(filter).await,
Self::Offline(s) => s.har(filter).await,
}
}
}
/// Calls the control API of the owner `control.json` names.
pub struct ControlClient {
info: ControlInfo,
http: reqwest::Client,
}
impl ControlClient {
pub fn new(info: ControlInfo) -> Self {
// Never through the system proxy, which may be the engine itself.
let http = reqwest::Client::builder().no_proxy().build().unwrap_or_default();
Self { info, http }
}
pub fn info(&self) -> &ControlInfo {
&self.info
}
pub async fn status(&self) -> Result<ControlStatus, String> {
self.get("/status").await
}
pub async fn start(&self, port: Option<u16>) -> Result<ControlStatus, String> {
self.post("/start", &StartRequest { port }).await
}
pub async fn stop(&self) -> Result<ControlStatus, String> {
self.post("/stop", &serde_json::json!({})).await
}
pub async fn flows(&self, filter: &FlowFilter) -> Result<Vec<Flow>, String> {
self.get(&format!("/flows?{}", filter.to_query())).await
}
pub async fn flow(&self, id: &str, body_limit: usize) -> Result<FlowDetail, String> {
self.get(&format!("/flows/{}?bodyLimit={body_limit}", encode(id))).await
}
pub async fn body(&self, id: &str, which: FlowBody) -> Result<DecodedBody, String> {
let side = match which {
FlowBody::Request => "request-body",
FlowBody::Response => "response-body",
};
let resp =
self.send(self.http.get(self.url(&format!("/flows/{}/{side}", encode(id))))).await?;
let truncated = resp.headers().get(TRUNCATED_HEADER).is_some_and(|v| v == "true");
let data = resp.bytes().await.map_err(|e| e.to_string())?.to_vec();
Ok(DecodedBody { stored_size: data.len(), data, truncated, decode_error: None })
}
pub async fn har(&self, filter: &FlowFilter) -> Result<Har, String> {
self.get(&format!("/har?{}", filter.to_query())).await
}
pub async fn ca_pem(&self) -> Result<String, String> {
let resp = self.send(self.http.get(self.url("/ca"))).await?;
resp.text().await.map_err(|e| e.to_string())
}
pub async fn trust_ca(&self) -> Result<CaStatus, String> {
self.post("/ca/trust", &serde_json::json!({})).await
}
pub async fn system_proxy(&self) -> Result<SystemProxyStatus, String> {
self.get("/system-proxy").await
}
pub async fn set_system_proxy(&self, enabled: bool) -> Result<SystemProxyStatus, String> {
self.post("/system-proxy", &SystemProxyRequest { enabled }).await
}
pub async fn workspaces(&self) -> Result<Vec<WorkspaceSummary>, String> {
self.get("/workspaces").await
}
pub async fn save(&self, id: &str, req: &SaveFlowRequest) -> Result<SaveFlowResponse, String> {
self.post(&format!("/flows/{}/save", encode(id)), req).await
}
fn url(&self, path: &str) -> String {
format!("{}{path}", self.info.base_url())
}
async fn get<T: DeserializeOwned>(&self, path: &str) -> Result<T, String> {
let resp = self.send(self.http.get(self.url(path))).await?;
resp.json().await.map_err(|e| format!("Invalid response from the proxy: {e}"))
}
async fn post<B: Serialize, T: DeserializeOwned>(
&self,
path: &str,
body: &B,
) -> Result<T, String> {
let body = serde_json::to_vec(body).map_err(|e| e.to_string())?;
let req =
self.http.post(self.url(path)).header("content-type", "application/json").body(body);
let resp = self.send(req).await?;
resp.json().await.map_err(|e| format!("Invalid response from the proxy: {e}"))
}
async fn send(&self, req: reqwest::RequestBuilder) -> Result<reqwest::Response, String> {
let resp = req
.bearer_auth(&self.info.token)
.send()
.await
.map_err(|e| format!("Failed to reach the proxy's control API: {e}"))?;
if resp.status().is_success() {
return Ok(resp);
}
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
Err(serde_json::from_str::<ErrorBody>(&text)
.map(|e| e.error)
.unwrap_or_else(|_| format!("The proxy answered {status}: {text}")))
}
}
fn encode(segment: &str) -> String {
url::form_urlencoded::byte_serialize(segment.as_bytes()).collect()
}
/// `proxy.db` and `blobs.sqlite`, opened read-only while no process owns the engine.
pub struct OfflineStore {
db: ProxyQueryManager,
blobs: BlobManager,
capture_body_limit: usize,
}
impl OfflineStore {
pub fn open(data_dir: &Path) -> Result<Self, String> {
let db_path = data_dir.join("proxy.db");
if !db_path.exists() {
return Err(format!(
"No captured traffic yet: {} does not exist. Start Yaak Proxy to capture some.",
db_path.display()
));
}
let db = ProxyQueryManager::open_read_only(&db_path)
.map_err(|e| format!("Failed to open {}: {e}", db_path.display()))?;
let blobs = crate::db::open_blobs_read_only(&data_dir.join("blobs.sqlite"))
.map_err(|e| format!("Failed to open the proxy's blob store: {e}"))?;
let settings = ProxySettings::load(&data_dir.join("settings.json"))?;
Ok(Self { db, blobs, capture_body_limit: settings.capture_body_limit })
}
pub fn flows(&self, filter: &FlowFilter) -> Result<Vec<Flow>, String> {
flows::list(&self.db, filter)
}
pub async fn flow(&self, id: &str, body_limit: usize) -> Result<FlowDetail, String> {
flows::detail(&self.db, &self.blobs, id, body_limit).await
}
pub async fn body(&self, id: &str, which: FlowBody) -> Result<DecodedBody, String> {
let flow = flows::get(&self.db, id)?;
bodies::read_body(&self.blobs, &flow, which, self.capture_body_limit).await
}
pub async fn har(&self, filter: &FlowFilter) -> Result<Har, String> {
let flows = flows::list(&self.db, filter)?;
har::export(&self.blobs, &flows, self.capture_body_limit).await
}
}
+605
View File
@@ -0,0 +1,605 @@
//! The control API: how the CLI, MCP tools and a second app launch reach the process that owns
//! the engine. The owner serves JSON on `127.0.0.1:<random>` behind a bearer token and writes
//! where to find it to `control.json` in the proxy data directory (mode 0600). One owner at a
//! time: a file whose process is gone, or whose port no longer answers, is stale and replaced.
use crate::ca::{self, CaStatus};
use crate::flows::{self, FlowFilter};
use crate::har;
use crate::macos::SystemProxyStatus;
use crate::models::{FlowBody, find_header};
use crate::save::SaveFlowRequest;
use crate::{ProxyCtx, ProxyState, bodies};
use bytes::Bytes;
use chrono::{DateTime, Utc};
use http_body_util::{BodyExt, Full, Limited};
use hyper::body::Incoming;
use hyper::header::{AUTHORIZATION, CONTENT_TYPE};
use hyper::{Method, Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use hyper_util::server::graceful::GracefulShutdown;
use log::{info, warn};
use serde::{Deserialize, Serialize};
use std::io::Write;
use std::net::{SocketAddr, TcpStream};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::mpsc as std_mpsc;
use std::time::Duration;
use tokio::net::TcpListener;
use ts_rs::TS;
use yaak_proxy::CaStore;
pub const CONTROL_FILE: &str = "control.json";
const LOCK_FILE: &str = "control.lock";
/// Body bytes `GET /flows/{id}` returns per side unless `bodyLimit` says otherwise.
pub const DEFAULT_BODY_PREVIEW_LIMIT: usize = 1024 * 1024;
const MAX_REQUEST_BODY: usize = 64 * 1024;
const LIVENESS_TIMEOUT: Duration = Duration::from_millis(500);
const SHUTDOWN_GRACE: Duration = Duration::from_secs(5);
/// The process that runs the engine.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export, export_to = "gen_rpc.ts")]
pub enum Owner {
App,
Cli,
}
/// `control.json`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ControlInfo {
pub pid: u32,
pub control_port: u16,
/// The engine's port while it runs.
pub proxy_port: Option<u16>,
pub token: String,
pub owner: Owner,
pub started_at: DateTime<Utc>,
}
impl ControlInfo {
pub fn base_url(&self) -> String {
format!("http://127.0.0.1:{}", self.control_port)
}
}
pub fn control_path(data_dir: &Path) -> PathBuf {
data_dir.join(CONTROL_FILE)
}
/// `control.json` as it is, live or stale. `None` when there is none or it does not parse.
pub fn read(data_dir: &Path) -> Option<ControlInfo> {
let contents = std::fs::read_to_string(control_path(data_dir)).ok()?;
serde_json::from_str(&contents).ok()
}
/// Held while `control.json` is checked and then removed or claimed, across processes, so two
/// of them cleaning up a stale file cannot delete each other's fresh claim. Released on drop.
fn claim_lock(data_dir: &Path) -> std::io::Result<std::fs::File> {
std::fs::create_dir_all(data_dir)?;
let file = std::fs::OpenOptions::new()
.create(true)
.truncate(false)
.write(true)
.open(data_dir.join(LOCK_FILE))?;
file.lock()?;
Ok(file)
}
/// The owner running the engine now. A stale `control.json` is removed on the way, as is one
/// that does not parse: it is only ever published whole, so that is never a claim in progress.
pub fn live_owner(data_dir: &Path) -> Option<ControlInfo> {
let _lock = claim_lock(data_dir)
.inspect_err(|e| warn!("Failed to lock {LOCK_FILE}, checking without it: {e}"))
.ok();
live_owner_locked(data_dir)
}
/// [`live_owner`] for a caller holding the claim lock.
fn live_owner_locked(data_dir: &Path) -> Option<ControlInfo> {
let path = control_path(data_dir);
let contents = std::fs::read_to_string(&path).ok()?;
let Ok(info) = serde_json::from_str::<ControlInfo>(&contents) else {
warn!("Removing unreadable {}", path.display());
let _ = std::fs::remove_file(&path);
return None;
};
if is_live(&info) {
return Some(info);
}
info!("Removing stale {CONTROL_FILE} of {:?} pid {}", info.owner, info.pid);
remove_if_token(data_dir, &info.token);
None
}
/// The process is alive and its control port answers. The port check catches a pid the OS
/// has handed to some other process since the owner died.
fn is_live(info: &ControlInfo) -> bool {
pid_alive(info.pid)
&& TcpStream::connect_timeout(
&SocketAddr::from(([127, 0, 0, 1], info.control_port)),
LIVENESS_TIMEOUT,
)
.is_ok()
}
#[cfg(unix)]
fn pid_alive(pid: u32) -> bool {
let Ok(pid) = libc::pid_t::try_from(pid) else {
return false;
};
// Signal 0 checks the process exists without signalling it; EPERM means it exists but is
// someone else's.
let alive = unsafe { libc::kill(pid, 0) } == 0;
alive || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
}
#[cfg(not(unix))]
fn pid_alive(_pid: u32) -> bool {
// The control port check decides.
true
}
/// Removes `control.json` if it is still the one with `token`, so an owner never removes its
/// successor's file. The caller holds the claim lock.
fn remove_if_token(data_dir: &Path, token: &str) {
if read(data_dir).is_some_and(|info| info.token == token) {
let _ = std::fs::remove_file(control_path(data_dir));
}
}
/// Writes `info` to `path`, readable by its owner alone and never seen half-written. With
/// `create_new`, fails when a file is already there, so two processes cannot both claim the
/// engine.
fn write_file(path: &Path, info: &ControlInfo, create_new: bool) -> std::io::Result<()> {
let contents = serde_json::to_vec_pretty(info)?;
let tmp = path.with_extension(format!("{}.tmp", std::process::id()));
let _ = std::fs::remove_file(&tmp);
let mut options = std::fs::OpenOptions::new();
options.write(true).create_new(true);
#[cfg(unix)]
std::os::unix::fs::OpenOptionsExt::mode(&mut options, 0o600);
options.open(&tmp)?.write_all(&contents)?;
let published = if create_new {
// A link, unlike a rename, refuses to replace an existing file.
std::fs::hard_link(&tmp, path).and_then(|_| std::fs::remove_file(&tmp))
} else {
std::fs::rename(&tmp, path)
};
if published.is_err() {
let _ = std::fs::remove_file(&tmp);
}
published
}
/// The owner's `control.json`, rewritten as the engine starts and stops.
pub(crate) struct ControlFile {
path: PathBuf,
pub(crate) info: ControlInfo,
}
impl ControlFile {
pub(crate) fn set_proxy_port(&mut self, port: Option<u16>) {
if self.info.proxy_port == port {
return;
}
self.info.proxy_port = port;
if let Err(e) = write_file(&self.path, &self.info, false) {
warn!("Failed to update {}: {e}", self.path.display());
}
}
}
// -- Responses --
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ControlStatus {
pub owner: Owner,
pub pid: u32,
pub started_at: DateTime<Utc>,
pub state: ProxyState,
/// The engine's port while it runs.
pub proxy_port: Option<u16>,
/// The port the engine starts on.
pub configured_port: u16,
pub flow_count: usize,
pub ca: CaStatus,
/// `None` when it could not be read.
pub system_proxy: Option<SystemProxyStatus>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct StartRequest {
/// For this run only; `settings.json` keeps its port.
pub port: Option<u16>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemProxyRequest {
pub enabled: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ErrorBody {
pub error: String,
}
/// Raw body headers: whether the body is cut short, at the capture limit or on the wire.
pub const TRUNCATED_HEADER: &str = "x-yaak-truncated";
// -- Server --
/// The control server of the process that owns the engine. Dropping it stops serving and
/// removes `control.json`.
pub struct ControlServer {
info: ControlInfo,
data_dir: PathBuf,
ctx: ProxyCtx,
shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
thread: Option<std::thread::JoinHandle<()>>,
stop_rx: Option<std_mpsc::Receiver<()>>,
}
impl ControlServer {
/// Claims the engine for this process and serves the control API for `ctx`. Fails when
/// another live owner has it.
pub fn start(ctx: ProxyCtx) -> Result<Self, String> {
let data_dir = ctx.data_dir().to_path_buf();
let path = control_path(&data_dir);
if let Some(owner) = live_owner(&data_dir) {
return Err(already_owned(&owner));
}
let (ready_tx, ready_rx) = std_mpsc::channel();
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
let (stop_tx, stop_rx) = std_mpsc::channel();
let token = yaak_database::generate_id_of_length(40);
let state = Arc::new(ServerState { ctx: ctx.clone(), token: token.clone(), stop_tx });
let thread = std::thread::spawn(move || {
let rt = match tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.thread_name("yaak-proxy-control")
.build()
{
Ok(rt) => rt,
Err(e) => {
let _ = ready_tx.send(Err(format!("Failed to create runtime: {e}")));
return;
}
};
rt.block_on(async move {
let listener = match TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0))).await
{
Ok(l) => l,
Err(e) => {
let _ = ready_tx.send(Err(format!("Failed to bind control port: {e}")));
return;
}
};
let _ = ready_tx
.send(listener.local_addr().map(|a| a.port()).map_err(|e| e.to_string()));
serve(listener, state, shutdown_rx).await;
});
});
let control_port = match ready_rx.recv() {
Ok(Ok(port)) => port,
Ok(Err(e)) => return Err(e),
Err(_) => return Err("Control thread died before binding".into()),
};
let info = ControlInfo {
pid: std::process::id(),
control_port,
proxy_port: ctx.port(),
token,
owner: ctx.owner(),
started_at: Utc::now(),
};
let mut server = Self {
info: info.clone(),
data_dir: data_dir.clone(),
ctx: ctx.clone(),
shutdown_tx: Some(shutdown_tx),
thread: Some(thread),
stop_rx: Some(stop_rx),
};
// Checked again under the lock: another process may have claimed it since.
let claimed = claim_lock(&data_dir)
.map_err(|e| format!("Failed to lock {LOCK_FILE}: {e}"))
.and_then(|_lock| match live_owner_locked(&data_dir) {
Some(owner) => Err(already_owned(&owner)),
None => write_file(&path, &info, true)
.map_err(|e| format!("Failed to write {}: {e}", path.display())),
});
if let Err(e) = claimed {
// Nothing of this server's to remove.
server.info.token.clear();
return Err(e);
}
ctx.set_control_file(Some(ControlFile { path, info }));
info!("Control API listening on 127.0.0.1:{control_port}");
Ok(server)
}
pub fn info(&self) -> &ControlInfo {
&self.info
}
/// Receives once per `POST /stop`. A CLI daemon exits on it; the app only stops its engine.
pub fn take_stop_rx(&mut self) -> Option<std_mpsc::Receiver<()>> {
self.stop_rx.take()
}
}
pub(crate) fn already_owned(owner: &ControlInfo) -> String {
let who = match owner.owner {
Owner::App => "Yaak Proxy",
Owner::Cli => "the yaak CLI",
};
format!("The proxy is already running in {who} (pid {})", owner.pid)
}
impl Drop for ControlServer {
fn drop(&mut self) {
if let Some(tx) = self.shutdown_tx.take() {
let _ = tx.send(());
}
if let Some(thread) = self.thread.take() {
let _ = thread.join();
}
if !self.info.token.is_empty() {
self.ctx.set_control_file(None);
let _lock = claim_lock(&self.data_dir).ok();
remove_if_token(&self.data_dir, &self.info.token);
}
}
}
struct ServerState {
ctx: ProxyCtx,
token: String,
stop_tx: std_mpsc::Sender<()>,
}
async fn serve(
listener: TcpListener,
state: Arc<ServerState>,
mut shutdown_rx: tokio::sync::oneshot::Receiver<()>,
) {
// Lets a request in flight at shutdown finish, such as the `POST /stop` a daemon exits on.
let graceful = GracefulShutdown::new();
loop {
tokio::select! {
result = listener.accept() => match result {
Ok((stream, _)) => {
let state = state.clone();
let service = hyper::service::service_fn(move |req| {
let state = state.clone();
async move { Ok::<_, std::convert::Infallible>(handle(&state, req).await) }
});
let conn = hyper::server::conn::http1::Builder::new()
.serve_connection(TokioIo::new(stream), service);
let conn = graceful.watch(conn);
tokio::spawn(async move {
let _ = conn.await;
});
}
Err(e) => warn!("Control accept error: {e}"),
},
_ = &mut shutdown_rx => break,
}
}
drop(listener);
if tokio::time::timeout(SHUTDOWN_GRACE, graceful.shutdown()).await.is_err() {
warn!("Control connections were still open after {SHUTDOWN_GRACE:?}");
}
}
type ControlResponse = Response<Full<Bytes>>;
struct ApiError {
status: StatusCode,
message: String,
}
impl ApiError {
fn new(status: StatusCode, message: impl Into<String>) -> Self {
Self { status, message: message.into() }
}
}
impl From<String> for ApiError {
fn from(message: String) -> Self {
Self::new(StatusCode::INTERNAL_SERVER_ERROR, message)
}
}
async fn handle(state: &ServerState, req: Request<Incoming>) -> ControlResponse {
if !authorized(&req, &state.token) {
return error_response(ApiError::new(
StatusCode::UNAUTHORIZED,
"Missing or wrong bearer token; it is in control.json",
));
}
match route(state, req).await {
Ok(resp) => resp,
Err(e) => error_response(e),
}
}
fn authorized(req: &Request<Incoming>, token: &str) -> bool {
let Some(given) = req
.headers()
.get(AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
else {
return false;
};
// Compares every byte, so the time taken says nothing about how much matched.
given.len() == token.len()
&& given.bytes().zip(token.bytes()).fold(0u8, |acc, (a, b)| acc | (a ^ b)) == 0
}
async fn route(state: &ServerState, req: Request<Incoming>) -> Result<ControlResponse, ApiError> {
let ctx = &state.ctx;
let method = req.method().clone();
let path = req.uri().path().to_string();
let query = req.uri().query().unwrap_or_default().to_string();
let segments: Vec<&str> = path.trim_matches('/').split('/').collect();
match (&method, segments.as_slice()) {
(&Method::GET, ["status"]) => json(&status(ctx).await?),
(&Method::POST, ["start"]) => {
let body: StartRequest = read_json(req).await?;
ctx.start_on(body.port)?;
json(&status(ctx).await?)
}
(&Method::POST, ["stop"]) => {
ctx.stop()?;
let _ = state.stop_tx.send(());
json(&status(ctx).await?)
}
(&Method::GET, ["flows"]) => {
let filter = FlowFilter::from_query(&query).map_err(bad_request)?;
json(&flows::list(&ctx.db, &filter)?)
}
(&Method::GET, ["flows", id]) => {
found(ctx, id)?;
let limit = body_limit(&query)?;
json(&flows::detail(&ctx.db, &ctx.blobs, id, limit).await?)
}
(&Method::GET, ["flows", id, "request-body"]) => raw_body(ctx, id, FlowBody::Request).await,
(&Method::GET, ["flows", id, "response-body"]) => {
raw_body(ctx, id, FlowBody::Response).await
}
(&Method::POST, ["flows", id, "save"]) => {
found(ctx, id)?;
let id = id.to_string();
let body: SaveFlowRequest = read_json(req).await?;
json(&ctx.save_flow(&id, &body).await.map_err(bad_request)?)
}
(&Method::GET, ["har"]) => {
let filter = FlowFilter::from_query(&query).map_err(bad_request)?;
let flows = flows::list(&ctx.db, &filter)?;
json(&har::export(&ctx.blobs, &flows, ctx.capture_body_limit()?).await?)
}
(&Method::GET, ["ca"]) => {
let ca = CaStore::load_or_create(&ctx.ca_dir())
.map_err(|e| format!("Failed to load proxy CA: {e}"))?;
Ok(response(StatusCode::OK, "application/x-pem-file", ca.pem().to_string()))
}
(&Method::POST, ["ca", "trust"]) => {
let ca = CaStore::load_or_create(&ctx.ca_dir())
.map_err(|e| format!("Failed to load proxy CA: {e}"))?;
ca::trust(ca.cert_path()).await.map_err(bad_request)?;
json(&ca::status(ca.cert_path()).await?)
}
(&Method::GET, ["system-proxy"]) => {
let ctx = ctx.clone();
json(&blocking(move || ctx.system_proxy_status()).await?)
}
(&Method::POST, ["system-proxy"]) => {
let body: SystemProxyRequest = read_json(req).await?;
let ctx = ctx.clone();
json(&blocking(move || ctx.set_system_proxy(body.enabled)).await.map_err(bad_request)?)
}
(&Method::GET, ["workspaces"]) => json(&ctx.workspaces().map_err(bad_request)?),
_ => Err(ApiError::new(StatusCode::NOT_FOUND, format!("No route {method} {path}"))),
}
}
async fn status(ctx: &ProxyCtx) -> Result<ControlStatus, ApiError> {
let control =
ctx.control_info().ok_or_else(|| "The control API is shutting down".to_string())?;
let proxy_port = ctx.port();
let ca = ca::status(&ctx.ca_dir().join("ca.pem")).await?;
let system_proxy = {
let ctx = ctx.clone();
blocking(move || ctx.system_proxy_status()).await.ok()
};
Ok(ControlStatus {
owner: control.owner,
pid: control.pid,
started_at: control.started_at,
state: if proxy_port.is_some() { ProxyState::Running } else { ProxyState::Stopped },
proxy_port,
configured_port: ctx.settings()?.port,
flow_count: flows::count(&ctx.db)?,
ca,
system_proxy,
})
}
fn found(ctx: &ProxyCtx, id: &str) -> Result<(), ApiError> {
match flows::find(&ctx.db, id) {
Some(_) => Ok(()),
None => Err(ApiError::new(StatusCode::NOT_FOUND, format!("No flow {id}"))),
}
}
async fn raw_body(ctx: &ProxyCtx, id: &str, which: FlowBody) -> Result<ControlResponse, ApiError> {
let flow = flows::find(&ctx.db, id)
.ok_or_else(|| ApiError::new(StatusCode::NOT_FOUND, format!("No flow {id}")))?;
let body = bodies::read_body(&ctx.blobs, &flow, which, ctx.capture_body_limit()?).await?;
let content_type =
find_header(flow.headers(which), "content-type").unwrap_or("application/octet-stream");
let mut resp = response(StatusCode::OK, content_type, body.data);
resp.headers_mut().insert(TRUNCATED_HEADER, body.truncated.to_string().parse().unwrap());
Ok(resp)
}
fn body_limit(query: &str) -> Result<usize, ApiError> {
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
if key == "bodyLimit" {
return value.parse().map_err(|_| bad_request(format!("Invalid bodyLimit {value:?}")));
}
}
Ok(DEFAULT_BODY_PREVIEW_LIMIT)
}
async fn blocking<T: Send + 'static>(
f: impl FnOnce() -> Result<T, String> + Send + 'static,
) -> Result<T, String> {
tokio::task::spawn_blocking(f).await.map_err(|e| e.to_string())?
}
/// The JSON request body; an empty one reads as `{}`.
async fn read_json<T: serde::de::DeserializeOwned>(req: Request<Incoming>) -> Result<T, ApiError> {
let bytes = Limited::new(req.into_body(), MAX_REQUEST_BODY)
.collect()
.await
.map_err(|e| bad_request(format!("Failed to read the request body: {e}")))?
.to_bytes();
let bytes: &[u8] = if bytes.iter().all(u8::is_ascii_whitespace) { b"{}" } else { &bytes };
serde_json::from_slice(bytes).map_err(|e| bad_request(format!("Invalid JSON body: {e}")))
}
fn bad_request(message: String) -> ApiError {
ApiError::new(StatusCode::BAD_REQUEST, message)
}
fn json<T: Serialize>(value: &T) -> Result<ControlResponse, ApiError> {
let body = serde_json::to_vec(value).map_err(|e| e.to_string())?;
Ok(response(StatusCode::OK, "application/json", body))
}
fn error_response(e: ApiError) -> ControlResponse {
let body = serde_json::to_vec(&ErrorBody { error: e.message }).unwrap_or_default();
response(e.status, "application/json", body)
}
fn response(status: StatusCode, content_type: &str, body: impl Into<Bytes>) -> ControlResponse {
let mut resp = Response::new(Full::new(body.into()));
*resp.status_mut() = status;
if let Ok(value) = content_type.parse() {
resp.headers_mut().insert(CONTENT_TYPE, value);
}
resp
}
+469
View File
@@ -0,0 +1,469 @@
//! A captured flow as a Yaak request. Headers Yaak sets itself when sending, and hop-by-hop
//! headers that only meant something between the app and the proxy, are dropped and listed so
//! the save dialog can say why. Authorization is kept as the plain header it was captured as.
use crate::bodies::DecodedBody;
use crate::models::{Flow, find_header};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::collections::BTreeMap;
use yaak_models::models::{HttpRequest, HttpRequestHeader, HttpUrlParameter};
// The client's body types (`apps/yaak-client/lib/model_util.ts`).
const BODY_TYPE_JSON: &str = "application/json";
const BODY_TYPE_XML: &str = "text/xml";
const BODY_TYPE_GRAPHQL: &str = "graphql";
const BODY_TYPE_FORM_URLENCODED: &str = "application/x-www-form-urlencoded";
const BODY_TYPE_OTHER: &str = "other";
/// Removed between the app and the proxy; never forwarded as they are.
const HOP_BY_HOP: &[&str] = &[
"connection",
"keep-alive",
"te",
"trailer",
"transfer-encoding",
"upgrade",
];
/// Yaak works these out itself when it sends the request.
const AUTOMATIC: &[&str] = &["host", "content-length", "accept-encoding"];
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct DroppedHeader {
pub name: String,
pub value: String,
pub reason: DropReason,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DropReason {
/// Only meant for the hop between the app and the proxy.
HopByHop,
/// Yaak sets it when it sends the request.
SetBySending,
/// The body is saved decoded, so its encoding no longer applies.
BodyDecoded,
}
#[derive(Debug, Clone)]
pub struct Conversion {
pub request: HttpRequest,
pub dropped_headers: Vec<DroppedHeader>,
/// What the saved request does not carry over exactly, one sentence each.
pub warnings: Vec<String>,
}
/// `flow` as a request in no workspace yet, with `body` its request body decoded.
pub fn flow_to_request(flow: &Flow, body: &DecodedBody) -> Conversion {
let mut warnings = Vec::new();
let (url, url_parameters) = split_query(&flow.url);
let (headers, dropped_headers) = split_headers(flow, body);
let has_authorization = headers.iter().any(|h| h.name.eq_ignore_ascii_case("authorization"));
if has_authorization {
warnings.push("The Authorization header is saved as captured, in plain text.".into());
}
let (body_type, body_fields) = request_body(flow, body, &mut warnings);
if body.truncated {
warnings.push(format!(
"The body was cut at the capture limit, so only its first {} bytes are saved.",
body.data.len()
));
}
let request = HttpRequest {
name: default_name(&url),
method: flow.method.clone(),
url,
url_parameters,
headers,
body_type,
body: body_fields,
// No auth rather than inherited, so a folder's or workspace's auth cannot replace the
// captured header when the request is sent.
authentication_type: has_authorization.then(|| "none".to_string()),
..Default::default()
};
Conversion { request, dropped_headers, warnings }
}
/// The URL without its query, and the query as parameters.
fn split_query(url: &str) -> (String, Vec<HttpUrlParameter>) {
let (url, fragment) = match url.split_once('#') {
Some((url, fragment)) => (url, Some(fragment)),
None => (url, None),
};
let Some((base, query)) = url.split_once('?') else {
return (url.to_string() + &fragment.map(|f| format!("#{f}")).unwrap_or_default(), vec![]);
};
let params = url::form_urlencoded::parse(query.as_bytes())
.map(|(name, value)| HttpUrlParameter {
enabled: true,
name: name.into_owned(),
value: value.into_owned(),
id: None,
})
.collect();
(base.to_string() + &fragment.map(|f| format!("#{f}")).unwrap_or_default(), params)
}
fn split_headers(flow: &Flow, body: &DecodedBody) -> (Vec<HttpRequestHeader>, Vec<DroppedHeader>) {
// Connection names further headers that are hop-by-hop for this request alone.
let listed: Vec<String> = flow
.request_headers
.iter()
.filter(|h| h.name.eq_ignore_ascii_case("connection"))
.flat_map(|h| h.value.split(','))
.map(|n| n.trim().to_ascii_lowercase())
.filter(|n| !n.is_empty())
.collect();
let encoding = find_header(&flow.request_headers, "content-encoding");
// A body that decodes to nothing was still decoded.
let decoded = body.stored_size > 0
&& encoding.is_some_and(|e| !e.trim().eq_ignore_ascii_case("identity"));
let mut kept = Vec::new();
let mut dropped = Vec::new();
for h in &flow.request_headers {
let name = h.name.to_ascii_lowercase();
let reason = if HOP_BY_HOP.contains(&name.as_str())
|| name.starts_with("proxy-")
|| name.starts_with(':')
|| listed.contains(&name)
{
Some(DropReason::HopByHop)
} else if AUTOMATIC.contains(&name.as_str()) {
Some(DropReason::SetBySending)
} else if name == "content-encoding" && decoded {
Some(DropReason::BodyDecoded)
} else {
None
};
match reason {
Some(reason) => {
dropped.push(DroppedHeader { name: h.name.clone(), value: h.value.clone(), reason })
}
None => kept.push(HttpRequestHeader {
enabled: true,
name: h.name.clone(),
value: h.value.clone(),
id: None,
}),
}
}
(kept, dropped)
}
fn request_body(
flow: &Flow,
body: &DecodedBody,
warnings: &mut Vec<String>,
) -> (Option<String>, BTreeMap<String, Value>) {
if body.data.is_empty() {
return (None, BTreeMap::new());
}
let content_type = find_header(&flow.request_headers, "content-type").unwrap_or_default();
let mime = content_type.split(';').next().unwrap_or_default().trim().to_ascii_lowercase();
let Ok(text) = std::str::from_utf8(&body.data) else {
warnings.push(
"The body is binary and is saved as text, so bytes that are not UTF-8 are replaced."
.into(),
);
return text_body(BODY_TYPE_OTHER, &String::from_utf8_lossy(&body.data));
};
if mime == "application/json" || mime.ends_with("+json") {
if let Some(graphql) = graphql_body(flow, text) {
return (Some(BODY_TYPE_GRAPHQL.into()), graphql);
}
return text_body(BODY_TYPE_JSON, text);
}
if mime == "application/xml" || mime == "text/xml" || mime.ends_with("+xml") {
return text_body(BODY_TYPE_XML, text);
}
if mime == BODY_TYPE_FORM_URLENCODED {
let form: Vec<Value> = url::form_urlencoded::parse(text.as_bytes())
.map(|(name, value)| json!({ "enabled": true, "name": name, "value": value }))
.collect();
return (
Some(BODY_TYPE_FORM_URLENCODED.into()),
BTreeMap::from([("form".into(), json!(form))]),
);
}
if mime == "multipart/form-data" {
warnings.push(
"The multipart body is saved as text with its Content-Type, not as form fields.".into(),
);
}
text_body(BODY_TYPE_OTHER, text)
}
fn text_body(body_type: &str, text: &str) -> (Option<String>, BTreeMap<String, Value>) {
(Some(body_type.into()), BTreeMap::from([("text".into(), json!(text))]))
}
/// A JSON body posted to a GraphQL endpoint, as the client's GraphQL body.
fn graphql_body(flow: &Flow, text: &str) -> Option<BTreeMap<String, Value>> {
let path = flow.url.split('?').next().unwrap_or_default();
if !path.to_ascii_lowercase().ends_with("/graphql") {
return None;
}
let Value::Object(body) = serde_json::from_str::<Value>(text).ok()? else {
return None;
};
// Anything else, such as persisted-query `extensions`, would be lost in a GraphQL body.
if body.keys().any(|k| !matches!(k.as_str(), "query" | "variables" | "operationName")) {
return None;
}
let query = body.get("query")?.as_str()?;
let mut fields = BTreeMap::from([("query".to_string(), json!(query))]);
if let Some(variables) = body.get("variables").filter(|v| !v.is_null()) {
let pretty = serde_json::to_string_pretty(variables).ok()?;
fields.insert("variables".into(), json!(pretty));
}
if let Some(name) = body.get("operationName").and_then(Value::as_str) {
fields.insert("operationName".into(), json!(name));
}
Some(fields)
}
/// The path, or the host for a request to the root.
fn default_name(url: &str) -> String {
let rest = url.split_once("://").map_or(url, |(_, rest)| rest);
match rest.split_once('/') {
Some((_, path)) if !path.is_empty() => format!("/{path}"),
Some((host, _)) => host.to_string(),
None => rest.to_string(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::models::ProxyHeader;
fn flow(method: &str, url: &str, headers: &[(&str, &str)]) -> Flow {
Flow {
method: method.into(),
url: url.into(),
request_headers: headers
.iter()
.map(|(n, v)| ProxyHeader { name: (*n).into(), value: (*v).into() })
.collect(),
..Default::default()
}
}
fn body(data: &str) -> DecodedBody {
DecodedBody {
data: data.as_bytes().to_vec(),
stored_size: data.len(),
truncated: false,
decode_error: None,
}
}
fn names(request: &HttpRequest) -> Vec<&str> {
request.headers.iter().map(|h| h.name.as_str()).collect()
}
#[test]
fn get_without_body() {
let f = flow(
"GET",
"https://api.github.com/search/issues?q=is%3Aopen+label%3Abug&per_page=5",
&[
("Host", "api.github.com"),
("Accept", "application/json"),
("Accept-Encoding", "gzip, br"),
("Connection", "keep-alive, X-Debug"),
("X-Debug", "1"),
("Proxy-Authorization", "Basic Zm9v"),
("Authorization", "Bearer ghp_secret"),
],
);
let c = flow_to_request(&f, &body(""));
let r = &c.request;
assert_eq!(r.method, "GET");
assert_eq!(r.url, "https://api.github.com/search/issues");
assert_eq!(r.name, "/search/issues");
let params: Vec<(&str, &str)> =
r.url_parameters.iter().map(|p| (p.name.as_str(), p.value.as_str())).collect();
assert_eq!(params, vec![("q", "is:open label:bug"), ("per_page", "5")]);
assert_eq!(names(r), vec!["Accept", "Authorization"]);
assert_eq!(r.headers[1].value, "Bearer ghp_secret");
assert_eq!(r.authentication_type.as_deref(), Some("none"));
assert_eq!(r.body_type, None);
assert!(r.body.is_empty());
let dropped: Vec<(&str, DropReason)> =
c.dropped_headers.iter().map(|d| (d.name.as_str(), d.reason)).collect();
assert_eq!(
dropped,
vec![
("Host", DropReason::SetBySending),
("Accept-Encoding", DropReason::SetBySending),
("Connection", DropReason::HopByHop),
("X-Debug", DropReason::HopByHop),
("Proxy-Authorization", DropReason::HopByHop),
]
);
assert_eq!(
c.warnings,
vec!["The Authorization header is saved as captured, in plain text."]
);
}
#[test]
fn json_body() {
let f = flow(
"POST",
"https://api.stripe.com/v1/charges",
&[
("Content-Type", "application/json; charset=utf-8"),
("Content-Length", "13"),
],
);
let c = flow_to_request(&f, &body(r#"{"amount":42}"#));
assert_eq!(c.request.body_type.as_deref(), Some("application/json"));
assert_eq!(c.request.body["text"], r#"{"amount":42}"#);
assert_eq!(names(&c.request), vec!["Content-Type"]);
assert!(c.warnings.is_empty());
// Without a captured Authorization, the request inherits its folder's auth.
assert_eq!(c.request.authentication_type, None);
}
#[test]
fn vendor_json_body() {
let f = flow("PATCH", "https://x.test/a", &[("Content-Type", "application/vnd.api+json")]);
let c = flow_to_request(&f, &body("{}"));
assert_eq!(c.request.body_type.as_deref(), Some("application/json"));
}
#[test]
fn graphql_body() {
let f =
flow("POST", "https://api.github.com/graphql", &[("Content-Type", "application/json")]);
let c = flow_to_request(
&f,
&body(
r#"{"query":"query($n:Int){viewer{login}}","variables":{"n":1},"operationName":"Me"}"#,
),
);
assert_eq!(c.request.body_type.as_deref(), Some("graphql"));
assert_eq!(c.request.body["query"], "query($n:Int){viewer{login}}");
assert_eq!(c.request.body["variables"], "{\n \"n\": 1\n}");
assert_eq!(c.request.body["operationName"], "Me");
// Fields a GraphQL body cannot hold keep it plain JSON.
let persisted = r#"{"query":"q","extensions":{"persistedQuery":{"version":1}}}"#;
let c = flow_to_request(&f, &body(persisted));
assert_eq!(c.request.body_type.as_deref(), Some("application/json"));
// The same body anywhere else is plain JSON.
let f = flow("POST", "https://x.test/api", &[("Content-Type", "application/json")]);
let c = flow_to_request(&f, &body(r#"{"query":"q"}"#));
assert_eq!(c.request.body_type.as_deref(), Some("application/json"));
}
#[test]
fn xml_body() {
let f = flow("POST", "https://x.test/soap", &[("Content-Type", "application/soap+xml")]);
let c = flow_to_request(&f, &body("<a/>"));
assert_eq!(c.request.body_type.as_deref(), Some("text/xml"));
assert_eq!(c.request.body["text"], "<a/>");
}
#[test]
fn form_body() {
let f = flow(
"POST",
"http://localhost:3000/login",
&[("Content-Type", "application/x-www-form-urlencoded")],
);
let c = flow_to_request(&f, &body("user=greg&note=hello+world%21"));
assert_eq!(c.request.body_type.as_deref(), Some("application/x-www-form-urlencoded"));
assert_eq!(
c.request.body["form"],
json!([
{ "enabled": true, "name": "user", "value": "greg" },
{ "enabled": true, "name": "note", "value": "hello world!" },
])
);
assert_eq!(c.request.name, "/login");
}
#[test]
fn multipart_body() {
let content_type = "multipart/form-data; boundary=XyZ";
let f = flow("POST", "https://x.test/upload", &[("Content-Type", content_type)]);
let text = "--XyZ\r\nContent-Disposition: form-data; name=\"a\"\r\n\r\n1\r\n--XyZ--\r\n";
let c = flow_to_request(&f, &body(text));
assert_eq!(c.request.body_type.as_deref(), Some("other"));
assert_eq!(c.request.body["text"], text);
assert_eq!(c.request.headers[0].value, content_type);
assert_eq!(c.warnings.len(), 1);
}
#[test]
fn text_body() {
let f = flow("PUT", "https://x.test/", &[("Content-Type", "text/plain")]);
let c = flow_to_request(&f, &body("hello"));
assert_eq!(c.request.body_type.as_deref(), Some("other"));
assert_eq!(c.request.body["text"], "hello");
assert_eq!(c.request.name, "x.test");
}
#[test]
fn binary_body() {
let f = flow("POST", "https://x.test/bin", &[("Content-Type", "application/octet-stream")]);
let data = DecodedBody {
data: vec![b'a', 0xff, b'b'],
stored_size: 3,
truncated: false,
decode_error: None,
};
let c = flow_to_request(&f, &data);
assert_eq!(c.request.body_type.as_deref(), Some("other"));
assert_eq!(c.request.body["text"], "a\u{fffd}b");
assert_eq!(c.warnings.len(), 1);
}
#[test]
fn decoded_empty_body() {
let f = flow("POST", "https://x.test/", &[("Content-Encoding", "gzip")]);
// A gzip stream of nothing: stored bytes, but an empty body once decoded.
let data =
DecodedBody { data: Vec::new(), stored_size: 20, truncated: false, decode_error: None };
let c = flow_to_request(&f, &data);
assert!(c.request.headers.is_empty());
assert_eq!(c.dropped_headers[0].reason, DropReason::BodyDecoded);
assert_eq!(c.request.body_type, None);
}
#[test]
fn decoded_and_truncated_body() {
let f = flow(
"POST",
"https://x.test/logs",
&[
("Content-Type", "application/json"),
("Content-Encoding", "gzip"),
],
);
let data = DecodedBody {
data: b"[1,2".to_vec(),
stored_size: 4,
truncated: true,
decode_error: None,
};
let c = flow_to_request(&f, &data);
assert_eq!(names(&c.request), vec!["Content-Type"]);
assert_eq!(c.dropped_headers[0].reason, DropReason::BodyDecoded);
assert_eq!(
c.warnings,
vec!["The body was cut at the capture limit, so only its first 4 bytes are saved."]
);
}
}
+49 -3
View File
@@ -1,10 +1,11 @@
use include_dir::{Dir, include_dir};
use r2d2::Pool;
use r2d2_sqlite::SqliteConnectionManager;
use rusqlite::{Transaction, TransactionBehavior};
use rusqlite::{OpenFlags, Transaction, TransactionBehavior};
use std::path::Path;
use std::time::Duration;
use yaak_database::{ConnectionOrTx, DbContext, Error, Result, SqlitePool, run_migrations};
use yaak_models::blob_manager::BlobManager;
static MIGRATIONS: Dir<'static> = include_dir!("$CARGO_MANIFEST_DIR/migrations");
@@ -26,6 +27,13 @@ impl ProxyQueryManager {
Ok(Self { readers: file_pool(db_path, 5)?, writer })
}
/// `proxy.db` for reading while another process may be writing it, or none is. Writes
/// fail. No migrations run, so an older database reads as it is.
pub fn open_read_only(db_path: &Path) -> Result<Self> {
let readers = read_only_pool(db_path)?;
Ok(Self { writer: readers.clone(), readers })
}
pub fn with_conn<F, T>(&self, func: F) -> T
where
F: FnOnce(&DbContext) -> T,
@@ -62,6 +70,23 @@ impl ProxyQueryManager {
}
}
/// `blobs.sqlite` for reading, alongside [`ProxyQueryManager::open_read_only`].
pub fn open_blobs_read_only(path: &Path) -> Result<BlobManager> {
let readers = read_only_pool(path)?;
Ok(BlobManager::new(readers.clone(), readers))
}
fn read_only_pool(path: &Path) -> Result<SqlitePool> {
let manager = SqliteConnectionManager::file(path)
.with_flags(OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX)
.with_init(|conn| conn.busy_timeout(Duration::from_millis(5000)));
Pool::builder()
.max_size(2)
.connection_timeout(Duration::from_secs(10))
.build(manager)
.map_err(|e| Error::Database(e.to_string()))
}
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")?;
@@ -85,6 +110,28 @@ mod tests {
stmt.query_map([], |r| r.get::<_, String>("name")).unwrap().map(|c| c.unwrap()).collect()
}
#[test]
fn reads_flows_read_only_before_migrations() {
use crate::models::Flow;
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("proxy.db");
drop(ProxyQueryManager::new(&path).unwrap());
// As a build from before the saved request column left it.
let conn = Connection::open(&path).unwrap();
conn.execute_batch(
"ALTER TABLE flows DROP COLUMN saved_request_id;
INSERT INTO flows (id, url, method) VALUES ('fl_old', 'https://example.com', 'GET');",
)
.unwrap();
drop(conn);
let db = ProxyQueryManager::open_read_only(&path).unwrap();
let flows = db.with_conn(|c| c.find_all::<Flow>()).unwrap();
assert_eq!(flows.len(), 1);
assert_eq!((flows[0].id.as_str(), flows[0].saved_request_id.as_deref()), ("fl_old", None));
}
#[test]
fn migrates_an_old_schema_db() {
let dir = tempfile::tempdir().unwrap();
@@ -124,8 +171,7 @@ mod tests {
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();
let journal_mode: String = conn.query_row("PRAGMA journal_mode", [], |r| r.get(0)).unwrap();
assert_eq!(journal_mode, "wal");
}
}
+305
View File
@@ -0,0 +1,305 @@
use crate::bodies;
use crate::db::ProxyQueryManager;
use crate::models::{Flow, FlowBody, FlowIden};
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
use serde::{Deserialize, Serialize};
use yaak_models::blob_manager::BlobManager;
/// Which flows a listing or a HAR export covers. Every field narrows it; none is all of them.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default, rename_all = "camelCase")]
pub struct FlowFilter {
/// Part of the host, any case: `stripe` matches `api.stripe.com`.
pub host: Option<String>,
/// The method, any case.
pub method: Option<String>,
/// A status (`404`), a class of them (`4xx`), or `error` for flows that failed.
pub status: Option<String>,
/// At most this many flows, the newest.
pub limit: Option<usize>,
/// Only flows recorded after the flow with this id, for following new traffic. Flows are
/// recorded when they end, so a slow one that started earlier still comes through.
pub after: Option<String>,
}
impl FlowFilter {
/// The filter from a query string's `host`, `method`, `status`, `limit` and `after`.
pub fn from_query(query: &str) -> Result<Self, String> {
let mut filter = Self::default();
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
let value = Some(value.into_owned()).filter(|v| !v.is_empty());
match key.as_ref() {
"host" => filter.host = value,
"method" => filter.method = value,
"status" => filter.status = value,
"after" => filter.after = value,
"limit" => {
filter.limit = value
.map(|v| v.parse().map_err(|_| format!("Invalid limit {v:?}")))
.transpose()?
}
_ => {}
}
}
Ok(filter)
}
pub fn to_query(&self) -> String {
let mut query = url::form_urlencoded::Serializer::new(String::new());
let limit = self.limit.map(|l| l.to_string());
for (key, value) in [
("host", &self.host),
("method", &self.method),
("status", &self.status),
("limit", &limit),
("after", &self.after),
] {
if let Some(value) = value {
query.append_pair(key, value);
}
}
query.finish()
}
pub fn matches(&self, flow: &Flow) -> bool {
if let Some(host) = &self.host
&& !host_of(&flow.url).to_ascii_lowercase().contains(&host.to_ascii_lowercase())
{
return false;
}
if let Some(method) = &self.method
&& !flow.method.eq_ignore_ascii_case(method)
{
return false;
}
match self.status.as_deref().map(str::trim) {
None | Some("") => true,
Some(s) if s.eq_ignore_ascii_case("error") => flow.error.is_some(),
Some(s) => match (flow.status, s.to_ascii_lowercase().strip_suffix("xx")) {
(Some(status), Some(class)) => class.parse() == Ok(status / 100),
(Some(status), None) => s.parse() == Ok(status),
(None, _) => false,
},
}
}
}
/// `host[:port]` of a captured URL.
fn host_of(url: &str) -> &str {
let rest = url.split_once("://").map_or(url, |(_, rest)| rest);
rest.split(['/', '?', '#']).next().unwrap_or_default()
}
/// The flows `filter` covers, newest first.
pub fn list(db: &ProxyQueryManager, filter: &FlowFilter) -> Result<Vec<Flow>, String> {
db.with_conn(|c| {
let flows = c.find_all::<Flow>().map_err(|e| e.to_string())?;
let recorded_after = match &filter.after {
Some(id) => Some(recorded_after(c.conn().resolve(), id)?),
None => None,
};
Ok(flows
.into_iter()
.filter(|f| recorded_after.as_ref().is_none_or(|ids| ids.contains(&f.id)))
.filter(|f| filter.matches(f))
.take(filter.limit.unwrap_or(usize::MAX))
.collect())
})
}
/// Ids of the flows inserted after flow `id`, by rowid, which only grows as flows are written.
fn recorded_after(
conn: &rusqlite::Connection,
id: &str,
) -> Result<std::collections::HashSet<String>, String> {
let rowid: i64 = conn
.query_row("SELECT rowid FROM flows WHERE id = ?", [id], |r| r.get(0))
.map_err(|_| format!("No flow {id}"))?;
let mut stmt =
conn.prepare("SELECT id FROM flows WHERE rowid > ?").map_err(|e| e.to_string())?;
let ids = stmt.query_map([rowid], |r| r.get(0)).map_err(|e| e.to_string())?;
ids.collect::<Result<_, _>>().map_err(|e| e.to_string())
}
pub fn find(db: &ProxyQueryManager, flow_id: &str) -> Option<Flow> {
db.with_conn(|c| c.find_optional::<Flow>(FlowIden::Id, flow_id))
}
pub fn get(db: &ProxyQueryManager, flow_id: &str) -> Result<Flow, String> {
find(db, flow_id).ok_or_else(|| format!("No flow {flow_id}"))
}
pub fn count(db: &ProxyQueryManager) -> Result<usize, String> {
db.with_conn(|c| {
c.conn()
.resolve()
.query_row("SELECT COUNT(*) FROM flows", [], |r| r.get::<_, i64>(0))
.map(|n| n as usize)
.map_err(|e| e.to_string())
})
}
/// A body in JSON: text when it is UTF-8, base64 when it is not.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BodyPreview {
pub data: String,
pub encoding: BodyEncoding,
/// Bytes in `data` once decoded.
pub size: usize,
/// The body is longer than `data`: cut at the capture limit or at the preview cap.
pub truncated: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum BodyEncoding {
Utf8,
Base64,
}
impl BodyPreview {
pub fn new(data: &[u8], truncated: bool) -> Self {
let (data, encoding) = match std::str::from_utf8(data) {
Ok(text) => (text.to_string(), BodyEncoding::Utf8),
// A cut can split a character; that alone does not make the body binary.
Err(e) if truncated && e.error_len().is_none() => {
(String::from_utf8_lossy(&data[..e.valid_up_to()]).into_owned(), BodyEncoding::Utf8)
}
Err(_) => (BASE64.encode(data), BodyEncoding::Base64),
};
let size = match encoding {
BodyEncoding::Utf8 => data.len(),
BodyEncoding::Base64 => BASE64.decode(&data).map(|d| d.len()).unwrap_or_default(),
};
Self { data, encoding, size, truncated }
}
}
/// One flow with its bodies, each decoded and cut at `body_cap` bytes.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FlowDetail {
pub flow: Flow,
pub request_body: Option<BodyPreview>,
pub response_body: Option<BodyPreview>,
}
pub async fn detail(
db: &ProxyQueryManager,
blobs: &BlobManager,
flow_id: &str,
body_cap: usize,
) -> Result<FlowDetail, String> {
let flow = get(db, flow_id)?;
let request_body = preview(blobs, &flow, FlowBody::Request, body_cap).await?;
let response_body = preview(blobs, &flow, FlowBody::Response, body_cap).await?;
Ok(FlowDetail { flow, request_body, response_body })
}
async fn preview(
blobs: &BlobManager,
flow: &Flow,
which: FlowBody,
cap: usize,
) -> Result<Option<BodyPreview>, String> {
let body = bodies::read_body(blobs, flow, which, cap).await?;
if body.data.is_empty() && !body.truncated {
return Ok(None);
}
Ok(Some(BodyPreview::new(&body.data, body.truncated)))
}
#[cfg(test)]
mod tests {
use super::*;
fn flow(url: &str, method: &str, status: Option<i32>, error: Option<&str>) -> Flow {
Flow {
url: url.into(),
method: method.into(),
status,
error: error.map(String::from),
..Default::default()
}
}
#[test]
fn filters() {
let ok = flow("https://api.stripe.com:8443/v1/charges?x=1", "POST", Some(200), None);
let missing = flow("http://example.com/missing", "GET", Some(404), None);
let failed = flow("https://down.example.com/", "GET", None, Some("refused"));
let by = |filter: FlowFilter| {
[&ok, &missing, &failed].iter().filter(|f| filter.matches(f)).count()
};
assert_eq!(by(FlowFilter::default()), 3);
assert_eq!(by(FlowFilter { host: Some("STRIPE".into()), ..Default::default() }), 1);
assert_eq!(by(FlowFilter { host: Some("example.com".into()), ..Default::default() }), 2);
assert_eq!(by(FlowFilter { host: Some("charges".into()), ..Default::default() }), 0);
assert_eq!(by(FlowFilter { method: Some("get".into()), ..Default::default() }), 2);
assert_eq!(by(FlowFilter { status: Some("404".into()), ..Default::default() }), 1);
assert_eq!(by(FlowFilter { status: Some("2xx".into()), ..Default::default() }), 1);
assert_eq!(by(FlowFilter { status: Some("5XX".into()), ..Default::default() }), 0);
assert_eq!(by(FlowFilter { status: Some("error".into()), ..Default::default() }), 1);
}
#[test]
fn after_follows_flows_as_they_are_recorded() {
use yaak_database::UpdateSource;
let dir = tempfile::tempdir().unwrap();
let db = ProxyQueryManager::new(&dir.path().join("proxy.db")).unwrap();
let at = |s| chrono::DateTime::from_timestamp(s, 0).unwrap().naive_utc();
// A slow flow that started first but ended, and so was written, after a quick one.
for (id, started) in [("fl_quick", 20), ("fl_slow", 10), ("fl_next", 30)] {
let flow = Flow { id: id.into(), started_at: at(started), ..Default::default() };
db.with_tx(|tx| tx.upsert(&flow, &UpdateSource::Background)).unwrap();
}
let after = |id: &str| {
let filter = FlowFilter { after: Some(id.into()), ..Default::default() };
list(&db, &filter).unwrap().into_iter().map(|f| f.id).collect::<Vec<_>>()
};
assert_eq!(after("fl_quick"), vec!["fl_next", "fl_slow"]);
assert_eq!(after("fl_next"), Vec::<String>::new());
let missing = FlowFilter { after: Some("fl_gone".into()), ..Default::default() };
assert_eq!(list(&db, &missing).unwrap_err(), "No flow fl_gone");
}
#[test]
fn query_round_trip() {
let filter = FlowFilter {
host: Some("api.stripe.com".into()),
status: Some("4xx".into()),
limit: Some(20),
..Default::default()
};
assert_eq!(filter.to_query(), "host=api.stripe.com&status=4xx&limit=20");
assert_eq!(FlowFilter::from_query(&filter.to_query()).unwrap(), filter);
assert_eq!(FlowFilter::from_query("").unwrap(), FlowFilter::default());
assert!(FlowFilter::from_query("limit=lots").is_err());
}
#[test]
fn previews() {
let text = BodyPreview::new("héllo".as_bytes(), false);
assert_eq!(
(text.data.as_str(), text.encoding, text.size),
("héllo", BodyEncoding::Utf8, 6)
);
// Cut inside the "é".
let cut = BodyPreview::new(&"héllo".as_bytes()[..2], true);
assert_eq!(
(cut.data.as_str(), cut.encoding, cut.truncated),
("h", BodyEncoding::Utf8, true)
);
let binary = BodyPreview::new(&[0xff, 0x00, 0xfe], false);
assert_eq!(
(binary.data.as_str(), binary.encoding, binary.size),
("/wD+", BodyEncoding::Base64, 3)
);
}
}
+422
View File
@@ -0,0 +1,422 @@
//! HAR 1.2 export (http://www.softwareishard.com/blog/har-12-spec/). Bodies are decoded with
//! their Content-Encoding; one that is not UTF-8 is written in base64 with `encoding: "base64"`.
//! Fields the spec leaves to the exporter start with an underscore: `_truncated` on a body cut
//! at the capture limit and `_error` on a flow that failed.
use crate::bodies::{self, DecodedBody};
use crate::flows::{BodyEncoding, BodyPreview};
use crate::models::{Flow, FlowBody, ProxyHeader, find_header};
use serde::{Deserialize, Serialize};
use yaak_models::blob_manager::BlobManager;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Har {
pub log: HarLog,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HarLog {
pub version: String,
pub creator: HarCreator,
pub entries: Vec<HarEntry>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HarCreator {
pub name: String,
pub version: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HarEntry {
pub started_date_time: String,
pub time: i64,
pub request: HarRequest,
pub response: HarResponse,
pub cache: serde_json::Map<String, serde_json::Value>,
pub timings: HarTimings,
#[serde(rename = "serverIPAddress", skip_serializing_if = "Option::is_none")]
pub server_ip_address: Option<String>,
#[serde(rename = "_error", skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HarRequest {
pub method: String,
pub url: String,
pub http_version: String,
pub cookies: Vec<HarNameValue>,
pub headers: Vec<HarNameValue>,
pub query_string: Vec<HarNameValue>,
#[serde(skip_serializing_if = "Option::is_none")]
pub post_data: Option<HarPostData>,
pub headers_size: i64,
pub body_size: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HarResponse {
pub status: i32,
pub status_text: String,
pub http_version: String,
pub cookies: Vec<HarNameValue>,
pub headers: Vec<HarNameValue>,
pub content: HarContent,
#[serde(rename = "redirectURL")]
pub redirect_url: String,
pub headers_size: i64,
pub body_size: i64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct HarNameValue {
pub name: String,
pub value: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HarPostData {
pub mime_type: String,
pub text: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub encoding: Option<String>,
#[serde(
rename = "_truncated",
skip_serializing_if = "std::ops::Not::not",
default
)]
pub truncated: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HarContent {
pub size: i64,
pub mime_type: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub encoding: Option<String>,
#[serde(
rename = "_truncated",
skip_serializing_if = "std::ops::Not::not",
default
)]
pub truncated: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HarTimings {
pub send: i64,
pub wait: i64,
pub receive: i64,
}
/// A HAR of `flows`, reading each body from `blobs` and keeping at most `limit` bytes of it.
pub async fn export(blobs: &BlobManager, flows: &[Flow], limit: usize) -> Result<Har, String> {
let mut entries = Vec::with_capacity(flows.len());
for flow in flows {
let request = bodies::read_body(blobs, flow, FlowBody::Request, limit).await?;
let response = bodies::read_body(blobs, flow, FlowBody::Response, limit).await?;
entries.push(entry(flow, &request, &response));
}
Ok(har(entries))
}
pub fn har(entries: Vec<HarEntry>) -> Har {
Har {
log: HarLog {
version: "1.2".into(),
creator: HarCreator {
name: "Yaak Proxy".into(),
version: env!("CARGO_PKG_VERSION").into(),
},
entries,
},
}
}
pub fn entry(flow: &Flow, request_body: &DecodedBody, response_body: &DecodedBody) -> HarEntry {
let elapsed = flow.elapsed_ms.map_or(0, i64::from).max(0);
let wait = flow.elapsed_headers_ms.map_or(elapsed, i64::from).clamp(0, elapsed);
let responded = flow.status.is_some();
HarEntry {
started_date_time: flow.started_at.and_utc().format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(),
time: elapsed,
request: HarRequest {
method: flow.method.clone(),
url: flow.url.clone(),
http_version: flow.client_http_version.clone(),
cookies: Vec::new(),
headers: name_values(&flow.request_headers),
query_string: query_string(&flow.url),
post_data: post_data(flow, request_body),
headers_size: -1,
body_size: flow.request_size,
},
response: HarResponse {
status: flow.status.unwrap_or(0),
status_text: flow
.status
.and_then(|s| hyper::StatusCode::from_u16(u16::try_from(s).ok()?).ok())
.and_then(|s| s.canonical_reason())
.unwrap_or_default()
.to_string(),
http_version: flow.upstream_http_version.clone().unwrap_or_default(),
cookies: Vec::new(),
headers: name_values(&flow.response_headers),
content: content(flow, response_body),
redirect_url: find_header(&flow.response_headers, "location")
.unwrap_or_default()
.to_string(),
headers_size: -1,
body_size: if responded { flow.response_size } else { -1 },
},
cache: Default::default(),
timings: HarTimings { send: 0, wait, receive: elapsed - wait },
server_ip_address: flow.remote_addr.as_deref().map(ip_of),
error: flow
.error
.clone()
.or_else(|| response_body.decode_error.clone())
.or_else(|| request_body.decode_error.clone()),
}
}
fn name_values(headers: &[ProxyHeader]) -> Vec<HarNameValue> {
headers.iter().map(|h| HarNameValue { name: h.name.clone(), value: h.value.clone() }).collect()
}
fn query_string(url: &str) -> Vec<HarNameValue> {
let Some((_, query)) = url.split_once('?') else {
return Vec::new();
};
let query = query.split('#').next().unwrap_or_default();
url::form_urlencoded::parse(query.as_bytes())
.map(|(name, value)| HarNameValue { name: name.into_owned(), value: value.into_owned() })
.collect()
}
fn post_data(flow: &Flow, body: &DecodedBody) -> Option<HarPostData> {
if body.data.is_empty() && !body.truncated {
return None;
}
let preview = BodyPreview::new(&body.data, body.truncated);
Some(HarPostData {
mime_type: find_header(&flow.request_headers, "content-type").unwrap_or_default().into(),
encoding: base64_marker(preview.encoding),
text: preview.data,
truncated: body.truncated,
})
}
fn content(flow: &Flow, body: &DecodedBody) -> HarContent {
let mime_type = find_header(&flow.response_headers, "content-type").unwrap_or_default().into();
if body.data.is_empty() && !body.truncated {
return HarContent { size: 0, mime_type, text: None, encoding: None, truncated: false };
}
let preview = BodyPreview::new(&body.data, body.truncated);
HarContent {
size: body.data.len() as i64,
mime_type,
encoding: base64_marker(preview.encoding),
text: Some(preview.data),
truncated: body.truncated,
}
}
fn base64_marker(encoding: BodyEncoding) -> Option<String> {
(encoding == BodyEncoding::Base64).then(|| "base64".to_string())
}
/// The address without its port: `127.0.0.1:443` to `127.0.0.1`, `[::1]:443` to `::1`.
fn ip_of(addr: &str) -> String {
match addr.parse::<std::net::SocketAddr>() {
Ok(a) => a.ip().to_string(),
Err(_) => addr.to_string(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::NaiveDate;
const GOLDEN: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/export.har");
fn header(name: &str, value: &str) -> ProxyHeader {
ProxyHeader { name: name.into(), value: value.into() }
}
fn body(data: &[u8], truncated: bool) -> DecodedBody {
DecodedBody { data: data.to_vec(), stored_size: data.len(), truncated, decode_error: None }
}
fn at(second: u32) -> chrono::NaiveDateTime {
NaiveDate::from_ymd_opt(2026, 10, 26)
.unwrap()
.and_hms_milli_opt(9, 30, second, 250)
.unwrap()
}
fn flows() -> Vec<(Flow, DecodedBody, DecodedBody)> {
let json = Flow {
id: "fl_json".into(),
started_at: at(1),
method: "GET".into(),
url:
"https://api.github.com/repos/mountain-loop/yaak/issues?state=open&labels=bug%20fix"
.into(),
client_http_version: "HTTP/1.1".into(),
upstream_http_version: Some("HTTP/2".into()),
remote_addr: Some("140.82.112.6:443".into()),
status: Some(200),
content_type: Some("application/json; charset=utf-8".into()),
request_headers: vec![
header("Host", "api.github.com"),
header("Accept", "application/vnd.github+json"),
],
response_headers: vec![
header("content-type", "application/json; charset=utf-8"),
header("content-encoding", "gzip"),
],
response_size: 41,
elapsed_headers_ms: Some(120),
elapsed_ms: Some(134),
..Default::default()
};
let form = Flow {
id: "fl_form".into(),
started_at: at(2),
method: "POST".into(),
url: "http://localhost:3000/login".into(),
client_http_version: "HTTP/1.1".into(),
upstream_http_version: Some("HTTP/1.1".into()),
remote_addr: Some("[::1]:3000".into()),
status: Some(302),
request_headers: vec![header("Content-Type", "application/x-www-form-urlencoded")],
response_headers: vec![header("Location", "/dashboard")],
request_size: 26,
elapsed_headers_ms: Some(8),
elapsed_ms: Some(8),
..Default::default()
};
let image = Flow {
id: "fl_image".into(),
started_at: at(3),
method: "GET".into(),
url: "https://images.example.com/logo.png".into(),
client_http_version: "HTTP/1.1".into(),
upstream_http_version: Some("HTTP/2".into()),
status: Some(200),
response_headers: vec![header("Content-Type", "image/png")],
response_size: 8,
elapsed_headers_ms: Some(30),
elapsed_ms: Some(31),
..Default::default()
};
let cut = Flow {
id: "fl_cut".into(),
started_at: at(4),
method: "GET".into(),
url: "https://cdn.example.com/big.txt".into(),
client_http_version: "HTTP/1.1".into(),
upstream_http_version: Some("HTTP/1.1".into()),
status: Some(200),
response_headers: vec![header("Content-Type", "text/plain")],
response_size: 20_000_000,
response_truncated: true,
elapsed_headers_ms: Some(10),
elapsed_ms: Some(900),
error: Some("connection reset by peer".into()),
..Default::default()
};
let failed = Flow {
id: "fl_failed".into(),
started_at: at(5),
method: "GET".into(),
url: "https://down.example.com/".into(),
client_http_version: "HTTP/1.1".into(),
elapsed_ms: Some(3),
error: Some("Connection refused (os error 61)".into()),
..Default::default()
};
vec![
(json, body(b"", false), body(br#"[{"number":761,"title":"MCP"}]"#, false)),
(form, body(b"user=greg&password=hunter2", false), body(b"", false)),
(
image,
body(b"", false),
body(&[0x89, b'P', b'N', b'G', 0x0d, 0x0a, 0x1a, 0x0a], false),
),
(cut, body(b"", false), body(b"0123456789", true)),
(failed, body(b"", false), body(b"", false)),
]
}
#[tokio::test]
async fn export_keeps_a_body_that_fails_to_decode() {
use std::io::Write;
use yaak_models::blob_manager::BodyChunkWriter;
// Text that does not compress to nothing, gzipped, with the stream cut in half.
let mut seed = 7u32;
let text: Vec<u8> = (0..8000)
.map(|_| {
seed = seed.wrapping_mul(1_103_515_245).wrapping_add(12_345);
b'a' + (seed >> 16) as u8 % 26
})
.collect();
let mut gz = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
gz.write_all(&text).unwrap();
let gz = gz.finish().unwrap();
let (_db, blobs, _rx) = yaak_models::init_in_memory().unwrap();
let mut writer =
BodyChunkWriter::new(blobs.clone(), Flow::body_id("fl_bad", FlowBody::Response));
writer.write(&gz[..gz.len() / 2]).unwrap();
writer.finish().unwrap();
let flow = Flow {
id: "fl_bad".into(),
method: "GET".into(),
url: "https://cdn.example.com/a.txt".into(),
status: Some(200),
response_headers: vec![
header("Content-Type", "text/plain"),
header("Content-Encoding", "gzip"),
],
..Default::default()
};
let har = export(&blobs, &[flow], 1024 * 1024).await.unwrap();
let entry = &har.log.entries[0];
let text_kept = entry.response.content.text.as_deref().unwrap();
assert!(!text_kept.is_empty() && text.starts_with(text_kept.as_bytes()));
assert!(entry.response.content.truncated);
assert!(
entry.error.as_deref().unwrap().starts_with("Failed to decode"),
"{:?}",
entry.error
);
}
#[test]
fn golden_file() {
let entries = flows().iter().map(|(f, req, res)| entry(f, req, res)).collect();
let mut har = har(entries);
har.log.creator.version = "test".into();
let actual = serde_json::to_string_pretty(&har).unwrap() + "\n";
if std::env::var_os("UPDATE_GOLDEN").is_some() {
std::fs::write(GOLDEN, &actual).unwrap();
}
let expected = std::fs::read_to_string(GOLDEN).unwrap();
assert_eq!(actual, expected, "rerun with UPDATE_GOLDEN=1 to accept a change");
}
}
+331 -9
View File
@@ -1,23 +1,35 @@
pub mod actions;
pub mod bodies;
pub mod ca;
pub mod client;
pub mod control;
pub mod convert;
pub mod db;
pub mod flows;
pub mod har;
pub mod macos;
pub mod models;
pub mod save;
pub mod settings;
use crate::actions::{ActionInvocation, ActionMetadata, GlobalAction};
use crate::ca::CaStatus;
use crate::control::{ControlFile, ControlInfo, Owner};
use crate::db::ProxyQueryManager;
use crate::macos::{SystemProxyBackend, SystemProxyStatus};
use crate::models::{Flow, FlowBody, FlowIden, ModelPayload, ProxyHeader, find_header};
use crate::save::{ClientDb, SaveFlowRequest, SaveFlowResponse, WorkspaceSummary};
use crate::settings::ProxySettings;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
use chrono::{DateTime, Utc};
use log::warn;
use log::{debug, warn};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::net::{SocketAddr, TcpStream};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use ts_rs::TS;
use yaak_database::{ModelChangeEvent, UpdateSource, UpsertModelInfo};
use yaak_models::blob_manager::{BlobManager, BodyChunkWriter};
@@ -31,6 +43,13 @@ use yaak_rpc::{RpcError, RpcEventEmitter, define_rpc};
pub struct ProxyCtx {
handle: Arc<Mutex<Option<ProxyHandle>>>,
data_dir: PathBuf,
owner: Owner,
client_data_dir: PathBuf,
control: Arc<Mutex<Option<ControlFile>>>,
/// Held through each system proxy change, password dialog included, so they apply in turn.
system_proxy_turn: Arc<Mutex<()>>,
/// The port the engine last ran on, which the system proxy may still point at.
last_port: Arc<Mutex<Option<u16>>>,
pub db: ProxyQueryManager,
pub blobs: BlobManager,
pub events: RpcEventEmitter,
@@ -46,27 +65,72 @@ impl ProxyCtx {
Ok(Self {
handle: Arc::new(Mutex::new(None)),
data_dir: data_dir.to_path_buf(),
owner: Owner::App,
client_data_dir: save::client_data_dir(),
control: Arc::new(Mutex::new(None)),
system_proxy_turn: Arc::new(Mutex::new(())),
last_port: Arc::new(Mutex::new(None)),
db,
blobs,
events,
})
}
/// The process this runs in, which decides how the system proxy gets admin rights.
pub fn with_owner(mut self, owner: Owner) -> Self {
self.owner = owner;
self
}
/// Yaak's data directory, where flows are saved as requests.
pub fn with_client_data_dir(mut self, dir: impl Into<PathBuf>) -> Self {
self.client_data_dir = dir.into();
self
}
pub fn data_dir(&self) -> &Path {
&self.data_dir
}
pub fn owner(&self) -> Owner {
self.owner
}
pub fn settings_path(&self) -> PathBuf {
self.data_dir.join("settings.json")
}
pub fn settings(&self) -> Result<ProxySettings, String> {
ProxySettings::load(&self.settings_path())
}
pub fn capture_body_limit(&self) -> Result<usize, String> {
Ok(self.settings()?.capture_body_limit)
}
/// Starts the engine with the current `settings.json`. Does nothing if it is running.
pub fn start(&self) -> Result<(), String> {
self.start_on(None)
}
/// Starts the engine on `port`, or the one in `settings.json`. Does nothing if it is
/// running. Points the system proxy at it when the settings ask for that. Refused while
/// another process owns the engine, so there are never two.
pub fn start_on(&self, port: Option<u16>) -> Result<(), String> {
if let Some(owner) = control::live_owner(&self.data_dir)
&& owner.pid != std::process::id()
{
return Err(control::already_owned(&owner));
}
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"))
let settings = self.settings()?;
let ca = CaStore::load_or_create(&self.ca_dir())
.map_err(|e| format!("Failed to load proxy CA: {e}"))?;
let mut config = ProxyConfig::new(settings.port, Arc::new(ca));
let mut config = ProxyConfig::new(port.unwrap_or(settings.port), Arc::new(ca));
config.capture_body_limit = settings.capture_body_limit;
let mut proxy_handle = yaak_proxy::start_proxy(config)?;
@@ -77,18 +141,63 @@ impl ProxyCtx {
std::thread::spawn(move || run_event_loop(event_rx, db, blobs, events));
}
let bound = proxy_handle.port;
*handle = Some(proxy_handle);
drop(handle);
self.set_control_proxy_port(Some(bound));
if let Ok(mut last_port) = self.last_port.lock() {
*last_port = Some(bound);
}
if settings.system_proxy {
// The password dialog can sit open for a while; the engine is up regardless.
let ctx = self.clone();
std::thread::spawn(move || {
if let Err(e) = ctx.change_system_proxy(true, bound, true) {
warn!("Failed to turn on the system proxy: {e}");
}
});
}
self.events.emit("proxy_state_changed", &ProxyStatePayload { state: ProxyState::Running });
Ok(())
}
/// Stops the engine and turns off a system proxy that points at it.
pub fn stop(&self) -> Result<(), String> {
let mut handle = self.handle.lock().map_err(|_| "lock poisoned".to_string())?;
handle.take();
let port = self.take_handle()?;
if let Some(port) = port {
let ctx = self.clone();
std::thread::spawn(move || {
if let Err(e) = ctx.change_system_proxy(false, port, true) {
warn!("Failed to turn off the system proxy: {e}");
}
});
}
self.events.emit("proxy_state_changed", &ProxyStatePayload { state: ProxyState::Stopped });
Ok(())
}
/// [`Self::stop`] for a process about to exit: waits for the system proxy to be off, also
/// when a `stop` just before is still turning it off.
pub fn shutdown(&self) {
let running = self.take_handle().ok().flatten();
let last = self.last_port.lock().ok().and_then(|p| *p);
if let Some(port) = running.or(last)
&& let Err(e) = self.change_system_proxy(false, port, false)
{
warn!("Failed to turn off the system proxy: {e}");
}
}
/// Stops the engine, returning the port it was on.
fn take_handle(&self) -> Result<Option<u16>, String> {
let mut handle = self.handle.lock().map_err(|_| "lock poisoned".to_string())?;
// Dropping the handle stops the listener and joins its thread.
let port = handle.take().map(|h| h.port);
drop(handle);
self.set_control_proxy_port(None);
Ok(port)
}
/// 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)
@@ -97,6 +206,135 @@ impl ProxyCtx {
pub fn ca_dir(&self) -> PathBuf {
self.data_dir.join("ca")
}
pub(crate) fn set_control_file(&self, file: Option<ControlFile>) {
if let Ok(mut control) = self.control.lock() {
*control = file;
}
}
/// This process's `control.json`, while it serves the control API.
pub fn control_info(&self) -> Option<ControlInfo> {
self.control.lock().ok()?.as_ref().map(|c| c.info.clone())
}
fn set_control_proxy_port(&self, port: Option<u16>) {
if let Ok(mut control) = self.control.lock()
&& let Some(control) = control.as_mut()
{
control.set_proxy_port(port);
}
}
pub fn system_proxy_backend(&self) -> SystemProxyBackend {
match self.owner {
Owner::App => SystemProxyBackend::AdminPrompt,
Owner::Cli => SystemProxyBackend::Sudo,
}
}
/// Judged against the running engine's port, or the configured one while it is stopped.
/// Blocks.
pub fn system_proxy_status(&self) -> Result<SystemProxyStatus, String> {
let port = self.port().unwrap_or(self.settings()?.port);
macos::status(port, self.system_proxy_backend())
}
/// Turns the system proxy on or off and remembers the choice. Turned on while the engine
/// is stopped, it waits for the engine to start, so the Mac is never pointed at nothing.
/// Blocks, through the password dialog.
pub fn set_system_proxy(&self, enabled: bool) -> Result<SystemProxyStatus, String> {
let mut settings = self.settings()?;
let previous = settings.system_proxy;
settings.system_proxy = enabled;
settings.save(&self.settings_path())?;
let port = self.port().unwrap_or(settings.port);
if let Err(e) = self.change_system_proxy(enabled, port, enabled) {
settings.system_proxy = previous;
let _ = settings.save(&self.settings_path());
return Err(e);
}
self.system_proxy_status()
}
/// Points the system proxy at `port`, or turns it off there, once the changes before it
/// are done. With `follow_engine`, the change is dropped if the engine has moved on by its
/// turn: an enable once the engine is off `port` or the setting is off, a disable once the
/// engine is back on `port`. A start and a stop in quick succession then end in the state
/// the engine is in, whichever password dialog is answered first.
fn change_system_proxy(
&self,
enabled: bool,
port: u16,
follow_engine: bool,
) -> Result<(), String> {
let _turn = self.system_proxy_turn.lock().map_err(|_| "lock poisoned".to_string())?;
if follow_engine {
let on_port = self.port() == Some(port);
let current = if enabled { on_port && self.settings()?.system_proxy } else { !on_port };
if !current {
return Ok(());
}
}
macos::set(enabled, port, self.system_proxy_backend())
}
/// Turns off a system proxy left pointing at one of `ports` by an owner that crashed, so
/// the Mac is not left offline. Leaves it alone when something is serving on the port.
/// Blocks.
pub fn reset_stale_system_proxy(&self, ports: &[u16]) -> Result<(), String> {
let mut ports = ports.to_vec();
ports.sort_unstable();
ports.dedup();
for port in ports {
if port == 0 {
continue;
}
// Checked in its turn, so an engine started on the port meanwhile keeps its proxy.
let _turn = self.system_proxy_turn.lock().map_err(|_| "lock poisoned".to_string())?;
if Some(port) == self.port() || is_listening(port) {
continue;
}
// Covers disabled services too, and asks for nothing when none point at the port.
debug!("Turning off any stale system proxy on port {port}");
macos::set(false, port, self.system_proxy_backend())?;
}
Ok(())
}
/// Yaak's workspaces and their folders, to save a flow into.
pub fn workspaces(&self) -> Result<Vec<WorkspaceSummary>, String> {
ClientDb::open(&self.client_data_dir)?.workspaces()
}
/// Saves a flow as a request in Yaak and records the request on the flow.
pub async fn save_flow(
&self,
flow_id: &str,
req: &SaveFlowRequest,
) -> Result<SaveFlowResponse, String> {
let flow = flows::get(&self.db, flow_id)?;
let limit = self.capture_body_limit()?;
let body = bodies::read_body(&self.blobs, &flow, FlowBody::Request, limit).await?;
let saved = ClientDb::open(&self.client_data_dir)?.save(&flow, &body, req)?;
let flow = Flow { saved_request_id: Some(saved.request.id.clone()), ..flow };
let (flow, created) = self
.db
.with_tx(|tx| tx.upsert(&flow, &UpdateSource::Background))
.map_err(|e| format!("Failed to record the saved request on the flow: {e}"))?;
self.events.emit(
"model_write",
&ModelPayload { model: flow, change: ModelChangeEvent::Upsert { created } },
);
Ok(saved)
}
}
/// Something accepts connections on `127.0.0.1:<port>`.
fn is_listening(port: u16) -> bool {
let addr = SocketAddr::from(([127, 0, 0, 1], port));
TcpStream::connect_timeout(&addr, Duration::from_millis(200)).is_ok()
}
// -- Proxy state --
@@ -214,6 +452,16 @@ pub struct CaStatusRequest {}
#[ts(export, export_to = "gen_rpc.ts")]
pub struct CaTrustRequest {}
#[derive(Deserialize, TS)]
#[ts(export, export_to = "gen_rpc.ts")]
pub struct SystemProxyStatusRequest {}
#[derive(Deserialize, TS)]
#[ts(export, export_to = "gen_rpc.ts")]
pub struct SystemProxySetRequest {
pub enabled: bool,
}
// -- Handlers --
fn execute_action(ctx: &ProxyCtx, invocation: ActionInvocation) -> Result<bool, RpcError> {
@@ -245,7 +493,7 @@ async fn get_flow_body(
) -> Result<GetFlowBodyResponse, RpcError> {
let flow = find_flow(&ctx, &req.flow_id)?;
let limit = settings(&ctx)?.capture_body_limit;
let body = bodies::read_body(&ctx.blobs, &flow, req.which, limit).await.map_err(rpc_error)?;
let body = viewer_body(&ctx, &flow, req.which, limit).await?;
Ok(GetFlowBodyResponse {
data: BASE64.encode(&body.data),
truncated: body.truncated,
@@ -253,13 +501,28 @@ async fn get_flow_body(
})
}
/// A body for the app's viewer, which shows a body that failed to decode as its error, as it
/// always has, rather than as what decoded before the failure.
async fn viewer_body(
ctx: &ProxyCtx,
flow: &Flow,
which: FlowBody,
limit: usize,
) -> Result<bodies::DecodedBody, RpcError> {
let body = bodies::read_body(&ctx.blobs, flow, which, limit).await.map_err(rpc_error)?;
match body.decode_error {
Some(e) => Err(rpc_error(e)),
None => Ok(body),
}
}
async fn filter_body(
ctx: ProxyCtx,
req: FilterBodyRequest,
) -> Result<FilterBodyResponse, RpcError> {
let flow = find_flow(&ctx, &req.flow_id)?;
let limit = settings(&ctx)?.capture_body_limit;
let body = bodies::read_body(&ctx.blobs, &flow, req.which, limit).await.map_err(rpc_error)?;
let body = viewer_body(&ctx, &flow, req.which, limit).await?;
let text = String::from_utf8_lossy(&body.data);
let language = match req.language {
Some(FilterLanguage::JsonPath) => yaak_filter::Language::JsonPath,
@@ -285,8 +548,28 @@ async fn ca_trust(ctx: ProxyCtx, _req: CaTrustRequest) -> Result<CaStatus, RpcEr
ca::status(ca.cert_path()).await.map_err(rpc_error)
}
async fn system_proxy_status(
ctx: ProxyCtx,
_req: SystemProxyStatusRequest,
) -> Result<SystemProxyStatus, RpcError> {
blocking(move || ctx.system_proxy_status()).await
}
async fn system_proxy_set(
ctx: ProxyCtx,
req: SystemProxySetRequest,
) -> Result<SystemProxyStatus, RpcError> {
blocking(move || ctx.set_system_proxy(req.enabled)).await
}
async fn blocking<T: Send + 'static>(
f: impl FnOnce() -> Result<T, String> + Send + 'static,
) -> Result<T, RpcError> {
tokio::task::spawn_blocking(f).await.map_err(rpc_error)?.map_err(rpc_error)
}
fn settings(ctx: &ProxyCtx) -> Result<ProxySettings, RpcError> {
ProxySettings::load(&ctx.settings_path()).map_err(rpc_error)
ctx.settings().map_err(rpc_error)
}
fn find_flow(ctx: &ProxyCtx, flow_id: &str) -> Result<Flow, RpcError> {
@@ -470,9 +753,48 @@ define_rpc! {
ca_trust(CaTrustRequest) -> CaStatus,
filter_body(FilterBodyRequest) -> FilterBodyResponse,
get_flow_body(GetFlowBodyRequest) -> GetFlowBodyResponse,
system_proxy_set(SystemProxySetRequest) -> SystemProxyStatus,
system_proxy_status(SystemProxyStatusRequest) -> SystemProxyStatus,
}
events {
model_write(ModelPayload),
proxy_state_changed(ProxyStatePayload),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn the_app_still_hears_why_a_body_failed_to_decode() {
let dir = tempfile::tempdir().unwrap();
let (events, _rx) = RpcEventEmitter::new();
let ctx = ProxyCtx::new(dir.path(), events).unwrap();
let flow = Flow {
id: "fl_bad".into(),
status: Some(200),
response_headers: vec![ProxyHeader {
name: "Content-Encoding".into(),
value: "gzip".into(),
}],
..Default::default()
};
ctx.db.with_tx(|tx| tx.upsert(&flow, &UpdateSource::Background)).unwrap();
let mut writer =
BodyChunkWriter::new(ctx.blobs.clone(), Flow::body_id("fl_bad", FlowBody::Response));
writer.write(b"not a gzip stream").unwrap();
writer.finish().unwrap();
let req = GetFlowBodyRequest { flow_id: "fl_bad".into(), which: FlowBody::Response };
let err = get_flow_body(ctx.clone(), req).await.err().unwrap();
assert!(err.message.starts_with("Failed to decode"), "{}", err.message);
let req = FilterBodyRequest {
flow_id: "fl_bad".into(),
which: FlowBody::Response,
expression: "$".into(),
language: None,
};
assert!(filter_body(ctx, req).await.is_err());
}
}
+334
View File
@@ -0,0 +1,334 @@
//! The macOS system proxy: `networksetup` pointing the web and secure web proxies of every
//! enabled network service at `127.0.0.1:<port>`. Reading needs no rights; changing needs an
//! administrator, so changes run as one script through the backend in use. Everything here
//! blocks, the password dialog included, so async callers go through `spawn_blocking`.
use serde::{Deserialize, Serialize};
use ts_rs::TS;
/// How a change gets administrator rights.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export, export_to = "gen_rpc.ts")]
pub enum SystemProxyBackend {
/// The privileged helper, approved once in System Settings (not built yet).
Helper,
/// The standard macOS password dialog, through `osascript`.
AdminPrompt,
/// `sudo` in the terminal running the CLI daemon.
Sudo,
}
#[derive(Debug, Clone, Serialize, Deserialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export, export_to = "gen_rpc.ts")]
pub struct SystemProxyStatus {
/// Every enabled network service sends HTTP and HTTPS to `127.0.0.1:<port>`.
pub enabled: bool,
/// The port checked against.
pub port: u16,
pub services: Vec<ServiceProxy>,
pub backend: SystemProxyBackend,
/// This platform has a system proxy the app can set.
pub supported: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export, export_to = "gen_rpc.ts")]
pub struct ServiceProxy {
/// The network service, as System Settings names it: "Wi-Fi", "USB 10/100/1000 LAN".
pub name: String,
pub http: ProxySetting,
pub https: ProxySetting,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export, export_to = "gen_rpc.ts")]
pub struct ProxySetting {
pub enabled: bool,
pub server: String,
pub port: u16,
}
const HOST: &str = "127.0.0.1";
impl ProxySetting {
fn points_at(&self, port: u16) -> bool {
self.enabled && self.server == HOST && self.port == port
}
}
impl ServiceProxy {
/// Either proxy of the service sends traffic to `127.0.0.1:<port>`.
pub fn points_at(&self, port: u16) -> bool {
self.http.points_at(port) || self.https.points_at(port)
}
}
/// The services from `networksetup -listallnetworkservices`, the disabled ones too with
/// `include_disabled`. The first line is a legend; a leading `*` marks a disabled service.
#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
fn parse_services(output: &str, include_disabled: bool) -> Vec<String> {
output
.lines()
.skip(1)
.map(str::trim)
.filter(|l| !l.is_empty() && (include_disabled || !l.starts_with('*')))
.map(|l| l.trim_start_matches('*').to_string())
.collect()
}
/// The output of `networksetup -getwebproxy` or `-getsecurewebproxy`.
#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
fn parse_proxy(output: &str) -> ProxySetting {
let mut setting = ProxySetting::default();
for line in output.lines() {
let Some((key, value)) = line.split_once(':') else {
continue;
};
let value = value.trim();
match key.trim() {
"Enabled" => setting.enabled = value.eq_ignore_ascii_case("yes"),
"Server" => setting.server = value.to_string(),
"Port" => setting.port = value.parse().unwrap_or_default(),
_ => {}
}
}
setting
}
/// The `networksetup` calls that point `services` at `port`, or turn off those pointing at it.
#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
fn commands(services: &[ServiceProxy], enabled: bool, port: u16) -> Vec<Vec<String>> {
let mut commands = Vec::new();
for service in services {
let name = service.name.clone();
if enabled {
let port = port.to_string();
for (set, state) in [
("-setwebproxy", "-setwebproxystate"),
("-setsecurewebproxy", "-setsecurewebproxystate"),
] {
commands.push(vec![set.into(), name.clone(), HOST.into(), port.clone()]);
commands.push(vec![state.into(), name.clone(), "on".into()]);
}
} else {
if service.http.points_at(port) {
commands.push(vec!["-setwebproxystate".into(), name.clone(), "off".into()]);
}
if service.https.points_at(port) {
commands.push(vec!["-setsecurewebproxystate".into(), name, "off".into()]);
}
}
}
commands
}
/// The calls as one `sh` script that stops at the first failure.
#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
fn script(commands: &[Vec<String>]) -> String {
commands
.iter()
.map(|args| {
std::iter::once(NETWORKSETUP.to_string())
.chain(args.iter().map(|a| shell_quote(a)))
.collect::<Vec<_>>()
.join(" ")
})
.collect::<Vec<_>>()
.join(" && ")
}
#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
fn shell_quote(s: &str) -> String {
format!("'{}'", s.replace('\'', r"'\''"))
}
/// `script` inside an AppleScript string literal.
#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
fn applescript_string(s: &str) -> String {
format!("\"{}\"", s.replace('\\', r"\\").replace('"', "\\\""))
}
#[cfg_attr(not(target_os = "macos"), allow(dead_code))]
const NETWORKSETUP: &str = "/usr/sbin/networksetup";
#[cfg(target_os = "macos")]
mod imp {
use super::*;
use std::process::{Command, Output, Stdio};
fn networksetup(args: &[&str]) -> Result<String, String> {
let output = Command::new(NETWORKSETUP)
.args(args)
.output()
.map_err(|e| format!("Failed to run {NETWORKSETUP}: {e}"))?;
check(&output, NETWORKSETUP)?;
Ok(String::from_utf8_lossy(&output.stdout).into_owned())
}
fn check(output: &Output, program: &str) -> Result<(), String> {
if output.status.success() {
return Ok(());
}
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
Err(if stderr.is_empty() {
format!("{program} exited with {}", output.status)
} else {
stderr
})
}
pub fn services(include_disabled: bool) -> Result<Vec<ServiceProxy>, String> {
parse_services(&networksetup(&["-listallnetworkservices"])?, include_disabled)
.into_iter()
.map(|name| {
let http = parse_proxy(&networksetup(&["-getwebproxy", &name])?);
let https = parse_proxy(&networksetup(&["-getsecurewebproxy", &name])?);
Ok(ServiceProxy { name, http, https })
})
.collect()
}
pub fn run_privileged(script: &str, backend: SystemProxyBackend) -> Result<(), String> {
match backend {
SystemProxyBackend::Helper => {
Err("The privileged helper is not available yet".to_string())
}
SystemProxyBackend::AdminPrompt => {
let source = format!(
"do shell script {} with prompt {} with administrator privileges",
applescript_string(script),
applescript_string("Yaak Proxy wants to change your system proxy settings."),
);
let output = Command::new("/usr/bin/osascript")
.args(["-e", &source])
.output()
.map_err(|e| format!("Failed to run osascript: {e}"))?;
// -128 is the dialog's Cancel button.
if String::from_utf8_lossy(&output.stderr).contains("(-128)") {
return Err("Cancelled in the password dialog".to_string());
}
check(&output, "osascript")
}
SystemProxyBackend::Sudo => {
// sudo asks on the terminal, so the daemon's own stdin stays attached.
let output = Command::new("/usr/bin/sudo")
.args(["/bin/sh", "-c", script])
.stdin(Stdio::inherit())
.output()
.map_err(|e| format!("Failed to run sudo: {e}"))?;
check(&output, "sudo")
}
}
}
}
/// Which services point where, judged against `port`.
#[cfg(target_os = "macos")]
pub fn status(port: u16, backend: SystemProxyBackend) -> Result<SystemProxyStatus, String> {
let services = imp::services(false)?;
let enabled = !services.is_empty()
&& services.iter().all(|s| s.http.points_at(port) && s.https.points_at(port));
Ok(SystemProxyStatus { enabled, port, services, backend, supported: true })
}
/// Points every enabled service at `127.0.0.1:<port>`, or turns off the services that point
/// there, disabled ones included, leaving other proxies alone. Turning off with nothing
/// pointing at `port` does nothing.
#[cfg(target_os = "macos")]
pub fn set(enabled: bool, port: u16, backend: SystemProxyBackend) -> Result<(), String> {
if port == 0 {
return Err("The system proxy needs a port".to_string());
}
// A service disabled since it was pointed at us would point at us again once re-enabled.
let commands = commands(&imp::services(!enabled)?, enabled, port);
if commands.is_empty() {
return Ok(());
}
imp::run_privileged(&script(&commands), backend)
}
#[cfg(not(target_os = "macos"))]
pub fn status(port: u16, backend: SystemProxyBackend) -> Result<SystemProxyStatus, String> {
Ok(SystemProxyStatus { enabled: false, port, services: Vec::new(), backend, supported: false })
}
#[cfg(not(target_os = "macos"))]
pub fn set(enabled: bool, _port: u16, _backend: SystemProxyBackend) -> Result<(), String> {
if !enabled {
return Ok(());
}
Err("Setting the system proxy is only supported on macOS".to_string())
}
#[cfg(test)]
mod tests {
use super::*;
fn setting(enabled: bool, server: &str, port: u16) -> ProxySetting {
ProxySetting { enabled, server: server.into(), port }
}
#[test]
fn parses_services() {
let output = "An asterisk (*) denotes that a network service is disabled.\nWi-Fi\n*Thunderbolt Bridge\nUSB 10/100/1000 LAN\n\n";
assert_eq!(parse_services(output, false), vec!["Wi-Fi", "USB 10/100/1000 LAN"]);
// A disabled service can still hold a proxy setting pointing at us.
assert_eq!(
parse_services(output, true),
vec!["Wi-Fi", "Thunderbolt Bridge", "USB 10/100/1000 LAN"]
);
}
#[test]
fn parses_proxy() {
let on = "Enabled: Yes\nServer: 127.0.0.1\nPort: 9090\nAuthenticated Proxy Enabled: 0\n";
assert_eq!(parse_proxy(on), setting(true, "127.0.0.1", 9090));
let off = "Enabled: No\nServer: \nPort: 0\nAuthenticated Proxy Enabled: 0\n";
assert_eq!(parse_proxy(off), setting(false, "", 0));
}
#[test]
fn turning_off_touches_only_our_proxies() {
let services = vec![
ServiceProxy {
name: "Wi-Fi".into(),
http: setting(true, "127.0.0.1", 9090),
https: setting(true, "127.0.0.1", 9090),
},
// Someone else's proxy, and ours on another port.
ServiceProxy {
name: "Ethernet".into(),
http: setting(true, "corp-proxy", 9090),
https: setting(true, "127.0.0.1", 8080),
},
];
assert!(services[0].points_at(9090) && !services[1].points_at(9090));
assert_eq!(
commands(&services, false, 9090),
vec![
vec!["-setwebproxystate", "Wi-Fi", "off"],
vec!["-setsecurewebproxystate", "Wi-Fi", "off"],
]
);
assert_eq!(commands(&services, true, 9090).len(), 8);
}
#[test]
fn quotes_service_names() {
let commands = vec![vec![
"-setwebproxystate".to_string(),
"Bob's \"LAN\"".to_string(),
"off".into(),
]];
let script = script(&commands);
assert_eq!(script, r#"/usr/sbin/networksetup '-setwebproxystate' 'Bob'\''s "LAN"' 'off'"#);
assert_eq!(
applescript_string(&script),
r#""/usr/sbin/networksetup '-setwebproxystate' 'Bob'\\''s \"LAN\"' 'off'""#
);
}
}
@@ -53,6 +53,8 @@ pub struct Flow {
/// Milliseconds from `started_at` to the end of the response.
pub elapsed_ms: Option<i32>,
pub error: Option<String>,
/// The id of the Yaak request this flow was last saved as.
pub saved_request_id: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
@@ -145,6 +147,7 @@ impl UpsertModelInfo for Flow {
(ElapsedHeadersMs, self.elapsed_headers_ms.into()),
(ElapsedMs, self.elapsed_ms.into()),
(Error, self.error.into()),
(SavedRequestId, self.saved_request_id.into()),
])
}
@@ -168,6 +171,7 @@ impl UpsertModelInfo for Flow {
FlowIden::ElapsedHeadersMs,
FlowIden::ElapsedMs,
FlowIden::Error,
FlowIden::SavedRequestId,
]
}
@@ -198,6 +202,11 @@ impl UpsertModelInfo for Flow {
elapsed_headers_ms: r.get("elapsed_headers_ms")?,
elapsed_ms: r.get("elapsed_ms")?,
error: r.get("error")?,
// Absent from a database read before its migrations have run.
saved_request_id: match r.get("saved_request_id") {
Err(rusqlite::Error::InvalidColumnName(_)) => None,
v => v?,
},
})
}
}
+140
View File
@@ -0,0 +1,140 @@
//! Saving a flow into the Yaak client's database. It is opened as a lifecycle guest, as the CLI
//! opens it, since Yaak may have it open too, and the request is written with
//! `UpdateSource::Sync` so a running Yaak picks it up from its model-change poller.
use crate::bodies::DecodedBody;
use crate::convert::{self, DroppedHeader};
use crate::models::Flow;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use yaak_models::models::HttpRequest;
use yaak_models::query_manager::QueryManager;
use yaak_models::util::UpdateSource;
/// Yaak's data directory, where its `db.sqlite` is. A debug build pairs with a debug Yaak.
pub fn client_data_dir() -> PathBuf {
let app_id = if cfg!(debug_assertions) { "app.yaak.desktop.dev" } else { "app.yaak.desktop" };
yaak_common::data_dir::resolve_data_dir(app_id)
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceSummary {
pub id: String,
pub name: String,
pub folders: Vec<FolderSummary>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FolderSummary {
pub id: String,
pub name: String,
/// The folder this one is in, if any.
pub folder_id: Option<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SaveFlowRequest {
pub workspace_id: String,
pub folder_id: Option<String>,
/// Defaults to the URL's path.
pub name: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SaveFlowResponse {
pub request: HttpRequest,
pub dropped_headers: Vec<DroppedHeader>,
pub warnings: Vec<String>,
}
pub struct ClientDb {
query_manager: QueryManager,
}
impl ClientDb {
/// Opens `db.sqlite` in Yaak's data directory `data_dir`. Fails when there is none, rather
/// than creating a database for a Yaak that has never run.
pub fn open(data_dir: &Path) -> Result<Self, String> {
let db_path = data_dir.join("db.sqlite");
if !db_path.exists() {
return Err(format!(
"Yaak's database was not found at {}. Install Yaak and open it once.",
db_path.display()
));
}
let (query_manager, blob_manager, _rx) =
yaak_models::init_standalone(&db_path, data_dir.join("blobs.sqlite"))
.map_err(|e| format!("Failed to open Yaak's database: {e}"))?;
query_manager
.with_tx(|tx| {
yaak_lifecycle::on_launch(&yaak_lifecycle::Host::guest(), tx, &blob_manager)
})
.map_err(|e| format!("Failed to open Yaak's database: {e}"))?;
Ok(Self { query_manager })
}
pub fn workspaces(&self) -> Result<Vec<WorkspaceSummary>, String> {
let db = self.query_manager.connect();
let workspaces = db.list_workspaces().map_err(|e| e.to_string())?;
workspaces
.into_iter()
.map(|w| {
let folders = db.list_folders(&w.id).map_err(|e| e.to_string())?;
Ok(WorkspaceSummary {
id: w.id,
name: w.name,
folders: folders
.into_iter()
.map(|f| FolderSummary { id: f.id, name: f.name, folder_id: f.folder_id })
.collect(),
})
})
.collect()
}
/// Writes `flow`, with `body` its decoded request body, as a new request.
pub fn save(
&self,
flow: &Flow,
body: &DecodedBody,
target: &SaveFlowRequest,
) -> Result<SaveFlowResponse, String> {
let db = self.query_manager.connect();
db.get_workspace(&target.workspace_id)
.map_err(|_| format!("No workspace {} in Yaak", target.workspace_id))?;
if let Some(folder_id) = &target.folder_id {
let folder =
db.get_folder(folder_id).map_err(|_| format!("No folder {folder_id} in Yaak"))?;
if folder.workspace_id != target.workspace_id {
return Err(format!(
"Folder {folder_id} is not in workspace {}",
target.workspace_id
));
}
}
let conversion = convert::flow_to_request(flow, body);
let mut request = conversion.request;
request.workspace_id = target.workspace_id.clone();
request.folder_id = target.folder_id.clone();
if let Some(name) = target.name.as_deref().map(str::trim).filter(|n| !n.is_empty()) {
request.name = name.to_string();
}
// First in its folder, as the client places a new request.
request.sort_priority = -(chrono::Utc::now().timestamp_millis() as f64);
let request = self
.query_manager
.with_tx(|tx| tx.upsert_http_request(&request, &UpdateSource::Sync))
.map_err(|e| format!("Failed to save the request to Yaak: {e}"))?;
Ok(SaveFlowResponse {
request,
dropped_headers: conversion.dropped_headers,
warnings: conversion.warnings,
})
}
}
@@ -0,0 +1,416 @@
use std::convert::Infallible;
use std::net::SocketAddr;
use std::path::Path;
use std::time::{Duration, Instant};
use bytes::Bytes;
use http_body_util::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_models::models::{Folder, Workspace};
use yaak_models::util::UpdateSource;
use yaak_proxy_lib::client::{ControlClient, ProxyConnection};
use yaak_proxy_lib::control::{self, ControlInfo, ControlServer, Owner};
use yaak_proxy_lib::flows::{BodyEncoding, FlowFilter};
use yaak_proxy_lib::models::{Flow, FlowBody};
use yaak_proxy_lib::save::SaveFlowRequest;
use yaak_proxy_lib::settings::ProxySettings;
use yaak_proxy_lib::{ProxyCtx, ProxyState};
use yaak_rpc::RpcEventEmitter;
async fn route(_req: Request<Incoming>) -> Result<Response<Full<Bytes>>, Infallible> {
Ok(Response::builder()
.header("content-type", "application/json")
.body(Full::from(r#"{"id":"ch_1","amount":42}"#))
.unwrap())
}
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
}
/// A context over `dir` whose engine picks a free port, saving into Yaak's data in `client`.
fn ctx(dir: &Path, client: &Path) -> ProxyCtx {
let settings = ProxySettings { port: 0, ..Default::default() };
settings.save(&dir.join("settings.json")).unwrap();
let (events, _rx) = RpcEventEmitter::new();
ProxyCtx::new(dir, events).unwrap().with_client_data_dir(client)
}
/// Sends one POST through the engine and waits for its flow.
async fn capture_flow(ctx: &ProxyCtx) -> Flow {
let upstream = spawn_upstream().await;
let url = format!("http://{upstream}/v1/charges?expand=customer");
let client = reqwest::Client::builder()
.proxy(reqwest::Proxy::http(format!("http://127.0.0.1:{}", ctx.port().unwrap())).unwrap())
.build()
.unwrap();
let resp = client
.post(&url)
.header("content-type", "application/json")
.header("authorization", "Bearer sk_test_123")
.body(r#"{"amount":42}"#)
.send()
.await
.unwrap();
assert_eq!(resp.text().await.unwrap(), r#"{"id":"ch_1","amount":42}"#);
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;
}
}
/// A process id that is certainly not running.
fn dead_pid() -> u32 {
let mut child = if cfg!(windows) {
std::process::Command::new("cmd").args(["/C", "exit"]).spawn().unwrap()
} else {
std::process::Command::new("true").spawn().unwrap()
};
let pid = child.id();
child.wait().unwrap();
pid
}
#[tokio::test(flavor = "multi_thread")]
async fn round_trip() {
let dir = tempfile::tempdir().unwrap();
let client_dir = tempfile::tempdir().unwrap();
let ctx = ctx(dir.path(), client_dir.path());
let server = ControlServer::start(ctx.clone()).unwrap();
let info = control::read(dir.path()).unwrap();
assert_eq!(&info, server.info());
assert_eq!((info.pid, info.owner, info.proxy_port), (std::process::id(), Owner::App, None));
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = std::fs::metadata(dir.path().join("control.json")).unwrap().permissions();
assert_eq!(mode.mode() & 0o777, 0o600);
}
// The token is required.
let http = reqwest::Client::builder().no_proxy().build().unwrap();
let url = format!("{}/status", info.base_url());
assert_eq!(http.get(&url).send().await.unwrap().status(), 401);
let wrong = http.get(&url).bearer_auth("not-the-token").send().await.unwrap();
assert_eq!(wrong.status(), 401);
let ProxyConnection::Live(client) = ProxyConnection::connect(dir.path()).unwrap() else {
panic!("expected the live owner");
};
let status = client.status().await.unwrap();
assert_eq!(
(status.state, status.proxy_port, status.flow_count),
(ProxyState::Stopped, None, 0)
);
let status = client.start(None).await.unwrap();
assert_eq!(status.state, ProxyState::Running);
assert_eq!(status.proxy_port, ctx.port());
assert_eq!(control::read(dir.path()).unwrap().proxy_port, ctx.port());
let flow = capture_flow(&ctx).await;
let flows = client.flows(&FlowFilter::default()).await.unwrap();
assert_eq!(flows.iter().map(|f| &f.id).collect::<Vec<_>>(), vec![&flow.id]);
let following = FlowFilter { after: Some(flow.id.clone()), ..Default::default() };
assert!(client.flows(&following).await.unwrap().is_empty());
let filter = FlowFilter { status: Some("4xx".into()), ..Default::default() };
assert!(client.flows(&filter).await.unwrap().is_empty());
let detail = client.flow(&flow.id, 8).await.unwrap();
assert_eq!(detail.flow.method, "POST");
let request_body = detail.request_body.unwrap();
assert_eq!((request_body.data.as_str(), request_body.truncated), (r#"{"amount"#, true));
assert_eq!(request_body.encoding, BodyEncoding::Utf8);
let response = client.body(&flow.id, FlowBody::Response).await.unwrap();
assert_eq!(response.data, br#"{"id":"ch_1","amount":42}"#);
assert!(!response.truncated);
let missing = client.flow("fl_missing", 8).await.unwrap_err();
assert_eq!(missing, "No flow fl_missing");
let har = client.har(&FlowFilter::default()).await.unwrap();
assert_eq!(har.log.entries.len(), 1);
assert_eq!(har.log.entries[0].request.post_data.as_ref().unwrap().text, r#"{"amount":42}"#);
assert!(client.ca_pem().await.unwrap().starts_with("-----BEGIN CERTIFICATE-----"));
let status = client.stop().await.unwrap();
assert_eq!(
(status.state, status.proxy_port, status.flow_count),
(ProxyState::Stopped, None, 1)
);
assert_eq!(control::read(dir.path()).unwrap().proxy_port, None);
drop(server);
assert!(!dir.path().join("control.json").exists());
assert!(control::live_owner(dir.path()).is_none());
}
#[tokio::test(flavor = "multi_thread")]
async fn one_owner_at_a_time() {
let dir = tempfile::tempdir().unwrap();
let client_dir = tempfile::tempdir().unwrap();
let ctx = ctx(dir.path(), client_dir.path());
let first = ControlServer::start(ctx.clone()).unwrap();
let err = ControlServer::start(ctx.clone()).err().unwrap();
assert!(err.contains("already running in Yaak Proxy"), "{err}");
assert_eq!(control::read(dir.path()).unwrap(), *first.info());
}
#[tokio::test(flavor = "multi_thread")]
async fn stale_control_file_is_replaced() {
let dir = tempfile::tempdir().unwrap();
let client_dir = tempfile::tempdir().unwrap();
let ctx = ctx(dir.path(), client_dir.path()).with_owner(Owner::Cli);
// A daemon that died without cleaning up, and one whose pid is alive (ours) but whose
// control port answers nothing, as after the OS reuses a pid.
let unused_port = TcpListener::bind("127.0.0.1:0").await.unwrap().local_addr().unwrap().port();
for pid in [dead_pid(), std::process::id()] {
let stale = ControlInfo {
pid,
control_port: unused_port,
proxy_port: Some(9090),
token: "stale".into(),
owner: Owner::Cli,
started_at: chrono::Utc::now(),
};
std::fs::write(dir.path().join("control.json"), serde_json::to_vec(&stale).unwrap())
.unwrap();
assert!(control::live_owner(dir.path()).is_none());
assert!(!dir.path().join("control.json").exists(), "stale file left for pid {pid}");
std::fs::write(dir.path().join("control.json"), serde_json::to_vec(&stale).unwrap())
.unwrap();
let server = ControlServer::start(ctx.clone()).unwrap();
let info = control::read(dir.path()).unwrap();
assert_eq!((info.pid, info.owner), (std::process::id(), Owner::Cli));
assert_ne!(info.token, "stale");
drop(server);
}
// A file that does not parse is never a claim in progress, since claims are published whole.
std::fs::write(dir.path().join("control.json"), "{\"pid\": 12").unwrap();
let server = ControlServer::start(ctx.clone()).unwrap();
assert_eq!(control::read(dir.path()).unwrap(), *server.info());
}
#[tokio::test(flavor = "multi_thread")]
async fn stop_reaches_a_cli_owner() {
let dir = tempfile::tempdir().unwrap();
let client_dir = tempfile::tempdir().unwrap();
let ctx = ctx(dir.path(), client_dir.path()).with_owner(Owner::Cli);
let mut server = ControlServer::start(ctx.clone()).unwrap();
let stop_rx = server.take_stop_rx().unwrap();
// The daemon exits as soon as it hears the stop, and the stop is still answered.
let daemon = std::thread::spawn(move || {
stop_rx.recv().unwrap();
drop(server);
});
let client = ControlClient::new(control::read(dir.path()).unwrap());
client.start(None).await.unwrap();
let status = client.stop().await.unwrap();
assert_eq!(status.state, ProxyState::Stopped);
daemon.join().unwrap();
assert_eq!(ctx.port(), None);
assert!(!dir.path().join("control.json").exists());
}
#[tokio::test(flavor = "multi_thread")]
async fn offline_reads_proxy_db() {
let dir = tempfile::tempdir().unwrap();
let client_dir = tempfile::tempdir().unwrap();
let flow = {
let ctx = ctx(dir.path(), client_dir.path());
ctx.start().unwrap();
let flow = capture_flow(&ctx).await;
ctx.stop().unwrap();
flow
};
let connection = ProxyConnection::connect(dir.path()).unwrap();
assert!(matches!(connection, ProxyConnection::Offline(_)));
let flows = connection.flows(&FlowFilter::default()).await.unwrap();
assert_eq!(flows.iter().map(|f| &f.id).collect::<Vec<_>>(), vec![&flow.id]);
let detail = connection.flow(&flow.id, 1024).await.unwrap();
assert_eq!(detail.response_body.unwrap().data, r#"{"id":"ch_1","amount":42}"#);
let body = connection.body(&flow.id, FlowBody::Request).await.unwrap();
assert_eq!(body.data, br#"{"amount":42}"#);
let posts = FlowFilter { method: Some("post".into()), ..Default::default() };
assert_eq!(connection.har(&posts).await.unwrap().log.entries.len(), 1);
let empty = tempfile::tempdir().unwrap();
let err = ProxyConnection::connect(empty.path()).err().unwrap();
assert!(err.starts_with("No captured traffic yet"), "{err}");
}
#[tokio::test(flavor = "multi_thread")]
async fn save_writes_the_request_into_yaak() {
let dir = tempfile::tempdir().unwrap();
let client_dir = tempfile::tempdir().unwrap();
// Yaak's database, with a workspace and a folder in it.
let (query_manager, _blobs, _rx) = yaak_models::init_standalone(
client_dir.path().join("db.sqlite"),
client_dir.path().join("blobs.sqlite"),
)
.unwrap();
let (workspace, folder) = query_manager
.with_tx(|tx| {
let workspace = tx.upsert_workspace(
&Workspace { name: "Payments".into(), ..Default::default() },
&UpdateSource::Background,
)?;
let folder = tx.upsert_folder(
&Folder {
name: "Stripe".into(),
workspace_id: workspace.id.clone(),
..Default::default()
},
&UpdateSource::Background,
)?;
Ok::<_, yaak_models::error::Error>((workspace, folder))
})
.unwrap();
let ctx = ctx(dir.path(), client_dir.path());
let _server = ControlServer::start(ctx.clone()).unwrap();
ctx.start().unwrap();
let flow = capture_flow(&ctx).await;
let client = ControlClient::new(control::read(dir.path()).unwrap());
let workspaces = client.workspaces().await.unwrap();
assert_eq!(workspaces.len(), 1);
assert_eq!(
(workspaces[0].id.as_str(), workspaces[0].name.as_str()),
(workspace.id.as_str(), "Payments")
);
assert_eq!(workspaces[0].folders[0].id, folder.id);
let target = SaveFlowRequest {
workspace_id: workspace.id.clone(),
folder_id: Some(folder.id.clone()),
name: Some("Create charge".into()),
};
let saved = client.save(&flow.id, &target).await.unwrap();
let dropped: Vec<&str> = saved.dropped_headers.iter().map(|h| h.name.as_str()).collect();
assert!(dropped.contains(&"host") && dropped.contains(&"content-length"), "{dropped:?}");
assert_eq!(
saved.warnings,
vec!["The Authorization header is saved as captured, in plain text."]
);
// The row is in Yaak's database, written as a sync so a running Yaak picks it up.
let request = query_manager.connect().get_http_request(&saved.request.id).unwrap();
assert_eq!(request.name, "Create charge");
assert_eq!(
(request.workspace_id, request.folder_id),
(workspace.id.clone(), Some(folder.id.clone()))
);
assert_eq!(request.method, "POST");
assert!(request.url.ends_with("/v1/charges"), "{}", request.url);
assert_eq!(request.url_parameters[0].name, "expand");
assert_eq!(request.body_type.as_deref(), Some("application/json"));
assert_eq!(request.body["text"], r#"{"amount":42}"#);
let auth = request.headers.iter().find(|h| h.name == "authorization").unwrap();
assert_eq!(auth.value, "Bearer sk_test_123");
let db = rusqlite::Connection::open(client_dir.path().join("db.sqlite")).unwrap();
let source: String = db
.query_row(
"SELECT update_source FROM model_changes WHERE model_id = ?",
[&saved.request.id],
|r| r.get(0),
)
.unwrap();
assert!(source.contains("sync"), "{source}");
// The flow remembers what it was saved as.
let flow = yaak_proxy_lib::flows::get(&ctx.db, &flow.id).unwrap();
assert_eq!(flow.saved_request_id, Some(saved.request.id));
let wrong = SaveFlowRequest { workspace_id: "wk_missing".into(), ..Default::default() };
assert!(client.save(&flow.id, &wrong).await.unwrap_err().contains("No workspace"));
}
#[tokio::test(flavor = "multi_thread")]
async fn stale_removal_waits_for_the_claim_lock() {
let dir = tempfile::tempdir().unwrap();
let stale = ControlInfo {
pid: dead_pid(),
control_port: 1,
proxy_port: None,
token: "stale".into(),
owner: Owner::Cli,
started_at: chrono::Utc::now(),
};
let path = dir.path().join("control.json");
std::fs::write(&path, serde_json::to_vec(&stale).unwrap()).unwrap();
// Another process mid-claim holds the lock; it may be about to replace the stale file.
let lock = std::fs::File::create(dir.path().join("control.lock")).unwrap();
lock.lock().unwrap();
let data_dir = dir.path().to_path_buf();
let cleaner = std::thread::spawn(move || control::live_owner(&data_dir));
std::thread::sleep(Duration::from_millis(300));
assert!(path.exists(), "the stale file was removed under another process's claim lock");
lock.unlock().unwrap();
assert!(cleaner.join().unwrap().is_none());
assert!(!path.exists());
}
#[tokio::test(flavor = "multi_thread")]
async fn start_is_refused_while_another_process_owns_the_engine() {
let dir = tempfile::tempdir().unwrap();
let client_dir = tempfile::tempdir().unwrap();
let ctx = ctx(dir.path(), client_dir.path());
// A live CLI daemon: a running process, and a control port that answers.
let mut daemon = if cfg!(windows) {
std::process::Command::new("ping").args(["-n", "30", "127.0.0.1"]).spawn().unwrap()
} else {
std::process::Command::new("sleep").arg("30").spawn().unwrap()
};
let control_port = TcpListener::bind("127.0.0.1:0").await.unwrap();
let owner = ControlInfo {
pid: daemon.id(),
control_port: control_port.local_addr().unwrap().port(),
proxy_port: Some(9090),
token: "daemon".into(),
owner: Owner::Cli,
started_at: chrono::Utc::now(),
};
std::fs::write(dir.path().join("control.json"), serde_json::to_vec(&owner).unwrap()).unwrap();
let err = ctx.start().unwrap_err();
assert!(err.contains("already running in the yaak CLI"), "{err}");
assert_eq!(ctx.port(), None);
daemon.kill().unwrap();
daemon.wait().unwrap();
}
+241
View File
@@ -0,0 +1,241 @@
{
"log": {
"version": "1.2",
"creator": {
"name": "Yaak Proxy",
"version": "test"
},
"entries": [
{
"startedDateTime": "2026-10-26T09:30:01.250Z",
"time": 134,
"request": {
"method": "GET",
"url": "https://api.github.com/repos/mountain-loop/yaak/issues?state=open&labels=bug%20fix",
"httpVersion": "HTTP/1.1",
"cookies": [],
"headers": [
{
"name": "Host",
"value": "api.github.com"
},
{
"name": "Accept",
"value": "application/vnd.github+json"
}
],
"queryString": [
{
"name": "state",
"value": "open"
},
{
"name": "labels",
"value": "bug fix"
}
],
"headersSize": -1,
"bodySize": 0
},
"response": {
"status": 200,
"statusText": "OK",
"httpVersion": "HTTP/2",
"cookies": [],
"headers": [
{
"name": "content-type",
"value": "application/json; charset=utf-8"
},
{
"name": "content-encoding",
"value": "gzip"
}
],
"content": {
"size": 30,
"mimeType": "application/json; charset=utf-8",
"text": "[{\"number\":761,\"title\":\"MCP\"}]"
},
"redirectURL": "",
"headersSize": -1,
"bodySize": 41
},
"cache": {},
"timings": {
"send": 0,
"wait": 120,
"receive": 14
},
"serverIPAddress": "140.82.112.6"
},
{
"startedDateTime": "2026-10-26T09:30:02.250Z",
"time": 8,
"request": {
"method": "POST",
"url": "http://localhost:3000/login",
"httpVersion": "HTTP/1.1",
"cookies": [],
"headers": [
{
"name": "Content-Type",
"value": "application/x-www-form-urlencoded"
}
],
"queryString": [],
"postData": {
"mimeType": "application/x-www-form-urlencoded",
"text": "user=greg&password=hunter2"
},
"headersSize": -1,
"bodySize": 26
},
"response": {
"status": 302,
"statusText": "Found",
"httpVersion": "HTTP/1.1",
"cookies": [],
"headers": [
{
"name": "Location",
"value": "/dashboard"
}
],
"content": {
"size": 0,
"mimeType": ""
},
"redirectURL": "/dashboard",
"headersSize": -1,
"bodySize": 0
},
"cache": {},
"timings": {
"send": 0,
"wait": 8,
"receive": 0
},
"serverIPAddress": "::1"
},
{
"startedDateTime": "2026-10-26T09:30:03.250Z",
"time": 31,
"request": {
"method": "GET",
"url": "https://images.example.com/logo.png",
"httpVersion": "HTTP/1.1",
"cookies": [],
"headers": [],
"queryString": [],
"headersSize": -1,
"bodySize": 0
},
"response": {
"status": 200,
"statusText": "OK",
"httpVersion": "HTTP/2",
"cookies": [],
"headers": [
{
"name": "Content-Type",
"value": "image/png"
}
],
"content": {
"size": 8,
"mimeType": "image/png",
"text": "iVBORw0KGgo=",
"encoding": "base64"
},
"redirectURL": "",
"headersSize": -1,
"bodySize": 8
},
"cache": {},
"timings": {
"send": 0,
"wait": 30,
"receive": 1
}
},
{
"startedDateTime": "2026-10-26T09:30:04.250Z",
"time": 900,
"request": {
"method": "GET",
"url": "https://cdn.example.com/big.txt",
"httpVersion": "HTTP/1.1",
"cookies": [],
"headers": [],
"queryString": [],
"headersSize": -1,
"bodySize": 0
},
"response": {
"status": 200,
"statusText": "OK",
"httpVersion": "HTTP/1.1",
"cookies": [],
"headers": [
{
"name": "Content-Type",
"value": "text/plain"
}
],
"content": {
"size": 10,
"mimeType": "text/plain",
"text": "0123456789",
"_truncated": true
},
"redirectURL": "",
"headersSize": -1,
"bodySize": 20000000
},
"cache": {},
"timings": {
"send": 0,
"wait": 10,
"receive": 890
},
"_error": "connection reset by peer"
},
{
"startedDateTime": "2026-10-26T09:30:05.250Z",
"time": 3,
"request": {
"method": "GET",
"url": "https://down.example.com/",
"httpVersion": "HTTP/1.1",
"cookies": [],
"headers": [],
"queryString": [],
"headersSize": -1,
"bodySize": 0
},
"response": {
"status": 0,
"statusText": "",
"httpVersion": "",
"cookies": [],
"headers": [],
"content": {
"size": 0,
"mimeType": ""
},
"redirectURL": "",
"headersSize": -1,
"bodySize": -1
},
"cache": {},
"timings": {
"send": 0,
"wait": 3,
"receive": 0
},
"_error": "Connection refused (os error 61)"
}
]
}
}
@@ -151,6 +151,7 @@ async fn flow_is_stored_with_its_bodies() {
data: b"hello proxy".to_vec(),
stored_size: compressed.len(),
truncated: false,
decode_error: None,
};
assert_eq!(decoded, expected);
+47 -1
View File
@@ -1,7 +1,9 @@
use log::{error, info, warn};
use std::sync::Mutex;
use tauri::Runtime;
use tauri::{Emitter, Manager, RunEvent, State, WebviewWindow};
use yaak_proxy_lib::ProxyCtx;
use yaak_proxy_lib::control::{self, ControlServer};
use yaak_rpc::{RpcEventEmitter, RpcRouter};
use yaak_window::window::CreateWindowConfig;
@@ -48,6 +50,40 @@ fn setup_window_menu<R: Runtime>(win: &WebviewWindow<R>) {
});
}
/// The control API, served while this app owns the engine.
struct Control(Mutex<Option<ControlServer>>);
/// Claims the engine and serves the control API, unless a live owner (the CLI daemon) has it.
fn host_control_api(ctx: &ProxyCtx) -> Option<ControlServer> {
// Read first: a crashed owner's file names the port its system proxy may still point at.
let previous = control::read(ctx.data_dir());
if let Some(owner) = control::live_owner(ctx.data_dir()) {
// Starting the engine here is refused too, until the owner stops.
warn!(
"The proxy engine is owned by {:?} (pid {}), so this app does not serve the control API",
owner.owner, owner.pid
);
return None;
}
let server = match ControlServer::start(ctx.clone()) {
Ok(server) => server,
Err(e) => {
error!("Failed to start the proxy control API: {e}");
return None;
}
};
let mut ports: Vec<u16> = previous.and_then(|p| p.proxy_port).into_iter().collect();
ports.extend(ctx.settings().ok().map(|s| s.port));
let ctx = ctx.clone();
std::thread::spawn(move || {
if let Err(e) = ctx.reset_stale_system_proxy(&ports) {
warn!("Failed to reset a stale system proxy: {e}");
}
});
Some(server)
}
#[tauri::command]
async fn rpc(
router: State<'_, RpcRouter<ProxyCtx>>,
@@ -67,7 +103,9 @@ 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, emitter)?);
let ctx = ProxyCtx::new(&data_dir, emitter)?;
app.manage(Control(Mutex::new(host_control_api(&ctx))));
app.manage(ctx);
app.manage(yaak_proxy_lib::build_router());
// Drain RPC events and forward as Tauri events
@@ -86,6 +124,14 @@ pub fn run() {
.build(tauri::generate_context!())
.expect("error while building yaak proxy tauri application")
.run(|app_handle, event| {
if let RunEvent::Exit = event {
// Never leave the Mac pointed at a proxy that is gone.
app_handle.state::<ProxyCtx>().shutdown();
if let Ok(mut control) = app_handle.state::<Control>().0.lock() {
control.take();
}
return;
}
if let RunEvent::Ready = event {
let config = CreateWindowConfig {
url: "/",
+9 -2
View File
@@ -38,6 +38,13 @@ impl CaptureBody {
Self { inner, direction, id, start, engine, captured: 0, size: 0, finished: false }
}
/// Every byte the body declared has arrived. An HTTP/2 stream's end can come in a frame after
/// its last byte, and hyper stops polling a Content-Length body once it has written that
/// many bytes, so the stream's end alone would call a complete body cut short.
fn complete(&self) -> bool {
self.inner.is_end_stream() || self.inner.size_hint().exact() == Some(0)
}
fn finish(&mut self, error: Option<String>) {
if self.finished {
return;
@@ -84,7 +91,7 @@ impl Body for CaptureBody {
Direction::Response => ProxyEvent::ResponseBodyChunk { id, len, data },
});
}
if self.inner.is_end_stream() {
if self.complete() {
self.finish(None);
}
}
@@ -109,7 +116,7 @@ impl Body for CaptureBody {
impl Drop for CaptureBody {
fn drop(&mut self) {
let error = (self.direction == Direction::Response && !self.inner.is_end_stream())
let error = (self.direction == Direction::Response && !self.complete())
.then(|| "Client closed the connection before the response finished".to_string());
self.finish(error);
}
+71 -12
View File
@@ -6,9 +6,10 @@ use std::time::{Duration, SystemTime};
use bytes::Bytes;
use futures_util::{SinkExt, StreamExt};
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper::header::{CONNECTION, UPGRADE};
use http_body_util::combinators::BoxBody;
use http_body_util::{BodyExt, Full, StreamBody};
use hyper::body::{Frame, Incoming};
use hyper::header::{CONNECTION, CONTENT_LENGTH, UPGRADE};
use hyper::service::service_fn;
use hyper::{Request, Response, StatusCode};
use hyper_util::rt::{TokioExecutor, TokioIo};
@@ -76,19 +77,42 @@ fn test_pki() -> TestPki {
TestPki { ca_pem: ca.pem(), ca_der: ca.der().clone(), server: Arc::new(server) }
}
async fn route(req: Request<Incoming>) -> Result<Response<Full<Bytes>>, Infallible> {
type TestBody = BoxBody<Bytes, Infallible>;
fn full(data: impl Into<Bytes>) -> TestBody {
Full::new(data.into()).boxed()
}
async fn route(req: Request<Incoming>) -> Result<Response<TestBody>, Infallible> {
Ok(match req.uri().path() {
"/hello" => Response::new(Full::from("hello")),
"/echo" => Response::new(Full::new(req.into_body().collect().await.unwrap().to_bytes())),
"/big" => Response::new(Full::from(vec![b'x'; BIG_BODY])),
"/hello" => Response::new(full("hello")),
"/echo" => Response::new(full(req.into_body().collect().await.unwrap().to_bytes())),
"/big" => Response::new(full(vec![b'x'; BIG_BODY])),
"/late-end" => late_end(5),
// The stream ends with half the promised bytes, so the upstream resets it mid-body.
"/cut" => late_end(10),
"/ws" => websocket_echo(req),
_ => {
Response::builder().status(StatusCode::NOT_FOUND).body(Full::from("not found")).unwrap()
}
_ => Response::builder().status(StatusCode::NOT_FOUND).body(full("not found")).unwrap(),
})
}
fn websocket_echo(mut req: Request<Incoming>) -> Response<Full<Bytes>> {
/// "hello" under a Content-Length of `content_length`, its end of stream sent a while after the
/// last byte, as an HTTP/2 server does when its final DATA frame has no END_STREAM.
fn late_end(content_length: usize) -> Response<TestBody> {
let frames = futures_util::stream::unfold(false, |sent| async move {
if sent {
tokio::time::sleep(Duration::from_millis(300)).await;
return None;
}
Some((Ok(Frame::data(Bytes::from("hello"))), true))
});
Response::builder()
.header(CONTENT_LENGTH, content_length)
.body(BodyExt::boxed(StreamBody::new(frames)))
.unwrap()
}
fn websocket_echo(mut req: Request<Incoming>) -> Response<TestBody> {
let accept = derive_accept_key(req.headers()["sec-websocket-key"].as_bytes());
let on_upgrade = hyper::upgrade::on(&mut req);
tokio::spawn(async move {
@@ -105,7 +129,7 @@ fn websocket_echo(mut req: Request<Incoming>) -> Response<Full<Bytes>> {
.header(UPGRADE, "websocket")
.header(CONNECTION, "Upgrade")
.header("sec-websocket-accept", accept)
.body(Full::default())
.body(full(Bytes::new()))
.unwrap()
}
@@ -374,6 +398,41 @@ async fn https_not_found() {
assert_eq!(body(&events, false), (b"not found".to_vec(), 9));
}
/// The client is sent every byte the Content-Length promised, and hyper drops the body there,
/// before an HTTP/2 upstream's end of stream arrives. That is a complete response, not an error.
#[tokio::test(flavor = "multi_thread")]
async fn https_body_complete_before_end_of_stream() {
let h = Harness::start(1024);
let upstream = spawn_upstream(Some(h.pki.server.clone())).await;
let url = format!("https://localhost:{}/late-end", upstream.port());
let resp = h.client().get(&url).send().await.unwrap();
assert_eq!(resp.text().await.unwrap(), "hello");
let events = h.flow_events();
let version =
find!(events, ProxyEvent::ResponseStart { http_version, .. } => http_version.clone());
assert_eq!(version, "HTTP/2");
let (size, error) = find!(events,
ProxyEvent::ResponseBodyComplete { size, error, .. } => (*size, error.clone()));
assert_eq!((size, error), (5, None));
assert_eq!(body(&events, false), (b"hello".to_vec(), 5));
}
#[tokio::test(flavor = "multi_thread")]
async fn https_body_cut_short_is_an_error() {
let h = Harness::start(1024);
let upstream = spawn_upstream(Some(h.pki.server.clone())).await;
let url = format!("https://localhost:{}/cut", upstream.port());
let resp = h.client().get(&url).send().await.unwrap();
assert!(resp.bytes().await.is_err());
let events = h.flow_events();
let error = find!(events, ProxyEvent::ResponseBodyComplete { error, .. } => error.clone());
assert!(error.is_some(), "{events:#?}");
}
#[tokio::test(flavor = "multi_thread")]
async fn bodies_over_the_cap_are_truncated_but_delivered() {
let h = Harness::start(1024);