Capture WebSocket, SSE and gRPC (#814)

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-11-06 09:53:56 -07:00
1 parent 08a98d5f82
commit 1129f7aa93
32 files changed
+5228 -100

No files matched your search

Generated
+71
View File
@@ -6480,6 +6480,70 @@ dependencies = [
"prost",
]
[[package]]
name = "protoc-bin-vendored"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8760a25b6ff9c620324822737e468478fa092234190d2e449760344354896ed9"
dependencies = [
"protoc-bin-vendored-linux-aarch_64",
"protoc-bin-vendored-linux-ppcle_64",
"protoc-bin-vendored-linux-s390_64",
"protoc-bin-vendored-linux-x86_32",
"protoc-bin-vendored-linux-x86_64",
"protoc-bin-vendored-macos-aarch_64",
"protoc-bin-vendored-macos-x86_64",
"protoc-bin-vendored-win32",
]
[[package]]
name = "protoc-bin-vendored-linux-aarch_64"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "73fa2624782ca04cd44f51554566717377acd240e4c0016d757dd74fccc9324f"
[[package]]
name = "protoc-bin-vendored-linux-ppcle_64"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e2417e9817fa237dab803ad4dda7357a111656e242959cc6b8f9a1a583367d42"
[[package]]
name = "protoc-bin-vendored-linux-s390_64"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4d189c34636356a46a7ed3188233dc8a88c431278cc54d4a19b096a2d270e985"
[[package]]
name = "protoc-bin-vendored-linux-x86_32"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "171e39f1e846e5f322ced1ac3b8d4cd3a3833ca24b6e5d58b3632574fe6204fa"
[[package]]
name = "protoc-bin-vendored-linux-x86_64"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "873cdcc097593432086661aa432b8078f1cd87bfb02847c332e98ae2c119e966"
[[package]]
name = "protoc-bin-vendored-macos-aarch_64"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eeb72df001783b8297847fe8f5f874ee400fd742c843d60583e8c23d96977c7f"
[[package]]
name = "protoc-bin-vendored-macos-x86_64"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b04652167eca899dda05f32f5481adeaf25c623a98ce2fc146a001cc59a2add7"
[[package]]
name = "protoc-bin-vendored-win32"
version = "3.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "263a3f48f01e7309e857138bd47f785585b4a005e8e56c6d2824ce91195999c3"
[[package]]
name = "quick-error"
version = "1.2.3"
@@ -11813,7 +11877,9 @@ dependencies = [
name = "yaak-proxy"
version = "0.1.0"
dependencies = [
"base64 0.23.1",
"bytes",
"flate2",
"futures-util",
"http",
"http-body-util",
@@ -11854,6 +11920,7 @@ dependencies = [
"chrono",
"encoding_rs",
"flate2",
"futures-util",
"http-body-util",
"hyper",
"hyper-util",
@@ -11865,6 +11932,7 @@ dependencies = [
"objc2-foundation 0.3.2",
"objc2-service-management",
"plist",
"protoc-bin-vendored",
"r2d2",
"r2d2_sqlite",
"reqwest",
@@ -11874,11 +11942,14 @@ dependencies = [
"serde_json",
"tempfile",
"tokio",
"tonic",
"tonic-reflection",
"ts-rs",
"url",
"yaak-common",
"yaak-database",
"yaak-filter",
"yaak-grpc",
"yaak-http",
"yaak-lifecycle",
"yaak-models",
+6
View File
@@ -42,6 +42,12 @@ const base: Flow = {
clientName: "curl",
clientBundleId: null,
clientPath: "/usr/bin/curl",
streamKind: null,
messageCount: 0,
grpcService: null,
grpcMethod: null,
grpcStatus: null,
grpcMessage: null,
};
const jsonRequest = JSON.stringify({ amount: 4200, currency: "usd", capture: true });
+6
View File
@@ -50,6 +50,12 @@ const flow: Flow = {
clientName: null,
clientBundleId: null,
clientPath: null,
streamKind: null,
messageCount: 0,
grpcService: null,
grpcMethod: null,
grpcStatus: null,
grpcMessage: null,
};
function bytes(text: string): FlowBodyData {
+8 -3
View File
@@ -22,12 +22,13 @@ rusqlite = { version = "0.38", features = ["bundled", "chrono", "functions"] }
sea-query = { version = "1.0", features = ["with-chrono", "attr"] }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt-multi-thread"] }
tokio = { workspace = true, features = ["io-util", "macros", "net", "process", "rt-multi-thread", "sync", "time"] }
ts-rs = { workspace = true, features = ["chrono-impl"] }
url = "2"
yaak-common = { workspace = true }
yaak-database = { workspace = true }
yaak-filter = { workspace = true }
yaak-grpc = { workspace = true }
yaak-http = { workspace = true }
yaak-lifecycle = { workspace = true }
yaak-models = { workspace = true }
@@ -71,9 +72,13 @@ plist = "1.7"
base64 = "0.23.1"
bytes = "1"
flate2 = "1"
futures-util = "0.3"
http-body-util = "0.1"
hyper = { version = "1", features = ["http1", "server"] }
hyper-util = { version = "0.1", features = ["tokio"] }
hyper = { version = "1", features = ["http1", "http2", "server"] }
hyper-util = { version = "0.1", features = ["tokio", "server-auto"] }
protoc-bin-vendored = "3"
tonic = { version = "0.14.6", default-features = false }
tonic-reflection = "0.14.6"
reqwest = { workspace = true }
tempfile = "3"
tokio = { workspace = true, features = ["macros", "net", "rt-multi-thread", "time"] }
+136 -1
View File
@@ -65,10 +65,145 @@ clientBundleId: string | null,
/**
* Its app bundle's path, or its executable's outside an app.
*/
clientPath: string | null, };
clientPath: string | null,
/**
* The flow carries a stream of messages, listed with `list_messages`. It is written as the
* stream opens and again as messages land; it is open until `elapsedMs`, `error` or
* `cancelled` is set.
*/
streamKind: StreamKind | null,
/**
* Messages recorded so far.
*/
messageCount: number,
/**
* A gRPC call's service, from its path: `package.Service`.
*/
grpcService: string | null,
/**
* A gRPC call's method, from its path.
*/
grpcMethod: string | null,
/**
* A gRPC call's status, from its trailers. Not 0 is a failed call.
*/
grpcStatus: number | null,
/**
* A gRPC call's `grpc-message`, decoded.
*/
grpcMessage: string | null, };
export type FlowBody = "request" | "response";
/**
* One message of a streaming flow. Its payload is fetched on its own.
*/
export type FlowMessage = { id: string, flowId: string,
/**
* Its place in the flow, from 0, in the order messages passed the proxy.
*/
seq: number, direction: MessageDirection, kind: MessageKind,
/**
* When it had passed the proxy whole.
*/
timestamp: string,
/**
* Bytes it took on the wire: the frames' payloads, the event's lines, the gRPC message
* without its prefix.
*/
size: number,
/**
* Bytes of payload stored: unmasked, reassembled, inflated or decompressed, and cut at the
* capture limit.
*/
storedSize: number,
/**
* The stored payload is cut short at the capture limit.
*/
truncated: boolean,
/**
* It was compressed on the wire (permessage-deflate, gRPC's compressed flag); the payload is
* stored decompressed unless `error` says otherwise.
*/
compressed: boolean,
/**
* Why the payload could not be decoded or stored.
*/
error: string | null, meta: MessageMeta,
/**
* The payload's start as text, when it is text.
*/
preview: string | null, };
/**
* Descriptors for decoding the gRPC calls to hosts that match `hostPattern`.
*/
export type GrpcSchema = { id: string, createdAt: string, updatedAt: string,
/**
* A host (`api.example.com`), a host and port (`localhost:50051`), any subdomain of one
* (`*.example.com`), or every host (`*`).
*/
hostPattern: string, source: GrpcSchemaSource,
/**
* The `.proto` files and include directories compiled, for a `files` schema.
*/
protoPaths: Array<string>,
/**
* The services the descriptors define, by full name.
*/
services: Array<string>, };
/**
* Where a schema's descriptors came from.
*/
export type GrpcSchemaSource = "files" | "reflection";
export type MessageDirection = "client" | "server";
export type MessageKind = "websocket_text" | "websocket_binary" | "websocket_ping" | "websocket_pong" | "websocket_close" | "sse_event" | "grpc_message" | "grpc_trailers";
/**
* What a kind of message carries beside its payload. Fields that do not apply are left out.
*/
export type MessageMeta = {
/**
* A WebSocket close frame's status code.
*/
closeCode?: number,
/**
* A WebSocket close frame's reason.
*/
closeReason?: string,
/**
* An SSE event's `event` field.
*/
event?: string,
/**
* An SSE event's `id` field.
*/
eventId?: string,
/**
* An SSE event's `retry` field, in milliseconds.
*/
retry?: number,
/**
* An SSE event's lines as they came, comments included, up to 4 KB.
*/
raw?: string,
/**
* gRPC trailers' `grpc-status`.
*/
grpcStatus?: number,
/**
* gRPC trailers' `grpc-message`, decoded.
*/
grpcMessage?: string, };
export type ModelPayload = { model: Flow, change: ModelChangeEvent, };
export type ProxyHeader = { name: string, value: string, };
/**
* The streaming protocol a flow carries.
*/
export type StreamKind = "websocket" | "sse" | "grpc" | "grpc_web";
+88 -4
View File
@@ -1,10 +1,20 @@
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
import type { Flow, FlowBody, ModelPayload } from "./gen_models";
import type { Flow, FlowBody, FlowMessage, GrpcSchema, ModelPayload } from "./gen_models";
export type ActionInvocation = { "scope": "global", action: GlobalAction, };
export type ActionMetadata = { label: string, defaultHotkey: string | null, };
export type AddGrpcSchemaRequest = {
/**
* `api.example.com`, `localhost:50051`, `*.example.com` or `*`.
*/
hostPattern: string,
/**
* `.proto` files, and directories they import from.
*/
protoPaths: Array<string>, };
export type AppIconRequest = {
/**
* A flow's `clientPath`.
@@ -44,6 +54,31 @@ bundleId: string | null,
*/
path: string, };
export type DecodeGrpcMessageRequest = { flowId: string, seq: number, };
/**
* A captured gRPC message as JSON.
*/
export type DecodedGrpcMessage = {
/**
* Pretty JSON: the message's fields by name with a schema; without one, an array of its
* fields as `{ field, wireType, value }`, a length-delimited one with its guesses (`string`,
* `message`, `bytes` in base64).
*/
json: string,
/**
* The message type decoded with: `package.HelloReply`. None when schemaless.
*/
messageType: string | null,
/**
* The schema decoded with. None when schemaless.
*/
schema: GrpcSchema | null,
/**
* Why it decoded schemaless, or what to know about the decode.
*/
note: string | null, };
export type DropReason = "hop_by_hop" | "set_by_sending" | "body_decoded";
export type DroppedHeader = { name: string, value: string, reason: DropReason, };
@@ -100,7 +135,8 @@ app?: Array<string>,
*/
method?: string,
/**
* A status (`404`), a class of them (`4xx`), or `error` for flows that failed.
* A status (`404`), a class of them (`4xx`), or `error` for flows that failed, gRPC calls
* that ended with a status other than OK among them.
*/
status?: string,
/**
@@ -149,6 +185,18 @@ storedSize: number,
*/
decodeError?: string, };
export type GetMessagePayloadRequest = { flowId: string, seq: number, };
export type GetMessagePayloadResponse = {
/**
* The payload as stored (unmasked, inflated, decompressed), in base64.
*/
data: string,
/**
* It is cut at the capture limit.
*/
truncated: boolean, };
export type GetProxyStateRequest = Record<string, never>;
export type GetProxyStateResponse = { state: ProxyState,
@@ -214,6 +262,28 @@ total: number,
*/
matching: number | null, };
export type ListGrpcSchemasRequest = Record<string, never>;
export type ListMessagesRequest = { flowId: string,
/**
* Only messages after this `seq`, for following a stream as it grows.
*/
afterSeq?: number,
/**
* At most this many, the first after `afterSeq`.
*/
limit?: number, };
export type ListMessagesResponse = {
/**
* In order, by `seq`.
*/
messages: Array<FlowMessage>,
/**
* Messages the flow has in all.
*/
total: number, };
export type ListWorkspacesRequest = Record<string, never>;
export type ListWorkspacesResponse = {
@@ -222,6 +292,11 @@ export type ListWorkspacesResponse = {
*/
yaakFound: boolean, workspaces: Array<WorkspaceSummary>, };
/**
* A message of a streaming flow was recorded.
*/
export type MessagePayload = { message: FlowMessage, };
export type OtherOwner = { owner: Owner, pid: number, startedAt: string,
/**
* The engine's port while it runs.
@@ -241,9 +316,18 @@ export type ProxyState = "running" | "stopped";
export type ProxyStatePayload = { state: ProxyState, };
export type RpcEventSchema = { flows_cleared: FlowsClearedPayload, model_write: ModelPayload, proxy_state_changed: ProxyStatePayload, };
export type ReflectGrpcHostRequest = {
/**
* `host` or `host:port`, over TLS on 443 unless a port is given, or either after
* `http://` to ask without TLS.
*/
host: string, };
export type RpcSchema = { execute_action: [ActionInvocation, boolean], format_body: [FormatBodyRequest, string], list_actions: [ListActionsRequest, ListActionsResponse], list_flows: [FlowFilter, ListFlowsResponse], set_recording: [SetRecordingRequest, boolean], set_system_proxy_at_launch: [SetSystemProxyAtLaunchRequest, boolean], set_http_version: [SetHttpVersionRequest, HttpVersion], app_icon: [AppIconRequest, string | null], ca_status: [CaStatusRequest, CaStatus], ca_trust: [CaTrustRequest, CaStatus], clear_flows: [ClearFlowsRequest, number], export_har: [ExportHarRequest, number], filter_body: [FilterBodyRequest, FilterBodyResponse], get_flow_body: [GetFlowBodyRequest, GetFlowBodyResponse], get_proxy_state: [GetProxyStateRequest, GetProxyStateResponse], helper_register: [HelperRegisterRequest, SystemProxyStatus], helper_unregister: [HelperUnregisterRequest, SystemProxyStatus], list_workspaces: [ListWorkspacesRequest, ListWorkspacesResponse], preview_save: [PreviewSaveRequest, SavePreview], save_flow: [SaveFlowToYaakRequest, SavedFlow], save_flow_body: [SaveFlowBodyRequest, SaveFlowBodyResponse], system_proxy_set: [SystemProxySetRequest, SystemProxyStatus], system_proxy_status: [SystemProxyStatusRequest, SystemProxyStatus], take_over: [TakeOverRequest, boolean], };
export type RemoveGrpcSchemaRequest = { id: string, };
export type RpcEventSchema = { flows_cleared: FlowsClearedPayload, message_write: MessagePayload, model_write: ModelPayload, proxy_state_changed: ProxyStatePayload, };
export type RpcSchema = { execute_action: [ActionInvocation, boolean], format_body: [FormatBodyRequest, string], get_message_payload: [GetMessagePayloadRequest, GetMessagePayloadResponse], list_actions: [ListActionsRequest, ListActionsResponse], list_flows: [FlowFilter, ListFlowsResponse], list_grpc_schemas: [ListGrpcSchemasRequest, Array<GrpcSchema>], list_messages: [ListMessagesRequest, ListMessagesResponse], remove_grpc_schema: [RemoveGrpcSchemaRequest, boolean], set_recording: [SetRecordingRequest, boolean], set_system_proxy_at_launch: [SetSystemProxyAtLaunchRequest, boolean], set_http_version: [SetHttpVersionRequest, HttpVersion], add_grpc_schema: [AddGrpcSchemaRequest, GrpcSchema], app_icon: [AppIconRequest, string | null], ca_status: [CaStatusRequest, CaStatus], ca_trust: [CaTrustRequest, CaStatus], clear_flows: [ClearFlowsRequest, number], decode_grpc_message: [DecodeGrpcMessageRequest, DecodedGrpcMessage], export_har: [ExportHarRequest, number], filter_body: [FilterBodyRequest, FilterBodyResponse], get_flow_body: [GetFlowBodyRequest, GetFlowBodyResponse], get_proxy_state: [GetProxyStateRequest, GetProxyStateResponse], helper_register: [HelperRegisterRequest, SystemProxyStatus], helper_unregister: [HelperUnregisterRequest, SystemProxyStatus], list_workspaces: [ListWorkspacesRequest, ListWorkspacesResponse], preview_save: [PreviewSaveRequest, SavePreview], reflect_grpc_host: [ReflectGrpcHostRequest, GrpcSchema], save_flow: [SaveFlowToYaakRequest, SavedFlow], save_flow_body: [SaveFlowBodyRequest, SaveFlowBodyResponse], system_proxy_set: [SystemProxySetRequest, SystemProxyStatus], system_proxy_status: [SystemProxyStatusRequest, SystemProxyStatus], take_over: [TakeOverRequest, boolean], };
export type SaveFlowBodyRequest = { flowId: string, which: FlowBody,
/**
@@ -0,0 +1,44 @@
-- Streaming flows: WebSocket, server-sent events, gRPC and gRPC-Web.
ALTER TABLE flows ADD COLUMN stream_kind TEXT;
ALTER TABLE flows ADD COLUMN message_count INTEGER NOT NULL DEFAULT 0;
-- A gRPC call's service and method, from its path, and its status, from its trailers.
ALTER TABLE flows ADD COLUMN grpc_service TEXT;
ALTER TABLE flows ADD COLUMN grpc_method TEXT;
ALTER TABLE flows ADD COLUMN grpc_status INTEGER;
ALTER TABLE flows ADD COLUMN grpc_message TEXT;
-- The messages of streaming flows, in the order they passed. A payload over a few KB is in
-- blobs.sqlite, under `<flow id>.msg.<message id>`, and `data` is NULL.
CREATE TABLE flow_messages
(
id TEXT NOT NULL PRIMARY KEY,
flow_id TEXT NOT NULL,
seq INTEGER NOT NULL,
direction TEXT NOT NULL,
kind TEXT NOT NULL,
timestamp DATETIME NOT NULL,
size INTEGER NOT NULL DEFAULT 0,
stored_size INTEGER NOT NULL DEFAULT 0,
truncated BOOLEAN NOT NULL DEFAULT FALSE,
compressed BOOLEAN NOT NULL DEFAULT FALSE,
error TEXT,
meta TEXT NOT NULL DEFAULT '{}',
preview TEXT,
data BLOB
);
CREATE UNIQUE INDEX idx_flow_messages_flow_id_seq ON flow_messages (flow_id, seq);
-- Where gRPC messages are decoded from: protos attached to a host pattern, or a server's
-- reflection, with the descriptors compiled from them.
CREATE TABLE grpc_schemas
(
id TEXT NOT NULL PRIMARY KEY,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL,
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL,
host_pattern TEXT NOT NULL,
source TEXT NOT NULL,
proto_paths TEXT NOT NULL DEFAULT '[]',
services TEXT NOT NULL DEFAULT '[]',
descriptor BLOB NOT NULL
);
+1 -1
View File
@@ -230,6 +230,6 @@ impl OfflineStore {
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
har::export(&self.db, &self.blobs, &flows, self.capture_body_limit).await
}
}
+89 -3
View File
@@ -4,12 +4,17 @@
//! 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::flows::{self, BodyPreview, FlowFilter};
use crate::grpc::{self, DecodedGrpcMessage};
use crate::har;
use crate::macos::SystemProxyStatus;
use crate::messages::{self, FlowMessage, MessageKind};
use crate::models::{FlowBody, find_header};
use crate::save::SaveFlowRequest;
use crate::{ProxyCtx, ProxyState, bodies};
use crate::{
AddGrpcSchemaRequest, ListMessagesResponse, ProxyCtx, ProxyState, ReflectGrpcHostRequest,
bodies,
};
use bytes::Bytes;
use chrono::{DateTime, Utc};
use http_body_util::{BodyExt, Full, Limited};
@@ -240,6 +245,25 @@ pub struct ErrorBody {
pub error: String,
}
/// `GET /flows/{id}/messages/{seq}`: a message with its payload, and a gRPC message decoded.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MessageDetail {
pub message: FlowMessage,
pub payload: BodyPreview,
/// A gRPC message as JSON, with its schema or without one.
#[serde(skip_serializing_if = "Option::is_none")]
pub decoded: Option<DecodedGrpcMessage>,
/// Why a gRPC message could not be decoded.
#[serde(skip_serializing_if = "Option::is_none")]
pub decode_error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RemovedResponse {
pub removed: bool,
}
/// 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";
@@ -504,6 +528,47 @@ async fn route(state: &ServerState, req: Request<Incoming>) -> Result<ControlRes
(&Method::GET, ["flows", id, "response-body"]) => {
raw_body(ctx, id, FlowBody::Response).await
}
(&Method::GET, ["flows", id, "messages"]) => {
found(ctx, id)?;
let (after_seq, limit) = message_page(&query)?;
let messages = messages::list(&ctx.db, id, after_seq, limit)?;
let total = messages::count(&ctx.db, id)?;
json(&ListMessagesResponse { messages, total })
}
(&Method::GET, ["flows", id, "messages", seq]) => {
let flow = flows::find(&ctx.db, id)
.ok_or_else(|| ApiError::new(StatusCode::NOT_FOUND, format!("No flow {id}")))?;
let seq: i64 = seq.parse().map_err(|_| bad_request(format!("Invalid seq {seq:?}")))?;
let message = messages::find(&ctx.db, id, seq)?.ok_or_else(|| {
ApiError::new(StatusCode::NOT_FOUND, format!("No message {seq} in flow {id}"))
})?;
let payload = messages::payload(&ctx.db, &ctx.blobs, &message)?;
let (decoded, decode_error) = if message.kind == MessageKind::GrpcMessage {
match grpc::decode(&ctx.db, &ctx.blobs, &ctx.grpc, &flow, seq).await {
Ok(decoded) => (Some(decoded), None),
Err(e) => (None, Some(e)),
}
} else {
(None, None)
};
let payload = BodyPreview::new(&payload, message.truncated);
json(&MessageDetail { message, payload, decoded, decode_error })
}
(&Method::GET, ["grpc-schemas"]) => json(&grpc::list(&ctx.db)?),
(&Method::POST, ["grpc-schemas"]) => {
let body: AddGrpcSchemaRequest = read_json(req).await?;
let added = grpc::add(&ctx.db, &ctx.grpc, &body.host_pattern, &body.proto_paths);
json(&added.await.map_err(bad_request)?)
}
(&Method::POST, ["grpc-schemas", "reflect"]) => {
let body: ReflectGrpcHostRequest = read_json(req).await?;
let (host, plaintext) = grpc::reflection_host(&body.host);
let reflected = grpc::reflect(&ctx.db, &ctx.grpc, &host, plaintext).await;
json(&reflected.map_err(bad_request)?)
}
(&Method::DELETE, ["grpc-schemas", id]) => {
json(&RemovedResponse { removed: grpc::remove(&ctx.db, &ctx.grpc, id)? })
}
(&Method::POST, ["flows", id, "save"]) => {
found(ctx, id)?;
let id = id.to_string();
@@ -513,7 +578,7 @@ async fn route(state: &ServerState, req: Request<Incoming>) -> Result<ControlRes
(&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?)
json(&har::export(&ctx.db, &ctx.blobs, &flows, ctx.capture_body_limit()?).await?)
}
(&Method::GET, ["ca"]) => {
let ca = CaStore::load_or_create(&ctx.ca_dir())
@@ -582,6 +647,27 @@ async fn raw_body(ctx: &ProxyCtx, id: &str, which: FlowBody) -> Result<ControlRe
Ok(resp)
}
/// `after` (a seq) and `limit` from a messages query.
fn message_page(query: &str) -> Result<(Option<i64>, Option<usize>), ApiError> {
let (mut after, mut limit) = (None, None);
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
match key.as_ref() {
"after" => {
after = Some(
value.parse().map_err(|_| bad_request(format!("Invalid after {value:?}")))?,
)
}
"limit" => {
limit = Some(
value.parse().map_err(|_| bad_request(format!("Invalid limit {value:?}")))?,
)
}
_ => {}
}
}
Ok((after, limit))
}
fn body_limit(query: &str) -> Result<usize, ApiError> {
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
if key == "bodyLimit" {
+18 -1
View File
@@ -141,6 +141,12 @@ mod tests {
ALTER TABLE flows DROP COLUMN client_bundle_id;
ALTER TABLE flows DROP COLUMN client_path;
ALTER TABLE flows DROP COLUMN cancelled;
ALTER TABLE flows DROP COLUMN stream_kind;
ALTER TABLE flows DROP COLUMN message_count;
ALTER TABLE flows DROP COLUMN grpc_service;
ALTER TABLE flows DROP COLUMN grpc_method;
ALTER TABLE flows DROP COLUMN grpc_status;
ALTER TABLE flows DROP COLUMN grpc_message;
INSERT INTO flows (id, url, method) VALUES ('fl_old', 'https://example.com', 'GET');",
)
.unwrap();
@@ -152,6 +158,11 @@ mod tests {
assert_eq!((flows[0].id.as_str(), flows[0].saved_request_id.as_deref()), ("fl_old", None));
assert_eq!((flows[0].client_pid, flows[0].client_name.as_deref()), (None, None));
assert!(!flows[0].cancelled);
assert_eq!((flows[0].stream_kind, flows[0].message_count), (None, 0));
// Failed flows are found without the gRPC status column too.
let failed =
crate::flows::FlowFilter { status: Some("error".into()), ..Default::default() };
assert_eq!(crate::flows::list(&db, &failed).unwrap().len(), 0);
}
#[test]
@@ -187,7 +198,7 @@ mod tests {
.unwrap()
.map(|t| t.unwrap())
.collect();
assert_eq!(tables, vec!["_sqlx_migrations", "flows"]);
assert_eq!(tables, vec!["_sqlx_migrations", "flow_messages", "flows", "grpc_schemas"]);
let flow_columns = columns(&conn, "flows");
assert!(flow_columns.contains(&"response_truncated".to_string()));
@@ -197,6 +208,12 @@ mod tests {
"client_bundle_id",
"client_path",
"cancelled",
"stream_kind",
"message_count",
"grpc_service",
"grpc_method",
"grpc_status",
"grpc_message",
] {
assert!(flow_columns.contains(&column.to_string()), "{column} in {flow_columns:?}");
}
+19 -5
View File
@@ -32,7 +32,8 @@ pub struct FlowFilter {
/// The method, any case.
#[ts(optional)]
pub method: Option<String>,
/// A status (`404`), a class of them (`4xx`), or `error` for flows that failed.
/// A status (`404`), a class of them (`4xx`), or `error` for flows that failed, gRPC calls
/// that ended with a status other than OK among them.
#[ts(optional)]
pub status: Option<String>,
/// At most this many flows, the newest.
@@ -114,7 +115,9 @@ impl FlowFilter {
}
/// The SQL condition for every field but `after`, `before` and `limit`, and its parameters.
fn condition(&self) -> (String, Vec<Value>) {
/// `grpc` says the database has gRPC statuses: one opened read-only before its migrations
/// have run may not.
fn condition(&self, grpc: bool) -> (String, Vec<Value>) {
let mut sql = String::from("1");
let mut params: Vec<Value> = Vec::new();
for term in self.q.iter().flat_map(|q| q.split_whitespace()) {
@@ -165,6 +168,10 @@ impl FlowFilter {
}
match self.status.as_deref().map(str::trim) {
None | Some("") => {}
// A gRPC call that ended with a status other than OK failed too.
Some(s) if s.eq_ignore_ascii_case("error") && grpc => {
sql.push_str(" AND (error IS NOT NULL OR coalesce(grpc_status, 0) != 0)")
}
Some(s) if s.eq_ignore_ascii_case("error") => sql.push_str(" AND error IS NOT NULL"),
// A whole status or a class, not the start of one as a search term may be.
Some(s) if s.len() == 3 => match status_condition(&s.to_ascii_lowercase(), &mut params)
@@ -268,8 +275,13 @@ pub fn page(
})
}
/// The database has the gRPC status column, added by a later migration.
fn has_grpc_status(conn: &rusqlite::Connection) -> bool {
conn.prepare("SELECT grpc_status FROM flows LIMIT 0").is_ok()
}
fn list_in(conn: &rusqlite::Connection, filter: &FlowFilter) -> Result<Vec<Flow>, String> {
let (condition, mut params) = filter.condition();
let (condition, mut params) = filter.condition(has_grpc_status(conn));
let mut sql = format!("SELECT * FROM flows WHERE {condition}");
if let Some(id) = &filter.after {
// By rowid, which only grows as flows are written.
@@ -300,7 +312,7 @@ fn list_in(conn: &rusqlite::Connection, filter: &FlowFilter) -> Result<Vec<Flow>
}
fn count_matching(conn: &rusqlite::Connection, filter: &FlowFilter) -> Result<usize, String> {
let (condition, params) = filter.condition();
let (condition, params) = filter.condition(has_grpc_status(conn));
conn.query_row(
&format!("SELECT COUNT(*) FROM flows WHERE {condition}"),
rusqlite::params_from_iter(params),
@@ -319,12 +331,14 @@ pub fn clear(db: &ProxyQueryManager, blobs: &BlobManager) -> Result<usize, Strin
let mut stmt = conn.prepare("SELECT id FROM flows")?;
let ids = stmt.query_map([], |r| r.get(0))?.collect::<Result<Vec<String>, _>>()?;
conn.execute("DELETE FROM flows", [])?;
conn.execute("DELETE FROM flow_messages", [])?;
// Only the deleted flows' bodies, so a flow still in flight keeps what it has written.
blobs
.with_tx(|b| {
for id in &ids {
b.delete_chunks(&Flow::body_id(id, FlowBody::Request))?;
b.delete_chunks(&Flow::body_id(id, FlowBody::Response))?;
b.delete_chunks_like(&format!("{id}.msg.%"))?;
}
Ok::<_, yaak_models::error::Error>(())
})
@@ -634,7 +648,7 @@ mod tests {
let plan: Vec<String> = db.with_conn(|c| {
let conn = c.conn().resolve();
let (condition, params) =
FlowFilter { q: Some("stripe".into()), ..Default::default() }.condition();
FlowFilter { q: Some("stripe".into()), ..Default::default() }.condition(true);
let sql = format!(
"EXPLAIN QUERY PLAN SELECT * FROM flows WHERE {condition} \
ORDER BY started_at DESC, id DESC LIMIT 500"
+595
View File
@@ -0,0 +1,595 @@
//! Decoding captured gRPC messages. A message decodes to JSON with the schema of its call's
//! service, found first in the protos attached to a host pattern, then in the server's own
//! reflection, asked once per host. Without one it decodes schemaless: field numbers, wire
//! types and values, with nested messages guessed.
use crate::db::ProxyQueryManager;
use crate::messages::{self, FlowMessage, MessageDirection, MessageKind};
use crate::models::Flow;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
use chrono::NaiveDateTime;
use rusqlite::{Row, params};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use ts_rs::TS;
use yaak_database::generate_prefixed_id;
use yaak_grpc::GrpcConfig;
use yaak_grpc::prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor};
use yaak_models::blob_manager::BlobManager;
/// Where a schema's descriptors came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export, export_to = "gen_models.ts")]
pub enum GrpcSchemaSource {
/// `.proto` files the user attached to a host pattern.
Files,
/// The server's reflection service.
Reflection,
}
/// Descriptors for decoding the gRPC calls to hosts that match `hostPattern`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export, export_to = "gen_models.ts")]
pub struct GrpcSchema {
pub id: String,
pub created_at: NaiveDateTime,
pub updated_at: NaiveDateTime,
/// A host (`api.example.com`), a host and port (`localhost:50051`), any subdomain of one
/// (`*.example.com`), or every host (`*`).
pub host_pattern: String,
pub source: GrpcSchemaSource,
/// The `.proto` files and include directories compiled, for a `files` schema.
pub proto_paths: Vec<String>,
/// The services the descriptors define, by full name.
pub services: Vec<String>,
}
/// A captured gRPC message as JSON.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export, export_to = "gen_rpc.ts")]
pub struct DecodedGrpcMessage {
/// Pretty JSON: the message's fields by name with a schema; without one, an array of its
/// fields as `{ field, wireType, value }`, a length-delimited one with its guesses (`string`,
/// `message`, `bytes` in base64).
pub json: String,
/// The message type decoded with: `package.HelloReply`. None when schemaless.
pub message_type: Option<String>,
/// The schema decoded with. None when schemaless.
pub schema: Option<GrpcSchema>,
/// Why it decoded schemaless, or what to know about the decode.
pub note: Option<String>,
}
/// What a host's reflection answered, once it has; decodes that need it while it is being asked
/// wait for the same answer.
type Lookup = Arc<tokio::sync::OnceCell<Result<GrpcSchema, String>>>;
/// The decoded descriptor pools, by schema id, and the hosts reflection was asked of.
#[derive(Clone, Default)]
pub struct GrpcState {
protoc: Arc<Mutex<Option<GrpcConfig>>>,
pools: Arc<Mutex<HashMap<String, DescriptorPool>>>,
/// By `host:port`.
lookups: Arc<Mutex<HashMap<String, Lookup>>>,
}
impl GrpcState {
pub fn set_protoc(&self, config: GrpcConfig) {
if let Ok(mut protoc) = self.protoc.lock() {
*protoc = Some(config);
}
}
/// The protoc set with [`Self::set_protoc`], or else one on `PATH`, with the `include`
/// directory installed beside it.
fn protoc(&self) -> Result<GrpcConfig, String> {
if let Some(config) = self.protoc.lock().ok().and_then(|p| p.clone()) {
return Ok(config);
}
protoc_on_path(std::env::var_os("PATH"))
}
}
/// What to do when there is no protoc to compile attached `.proto` files with.
pub const NO_PROTOC: &str = "Yaak Proxy has no protoc to compile .proto files with. A packaged \
Yaak Proxy includes one; otherwise install protoc (on macOS, `brew install protobuf`) so it \
is on PATH, or let the server's reflection supply the schema.";
/// A `protoc` in one of the directories of `path`, a `PATH` value, with the `include` directory
/// installed beside it.
fn protoc_on_path(path: Option<std::ffi::OsString>) -> Result<GrpcConfig, String> {
let name = if cfg!(windows) { "protoc.exe" } else { "protoc" };
let bin = path
.iter()
.flat_map(std::env::split_paths)
.map(|dir| dir.join(name))
.find(|p| p.is_file())
.ok_or(NO_PROTOC)?;
let include = bin.parent().and_then(|d| d.parent()).map(|d| d.join("include"));
Ok(GrpcConfig { protoc_include_dir: include.unwrap_or_default(), protoc_bin_path: bin })
}
const COLUMNS: &str = "id, created_at, updated_at, host_pattern, source, proto_paths, services";
fn schema_from_row(r: &Row) -> rusqlite::Result<GrpcSchema> {
let source: String = r.get("source")?;
let proto_paths: String = r.get("proto_paths")?;
let services: String = r.get("services")?;
Ok(GrpcSchema {
id: r.get("id")?,
created_at: r.get("created_at")?,
updated_at: r.get("updated_at")?,
host_pattern: r.get("host_pattern")?,
source: if source == "reflection" {
GrpcSchemaSource::Reflection
} else {
GrpcSchemaSource::Files
},
proto_paths: serde_json::from_str(&proto_paths).unwrap_or_default(),
services: serde_json::from_str(&services).unwrap_or_default(),
})
}
pub fn list(db: &ProxyQueryManager) -> Result<Vec<GrpcSchema>, String> {
db.with_conn(|c| {
let conn = c.conn().resolve();
let mut stmt = conn
.prepare(&format!("SELECT {COLUMNS} FROM grpc_schemas ORDER BY created_at, id"))
.map_err(|e| e.to_string())?;
let rows = stmt.query_map([], schema_from_row).map_err(|e| e.to_string())?;
rows.collect::<Result<Vec<_>, _>>().map_err(|e| e.to_string())
})
}
fn get(db: &ProxyQueryManager, id: &str) -> Result<GrpcSchema, String> {
db.with_conn(|c| {
c.conn()
.resolve()
.query_row(
&format!("SELECT {COLUMNS} FROM grpc_schemas WHERE id = ?1"),
[id],
schema_from_row,
)
.map_err(|_| format!("No gRPC schema {id}"))
})
}
fn insert(
db: &ProxyQueryManager,
host_pattern: &str,
source: GrpcSchemaSource,
proto_paths: &[String],
pool: &DescriptorPool,
) -> Result<GrpcSchema, String> {
let id = generate_prefixed_id("gs");
let services: Vec<String> = pool.services().map(|s| s.full_name().to_string()).collect();
let source_name = match source {
GrpcSchemaSource::Files => "files",
GrpcSchemaSource::Reflection => "reflection",
};
db.with_tx(|c| {
let conn = c.conn().resolve();
// A host's reflection is kept once, the latest.
if source == GrpcSchemaSource::Reflection {
conn.execute(
"DELETE FROM grpc_schemas WHERE source = 'reflection' AND host_pattern = ?1",
[host_pattern],
)?;
}
conn.execute(
"INSERT INTO grpc_schemas (id, host_pattern, source, proto_paths, services, descriptor) \
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
params![
id,
host_pattern,
source_name,
serde_json::to_string(proto_paths).unwrap_or_default(),
serde_json::to_string(&services).unwrap_or_default(),
pool.encode_to_vec(),
],
)?;
Ok::<_, yaak_database::Error>(())
})
.map_err(|e| format!("Failed to save the gRPC schema: {e}"))?;
get(db, &id)
}
/// Compiles `proto_paths` (files, and directories to import from) and attaches them to
/// `host_pattern`.
pub async fn add(
db: &ProxyQueryManager,
state: &GrpcState,
host_pattern: &str,
proto_paths: &[String],
) -> Result<GrpcSchema, String> {
let host_pattern = host_pattern.trim().to_ascii_lowercase();
if host_pattern.is_empty() {
return Err("A host pattern is required".into());
}
if proto_paths.is_empty() {
return Err("At least one .proto file is required".into());
}
let paths: Vec<PathBuf> = proto_paths.iter().map(PathBuf::from).collect();
if let Some(missing) = paths.iter().find(|p| !p.exists()) {
return Err(format!("{} does not exist", missing.display()));
}
let config = state.protoc()?;
let pool = yaak_grpc::fill_pool_from_files(&config, &paths)
.await
.map_err(|e| format!("Failed to compile the protos: {e}"))?;
insert(db, &host_pattern, GrpcSchemaSource::Files, proto_paths, &pool)
}
pub fn remove(db: &ProxyQueryManager, state: &GrpcState, id: &str) -> Result<bool, String> {
let removed = db
.with_tx(|c| {
let removed =
c.conn().resolve().execute("DELETE FROM grpc_schemas WHERE id = ?1", [id])?;
Ok::<_, yaak_database::Error>(removed)
})
.map_err(|e| e.to_string())?;
if let Ok(mut pools) = state.pools.lock() {
pools.remove(id);
}
Ok(removed > 0)
}
/// How long a server has to answer reflection, all of it.
const REFLECTION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15);
/// A host to ask for reflection, as given to `reflect_grpc_host`: `host`, `host:port`, or
/// either after `http://` or `https://`. Returns `host:port` and whether to ask without TLS;
/// without a scheme, over TLS on 443.
pub fn reflection_host(host: &str) -> (String, bool) {
let host = host.trim().trim_end_matches('/').to_ascii_lowercase();
let (host, plaintext) = match host.strip_prefix("http://") {
Some(rest) => (rest.to_string(), true),
None => (host.strip_prefix("https://").unwrap_or(&host).to_string(), false),
};
let has_port = host.rsplit_once(':').is_some_and(|(_, port)| port.parse::<u16>().is_ok());
if has_port {
return (host, plaintext);
}
(format!("{host}:{}", if plaintext { 80 } else { 443 }), plaintext)
}
/// Asks the server at `authority` (`host:port`) for its services by reflection, without TLS
/// when `plaintext`, and keeps them as the host's schema, also for decodes to come. Fails when
/// it offers none.
pub async fn reflect(
db: &ProxyQueryManager,
state: &GrpcState,
authority: &str,
plaintext: bool,
) -> Result<GrpcSchema, String> {
let authority = authority.trim().to_ascii_lowercase();
let reflected = reflect_uncached(db, &authority, plaintext).await;
if let Ok(mut lookups) = state.lookups.lock() {
let answered = tokio::sync::OnceCell::new_with(Some(reflected.clone()));
lookups.insert(authority, Arc::new(answered));
}
reflected
}
async fn reflect_uncached(
db: &ProxyQueryManager,
authority: &str,
plaintext: bool,
) -> Result<GrpcSchema, String> {
let scheme = if plaintext { "http" } else { "https" };
let uri: hyper::Uri = format!("{scheme}://{authority}")
.parse()
.map_err(|e| format!("Invalid host {authority}: {e}"))?;
// Reflection's client panics on an answer it does not expect; that ends this task alone.
// Timed out inside the task, so a server that stalls leaves nothing running.
let timeout_error = format!("{authority} did not answer reflection in time");
let reflected = tokio::spawn(async move {
let metadata = Default::default();
let reflecting =
yaak_grpc::fill_pool_from_reflection(&uri, &metadata, true, None, 64 << 20);
match tokio::time::timeout(REFLECTION_TIMEOUT, reflecting).await {
Ok(reflected) => reflected.map_err(|e| e.to_string()),
Err(_) => Err(timeout_error),
}
})
.await
.map_err(|_| format!("{authority} answered reflection with something unexpected"))??;
if reflected.services().next().is_none() {
return Err(format!("{authority} offers no services by reflection"));
}
insert(db, authority, GrpcSchemaSource::Reflection, &[], &reflected)
}
/// Whether `pattern` covers `host` (lowercase, no port) at `authority` (`host:port`).
fn matches(pattern: &str, host: &str, authority: &str) -> bool {
if pattern == "*" {
return true;
}
let target = if pattern.contains(':') { authority } else { host };
match pattern.strip_prefix("*.") {
Some(domain) => target.strip_suffix(domain).is_some_and(|rest| rest.ends_with('.')),
None => pattern == target,
}
}
impl GrpcState {
fn pool(&self, db: &ProxyQueryManager, id: &str) -> Result<DescriptorPool, String> {
if let Some(pool) = self.pools.lock().ok().and_then(|p| p.get(id).cloned()) {
return Ok(pool);
}
let bytes: Vec<u8> = db.with_conn(|c| {
c.conn()
.resolve()
.query_row("SELECT descriptor FROM grpc_schemas WHERE id = ?1", [id], |r| r.get(0))
.map_err(|e| e.to_string())
})?;
let pool = DescriptorPool::decode(bytes.as_slice()).map_err(|e| e.to_string())?;
if let Ok(mut pools) = self.pools.lock() {
pools.insert(id.to_string(), pool.clone());
}
Ok(pool)
}
}
/// The schema that defines `service` for the host at `authority`: attached protos first, the
/// most specific pattern first, then the host's reflection, asked of the server once, without
/// TLS when `plaintext`, as the call was made.
async fn find_schema(
db: &ProxyQueryManager,
state: &GrpcState,
authority: &str,
plaintext: bool,
service: &str,
) -> Result<(Option<(GrpcSchema, DescriptorPool)>, Option<String>), String> {
let host = authority.rsplit_once(':').map_or(authority, |(h, _)| h);
let mut schemas: Vec<GrpcSchema> =
list(db)?.into_iter().filter(|s| matches(&s.host_pattern, host, authority)).collect();
// Files before reflection; then exact before wildcards, the narrower first.
schemas.sort_by_key(|s| {
let wildcard = s.host_pattern.starts_with('*');
(
s.source == GrpcSchemaSource::Reflection,
wildcard,
std::cmp::Reverse(s.host_pattern.len()),
)
});
for schema in &schemas {
let pool = state.pool(db, &schema.id)?;
if pool.get_service_by_name(service).is_some() {
return Ok((Some((schema.clone(), pool)), None));
}
}
if schemas.iter().any(|s| s.source == GrpcSchemaSource::Reflection) {
return Ok((None, Some(format!("No schema for {service} on {authority}"))));
}
// Asked once: decodes that come while it is being asked wait for the same answer.
let Some(lookup) =
state.lookups.lock().ok().map(|mut l| l.entry(authority.into()).or_default().clone())
else {
return Ok((None, Some(format!("No schema for {service} on {authority}"))));
};
let answer = lookup.get_or_init(|| reflect_uncached(db, authority, plaintext)).await.clone();
match answer {
// Unless it was removed since.
Ok(schema) => match state.pool(db, &schema.id) {
Ok(pool) if pool.get_service_by_name(service).is_some() => {
Ok((Some((schema, pool)), None))
}
Ok(_) => {
Ok((None, Some(format!("{authority}'s reflection does not include {service}"))))
}
Err(_) => Ok((None, Some(format!("No schema for {service} on {authority}")))),
},
Err(e) => Ok((None, Some(format!("No schema for {service}; reflection failed: {e}")))),
}
}
/// Decodes a captured gRPC message.
pub async fn decode(
db: &ProxyQueryManager,
blobs: &BlobManager,
state: &GrpcState,
flow: &Flow,
seq: i64,
) -> Result<DecodedGrpcMessage, String> {
let message = messages::get(db, &flow.id, seq)?;
if message.kind != MessageKind::GrpcMessage {
return Err(format!("Message {seq} is not a gRPC message"));
}
let payload = messages::payload(db, blobs, &message)?;
if message.compressed && message.error.is_some() {
return Err(message.error.unwrap_or_default());
}
let (Some(service), Some(method)) = (&flow.grpc_service, &flow.grpc_method) else {
return Ok(schemaless(&payload, &message, Some("The call's path names no service".into())));
};
let url = url::Url::parse(&flow.url).map_err(|e| e.to_string())?;
let host = url.host_str().unwrap_or_default().to_ascii_lowercase();
let authority = format!("{host}:{}", url.port_or_known_default().unwrap_or(443));
let plaintext = url.scheme() == "http";
let (found, note) = find_schema(db, state, &authority, plaintext, service).await?;
let Some((schema, pool)) = found else {
return Ok(schemaless(&payload, &message, note));
};
let Some(descriptor) = pool
.get_service_by_name(service)
.and_then(|s| s.methods().find(|m| m.name() == method))
.map(|m| match message.direction {
MessageDirection::Client => m.input(),
MessageDirection::Server => m.output(),
})
else {
let note = format!("{} has no method {method}", service);
return Ok(schemaless(&payload, &message, Some(note)));
};
decode_with(&payload, &message, descriptor, schema)
}
fn decode_with(
payload: &[u8],
message: &FlowMessage,
descriptor: MessageDescriptor,
schema: GrpcSchema,
) -> Result<DecodedGrpcMessage, String> {
let message_type = descriptor.full_name().to_string();
match DynamicMessage::decode(descriptor, payload) {
Ok(decoded) => Ok(DecodedGrpcMessage {
json: yaak_grpc::message_to_json(&decoded)?,
message_type: Some(message_type),
schema: Some(schema),
note: message.truncated.then(|| "The message is cut at the capture limit".into()),
}),
Err(e) => {
Ok(schemaless(payload, message, Some(format!("Did not decode as {message_type}: {e}"))))
}
}
}
fn schemaless(payload: &[u8], message: &FlowMessage, reason: Option<String>) -> DecodedGrpcMessage {
let mut notes: Vec<String> = reason.into_iter().collect();
notes.push("Decoded without a schema: field names and types are guesses".into());
if message.truncated {
notes.push("The message is cut at the capture limit".into());
}
let fields = decode_fields(payload, 0).unwrap_or_else(|| {
notes.push("Not a valid protobuf message: shown as bytes".into());
vec![json!({ "bytes": BASE64.encode(payload) })]
});
DecodedGrpcMessage {
json: serde_json::to_string_pretty(&fields).unwrap_or_default(),
message_type: None,
schema: None,
note: Some(notes.join(". ")),
}
}
const MAX_DEPTH: usize = 16;
/// A message's fields from the protobuf wire format, None when the bytes are not one.
pub fn decode_fields(mut bytes: &[u8], depth: usize) -> Option<Vec<Value>> {
let mut fields = Vec::new();
while !bytes.is_empty() {
let key = read_varint(&mut bytes)?;
let field = key >> 3;
if field == 0 || field > 536_870_911 {
return None;
}
let value = match key & 7 {
0 => {
let v = read_varint(&mut bytes)?;
json!({ "field": field, "wireType": "varint", "value": number(v) })
}
1 => {
let raw: [u8; 8] = bytes.get(..8)?.try_into().ok()?;
bytes = &bytes[8..];
let v = u64::from_le_bytes(raw);
json!({ "field": field, "wireType": "fixed64", "value": number(v), "double": f64::from_le_bytes(raw) })
}
2 => {
let len = usize::try_from(read_varint(&mut bytes)?).ok()?;
let data = bytes.get(..len)?;
bytes = &bytes[len..];
let mut value = json!({ "field": field, "wireType": "len" });
let text = std::str::from_utf8(data)
.ok()
.filter(|s| s.chars().all(|c| !c.is_control() || c.is_whitespace()));
if let Some(text) = text {
value["string"] = json!(text);
}
if depth < MAX_DEPTH
&& !data.is_empty()
&& let Some(nested) = decode_fields(data, depth + 1)
{
value["message"] = json!(nested);
}
if text.is_none() {
value["bytes"] = json!(BASE64.encode(data));
}
value
}
5 => {
let raw: [u8; 4] = bytes.get(..4)?.try_into().ok()?;
bytes = &bytes[4..];
let v = u32::from_le_bytes(raw);
json!({ "field": field, "wireType": "fixed32", "value": v, "float": f32::from_le_bytes(raw) })
}
// Groups are long deprecated; 6 and 7 are no wire type.
_ => return None,
};
fields.push(value);
}
Some(fields)
}
/// A u64 as a JSON number when JavaScript can hold it exactly, else as a string.
fn number(v: u64) -> Value {
if v <= (1 << 53) { json!(v) } else { json!(v.to_string()) }
}
fn read_varint(bytes: &mut &[u8]) -> Option<u64> {
let mut value = 0u64;
for i in 0..10 {
let byte = *bytes.get(i)?;
value |= u64::from(byte & 0x7f) << (7 * i);
if byte & 0x80 == 0 {
*bytes = &bytes[i + 1..];
return Some(value);
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn without_protoc_the_error_says_what_to_install() {
let empty = tempfile::tempdir().unwrap();
let err = protoc_on_path(Some(empty.path().as_os_str().to_owned())).err().unwrap();
assert_eq!(err, NO_PROTOC);
assert!(err.contains("brew install protobuf"), "{err}");
assert!(protoc_on_path(None).is_err());
}
#[test]
fn reflection_hosts() {
assert_eq!(reflection_host("api.example.com"), ("api.example.com:443".into(), false));
assert_eq!(reflection_host("localhost:50051"), ("localhost:50051".into(), false));
assert_eq!(reflection_host("http://localhost:50051"), ("localhost:50051".into(), true));
assert_eq!(reflection_host("http://Example.com/"), ("example.com:80".into(), true));
assert_eq!(reflection_host("https://example.com"), ("example.com:443".into(), false));
}
#[test]
fn host_patterns() {
assert!(matches("api.example.com", "api.example.com", "api.example.com:443"));
assert!(!matches("api.example.com", "www.example.com", "www.example.com:443"));
assert!(matches("*.example.com", "a.b.example.com", "a.b.example.com:443"));
assert!(!matches("*.example.com", "example.com", "example.com:443"));
assert!(!matches("*.example.com", "badexample.com", "badexample.com:443"));
assert!(matches("localhost:50051", "localhost", "localhost:50051"));
assert!(!matches("localhost:50051", "localhost", "localhost:50052"));
assert!(matches("*", "anything", "anything:1"));
}
#[test]
fn decodes_the_wire_format_without_a_schema() {
// field 1: varint 150; field 2: "hi"; field 3: nested { field 1: 1 }
let bytes = [
0x08, 0x96, 0x01, 0x12, 0x02, b'h', b'i', 0x1a, 0x02, 0x08, 0x01,
];
let fields = decode_fields(&bytes, 0).unwrap();
assert_eq!(fields[0], json!({ "field": 1, "wireType": "varint", "value": 150 }));
assert_eq!(fields[1]["string"], "hi");
assert_eq!(fields[2]["message"][0]["value"], 1);
assert!(decode_fields(&[0x08], 0).is_none(), "a varint cut short");
}
}
+74 -4
View File
@@ -4,8 +4,12 @@
//! at the capture limit and `_error` on a flow that failed.
use crate::bodies::{self, DecodedBody};
use crate::db::ProxyQueryManager;
use crate::flows::{BodyEncoding, BodyPreview};
use crate::models::{Flow, FlowBody, ProxyHeader, find_header};
use crate::messages::{self, MessageDirection, MessageKind};
use crate::models::{Flow, FlowBody, ProxyHeader, StreamKind, find_header};
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
use serde::{Deserialize, Serialize};
use yaak_models::blob_manager::BlobManager;
@@ -40,6 +44,26 @@ pub struct HarEntry {
pub server_ip_address: Option<String>,
#[serde(rename = "_error", skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
/// A WebSocket's data messages, as Chrome's devtools export them.
#[serde(
rename = "_webSocketMessages",
skip_serializing_if = "Option::is_none",
default
)]
pub web_socket_messages: Option<Vec<HarWebSocketMessage>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HarWebSocketMessage {
/// `send` from the client, `receive` from the server.
#[serde(rename = "type")]
pub kind: String,
/// Seconds since the epoch.
pub time: f64,
/// 1 for text, 2 for binary.
pub opcode: u8,
/// The text, or a binary message in base64.
pub data: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -118,16 +142,59 @@ pub struct HarTimings {
}
/// 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> {
/// A WebSocket's messages are in its entry's `_webSocketMessages`; an event stream's events are
/// its response body, as they came.
pub async fn export(
db: &ProxyQueryManager,
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));
let mut entry = entry(flow, &request, &response);
if flow.stream_kind == Some(StreamKind::Websocket) {
entry.web_socket_messages = Some(websocket_messages(db, blobs, flow)?);
}
entries.push(entry);
}
Ok(har(entries))
}
fn websocket_messages(
db: &ProxyQueryManager,
blobs: &BlobManager,
flow: &Flow,
) -> Result<Vec<HarWebSocketMessage>, String> {
let mut out = Vec::new();
for message in messages::list(db, &flow.id, None, None)? {
let opcode = match message.kind {
MessageKind::WebsocketText => 1,
MessageKind::WebsocketBinary => 2,
_ => continue,
};
let payload = messages::payload(db, blobs, &message)?;
let data = if opcode == 1 {
String::from_utf8_lossy(&payload).into_owned()
} else {
BASE64.encode(&payload)
};
out.push(HarWebSocketMessage {
kind: match message.direction {
MessageDirection::Client => "send",
MessageDirection::Server => "receive",
}
.into(),
time: message.timestamp.and_utc().timestamp_micros() as f64 / 1_000_000.0,
opcode,
data,
});
}
Ok(out)
}
pub fn har(entries: Vec<HarEntry>) -> Har {
Har {
log: HarLog {
@@ -186,6 +253,7 @@ pub fn entry(flow: &Flow, request_body: &DecodedBody, response_body: &DecodedBod
.clone()
.or_else(|| response_body.decode_error.clone())
.or_else(|| request_body.decode_error.clone()),
web_socket_messages: None,
}
}
@@ -394,7 +462,9 @@ mod tests {
..Default::default()
};
let har = export(&blobs, &[flow], 1024 * 1024).await.unwrap();
let dir = tempfile::tempdir().unwrap();
let db = ProxyQueryManager::new(&dir.path().join("proxy.db")).unwrap();
let har = export(&db, &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()));
File diff suppressed because it is too large. Load diff
+368
View File
@@ -0,0 +1,368 @@
//! The messages of streaming flows: WebSocket frames, server-sent events, and gRPC messages and
//! trailers. Each is a row of `flow_messages`, its payload inline when small, else in
//! `blobs.sqlite`. Like flows, they are never pruned; clearing removes them with their flow.
use crate::db::ProxyQueryManager;
use chrono::{DateTime, NaiveDateTime, Utc};
use rusqlite::{Row, params};
use serde::{Deserialize, Serialize};
use ts_rs::TS;
use yaak_database::generate_prefixed_id;
use yaak_models::blob_manager::BlobManager;
/// Payloads up to this many bytes are kept in the row; larger ones in `blobs.sqlite`.
pub const INLINE_LIMIT: usize = 4096;
/// Characters of a text payload kept as its preview.
pub const PREVIEW_CHARS: usize = 200;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export, export_to = "gen_models.ts")]
pub enum MessageDirection {
Client,
Server,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export, export_to = "gen_models.ts")]
pub enum MessageKind {
WebsocketText,
WebsocketBinary,
WebsocketPing,
WebsocketPong,
WebsocketClose,
SseEvent,
GrpcMessage,
GrpcTrailers,
}
const KINDS: [MessageKind; 8] = [
MessageKind::WebsocketText,
MessageKind::WebsocketBinary,
MessageKind::WebsocketPing,
MessageKind::WebsocketPong,
MessageKind::WebsocketClose,
MessageKind::SseEvent,
MessageKind::GrpcMessage,
MessageKind::GrpcTrailers,
];
impl MessageKind {
pub fn as_str(self) -> &'static str {
match self {
Self::WebsocketText => "websocket_text",
Self::WebsocketBinary => "websocket_binary",
Self::WebsocketPing => "websocket_ping",
Self::WebsocketPong => "websocket_pong",
Self::WebsocketClose => "websocket_close",
Self::SseEvent => "sse_event",
Self::GrpcMessage => "grpc_message",
Self::GrpcTrailers => "grpc_trailers",
}
}
pub fn parse(s: &str) -> Option<Self> {
KINDS.into_iter().find(|k| k.as_str() == s)
}
/// The WebSocket opcode, for HAR's `_webSocketMessages`.
pub fn websocket_opcode(self) -> Option<u8> {
match self {
Self::WebsocketText => Some(1),
Self::WebsocketBinary => Some(2),
Self::WebsocketClose => Some(8),
Self::WebsocketPing => Some(9),
Self::WebsocketPong => Some(10),
_ => None,
}
}
}
/// What a kind of message carries beside its payload. Fields that do not apply are left out.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(default, rename_all = "camelCase")]
#[ts(export, export_to = "gen_models.ts")]
pub struct MessageMeta {
/// A WebSocket close frame's status code.
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub close_code: Option<u16>,
/// A WebSocket close frame's reason.
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub close_reason: Option<String>,
/// An SSE event's `event` field.
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub event: Option<String>,
/// An SSE event's `id` field.
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub event_id: Option<String>,
/// An SSE event's `retry` field, in milliseconds.
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional, type = "number")]
pub retry: Option<u64>,
/// An SSE event's lines as they came, comments included, up to 4 KB.
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub raw: Option<String>,
/// gRPC trailers' `grpc-status`.
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub grpc_status: Option<i32>,
/// gRPC trailers' `grpc-message`, decoded.
#[serde(skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub grpc_message: Option<String>,
}
/// One message of a streaming flow. Its payload is fetched on its own.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export, export_to = "gen_models.ts")]
pub struct FlowMessage {
pub id: String,
pub flow_id: String,
/// Its place in the flow, from 0, in the order messages passed the proxy.
#[ts(type = "number")]
pub seq: i64,
pub direction: MessageDirection,
pub kind: MessageKind,
/// When it had passed the proxy whole.
pub timestamp: NaiveDateTime,
/// Bytes it took on the wire: the frames' payloads, the event's lines, the gRPC message
/// without its prefix.
#[ts(type = "number")]
pub size: i64,
/// Bytes of payload stored: unmasked, reassembled, inflated or decompressed, and cut at the
/// capture limit.
#[ts(type = "number")]
pub stored_size: i64,
/// The stored payload is cut short at the capture limit.
pub truncated: bool,
/// It was compressed on the wire (permessage-deflate, gRPC's compressed flag); the payload is
/// stored decompressed unless `error` says otherwise.
pub compressed: bool,
/// Why the payload could not be decoded or stored.
pub error: Option<String>,
pub meta: MessageMeta,
/// The payload's start as text, when it is text.
pub preview: Option<String>,
}
impl FlowMessage {
/// The message as it is stored, from the engine's.
pub fn from_engine(flow_id: &str, seq: i64, message: &yaak_proxy::StreamMessage) -> Self {
use yaak_proxy::MessageKind as K;
let mut meta = MessageMeta::default();
let kind = match &message.kind {
K::WebSocketText => MessageKind::WebsocketText,
K::WebSocketBinary => MessageKind::WebsocketBinary,
K::WebSocketPing => MessageKind::WebsocketPing,
K::WebSocketPong => MessageKind::WebsocketPong,
K::WebSocketClose { code, reason } => {
meta.close_code = *code;
meta.close_reason = Some(reason.clone()).filter(|r| !r.is_empty());
MessageKind::WebsocketClose
}
K::SseEvent { event, id, retry, raw } => {
meta.event = event.clone();
meta.event_id = id.clone();
meta.retry = *retry;
meta.raw = Some(raw.clone());
MessageKind::SseEvent
}
K::GrpcMessage => MessageKind::GrpcMessage,
K::GrpcTrailers { status, message } => {
meta.grpc_status = *status;
meta.grpc_message = message.clone();
MessageKind::GrpcTrailers
}
};
let direction = match message.direction {
yaak_proxy::MessageDirection::Client => MessageDirection::Client,
yaak_proxy::MessageDirection::Server => MessageDirection::Server,
};
let preview = match kind {
// Binary by nature, however it happens to read.
MessageKind::WebsocketBinary | MessageKind::GrpcMessage => None,
MessageKind::WebsocketClose => meta.close_reason.clone(),
_ => text_preview(&message.data, message.truncated),
};
Self {
id: generate_prefixed_id("fm"),
flow_id: flow_id.to_string(),
seq,
direction,
kind,
timestamp: DateTime::<Utc>::from(message.at).naive_utc(),
size: i64::try_from(message.size).unwrap_or(i64::MAX),
stored_size: message.data.len() as i64,
truncated: message.truncated,
compressed: message.compressed,
error: message.error.clone(),
meta,
preview,
}
}
/// Where a payload too large for the row is kept in `blobs.sqlite`.
pub fn blob_id(&self) -> String {
format!("{}.msg.{}", self.flow_id, self.id)
}
fn from_row(r: &Row) -> rusqlite::Result<Self> {
let direction: String = r.get("direction")?;
let kind: String = r.get("kind")?;
let meta: String = r.get("meta")?;
Ok(Self {
id: r.get("id")?,
flow_id: r.get("flow_id")?,
seq: r.get("seq")?,
direction: if direction == "client" {
MessageDirection::Client
} else {
MessageDirection::Server
},
kind: MessageKind::parse(&kind).unwrap_or(MessageKind::WebsocketBinary),
timestamp: r.get("timestamp")?,
size: r.get("size")?,
stored_size: r.get("stored_size")?,
truncated: r.get("truncated")?,
compressed: r.get("compressed")?,
error: r.get("error")?,
meta: serde_json::from_str(&meta).unwrap_or_default(),
preview: r.get("preview")?,
})
}
}
/// The start of `data` as text, if it is UTF-8: a cut can split a character at the end.
fn text_preview(data: &[u8], truncated: bool) -> Option<String> {
let text = match std::str::from_utf8(data) {
Ok(text) => text,
Err(e) if truncated && e.error_len().is_none() => {
std::str::from_utf8(&data[..e.valid_up_to()]).ok()?
}
Err(_) => return None,
};
Some(text.chars().take(PREVIEW_CHARS).collect())
}
/// Writes `message`, with its payload in the row when `inline` has it.
pub fn insert(
conn: &rusqlite::Connection,
message: &FlowMessage,
inline: Option<&[u8]>,
) -> rusqlite::Result<()> {
let meta = serde_json::to_string(&message.meta).unwrap_or_else(|_| "{}".into());
let direction = match message.direction {
MessageDirection::Client => "client",
MessageDirection::Server => "server",
};
conn.execute(
"INSERT INTO flow_messages (id, flow_id, seq, direction, kind, timestamp, size, \
stored_size, truncated, compressed, error, meta, preview, data) \
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)",
params![
message.id,
message.flow_id,
message.seq,
direction,
message.kind.as_str(),
message.timestamp,
message.size,
message.stored_size,
message.truncated,
message.compressed,
message.error,
meta,
message.preview,
inline,
],
)?;
Ok(())
}
const COLUMNS: &str = "id, flow_id, seq, direction, kind, timestamp, size, stored_size, \
truncated, compressed, error, meta, preview";
/// A flow's messages in order, after `after_seq` when given, at most `limit` of them.
pub fn list(
db: &ProxyQueryManager,
flow_id: &str,
after_seq: Option<i64>,
limit: Option<usize>,
) -> Result<Vec<FlowMessage>, String> {
db.with_conn(|c| {
let conn = c.conn().resolve();
let limit = limit.map_or(-1, |l| l.min(i64::MAX as usize) as i64);
let mut stmt = conn
.prepare(&format!(
"SELECT {COLUMNS} FROM flow_messages WHERE flow_id = ?1 AND seq > ?2 \
ORDER BY seq LIMIT ?3"
))
.map_err(|e| e.to_string())?;
let rows = stmt
.query_map(params![flow_id, after_seq.unwrap_or(-1), limit], FlowMessage::from_row)
.map_err(|e| e.to_string())?;
rows.collect::<Result<Vec<_>, _>>().map_err(|e| e.to_string())
})
}
/// Messages recorded for a flow.
pub fn count(db: &ProxyQueryManager, flow_id: &str) -> Result<i64, String> {
db.with_conn(|c| {
c.conn()
.resolve()
.query_row("SELECT COUNT(*) FROM flow_messages WHERE flow_id = ?1", [flow_id], |r| {
r.get(0)
})
.map_err(|e| e.to_string())
})
}
pub fn find(
db: &ProxyQueryManager,
flow_id: &str,
seq: i64,
) -> Result<Option<FlowMessage>, String> {
db.with_conn(|c| {
let conn = c.conn().resolve();
let found = conn.query_row(
&format!("SELECT {COLUMNS} FROM flow_messages WHERE flow_id = ?1 AND seq = ?2"),
params![flow_id, seq],
FlowMessage::from_row,
);
match found {
Ok(message) => Ok(Some(message)),
Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
Err(e) => Err(e.to_string()),
}
})
}
pub fn get(db: &ProxyQueryManager, flow_id: &str, seq: i64) -> Result<FlowMessage, String> {
find(db, flow_id, seq)?.ok_or_else(|| format!("No message {seq} in flow {flow_id}"))
}
/// A message's stored payload: decoded, and cut at the capture limit.
pub fn payload(
db: &ProxyQueryManager,
blobs: &BlobManager,
message: &FlowMessage,
) -> Result<Vec<u8>, String> {
let inline: Option<Vec<u8>> = db.with_conn(|c| {
c.conn()
.resolve()
.query_row("SELECT data FROM flow_messages WHERE id = ?1", [&message.id], |r| r.get(0))
.map_err(|e| e.to_string())
})?;
if let Some(data) = inline {
return Ok(data);
}
let chunks = blobs.connect().get_chunks(&message.blob_id()).map_err(|e| e.to_string())?;
Ok(chunks.into_iter().flat_map(|c| c.data).collect())
}
+77
View File
@@ -66,6 +66,64 @@ pub struct Flow {
pub client_bundle_id: Option<String>,
/// Its app bundle's path, or its executable's outside an app.
pub client_path: Option<String>,
/// The flow carries a stream of messages, listed with `list_messages`. It is written as the
/// stream opens and again as messages land; it is open until `elapsedMs`, `error` or
/// `cancelled` is set.
pub stream_kind: Option<StreamKind>,
/// Messages recorded so far.
#[ts(type = "number")]
pub message_count: i64,
/// A gRPC call's service, from its path: `package.Service`.
pub grpc_service: Option<String>,
/// A gRPC call's method, from its path.
pub grpc_method: Option<String>,
/// A gRPC call's status, from its trailers. Not 0 is a failed call.
pub grpc_status: Option<i32>,
/// A gRPC call's `grpc-message`, decoded.
pub grpc_message: Option<String>,
}
/// The streaming protocol a flow carries.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
#[serde(rename_all = "snake_case")]
#[ts(export, export_to = "gen_models.ts")]
pub enum StreamKind {
Websocket,
Sse,
Grpc,
GrpcWeb,
}
impl StreamKind {
pub fn as_str(self) -> &'static str {
match self {
Self::Websocket => "websocket",
Self::Sse => "sse",
Self::Grpc => "grpc",
Self::GrpcWeb => "grpc_web",
}
}
pub fn parse(s: &str) -> Option<Self> {
[Self::Websocket, Self::Sse, Self::Grpc, Self::GrpcWeb]
.into_iter()
.find(|k| k.as_str() == s)
}
pub fn is_grpc(self) -> bool {
matches!(self, Self::Grpc | Self::GrpcWeb)
}
}
impl From<yaak_proxy::StreamKind> for StreamKind {
fn from(kind: yaak_proxy::StreamKind) -> Self {
match kind {
yaak_proxy::StreamKind::WebSocket => Self::Websocket,
yaak_proxy::StreamKind::Sse => Self::Sse,
yaak_proxy::StreamKind::Grpc => Self::Grpc,
yaak_proxy::StreamKind::GrpcWeb => Self::GrpcWeb,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)]
@@ -164,6 +222,12 @@ impl UpsertModelInfo for Flow {
(ClientName, self.client_name.into()),
(ClientBundleId, self.client_bundle_id.into()),
(ClientPath, self.client_path.into()),
(FlowIden::StreamKind, self.stream_kind.map(|k| k.as_str()).into()),
(MessageCount, self.message_count.into()),
(GrpcService, self.grpc_service.into()),
(GrpcMethod, self.grpc_method.into()),
(GrpcStatus, self.grpc_status.into()),
(GrpcMessage, self.grpc_message.into()),
])
}
@@ -193,6 +257,12 @@ impl UpsertModelInfo for Flow {
FlowIden::ClientName,
FlowIden::ClientBundleId,
FlowIden::ClientPath,
FlowIden::StreamKind,
FlowIden::MessageCount,
FlowIden::GrpcService,
FlowIden::GrpcMethod,
FlowIden::GrpcStatus,
FlowIden::GrpcMessage,
]
}
@@ -229,6 +299,13 @@ impl UpsertModelInfo for Flow {
client_name: added_column(r, "client_name")?,
client_bundle_id: added_column(r, "client_bundle_id")?,
client_path: added_column(r, "client_path")?,
stream_kind: added_column::<String>(r, "stream_kind")?
.and_then(|k| StreamKind::parse(&k)),
message_count: added_column(r, "message_count")?.unwrap_or_default(),
grpc_service: added_column(r, "grpc_service")?,
grpc_method: added_column(r, "grpc_method")?,
grpc_status: added_column(r, "grpc_status")?,
grpc_message: added_column(r, "grpc_message")?,
})
}
}
@@ -0,0 +1,437 @@
//! Streaming flows through a running engine: server-sent events and gRPC-Web over plain
//! HTTP, stored as messages while they stream, and gRPC messages decoded with and without a
//! schema.
use std::convert::Infallible;
use std::net::SocketAddr;
use std::path::Path;
use std::time::{Duration, Instant};
use bytes::Bytes;
use http_body_util::combinators::UnsyncBoxBody;
use http_body_util::{BodyExt, Full, StreamBody};
use hyper::body::{Frame, Incoming};
use hyper::service::service_fn;
use hyper::{Request, Response};
use hyper_util::rt::{TokioExecutor, TokioIo};
use hyper_util::server::conn::auto;
use serde_json::{Value, json};
use tokio::net::TcpListener;
use yaak_proxy_lib::models::Flow;
use yaak_proxy_lib::settings::ProxySettings;
use yaak_proxy_lib::{GrpcConfig, ProxyCtx};
use yaak_rpc::RpcEventEmitter;
type BoxError = Box<dyn std::error::Error + Send + Sync>;
type TestBody = UnsyncBoxBody<Bytes, BoxError>;
const PROTO: &str = r#"
syntax = "proto3";
package pkg;
service Greeter {
rpc SayHello (HelloRequest) returns (HelloReply);
}
message HelloRequest { string name = 1; }
message HelloReply { string message = 1; int32 count = 2; }
"#;
/// The longest a test here may run: one that hangs fails with a message instead of holding CI
/// until its job is cancelled.
const TEST_TIMEOUT: Duration = Duration::from_secs(60);
async fn within<T>(test: impl Future<Output = T>) -> T {
match tokio::time::timeout(TEST_TIMEOUT, test).await {
Ok(result) => result,
Err(_) => panic!("the test did not finish within {TEST_TIMEOUT:?}"),
}
}
fn full(data: impl Into<Bytes>) -> TestBody {
Full::new(data.into()).map_err(|never| match never {}).boxed_unsync()
}
fn grpc_frame(flag: u8, payload: &[u8]) -> Vec<u8> {
let mut out = vec![flag];
out.extend((payload.len() as u32).to_be_bytes());
out.extend(payload);
out
}
/// A protobuf message with a string in field 1, and a varint in field 2 when given.
fn proto(text: &str, count: Option<u8>) -> Vec<u8> {
let mut out = vec![0x0a, text.len() as u8];
out.extend(text.as_bytes());
if let Some(count) = count {
out.extend([0x10, count]);
}
out
}
fn protoc() -> GrpcConfig {
GrpcConfig {
protoc_bin_path: protoc_bin_vendored::protoc_bin_path().unwrap(),
protoc_include_dir: protoc_bin_vendored::include_path().unwrap(),
}
}
fn proto_file(dir: &Path) -> String {
let path = dir.join("greeter.proto");
std::fs::write(&path, PROTO).unwrap();
path.to_string_lossy().into_owned()
}
type ReflectionFuture = std::pin::Pin<Box<dyn Future<Output = Response<tonic::body::Body>> + Send>>;
/// A gRPC reflection service, its type erased.
type Reflection =
std::sync::Arc<dyn Fn(Request<tonic::body::Body>) -> ReflectionFuture + Send + Sync>;
fn reflection_service<S>(service: S) -> Reflection
where
S: tonic::codegen::Service<
Request<tonic::body::Body>,
Response = Response<tonic::body::Body>,
Error = Infallible,
> + Clone
+ Send
+ Sync
+ 'static,
S::Future: Send,
{
std::sync::Arc::new(move |req| {
let mut service = service.clone();
Box::pin(async move {
std::future::poll_fn(|cx| service.poll_ready(cx)).await.unwrap();
service.call(req).await.unwrap()
})
})
}
/// Serves the upstream's routes, and gRPC reflection for `pkg.Greeter` over HTTP/2 when given
/// one.
async fn route(
req: Request<Incoming>,
reflection: Option<Reflection>,
) -> Result<Response<TestBody>, Infallible> {
let path = req.uri().path().to_string();
if path.starts_with("/grpc.reflection.v1.ServerReflection/") {
let Some(reflection) = reflection else {
return Ok(Response::builder().status(404).body(full("")).unwrap());
};
let resp = reflection(req.map(tonic::body::Body::new)).await;
return Ok(resp.map(|b| b.map_err(|e| Box::new(e) as BoxError).boxed_unsync()));
}
Ok(match path.as_str() {
// Two events a while apart, so the first is stored while the stream is open.
"/events" => {
let frames = futures_util::stream::unfold(0, |i| async move {
let chunk = match i {
0 => "event: tick\nid: 1\ndata: first\n\n",
1 => {
tokio::time::sleep(Duration::from_millis(1500)).await;
"data: second\n\n"
}
_ => return None,
};
Some((Ok::<_, BoxError>(Frame::data(Bytes::from(chunk))), i + 1))
});
Response::builder()
.header("content-type", "text/event-stream")
.body(StreamBody::new(frames).boxed_unsync())
.unwrap()
}
"/pkg.Greeter/SayHello" => {
let body = req.into_body().collect().await.unwrap().to_bytes();
let name = String::from_utf8_lossy(&body[7..]).into_owned();
let mut out = grpc_frame(0, &proto(&format!("Hello {name}"), Some(3)));
out.extend(grpc_frame(0x80, b"grpc-status: 0\r\n"));
Response::builder()
.header("content-type", "application/grpc-web+proto")
.body(full(out))
.unwrap()
}
_ => Response::builder()
.header("content-type", "application/grpc-web+proto")
.body(full(grpc_frame(0x80, b"grpc-status: 12\r\ngrpc-message: not%20here\r\n")))
.unwrap(),
})
}
async fn spawn_upstream(reflection: Option<Reflection>) -> 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();
let reflection = reflection.clone();
tokio::spawn(async move {
let svc = service_fn(move |req| route(req, reflection.clone()));
let _ = auto::Builder::new(TokioExecutor::new())
.serve_connection(TokioIo::new(tcp), svc)
.await;
});
}
});
addr
}
fn start(dir: &Path) -> (ProxyCtx, std::sync::mpsc::Receiver<yaak_rpc::RpcEvent>) {
let settings = ProxySettings {
port: 0,
capture_body_limit: 1024 * 1024,
system_proxy_at_launch: false,
..Default::default()
};
settings.save(&dir.join("settings.json")).unwrap();
let (events, event_rx) = RpcEventEmitter::new();
let ctx = ProxyCtx::new(dir, events).unwrap().with_protoc(protoc());
ctx.start().unwrap();
(ctx, event_rx)
}
fn client(ctx: &ProxyCtx) -> reqwest::Client {
let proxy = format!("http://127.0.0.1:{}", ctx.port().unwrap());
reqwest::Client::builder().proxy(reqwest::Proxy::all(proxy).unwrap()).build().unwrap()
}
async fn rpc(ctx: &ProxyCtx, cmd: &str, payload: Value) -> Result<Value, String> {
yaak_proxy_lib::build_router().dispatch(cmd, payload, ctx).await.map_err(|e| e.message)
}
/// Waits until `check` holds for the flow of `url`, as stored.
async fn wait_for(ctx: &ProxyCtx, url: &str, check: impl Fn(&Flow) -> bool) -> Flow {
let deadline = Instant::now() + Duration::from_secs(10);
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 && check(f)) {
return flow;
}
assert!(Instant::now() < deadline, "no flow for {url} as expected");
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
#[tokio::test(flavor = "multi_thread")]
async fn server_sent_events_are_stored_while_the_stream_is_open() {
within(async {
let dir = tempfile::tempdir().unwrap();
let (ctx, event_rx) = start(dir.path());
let upstream = spawn_upstream(None).await;
let url = format!("http://{upstream}/events");
let reading = tokio::spawn({
let client = client(&ctx);
let url = url.clone();
async move { client.get(&url).send().await.unwrap().text().await.unwrap() }
});
// The flow is in the database with its first event while the second is still to come.
let open = wait_for(&ctx, &url, |f| f.message_count == 1).await;
assert_eq!(open.stream_kind.map(|k| k.as_str()), Some("sse"));
assert_eq!(open.elapsed_ms, None, "still streaming");
let listed = rpc(&ctx, "list_messages", json!({ "flowId": open.id })).await.unwrap();
assert_eq!(listed["total"], 1);
let first = &listed["messages"][0];
assert_eq!(
(first["kind"].as_str(), first["direction"].as_str()),
(Some("sse_event"), Some("server"))
);
assert_eq!(first["preview"], "first");
assert_eq!(
first["meta"],
json!({ "event": "tick", "eventId": "1", "raw": "event: tick\nid: 1\ndata: first\n\n" })
);
assert_eq!(reading.await.unwrap(), "event: tick\nid: 1\ndata: first\n\ndata: second\n\n");
let ended = wait_for(&ctx, &url, |f| f.elapsed_ms.is_some()).await;
assert_eq!((ended.message_count, ended.error), (2, None));
let after =
rpc(&ctx, "list_messages", json!({ "flowId": ended.id, "afterSeq": 0 })).await.unwrap();
assert_eq!(after["messages"].as_array().unwrap().len(), 1);
let payload = rpc(&ctx, "get_message_payload", json!({ "flowId": ended.id, "seq": 1 }))
.await
.unwrap();
assert_eq!(payload["data"], "c2Vjb25k", "\"second\" in base64");
// The app heard of the open flow, of each message, and of the end.
let events: Vec<_> = event_rx.try_iter().map(|e| e.event).collect();
let writes = |name: &str| events.iter().filter(|e| **e == name).count();
assert_eq!(writes("message_write"), 2, "{events:?}");
assert!(writes("model_write") >= 2, "{events:?}");
})
.await
}
#[tokio::test(flavor = "multi_thread")]
async fn grpc_web_calls_are_stored_and_decoded() {
within(async {
let dir = tempfile::tempdir().unwrap();
let (ctx, _event_rx) = start(dir.path());
let upstream = spawn_upstream(None).await;
let client = client(&ctx);
let url = format!("http://{upstream}/pkg.Greeter/SayHello");
let body = client
.post(&url)
.header("content-type", "application/grpc-web+proto")
.body(grpc_frame(0, &proto("Ada", None)))
.send()
.await
.unwrap()
.bytes()
.await
.unwrap();
assert!(body.starts_with(&grpc_frame(0, &proto("Hello Ada", Some(3)))));
let flow = wait_for(&ctx, &url, |f| f.elapsed_ms.is_some()).await;
assert_eq!(flow.stream_kind.map(|k| k.as_str()), Some("grpc_web"));
assert_eq!(
(flow.grpc_service.as_deref(), flow.grpc_method.as_deref()),
(Some("pkg.Greeter"), Some("SayHello"))
);
assert_eq!((flow.grpc_status, flow.message_count), (Some(0), 3));
let failing = format!("http://{upstream}/pkg.Greeter/Missing");
client
.post(&failing)
.header("content-type", "application/grpc-web+proto")
.body(grpc_frame(0, &proto("Bo", None)))
.send()
.await
.unwrap();
let failed = wait_for(&ctx, &failing, |f| f.elapsed_ms.is_some()).await;
assert_eq!(
(failed.grpc_status, failed.grpc_message.as_deref()),
(Some(12), Some("not here"))
);
let errors = rpc(&ctx, "list_flows", json!({ "status": "error" })).await.unwrap();
let ids: Vec<_> =
errors["flows"].as_array().unwrap().iter().map(|f| f["id"].clone()).collect();
assert_eq!(ids, vec![json!(failed.id)], "a non-zero status is a failed call");
// No schema, and reflection against the server fails: decoded by wire type alone.
let decode = json!({ "flowId": flow.id, "seq": 1 });
let schemaless = rpc(&ctx, "decode_grpc_message", decode.clone()).await.unwrap();
assert_eq!(schemaless["messageType"], Value::Null);
assert!(schemaless["note"].as_str().unwrap().contains("reflection failed"), "{schemaless}");
let fields: Value = serde_json::from_str(schemaless["json"].as_str().unwrap()).unwrap();
assert_eq!(fields[0]["string"], "Hello Ada");
assert_eq!(fields[1], json!({ "field": 2, "wireType": "varint", "value": 3 }));
// With the proto attached to the host, by name.
let protos = tempfile::tempdir().unwrap();
let added = rpc(
&ctx,
"add_grpc_schema",
json!({ "hostPattern": "127.0.0.1", "protoPaths": [proto_file(protos.path())] }),
)
.await
.unwrap();
assert_eq!(added["services"], json!(["pkg.Greeter"]));
assert_eq!(added["source"], "files");
let decoded = rpc(&ctx, "decode_grpc_message", decode.clone()).await.unwrap();
assert_eq!(decoded["messageType"], "pkg.HelloReply");
let message: Value = serde_json::from_str(decoded["json"].as_str().unwrap()).unwrap();
assert_eq!(message, json!({ "message": "Hello Ada", "count": 3 }));
let request = json!({ "flowId": flow.id, "seq": 0 });
let decoded = rpc(&ctx, "decode_grpc_message", request).await.unwrap();
assert_eq!(decoded["messageType"], "pkg.HelloRequest");
let trailers = json!({ "flowId": flow.id, "seq": 2 });
assert!(
rpc(&ctx, "decode_grpc_message", trailers).await.is_err(),
"trailers are no message"
);
// The control API returns the same, and lists the messages.
ctx.serve_control_api().unwrap();
let info = ctx.control_info().unwrap();
let api = reqwest::Client::new();
let get = |path: String| {
api.get(format!("{}{path}", info.base_url())).bearer_auth(&info.token).send()
};
let listed: Value = get(format!("/flows/{}/messages?after=0&limit=1", flow.id))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(
(listed["total"].as_i64(), listed["messages"][0]["seq"].as_i64()),
(Some(3), Some(1))
);
let detail: Value =
get(format!("/flows/{}/messages/1", flow.id)).await.unwrap().json().await.unwrap();
assert_eq!(detail["decoded"]["messageType"], "pkg.HelloReply");
assert_eq!(detail["payload"]["encoding"], "utf8");
let schemas: Value = get("/grpc-schemas".into()).await.unwrap().json().await.unwrap();
assert_eq!(schemas.as_array().unwrap().len(), 1);
let removed = rpc(&ctx, "remove_grpc_schema", json!({ "id": added["id"] })).await.unwrap();
assert_eq!(removed, true);
let after = rpc(&ctx, "decode_grpc_message", decode).await.unwrap();
assert_eq!(after["messageType"], Value::Null);
// Clearing removes the messages with their flows.
rpc(&ctx, "clear_flows", json!({})).await.unwrap();
assert!(rpc(&ctx, "list_messages", json!({ "flowId": flow.id })).await.is_err());
let left: i64 = ctx.db.with_conn(|c| {
c.conn()
.resolve()
.query_row("SELECT COUNT(*) FROM flow_messages", [], |r| r.get(0))
.unwrap()
});
assert_eq!(left, 0);
ctx.stop_control_api();
})
.await
}
#[tokio::test(flavor = "multi_thread")]
async fn grpc_schemas_come_from_the_servers_reflection() {
within(async {
let protos = tempfile::tempdir().unwrap();
let pool =
yaak_grpc::fill_pool_from_files(&protoc(), &vec![proto_file(protos.path()).into()])
.await
.unwrap();
let descriptors = pool.encode_to_vec();
let reflection = tonic_reflection::server::Builder::configure()
.register_encoded_file_descriptor_set(&descriptors)
.build_v1()
.unwrap();
let upstream = spawn_upstream(Some(reflection_service(reflection))).await;
let dir = tempfile::tempdir().unwrap();
let (ctx, _event_rx) = start(dir.path());
// Plain HTTP, so reflection is asked without TLS, as the call was made.
let url = format!("http://{upstream}/pkg.Greeter/SayHello");
client(&ctx)
.post(&url)
.header("content-type", "application/grpc-web+proto")
.body(grpc_frame(0, &proto("Cy", None)))
.send()
.await
.unwrap();
let flow = wait_for(&ctx, &url, |f| f.elapsed_ms.is_some()).await;
// Nothing attached: the server is asked once, and its answer kept for the host. Two
// decodes at once both wait for that answer.
let request = json!({ "flowId": flow.id, "seq": 1 });
let (decoded, alongside) = tokio::join!(
rpc(&ctx, "decode_grpc_message", request.clone()),
rpc(&ctx, "decode_grpc_message", request)
);
let (decoded, alongside) = (decoded.unwrap(), alongside.unwrap());
assert_eq!(decoded["messageType"], "pkg.HelloReply", "{decoded}");
assert_eq!(alongside["messageType"], "pkg.HelloReply", "{alongside}");
assert_eq!(decoded["schema"]["source"], "reflection");
assert_eq!(decoded["schema"]["hostPattern"], format!("127.0.0.1:{}", upstream.port()));
let schemas = rpc(&ctx, "list_grpc_schemas", json!({})).await.unwrap();
assert_eq!(schemas.as_array().unwrap().len(), 1);
// Asked again on demand, it replaces the host's earlier answer.
let host = format!("http://127.0.0.1:{}", upstream.port());
let again = rpc(&ctx, "reflect_grpc_host", json!({ "host": host })).await.unwrap();
assert_eq!(again["services"], json!(["pkg.Greeter"]));
let schemas = rpc(&ctx, "list_grpc_schemas", json!({})).await.unwrap();
assert_eq!(schemas.as_array().unwrap().len(), 1);
})
.await
}
+45 -1
View File
@@ -1,5 +1,6 @@
use log::{error, info, warn};
use tauri::Runtime;
use tauri::path::BaseDirectory;
use tauri::{Emitter, Manager, RunEvent, State, WebviewWindow};
use yaak_proxy_lib::ProxyCtx;
use yaak_rpc::{RpcEventEmitter, RpcRouter};
@@ -48,6 +49,24 @@ fn setup_window_menu<R: Runtime>(win: &WebviewWindow<R>) {
});
}
/// Where `tauri.bundle.conf.json` puts the `yaakprotoc` that Yaak's vendor script fetches, in
/// the app's resources.
const PROTOC_DIR: &str = "vendored/protoc";
/// Where it puts protoc's include files.
const PROTOC_INCLUDE_DIR: &str = "vendored/protoc/include";
/// The `yaakprotoc` and its includes in the app's resources, if the bundle has them, as a
/// packaged app's does; a development build has none.
fn bundled_protoc<R: Runtime>(app: &tauri::AppHandle<R>) -> Option<yaak_proxy_lib::GrpcConfig> {
let bin_name = if cfg!(windows) { "yaakprotoc.exe" } else { "yaakprotoc" };
let resource = |path: &str| app.path().resolve(path, BaseDirectory::Resource).ok();
let protoc_bin_path = resource(&format!("{PROTOC_DIR}/{bin_name}"))?;
let protoc_include_dir = resource(PROTOC_INCLUDE_DIR)?;
protoc_bin_path
.is_file()
.then_some(yaak_proxy_lib::GrpcConfig { protoc_include_dir, protoc_bin_path })
}
#[tauri::command]
async fn rpc(
router: State<'_, RpcRouter<ProxyCtx>>,
@@ -69,7 +88,12 @@ pub fn run() {
std::fs::create_dir_all(&data_dir).expect("failed to create app data dir");
let (emitter, event_rx) = RpcEventEmitter::new();
let ctx = ProxyCtx::new(&data_dir, emitter)?;
let mut ctx = ProxyCtx::new(&data_dir, emitter)?;
// Compiles protos attached for gRPC decoding when the bundle has protoc, as Yaak's
// does; otherwise one on PATH is used.
if let Some(config) = bundled_protoc(app.handle()) {
ctx = ctx.with_protoc(config);
}
// Not while the CLI daemon owns the engine: the window offers to take it over.
match ctx.serve_control_api() {
// The engine runs for as long as the app does. A failure, such as the port
@@ -127,3 +151,23 @@ pub fn run() {
}
});
}
#[cfg(test)]
mod tests {
use super::*;
/// The bundle puts protoc, and its includes, where the app looks for them.
#[test]
fn the_bundle_ships_protoc_where_the_app_looks() {
let config: serde_json::Value =
serde_json::from_str(include_str!("../tauri.bundle.conf.json")).unwrap();
let resources = config["bundle"]["resources"].as_object().unwrap();
let target = |source: &str| resources.get(source).and_then(|t| t.as_str());
assert_eq!(
target("../yaak-app-client/vendored/protoc/yaakprotoc*")
.map(|t| t.trim_end_matches('/')),
Some(PROTOC_DIR)
);
assert_eq!(target("../yaak-app-client/vendored/protoc/include"), Some(PROTOC_INCLUDE_DIR));
}
}
@@ -4,6 +4,10 @@
},
"bundle": {
"externalBin": ["vendored/helper/yaak-proxy-helper"],
"resources": {
"../yaak-app-client/vendored/protoc/include": "vendored/protoc/include",
"../yaak-app-client/vendored/protoc/yaakprotoc*": "vendored/protoc/"
},
"macOS": {
"minimumSystemVersion": "13.0",
"hardenedRuntime": true,
+8
View File
@@ -11,9 +11,17 @@ pub mod manager;
mod reflection;
mod transport;
pub use manager::GrpcConfig;
pub use prost_reflect;
pub use reflection::{fill_pool_from_files, fill_pool_from_reflection};
pub use tonic::Code;
pub use tonic::metadata::*;
/// A message as pretty JSON, its default fields included, as Yaak shows gRPC messages.
pub fn message_to_json(msg: &DynamicMessage) -> Result<String, String> {
serialize_dynamic_message_json(msg)
}
pub fn serialize_options() -> SerializeOptions {
SerializeOptions::new().skip_default_fields(false)
}
+3 -1
View File
@@ -5,7 +5,9 @@ edition = "2024"
publish = false
[dependencies]
hyper = { version = "1", features = ["http1", "http2", "server", "client"] }
base64 = "0.23.1"
flate2 = "1"
hyper ={ version = "1", features = ["http1", "http2", "server", "client"] }
hyper-util = { version = "0.1", features = ["tokio", "http1", "client-legacy", "server-auto"] }
http-body-util = "0.1"
http = "1"
+65 -1
View File
@@ -7,6 +7,9 @@ use std::time::Instant;
use bytes::Bytes;
use hyper::body::{Body, Frame, Incoming, SizeHint};
use crate::grpc::{self, GrpcParser};
use crate::messages::StreamMessage;
use crate::sse::SseParser;
use crate::{Engine, ProxyEvent};
#[derive(Clone, Copy, PartialEq)]
@@ -15,6 +18,36 @@ pub(crate) enum Direction {
Response,
}
/// Reads a streaming body's messages as the body passes.
pub(crate) enum BodyParser {
Sse(Box<SseParser>),
Grpc(GrpcParser),
}
impl BodyParser {
fn feed(&mut self, data: &[u8]) -> Vec<StreamMessage> {
match self {
Self::Sse(p) => p.feed(data),
Self::Grpc(p) => p.feed(data),
}
}
fn trailers(&mut self, trailers: &hyper::HeaderMap) -> Option<StreamMessage> {
match self {
Self::Sse(_) => None,
Self::Grpc(_) => grpc::trailers_from_headers(trailers),
}
}
/// The body ended: messages held back for what might have followed.
fn finish(&mut self) -> Vec<StreamMessage> {
match self {
Self::Sse(p) => p.finish(),
Self::Grpc(_) => Vec::new(),
}
}
}
/// Streams a body through untouched, reporting each frame as a chunk event (with data up to
/// `capture_body_limit`) and the end as a completion event.
pub(crate) struct CaptureBody {
@@ -28,6 +61,7 @@ pub(crate) struct CaptureBody {
finished: bool,
/// Set when reading the body fails.
failed: Option<Arc<AtomicBool>>,
parser: Option<BodyParser>,
}
impl CaptureBody {
@@ -48,6 +82,19 @@ impl CaptureBody {
size: 0,
finished: false,
failed: None,
parser: None,
}
}
/// Records the messages `parser` reads from the body, all of it, past the capture limit too.
pub(crate) fn with_parser(mut self, parser: Option<BodyParser>) -> Self {
self.parser = parser;
self
}
fn emit_messages(&self, messages: impl IntoIterator<Item = StreamMessage>) {
for message in messages {
self.engine.emit(ProxyEvent::Message { id: self.id, message });
}
}
@@ -69,6 +116,11 @@ impl CaptureBody {
return;
}
self.finished = true;
// Before the body's end, so they land while the flow is still open.
if let Some(parser) = self.parser.as_mut() {
let messages = parser.finish();
self.emit_messages(messages);
}
let truncated = self.size > self.captured as u64;
let event = match self.direction {
Direction::Request if self.size == 0 => return,
@@ -110,8 +162,20 @@ impl Body for CaptureBody {
Direction::Request => ProxyEvent::RequestBodyChunk { id, len, data },
Direction::Response => ProxyEvent::ResponseBodyChunk { id, len, data },
});
if let Some(parser) = self.parser.as_mut() {
let messages = parser.feed(frame_data);
self.emit_messages(messages);
}
} else if let Some(trailers) = frame.trailers_ref()
&& let Some(parser) = self.parser.as_mut()
{
let message = parser.trailers(trailers);
self.emit_messages(message);
}
if self.complete() {
// A gRPC call's status comes in trailers after its last byte, so its body is
// done only at the stream's end.
let grpc = matches!(self.parser, Some(BodyParser::Grpc(_)));
if if grpc { self.inner.is_end_stream() } else { self.complete() } {
self.finish(None, false);
}
}
+378
View File
@@ -0,0 +1,378 @@
//! A response body's `Content-Encoding` undone as it streams, for parsers that read messages
//! from the body. Output comes in pieces of at most [`PIECE`] bytes, so a highly compressed
//! body never expands in memory all at once.
use flate2::{Decompress, FlushDecompress, Status};
/// The most decompressed bytes handed on at once.
const PIECE: usize = 16 * 1024;
/// The longest gzip header waited for, its optional name, comment and extra field included.
const MAX_GZIP_HEADER: usize = 128 * 1024;
const GZIP_FHCRC: u8 = 0x02;
const GZIP_FEXTRA: u8 = 0x04;
const GZIP_FNAME: u8 = 0x08;
const GZIP_FCOMMENT: u8 = 0x10;
/// How a body is encoded, from its `Content-Encoding`.
pub(crate) enum Decoding {
Identity,
Decoder(Box<ContentDecoder>),
/// One this cannot undo.
Unsupported,
}
impl Decoding {
/// `limit` bounds what decoding may hold on to: the input of a deflate body kept until it
/// shows whether it is zlib-framed.
pub(crate) fn new(content_encoding: Option<&str>, limit: usize) -> Self {
match content_encoding.map(|e| e.trim().to_ascii_lowercase()).as_deref() {
None | Some("") | Some("identity") => Self::Identity,
Some("gzip") | Some("x-gzip") => Self::Decoder(Box::new(ContentDecoder::gzip(limit))),
Some("deflate") => Self::Decoder(Box::new(ContentDecoder::deflate(limit))),
Some(_) => Self::Unsupported,
}
}
}
pub(crate) struct ContentDecoder {
decompress: Decompress,
/// gzip, which comes in members, each with a header and a trailer.
gzip: bool,
/// A gzip member's header still to be read, as much of it as has come; None while a member
/// is being inflated, and for deflate, which zlib's own header covers.
gzip_header: Option<Vec<u8>>,
/// How far the gzip header's name and comment have been searched for their ends.
header_scanned: [usize; 2],
/// Bytes of a gzip member's trailer (its CRC and length, not checked) still to skip.
trailer_left: usize,
/// gzip members read whole.
members: usize,
/// The compressed stream ended; anything after it is not read.
ended: bool,
/// Deflate's `zlib` framing, which some servers leave out.
zlib: bool,
/// No output yet, so a deflate body without zlib framing can still be tried raw.
started: bool,
/// The input until the first output, while it may yet turn out not to be zlib-framed.
unframed: Vec<u8>,
/// The most input kept in `unframed`: a body that has given no output by then is not
/// decoded further.
unframed_limit: usize,
}
impl ContentDecoder {
fn gzip(limit: usize) -> Self {
Self {
decompress: Decompress::new(false),
gzip: true,
gzip_header: Some(Vec::new()),
header_scanned: [0; 2],
trailer_left: 0,
members: 0,
ended: false,
zlib: false,
started: false,
unframed: Vec::new(),
unframed_limit: limit,
}
}
fn deflate(limit: usize) -> Self {
Self {
decompress: Decompress::new(true),
gzip: false,
gzip_header: None,
header_scanned: [0; 2],
trailer_left: 0,
members: 0,
ended: false,
zlib: true,
started: false,
unframed: Vec::new(),
unframed_limit: limit,
}
}
/// Decompresses `input`, handing the output on to `out` a piece at a time.
pub(crate) fn feed(&mut self, input: &[u8], out: &mut impl FnMut(&[u8])) -> Result<(), String> {
self.feed_inner(input, out)?;
// Kept only until the body gives output; empty deflate blocks can go on without end,
// and what is kept for them cannot.
if self.zlib && !self.started && self.unframed.len() > self.unframed_limit {
self.unframed = Vec::new();
return Err("No output from the deflate body within the capture limit".into());
}
Ok(())
}
fn feed_inner(&mut self, mut input: &[u8], out: &mut impl FnMut(&[u8])) -> Result<(), String> {
if self.zlib && !self.started {
self.unframed.extend_from_slice(input);
}
let mut buf = [0u8; PIECE];
loop {
if self.ended {
return Ok(());
}
if self.trailer_left > 0 {
let skip = self.trailer_left.min(input.len());
input = &input[skip..];
self.trailer_left -= skip;
if self.trailer_left > 0 {
return Ok(());
}
}
if let Some(header) = self.gzip_header.as_mut() {
if input.is_empty() {
return Ok(());
}
let before = header.len();
header.extend_from_slice(input);
match gzip_header_len(header, &mut self.header_scanned) {
// A name or comment that never ends is no header to wait for.
Ok(None) if header.len() > MAX_GZIP_HEADER => {
return Err("The gzip header is too long".into());
}
Ok(None) => return Ok(()),
Ok(Some(len)) => {
input = &input[len - before..];
self.gzip_header = None;
self.header_scanned = [0; 2];
}
// What follows a whole member, padding say, is not another one.
Err(_) if self.members > 0 => {
self.ended = true;
return Ok(());
}
Err(e) => return Err(e),
}
}
let (read_before, written_before) =
(self.decompress.total_in(), self.decompress.total_out());
let status = match self.decompress.decompress(input, &mut buf, FlushDecompress::None) {
Ok(status) => status,
// "deflate" is meant to be zlib-framed, but some servers send it raw: read again
// from the start as raw.
Err(_) if self.zlib && !self.started => {
self.decompress = Decompress::new(false);
self.zlib = false;
let replay = std::mem::take(&mut self.unframed);
return self.feed_inner(&replay, out);
}
Err(e) => return Err(e.to_string()),
};
let read = (self.decompress.total_in() - read_before) as usize;
let written = (self.decompress.total_out() - written_before) as usize;
input = &input[read..];
if written > 0 {
self.started = true;
self.unframed = Vec::new();
out(&buf[..written]);
}
if status == Status::StreamEnd {
// A gzip body may hold more members, each read in turn.
if self.gzip {
self.members += 1;
self.decompress.reset(false);
self.trailer_left = 8;
self.gzip_header = Some(Vec::new());
continue;
}
self.ended = true;
return Ok(());
}
// A full piece may have more output behind it, even once the input is used up.
if (input.is_empty() && written < PIECE) || (read == 0 && written == 0) {
return Ok(());
}
}
}
}
/// The length of the gzip header `bytes` starts with (RFC 1952 2.3), None while more of it is
/// still to come.
/// `scanned` keeps how far the name and the comment have been searched for their ends, so a
/// header that comes a byte at a time is not searched again from its start each time.
fn gzip_header_len(bytes: &[u8], scanned: &mut [usize; 2]) -> Result<Option<usize>, String> {
if bytes.len() < 10 {
return Ok(None);
}
if bytes[..3] != [0x1f, 0x8b, 0x08] {
return Err("Not a gzip stream".into());
}
let flags = bytes[3];
let mut len = 10;
if flags & GZIP_FEXTRA != 0 {
let Some(extra) = bytes.get(len..len + 2) else {
return Ok(None);
};
len += 2 + u16::from_le_bytes([extra[0], extra[1]]) as usize;
}
for (i, flag) in [GZIP_FNAME, GZIP_FCOMMENT].into_iter().enumerate() {
if flags & flag != 0 {
let from = len.max(scanned[i]);
match bytes.get(from..).and_then(|rest| rest.iter().position(|&b| b == 0)) {
Some(end) => len = from + end + 1,
None => {
scanned[i] = bytes.len().max(len);
return Ok(None);
}
}
}
}
if flags & GZIP_FHCRC != 0 {
len += 2;
}
Ok((bytes.len() >= len).then_some(len))
}
#[cfg(test)]
mod tests {
use super::*;
const TEST_LIMIT: usize = 1 << 20;
use std::io::Write;
fn decode_bytewise(encoding: &str, compressed: &[u8]) -> (Vec<u8>, usize) {
let Decoding::Decoder(mut decoder) = Decoding::new(Some(encoding), TEST_LIMIT) else {
panic!("no decoder for {encoding}");
};
let (mut out, mut largest) = (Vec::new(), 0);
for byte in compressed {
decoder
.feed(&[*byte], &mut |piece| {
largest = largest.max(piece.len());
out.extend_from_slice(piece);
})
.unwrap();
}
(out, largest)
}
#[test]
fn gzip_with_a_file_name_split_into_single_bytes() {
let text = b"data: hello\n\n".repeat(50);
let mut gz = flate2::GzBuilder::new()
.filename("events.txt")
.write(Vec::new(), flate2::Compression::default());
gz.write_all(&text).unwrap();
let (out, _) = decode_bytewise("gzip", &gz.finish().unwrap());
assert_eq!(out, text);
}
#[test]
fn deflate_framed_or_raw() {
let text = b"some text to squeeze".repeat(20);
let mut zlib = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default());
zlib.write_all(&text).unwrap();
assert_eq!(decode_bytewise("deflate", &zlib.finish().unwrap()).0, text);
let mut raw =
flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::default());
raw.write_all(&text).unwrap();
assert_eq!(decode_bytewise("deflate", &raw.finish().unwrap()).0, text);
}
#[test]
fn output_buffered_past_a_piece_is_drained_when_the_input_runs_out() {
// Raw deflate inflates more than a piece from the last of its input.
let text: Vec<u8> =
b"data: ".iter().chain(&[b'x'; 40_000]).chain(b"\n\n").copied().collect();
let mut raw = flate2::write::DeflateEncoder::new(Vec::new(), flate2::Compression::best());
raw.write_all(&text).unwrap();
let Decoding::Decoder(mut decoder) = Decoding::new(Some("deflate"), TEST_LIMIT) else {
panic!()
};
let mut out = Vec::new();
decoder.feed(&raw.finish().unwrap(), &mut |piece| out.extend_from_slice(piece)).unwrap();
assert_eq!(out, text);
}
#[test]
fn a_gzip_name_that_never_ends_is_refused() {
let Decoding::Decoder(mut decoder) = Decoding::new(Some("gzip"), TEST_LIMIT) else {
panic!()
};
let header = [0x1f, 0x8b, 0x08, GZIP_FNAME, 0, 0, 0, 0, 0, 0];
decoder.feed(&header, &mut |_| {}).unwrap();
let mut refused = false;
for _ in 0..MAX_GZIP_HEADER / 512 + 1 {
if decoder.feed(&[b'a'; 512], &mut |_| {}).is_err() {
refused = true;
break;
}
}
assert!(refused);
}
#[test]
fn deflate_without_output_is_not_kept_past_the_limit() {
// A zlib header, then empty stored blocks: valid deflate that never produces a byte.
let empty_block = [0x00, 0x00, 0x00, 0xff, 0xff];
let Decoding::Decoder(mut decoder) = Decoding::new(Some("deflate"), 1024) else {
panic!()
};
decoder.feed(&[0x78, 0x9c], &mut |_| {}).unwrap();
let mut result = Ok(());
for _ in 0..1024 {
result = decoder.feed(&empty_block, &mut |_| {});
if result.is_err() {
break;
}
}
assert!(result.is_err(), "stopped once past the limit");
assert!(decoder.unframed.len() <= 1024);
}
#[test]
fn a_chunk_past_the_limit_is_decoded_when_it_gives_output() {
let text = b"data: tick\n\n".repeat(400);
let mut zlib = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::none());
zlib.write_all(&text).unwrap();
let body = zlib.finish().unwrap();
assert!(body.len() > 1024);
let Decoding::Decoder(mut decoder) = Decoding::new(Some("deflate"), 1024) else {
panic!()
};
let mut out = Vec::new();
decoder.feed(&body, &mut |piece| out.extend_from_slice(piece)).unwrap();
assert_eq!(out, text);
}
#[test]
fn every_gzip_member_is_read() {
let mut body = Vec::new();
for part in [&b"data: one\n\n"[..], b"data: two\n\n"] {
let mut gz = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
gz.write_all(part).unwrap();
body.extend(gz.finish().unwrap());
}
assert_eq!(decode_bytewise("gzip", &body).0, b"data: one\n\ndata: two\n\n");
let Decoding::Decoder(mut decoder) = Decoding::new(Some("gzip"), TEST_LIMIT) else {
panic!()
};
let mut out = Vec::new();
decoder.feed(&body, &mut |piece| out.extend_from_slice(piece)).unwrap();
assert_eq!(out, b"data: one\n\ndata: two\n\n", "in one chunk");
}
#[test]
fn a_highly_compressed_body_comes_out_in_pieces() {
let mut gz = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::best());
gz.write_all(&vec![b'a'; 4 << 20]).unwrap();
let compressed = gz.finish().unwrap();
let Decoding::Decoder(mut decoder) = Decoding::new(Some("gzip"), TEST_LIMIT) else {
panic!()
};
let (mut total, mut largest) = (0, 0);
decoder
.feed(&compressed, &mut |piece| {
total += piece.len();
largest = largest.max(piece.len());
})
.unwrap();
assert_eq!(total, 4 << 20);
assert!(largest <= PIECE, "{largest}");
}
}
+346
View File
@@ -0,0 +1,346 @@
//! gRPC and gRPC-Web: the length-prefixed messages of each direction's body, and the call's
//! status from its trailers.
use std::io::Read;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64;
use bytes::Bytes;
use flate2::read::{GzDecoder, ZlibDecoder};
use hyper::header::HeaderMap;
use crate::decode::Decoding;
use crate::messages::{Capped, MessageDirection, MessageKind, StreamKind, StreamMessage};
/// gRPC's compressed flag, in the first byte of each message's prefix.
const FLAG_COMPRESSED: u8 = 0x01;
/// gRPC-Web's flag for the frame that carries the trailers.
const FLAG_TRAILERS: u8 = 0x80;
const PREFIX_LEN: usize = 5;
/// The gRPC flavor a request's `Content-Type` names, and whether it is gRPC-Web's base64 one.
pub(crate) fn detect(content_type: Option<&str>) -> Option<(StreamKind, bool)> {
let content_type = content_type?.trim().to_ascii_lowercase();
let essence = content_type.split(';').next().unwrap_or_default().trim();
if let Some(rest) = essence.strip_prefix("application/grpc-web") {
return match rest {
"" | "+proto" => Some((StreamKind::GrpcWeb, false)),
"-text" | "-text+proto" => Some((StreamKind::GrpcWeb, true)),
_ => None,
};
}
match essence.strip_prefix("application/grpc") {
Some("" | "+proto") => Some((StreamKind::Grpc, false)),
_ => None,
}
}
/// `/package.Service/Method` as its service and method.
pub fn service_and_method(path: &str) -> Option<(&str, &str)> {
let path = path.split('?').next()?;
let (service, method) = path.strip_prefix('/')?.split_once('/')?;
(!service.is_empty() && !method.is_empty() && !method.contains('/'))
.then_some((service, method))
}
pub(crate) struct GrpcParser {
direction: MessageDirection,
/// `grpc-encoding`, what a compressed message is compressed with.
encoding: Option<String>,
/// gRPC-Web's base64 flavor, with base64 not yet a whole quantum.
base64: Option<Vec<u8>>,
prefix: Vec<u8>,
/// The message being read: its flag, length, bytes left, and payload.
current: Option<(u8, u32, u32, Capped)>,
limit: usize,
broken: bool,
/// The body's `Content-Encoding`, which gRPC-Web may use, undone before anything else.
decoding: Decoding,
}
impl GrpcParser {
/// Reads a body sent with `content_encoding`, as a gRPC-Web one can be.
pub(crate) fn with_content_encoding(mut self, content_encoding: Option<&str>) -> Self {
self.decoding = Decoding::new(content_encoding, self.limit);
self
}
pub(crate) fn new(
direction: MessageDirection,
encoding: Option<&str>,
text: bool,
limit: usize,
) -> Self {
Self {
direction,
encoding: encoding.map(|e| e.trim().to_ascii_lowercase()),
base64: text.then(Vec::new),
prefix: Vec::with_capacity(PREFIX_LEN),
current: None,
limit,
broken: false,
decoding: Decoding::Identity,
}
}
pub(crate) fn feed(&mut self, bytes: &[u8]) -> Vec<StreamMessage> {
match std::mem::replace(&mut self.decoding, Decoding::Unsupported) {
Decoding::Identity => {
self.decoding = Decoding::Identity;
self.feed_decoded(bytes)
}
// An encoding this cannot undo leaves the frames unreadable.
Decoding::Unsupported => Vec::new(),
// A piece at a time, so a highly compressed body never expands all at once.
Decoding::Decoder(mut decoder) => {
let mut messages = Vec::new();
let decoded = decoder.feed(bytes, &mut |piece| {
messages.extend(self.feed_decoded(piece));
});
if decoded.is_ok() {
self.decoding = Decoding::Decoder(decoder);
}
messages
}
}
}
/// Reads bytes with any Content-Encoding undone: base64 first for gRPC-Web's text flavor,
/// then the frames.
fn feed_decoded(&mut self, bytes: &[u8]) -> Vec<StreamMessage> {
if self.broken {
return Vec::new();
}
let Some(pending) = self.base64.as_mut() else {
return self.read(bytes);
};
pending.extend(bytes.iter().filter(|b| !b.is_ascii_whitespace()));
// Each response chunk is base64 of its own, padded, so padding can come mid-stream:
// decoded up to and including each padded quantum.
let whole = pending.len() / 4 * 4;
let mut decoded = Vec::new();
let mut start = 0;
for end in (4..=whole).step_by(4) {
if pending[end - 1] == b'=' || end == whole {
match BASE64.decode(&pending[start..end]) {
Ok(bytes) => decoded.extend(bytes),
Err(_) => {
self.broken = true;
return Vec::new();
}
}
start = end;
}
}
pending.drain(..whole);
self.read(&decoded)
}
fn read(&mut self, mut bytes: &[u8]) -> Vec<StreamMessage> {
let mut messages = Vec::new();
while !bytes.is_empty() {
let Some((_, _, remaining, payload)) = self.current.as_mut() else {
let take = (PREFIX_LEN - self.prefix.len()).min(bytes.len());
self.prefix.extend_from_slice(&bytes[..take]);
bytes = &bytes[take..];
if self.prefix.len() == PREFIX_LEN {
let flag = self.prefix[0];
let length = u32::from_be_bytes(self.prefix[1..5].try_into().unwrap());
self.prefix.clear();
self.current = Some((flag, length, length, Capped::new(self.limit)));
if length == 0 {
messages.push(self.finish());
}
}
continue;
};
let n = (*remaining as usize).min(bytes.len());
payload.push(&bytes[..n]);
*remaining -= n as u32;
bytes = &bytes[n..];
if *remaining == 0 {
messages.push(self.finish());
}
}
messages
}
fn finish(&mut self) -> StreamMessage {
let (flag, length, _, mut payload) = self.current.take().expect("a message being read");
let (data, truncated) = payload.take();
let compressed = flag & FLAG_COMPRESSED != 0;
let (data, truncated, error) =
if compressed { self.decompress(data, truncated) } else { (data, truncated, None) };
// gRPC-Web's trailers, read once decompressed, as they can be.
let mut message = if flag & FLAG_TRAILERS != 0 {
// Trailers that would not decompress are kept as they came, with no status read.
let text = if error.is_none() { String::from_utf8_lossy(&data) } else { "".into() };
let mut message = trailers_message(self.direction, &text);
message.data = data;
message
} else {
let mut message = StreamMessage::new(self.direction, MessageKind::GrpcMessage);
message.data = data;
message
};
message.size = u64::from(length);
message.truncated = truncated;
message.compressed = compressed;
message.error = error;
message
}
/// A compressed message's payload decompressed with `grpc-encoding`, whether that is cut at
/// the limit, and why it could not be decompressed, when it could not: then the payload as
/// it came.
fn decompress(&self, data: Bytes, truncated: bool) -> (Bytes, bool, Option<String>) {
let mut decoded = Vec::new();
let take = self.limit as u64 + 1;
let result = match self.encoding.as_deref() {
Some("gzip") => GzDecoder::new(&data[..]).take(take).read_to_end(&mut decoded),
Some("deflate") => ZlibDecoder::new(&data[..]).take(take).read_to_end(&mut decoded),
other => {
let error =
format!("Compressed with {}, not decoded", other.unwrap_or("no encoding"));
return (data, truncated, Some(error));
}
};
match result {
Ok(_) if !truncated => {
let truncated = decoded.len() > self.limit;
decoded.truncate(self.limit);
(decoded.into(), truncated, None)
}
// A payload cut at the limit cannot be decompressed whole.
Ok(_) => (data, truncated, Some("Too large to decompress".into())),
Err(e) => (data, truncated, Some(format!("Failed to decompress the message: {e}"))),
}
}
}
/// The trailers message for `grpc-status` and `grpc-message` in `headers`, the HTTP/2 trailers
/// or a trailers-only response's headers. None when they carry no status.
pub(crate) fn trailers_from_headers(headers: &HeaderMap) -> Option<StreamMessage> {
headers.get("grpc-status")?;
let text: String = headers
.iter()
.map(|(name, value)| format!("{name}: {}\r\n", String::from_utf8_lossy(value.as_bytes())))
.collect();
let mut message = trailers_message(MessageDirection::Server, &text);
message.size = text.len() as u64;
Some(message)
}
/// Trailers written as `name: value` lines, as gRPC-Web frames them.
fn trailers_message(direction: MessageDirection, text: &str) -> StreamMessage {
let mut status = None;
let mut status_message = None;
for line in text.lines() {
let Some((name, value)) = line.split_once(':') else {
continue;
};
let value = value.trim();
if name.trim().eq_ignore_ascii_case("grpc-status") {
status = value.parse().ok();
} else if name.trim().eq_ignore_ascii_case("grpc-message") {
status_message = Some(percent_decode(value));
}
}
let kind = MessageKind::GrpcTrailers { status, message: status_message };
let mut message = StreamMessage::new(direction, kind);
message.data = text.to_string().into();
message
}
/// `grpc-message` is percent-encoded UTF-8.
fn percent_decode(value: &str) -> String {
let bytes = value.as_bytes();
let mut out = Vec::with_capacity(bytes.len());
let mut i = 0;
while i < bytes.len() {
if bytes[i] == b'%'
&& let Some(byte) =
value.get(i + 1..i + 3).and_then(|hex| u8::from_str_radix(hex, 16).ok())
{
out.push(byte);
i += 3;
continue;
}
out.push(bytes[i]);
i += 1;
}
String::from_utf8_lossy(&out).into_owned()
}
#[cfg(test)]
mod tests {
use super::*;
fn frame(flag: u8, payload: &[u8]) -> Vec<u8> {
let mut out = vec![flag];
out.extend((payload.len() as u32).to_be_bytes());
out.extend(payload);
out
}
#[test]
fn detects_flavors() {
assert_eq!(detect(Some("application/grpc")), Some((StreamKind::Grpc, false)));
assert_eq!(detect(Some("application/grpc+proto")), Some((StreamKind::Grpc, false)));
assert_eq!(detect(Some("application/grpc-web+proto")), Some((StreamKind::GrpcWeb, false)));
assert_eq!(detect(Some("application/grpc-web-text")), Some((StreamKind::GrpcWeb, true)));
assert_eq!(detect(Some("application/grpc+json")), None);
assert_eq!(detect(Some("application/json")), None);
assert_eq!(service_and_method("/pkg.Greeter/SayHello"), Some(("pkg.Greeter", "SayHello")));
assert_eq!(service_and_method("/"), None);
}
#[test]
fn splits_messages_across_chunks() {
let mut bytes = frame(0, b"hello");
bytes.extend(frame(0, b""));
bytes.extend(frame(0, b"world!"));
let mut parser = GrpcParser::new(MessageDirection::Client, None, false, 1024);
let mut messages = Vec::new();
for byte in bytes {
messages.extend(parser.feed(&[byte]));
}
let payloads: Vec<&[u8]> = messages.iter().map(|m| &m.data[..]).collect();
assert_eq!(payloads, vec![&b"hello"[..], b"", b"world!"]);
}
#[test]
fn decompresses_gzip_and_reads_web_trailers_in_base64() {
use flate2::write::GzEncoder;
use std::io::Write;
let mut gz = GzEncoder::new(Vec::new(), flate2::Compression::default());
gz.write_all(b"squeezed").unwrap();
let compressed = gz.finish().unwrap();
// Two separately padded base64 chunks, as gRPC-Web text responses come.
let first = BASE64.encode(frame(FLAG_COMPRESSED, &compressed));
let second =
BASE64.encode(frame(FLAG_TRAILERS, b"grpc-status: 5\r\ngrpc-message: not%20found\r\n"));
let mut parser = GrpcParser::new(MessageDirection::Server, Some("gzip"), true, 1024);
let mut messages = parser.feed(first.as_bytes());
messages.extend(parser.feed(second.as_bytes()));
assert_eq!(&messages[0].data[..], b"squeezed");
assert!(messages[0].compressed);
assert_eq!(
messages[1].kind,
MessageKind::GrpcTrailers { status: Some(5), message: Some("not found".into()) }
);
}
#[test]
fn compressed_web_trailers_are_read_once_decompressed() {
use flate2::write::GzEncoder;
use std::io::Write;
let mut gz = GzEncoder::new(Vec::new(), flate2::Compression::default());
gz.write_all(b"grpc-status: 5\r\n").unwrap();
let trailers = frame(FLAG_TRAILERS | FLAG_COMPRESSED, &gz.finish().unwrap());
let mut parser = GrpcParser::new(MessageDirection::Server, Some("gzip"), false, 1024);
let messages = parser.feed(&trailers);
assert_eq!(messages[0].kind, MessageKind::GrpcTrailers { status: Some(5), message: None });
assert!(messages[0].compressed && messages[0].error.is_none());
}
}
+30 -4
View File
@@ -1,8 +1,13 @@
mod body;
mod cert;
mod connection;
mod decode;
mod grpc;
mod messages;
mod request;
mod sse;
mod tunnel;
mod ws;
use std::collections::hash_map::RandomState;
use std::collections::{HashMap, HashSet};
@@ -24,6 +29,8 @@ use tokio::net::TcpListener;
pub use cert::CaStore;
use connection::handle_connection;
pub use grpc::service_and_method;
pub use messages::{MessageDirection, MessageKind, RAW_LIMIT, StreamKind, StreamMessage};
use request::BoxBody;
const PASSTHROUGH_AFTER_FAILURES: u32 = 2;
@@ -49,6 +56,9 @@ pub struct ProxyConfig {
/// Offer HTTP/2 to clients in intercepted TLS, beside HTTP/1.1. Off, they all speak
/// HTTP/1.1, for clients that misbehave on HTTP/2. See [`ProxyHandle::set_http2`].
pub http2: bool,
/// Record WebSocket frames as messages. Off, upgraded connections are spliced blind, as
/// other upgrades are.
pub websocket_capture: bool,
}
/// See [`ProxyConfig::on_accept`].
@@ -63,15 +73,17 @@ impl ProxyConfig {
upstream_root_certs: Vec::new(),
on_accept: None,
http2: true,
websocket_capture: true,
}
}
}
/// Lifecycle of each proxied exchange (a flow), keyed by `id`. A flow starts with
/// `RequestStart` and ends with exactly one `ResponseBodyComplete`, `Error` or `Cancelled`.
/// Each HTTP/2 stream is a flow of its own. Traffic the engine tunnels without decrypting
/// (pinned clients, non-TLS CONNECT payloads, the frames after a `101 Switching Protocols`)
/// produces no events.
/// `RequestStart` and ends with exactly one `ResponseBodyComplete`, `StreamEnd`, `Error` or
/// `Cancelled`. Each HTTP/2 stream is a flow of its own. A streaming flow (WebSocket,
/// server-sent events, gRPC) has a `StreamStart` and then a `Message` for each message as it
/// passes. Traffic the engine tunnels without decrypting (pinned clients, non-TLS CONNECT
/// payloads, upgrades to protocols other than WebSocket) produces no events.
#[derive(Debug, Clone)]
pub enum ProxyEvent {
/// Request headers arrived from the client.
@@ -154,6 +166,18 @@ pub enum ProxyEvent {
/// The client gave up before the response headers arrived, by resetting the stream or
/// closing the connection. The request upstream is cancelled with it.
Cancelled { id: u64, elapsed_ms: u64 },
/// The flow carries a stream of messages: from its request headers for gRPC, its response
/// headers for server-sent events, its `101` for a WebSocket.
StreamStart { id: u64, kind: StreamKind },
/// A message of a streaming flow passed through.
Message { id: u64, message: StreamMessage },
/// A WebSocket ended: both sides sent a close frame, or the connection ended, with `error`
/// when it failed. Its flow ends here, as the `101` has no body to complete.
StreamEnd {
id: u64,
elapsed_ms: u64,
error: Option<String>,
},
}
/// A running engine. Dropping it stops the listener and joins its thread.
@@ -225,6 +249,7 @@ pub(crate) struct Engine {
handshake_failures: Mutex<HashMap<String, u32>>,
/// Offer HTTP/2 to clients in intercepted TLS.
http2: AtomicBool,
websocket_capture: bool,
/// `Via` value marking requests whose URL the engine rebuilt from `Host`, unique to this run,
/// so one that comes back to the engine is refused instead of forwarded again.
via: HeaderValue,
@@ -308,6 +333,7 @@ pub fn start_proxy(config: ProxyConfig) -> Result<ProxyHandle, String> {
passthrough: Mutex::new(HashSet::new()),
handshake_failures: Mutex::new(HashMap::new()),
http2: AtomicBool::new(config.http2),
websocket_capture: config.websocket_capture,
via: via_marker(),
});
+113
View File
@@ -0,0 +1,113 @@
use std::time::SystemTime;
use bytes::{Bytes, BytesMut};
/// The streaming protocol a flow carries, recorded as messages beside its bodies.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StreamKind {
/// An HTTP/1.1 WebSocket, after its `101 Switching Protocols`.
WebSocket,
/// A `text/event-stream` response.
Sse,
/// `application/grpc`, over HTTP/2.
Grpc,
/// `application/grpc-web`, binary or base64 (`-text`).
GrpcWeb,
}
/// Which side sent a message.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageDirection {
Client,
Server,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MessageKind {
WebSocketText,
WebSocketBinary,
WebSocketPing,
WebSocketPong,
WebSocketClose {
code: Option<u16>,
reason: String,
},
SseEvent {
event: Option<String>,
id: Option<String>,
retry: Option<u64>,
/// The event's lines as they came, comments included, up to [`RAW_LIMIT`] bytes.
raw: String,
},
GrpcMessage,
/// The status of a gRPC call: its HTTP/2 trailers, the headers of a trailers-only response,
/// or a gRPC-Web trailer frame. The payload is the trailers as `name: value` lines.
GrpcTrailers {
status: Option<i32>,
message: Option<String>,
},
}
/// One message of a stream, as it passed through.
#[derive(Debug, Clone)]
pub struct StreamMessage {
pub direction: MessageDirection,
pub kind: MessageKind,
pub at: SystemTime,
/// Bytes it took on the wire: the frames' payloads, the event's lines, the length-prefixed
/// message without its prefix.
pub size: u64,
/// The payload decoded: unmasked, reassembled, inflated or decompressed. Cut at the
/// engine's `capture_body_limit`.
pub data: Bytes,
/// `data` is cut short.
pub truncated: bool,
/// It was compressed on the wire: permessage-deflate, or gRPC's compressed flag.
pub compressed: bool,
/// Why the payload could not be decoded; `data` is then what could be read of it.
pub error: Option<String>,
}
/// Bytes kept of an SSE event's raw lines.
pub const RAW_LIMIT: usize = 4096;
impl StreamMessage {
pub(crate) fn new(direction: MessageDirection, kind: MessageKind) -> Self {
Self {
direction,
kind,
at: SystemTime::now(),
size: 0,
data: Bytes::new(),
truncated: false,
compressed: false,
error: None,
}
}
}
/// A payload collected up to a limit, counting what is left out.
pub(crate) struct Capped {
pub(crate) data: BytesMut,
pub(crate) truncated: bool,
limit: usize,
}
impl Capped {
pub(crate) fn new(limit: usize) -> Self {
Self { data: BytesMut::new(), truncated: false, limit }
}
pub(crate) fn push(&mut self, bytes: &[u8]) {
let room = self.limit.saturating_sub(self.data.len());
if bytes.len() > room {
self.truncated = true;
}
self.data.extend_from_slice(&bytes[..bytes.len().min(room)]);
}
pub(crate) fn take(&mut self) -> (Bytes, bool) {
let truncated = std::mem::take(&mut self.truncated);
(self.data.split().freeze(), truncated)
}
}
+102 -8
View File
@@ -25,9 +25,13 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
use tokio::net::TcpStream;
use tokio_rustls::{LazyConfigAcceptor, TlsConnector};
use crate::body::{CaptureBody, Direction};
use crate::body::{BodyParser, CaptureBody, Direction};
use crate::connection::serve_tls;
use crate::grpc::{self, GrpcParser};
use crate::messages::{MessageDirection, StreamKind};
use crate::sse::SseParser;
use crate::tunnel::{Rewind, splice, tunnel};
use crate::ws::{self, WsFlow};
use crate::{ClientConn, Engine, ProxyEvent};
pub(crate) type BoxBody = http_body_util::combinators::BoxBody<Bytes, hyper::Error>;
@@ -65,6 +69,20 @@ fn error_chain(e: &(dyn Error + 'static)) -> String {
message
}
fn header_str<'a>(headers: &'a hyper::HeaderMap, name: &str) -> Option<&'a str> {
headers.get(name).and_then(|v| v.to_str().ok())
}
/// A gRPC call, from its request headers.
struct GrpcCall {
/// gRPC-Web's base64 flavor.
text: bool,
/// What the client's compressed messages are compressed with.
encoding: Option<String>,
/// The request body's `Content-Encoding`, which gRPC-Web may use.
content_encoding: Option<String>,
}
/// Per-flow bookkeeping shared by the plain HTTP and the intercepted HTTPS paths.
struct Flow {
id: u64,
@@ -74,6 +92,7 @@ struct Flow {
answered: bool,
/// Reading the client's upload failed: it reset the stream or closed the connection.
upload_failed: Arc<AtomicBool>,
grpc: Option<GrpcCall>,
}
impl Flow {
@@ -83,12 +102,13 @@ impl Flow {
req: &Request<Incoming>,
url: String,
) -> Self {
let flow = Self {
let mut flow = Self {
id: engine.next_flow_id(),
start: Instant::now(),
engine,
answered: false,
upload_failed: Arc::new(AtomicBool::new(false)),
grpc: None,
};
flow.engine.emit(ProxyEvent::RequestStart {
id: flow.id,
@@ -118,20 +138,57 @@ impl Flow {
value: String::from_utf8_lossy(value.as_bytes()).into_owned(),
});
}
if let Some((kind, text)) = grpc::detect(header_str(req.headers(), "content-type")) {
flow.engine.emit(ProxyEvent::StreamStart { id: flow.id, kind });
let encoding = header_str(req.headers(), "grpc-encoding").map(String::from);
let content_encoding = header_str(req.headers(), "content-encoding").map(String::from);
flow.grpc = Some(GrpcCall { text, encoding, content_encoding });
}
flow
}
fn capture_request(&self, body: Incoming) -> BoxBody {
let parser = self.grpc.as_ref().map(|call| {
let limit = self.engine.capture_body_limit;
let encoding = call.encoding.as_deref();
let parser = GrpcParser::new(MessageDirection::Client, encoding, call.text, limit)
.with_content_encoding(call.content_encoding.as_deref());
BodyParser::Grpc(parser)
});
CaptureBody::new(body, Direction::Request, self.id, self.start, self.engine.clone())
.flag_failure(self.upload_failed.clone())
.with_parser(parser)
.boxed()
}
fn respond(
mut self,
resp: Response<Incoming>,
remote_addr: Option<SocketAddr>,
) -> Response<BoxBody> {
/// How the response's body is read for messages, if it streams them, after recording that
/// it does.
fn response_parser(&self, resp: &Response<Incoming>) -> Option<BodyParser> {
let limit = self.engine.capture_body_limit;
let headers = resp.headers();
if self.grpc.is_some() {
// A call that fails at once answers with its status in the headers, and no body.
if let Some(message) = grpc::trailers_from_headers(headers) {
self.engine.emit(ProxyEvent::Message { id: self.id, message });
}
// A gateway's own error page is no gRPC body.
let (_, text) = grpc::detect(header_str(headers, "content-type"))?;
let encoding = header_str(headers, "grpc-encoding");
let parser = GrpcParser::new(MessageDirection::Server, encoding, text, limit)
.with_content_encoding(header_str(headers, "content-encoding"));
return Some(BodyParser::Grpc(parser));
}
let content_type = header_str(headers, "content-type").unwrap_or_default();
let essence = content_type.split(';').next().unwrap_or_default().trim();
if !essence.eq_ignore_ascii_case("text/event-stream") {
return None;
}
self.engine.emit(ProxyEvent::StreamStart { id: self.id, kind: StreamKind::Sse });
let encoding = header_str(headers, "content-encoding");
Some(BodyParser::Sse(Box::new(SseParser::new(encoding, limit))))
}
fn respond_head(&mut self, resp: &Response<Incoming>, remote_addr: Option<SocketAddr>) {
self.answered = true;
self.engine.emit(ProxyEvent::ResponseStart {
id: self.id,
@@ -147,12 +204,36 @@ impl Flow {
value: String::from_utf8_lossy(value.as_bytes()).into_owned(),
});
}
}
fn respond(
mut self,
resp: Response<Incoming>,
remote_addr: Option<SocketAddr>,
) -> Response<BoxBody> {
self.respond_head(&resp, remote_addr);
let parser = self.response_parser(&resp);
let (parts, body) = resp.into_parts();
let engine = self.engine.clone();
let body = CaptureBody::new(body, Direction::Response, self.id, self.start, engine);
let body = CaptureBody::new(body, Direction::Response, self.id, self.start, engine)
.with_parser(parser);
Response::from_parts(parts, body.boxed())
}
/// Answers with a WebSocket's `101`. The flow goes on in the returned [`WsFlow`], which
/// ends it, as the `101` has no body to.
fn respond_websocket(
mut self,
resp: Response<Incoming>,
remote_addr: Option<SocketAddr>,
) -> (Response<BoxBody>, Arc<WsFlow>) {
self.respond_head(&resp, remote_addr);
self.engine.emit(ProxyEvent::StreamStart { id: self.id, kind: StreamKind::WebSocket });
let ws = Arc::new(WsFlow::new(self.id, self.start, self.engine.clone()));
let (parts, body) = resp.into_parts();
(Response::from_parts(parts, body.boxed()), ws)
}
/// Answers the client with `status` and `error` without contacting any server.
fn reject(mut self, status: StatusCode, error: &str) -> Response<BoxBody> {
self.answered = true;
@@ -683,6 +764,19 @@ async fn forward_https(
{
let upstream_upgrade = hyper::upgrade::on(&mut resp);
let addr = target.addr.clone();
let websocket = header_str(resp.headers(), UPGRADE.as_str())
.is_some_and(|u| u.trim().eq_ignore_ascii_case("websocket"));
if websocket && engine.websocket_capture {
let deflate = ws::negotiated_deflate(resp.headers());
let (resp, ws_flow) = flow.respond_websocket(resp, remote_addr);
tokio::spawn(async move {
if let Err(e) = ws::relay(client_upgrade, upstream_upgrade, ws_flow, deflate).await
{
debug!("WebSocket to {addr} ended: {e}");
}
});
return Ok(resp);
}
tokio::spawn(async move {
if let Err(e) = splice(client_upgrade, upstream_upgrade).await {
debug!("Upgraded tunnel to {addr} ended: {e}");
+306
View File
@@ -0,0 +1,306 @@
//! Server-sent events (the WHATWG `text/event-stream` format), read from a response body as it
//! streams.
use crate::decode::Decoding;
use crate::messages::{Capped, MessageDirection, MessageKind, RAW_LIMIT, StreamMessage};
const BOM: &[u8] = b"\xEF\xBB\xBF";
/// Bytes a line may take beyond the capture limit, for its field name.
const LINE_SLACK: usize = 64;
pub(crate) struct SseParser {
/// The response's Content-Encoding, undone before the events are read.
decoding: Decoding,
line: Capped,
/// A chunk ended in a CR whose line is not ended yet: an LF starting the next chunk makes it
/// a CRLF, which stays with the event it ends.
pending_cr: bool,
/// The stream ends lines with a bare CR, so a CR ending a chunk ends its line at once, and
/// an LF that starts the next chunk is skipped as the end of a CRLF.
bare_cr: bool,
/// The last chunk ended in a CR read as a bare one.
after_cr: bool,
event: Option<String>,
id: Option<String>,
retry: Option<u64>,
data: Capped,
/// `data` fields since the last event, which `data` joins with line feeds.
data_lines: usize,
/// A `data` line was longer than could be kept.
data_truncated: bool,
/// The start of the stream read so far while it may be a byte order mark; None after.
bom: Option<Vec<u8>>,
raw: Capped,
size: u64,
}
impl SseParser {
pub(crate) fn new(content_encoding: Option<&str>, limit: usize) -> Self {
Self {
decoding: Decoding::new(content_encoding, limit),
// Room for a field name before a value up to the limit.
line: Capped::new(limit.saturating_add(LINE_SLACK)),
pending_cr: false,
bare_cr: false,
after_cr: false,
event: None,
id: None,
retry: None,
data: Capped::new(limit),
data_lines: 0,
data_truncated: false,
bom: Some(Vec::new()),
raw: Capped::new(RAW_LIMIT),
size: 0,
}
}
pub(crate) fn feed(&mut self, bytes: &[u8]) -> Vec<StreamMessage> {
// Left unsupported if decoding fails: what follows cannot be read.
match std::mem::replace(&mut self.decoding, Decoding::Unsupported) {
Decoding::Identity => {
self.decoding = Decoding::Identity;
self.read(bytes)
}
Decoding::Unsupported => Vec::new(),
// A piece at a time, so a highly compressed body never expands all at once.
Decoding::Decoder(mut decoder) => {
let mut events = Vec::new();
if decoder.feed(bytes, &mut |piece| events.extend(self.read(piece))).is_ok() {
self.decoding = Decoding::Decoder(decoder);
}
events
}
}
}
/// The body ended: a line ending held back for the LF that might have followed it ends its
/// line now.
pub(crate) fn finish(&mut self) -> Vec<StreamMessage> {
if std::mem::take(&mut self.pending_cr) {
return self.end_line().into_iter().collect();
}
Vec::new()
}
fn read(&mut self, mut bytes: &[u8]) -> Vec<StreamMessage> {
// One UTF-8 byte order mark at the start of the stream is no part of it, however the
// chunks split it.
if let Some(seen) = self.bom.as_mut() {
let take = (BOM.len() - seen.len()).min(bytes.len());
seen.extend_from_slice(&bytes[..take]);
if !BOM.starts_with(seen) {
let seen = self.bom.take().unwrap_or_default();
let mut events = self.read(&seen);
events.extend(self.read(&bytes[take..]));
return events;
}
bytes = &bytes[take..];
if seen.len() < BOM.len() {
return Vec::new();
}
self.bom = None;
}
let mut events = Vec::new();
let mut i = 0;
while i < bytes.len() {
let byte = bytes[i];
i += 1;
// A CR that ended the last chunk: an LF now makes it a CRLF, kept with its line.
if std::mem::take(&mut self.pending_cr) {
if byte == b'\n' {
self.size += 1;
self.raw.push(b"\n");
}
events.extend(self.end_line());
if byte == b'\n' {
continue;
}
}
// The LF of a CRLF split from a CR already read as a bare one.
if std::mem::take(&mut self.after_cr) && byte == b'\n' {
continue;
}
self.size += 1;
self.raw.push(&[byte]);
match byte {
b'\r' if bytes.get(i) == Some(&b'\n') => {
i += 1;
self.size += 1;
self.raw.push(b"\n");
events.extend(self.end_line());
}
// Whether an LF follows is for the next chunk to say, unless the stream has
// shown it ends lines with bare CRs; then the line ends now, not a chunk late.
b'\r' if i == bytes.len() => {
if self.bare_cr {
self.after_cr = true;
events.extend(self.end_line());
} else {
self.pending_cr = true;
}
}
b'\r' => {
self.bare_cr = true;
events.extend(self.end_line());
}
b'\n' => events.extend(self.end_line()),
_ => self.line.push(&[byte]),
}
}
events
}
fn end_line(&mut self) -> Option<StreamMessage> {
let (line, line_truncated) = self.line.take();
self.data_truncated |= line_truncated && line.starts_with(b"data");
if line.is_empty() {
return self.dispatch();
}
if line[0] == b':' {
return None;
}
let (name, value) = match line.iter().position(|&b| b == b':') {
Some(i) => {
let value = &line[i + 1..];
(&line[..i], value.strip_prefix(b" ").unwrap_or(value))
}
None => (&line[..], &b""[..]),
};
let value = String::from_utf8_lossy(value);
match name {
b"event" => self.event = Some(value.into_owned()),
b"data" => {
if self.data_lines > 0 {
self.data.push(b"\n");
}
self.data.push(value.as_bytes());
self.data_lines += 1;
}
b"id" if !value.contains('\0') => self.id = Some(value.into_owned()),
b"retry" if !value.is_empty() && value.bytes().all(|b| b.is_ascii_digit()) => {
self.retry = value.parse().ok();
}
_ => {}
}
None
}
/// Ends the event at a blank line. As a browser does, only a block with a `data` line is an
/// event; one with only `id`, `retry` or `event` lines dispatches nothing.
fn dispatch(&mut self) -> Option<StreamMessage> {
let (raw, _) = self.raw.take();
let size = std::mem::take(&mut self.size);
let (event, id, retry) = (self.event.take(), self.id.take(), self.retry.take());
let (data, truncated) = self.data.take();
let truncated = truncated || std::mem::take(&mut self.data_truncated);
if std::mem::take(&mut self.data_lines) == 0 {
return None;
}
let raw = String::from_utf8_lossy(&raw).into_owned();
let kind = MessageKind::SseEvent { event, id, retry, raw };
let mut message = StreamMessage::new(MessageDirection::Server, kind);
message.size = size;
message.data = data;
message.truncated = truncated;
Some(message)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn feed_bytewise(parser: &mut SseParser, stream: &[u8]) -> Vec<StreamMessage> {
let mut events: Vec<_> = stream.iter().flat_map(|b| parser.feed(&[*b])).collect();
events.extend(parser.finish());
events
}
fn raw(message: &StreamMessage) -> &str {
match &message.kind {
MessageKind::SseEvent { raw, .. } => raw,
_ => panic!("not an event"),
}
}
#[test]
fn a_byte_order_mark_split_across_chunks_is_skipped() {
let mut parser = SseParser::new(None, 1024);
let events =
feed_bytewise(&mut parser, b"\xEF\xBB\xBFdata: hello\n\ndata: \xEF\xBB\xBF\n\n");
assert_eq!(&events[0].data[..], b"hello");
assert_eq!(&events[1].data[..], b"\xEF\xBB\xBF", "only at the start of the stream");
}
#[test]
fn crlf_stays_with_the_event_it_ends_however_it_is_split() {
for split in [false, true] {
let stream = b"data: a\r\n\r\ndata: b\r\n\r\n";
let mut parser = SseParser::new(None, 1024);
let events =
if split { feed_bytewise(&mut parser, stream) } else { parser.feed(stream) };
assert_eq!(events.len(), 2, "split: {split}");
assert_eq!(raw(&events[0]), "data: a\r\n\r\n", "split: {split}");
assert_eq!(events[0].size, 11, "split: {split}");
assert_eq!(raw(&events[1]), "data: b\r\n\r\n", "split: {split}");
}
}
#[test]
fn a_cr_ending_a_chunk_waits_for_the_next_unless_the_stream_uses_bare_crs() {
// A CR ending the chunk might start a CRLF, so its line waits for the next byte.
let mut parser = SseParser::new(None, 1024);
assert!(parser.feed(b"data: x\n\r").is_empty());
let events = parser.feed(b"\n");
assert_eq!(raw(&events[0]), "data: x\n\r\n");
// Bare CRs seen mid-chunk: one ending a chunk ends its event at once.
let mut parser = SseParser::new(None, 1024);
let events = parser.feed(b"data: a\r\rdata: b\r\r");
assert_eq!(events.len(), 2);
assert_eq!((&events[0].data[..], &events[1].data[..]), (&b"a"[..], &b"b"[..]));
// An LF after it is the rest of a CRLF, no blank line of its own.
let events = parser.feed(b"\ndata: c\r\r");
assert_eq!(events.len(), 1);
assert_eq!(&events[0].data[..], b"c");
// A held CR at the body's end ends its line then.
let mut parser = SseParser::new(None, 1024);
assert!(parser.feed(b"data: y\n\r").is_empty());
assert_eq!(&parser.finish()[0].data[..], b"y");
}
#[test]
fn only_blocks_with_data_are_events() {
let mut parser = SseParser::new(None, 1024);
let events = parser.feed(b"retry: 1500\n\nid: 7\nevent: x\n\ndata:\n\n: comment\n\n");
assert_eq!(events.len(), 1, "only the empty data event");
assert_eq!(&events[0].data[..], b"");
assert_eq!(raw(&events[0]), "data:\n\n");
}
#[test]
fn a_data_line_past_the_limit_is_truncated() {
let mut parser = SseParser::new(None, 16);
let events = parser.feed(format!("data: {}\n\n", "x".repeat(100)).as_bytes());
assert_eq!(events[0].data.len(), 16);
assert!(events[0].truncated);
}
#[test]
fn gzip_events_are_read_as_they_decompress() {
use std::io::Write;
let mut gz = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
gz.write_all(b"data: one\n\n").unwrap();
gz.flush().unwrap();
let first = gz.get_ref().len();
gz.write_all(b"data: two\n\n").unwrap();
let compressed = gz.finish().unwrap();
let mut parser = SseParser::new(Some("gzip"), 1024);
let events = parser.feed(&compressed[..first]);
assert_eq!(&events[0].data[..], b"one", "before the rest of the body");
assert_eq!(&parser.feed(&compressed[first..])[0].data[..], b"two");
}
}
+419
View File
@@ -0,0 +1,419 @@
//! WebSocket capture: after the `101`, the upgraded connections are relayed byte for byte, and
//! a frame parser on each direction records every frame as it passes. The parser reads a copy
//! of what was already written on, so it never holds bytes back or changes them.
use std::io;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
use std::time::Instant;
use flate2::{Decompress, FlushDecompress, Status};
use hyper::header::{HeaderMap, SEC_WEBSOCKET_EXTENSIONS};
use hyper::upgrade::OnUpgrade;
use hyper_util::rt::TokioIo;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use crate::messages::{Capped, MessageDirection, MessageKind, StreamMessage};
use crate::{Engine, ProxyEvent};
const OP_CONTINUATION: u8 = 0x0;
const OP_TEXT: u8 = 0x1;
const OP_BINARY: u8 = 0x2;
const OP_CLOSE: u8 = 0x8;
const OP_PING: u8 = 0x9;
const OP_PONG: u8 = 0xA;
/// The most a control frame carries (RFC 6455 5.5).
const MAX_CONTROL_PAYLOAD: usize = 125;
/// The tail permessage-deflate strips from each compressed message (RFC 7692 7.2.2).
const DEFLATE_TAIL: [u8; 4] = [0x00, 0x00, 0xff, 0xff];
/// permessage-deflate as the server accepted it.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub(crate) struct Deflate {
client_no_context_takeover: bool,
server_no_context_takeover: bool,
}
/// The permessage-deflate parameters in the server's `Sec-WebSocket-Extensions`, if it accepted
/// the extension.
pub(crate) fn negotiated_deflate(headers: &HeaderMap) -> Option<Deflate> {
let extensions =
headers.get_all(SEC_WEBSOCKET_EXTENSIONS).iter().filter_map(|v| v.to_str().ok());
for extension in extensions.flat_map(|v| v.split(',')) {
let mut params = extension.split(';').map(str::trim);
if !params.next().is_some_and(|name| name.eq_ignore_ascii_case("permessage-deflate")) {
continue;
}
let mut deflate = Deflate::default();
for param in params {
let name = param.split('=').next().unwrap_or_default().trim();
if name.eq_ignore_ascii_case("client_no_context_takeover") {
deflate.client_no_context_takeover = true;
} else if name.eq_ignore_ascii_case("server_no_context_takeover") {
deflate.server_no_context_takeover = true;
}
}
return Some(deflate);
}
None
}
/// One WebSocket flow, shared by its two directions. It ends once both sides have sent a close
/// frame, or when the connection does, whichever is first.
pub(crate) struct WsFlow {
id: u64,
start: Instant,
engine: Arc<Engine>,
/// A bit per direction that has sent a close frame.
closes: AtomicU8,
ended: AtomicBool,
}
impl WsFlow {
pub(crate) fn new(id: u64, start: Instant, engine: Arc<Engine>) -> Self {
Self { id, start, engine, closes: AtomicU8::new(0), ended: AtomicBool::new(false) }
}
fn emit(&self, message: StreamMessage) {
if !self.ended.load(Ordering::Relaxed) {
self.engine.emit(ProxyEvent::Message { id: self.id, message });
}
}
fn closed(&self, direction: MessageDirection) {
let bit = match direction {
MessageDirection::Client => 1,
MessageDirection::Server => 2,
};
if self.closes.fetch_or(bit, Ordering::Relaxed) | bit == 3 {
self.end(None);
}
}
pub(crate) fn end(&self, error: Option<String>) {
if !self.ended.swap(true, Ordering::Relaxed) {
self.engine.emit(ProxyEvent::StreamEnd {
id: self.id,
elapsed_ms: self.start.elapsed().as_millis() as u64,
error,
});
}
}
}
/// Relays an upgraded WebSocket connection, recording its frames on `flow`, and ends the flow
/// when the connection ends.
pub(crate) async fn relay(
client: OnUpgrade,
upstream: OnUpgrade,
flow: Arc<WsFlow>,
deflate: Option<Deflate>,
) -> io::Result<()> {
let (client, upstream) = match tokio::try_join!(client, upstream) {
Ok(joined) => joined,
Err(e) => {
flow.end(Some(e.to_string()));
return Err(io::Error::other(e));
}
};
let (client_read, client_write) = tokio::io::split(TokioIo::new(client));
let (upstream_read, upstream_write) = tokio::io::split(TokioIo::new(upstream));
let limit = flow.engine.capture_body_limit;
let outbound = FrameParser::new(
MessageDirection::Client,
flow.clone(),
limit,
deflate.map(|d| d.client_no_context_takeover),
);
let inbound = FrameParser::new(
MessageDirection::Server,
flow.clone(),
limit,
deflate.map(|d| d.server_no_context_takeover),
);
let result = tokio::try_join!(
pump(client_read, upstream_write, outbound),
pump(upstream_read, client_write, inbound),
);
flow.end(result.as_ref().err().map(ToString::to_string));
result.map(|_| ())
}
/// Copies `from` to `to`, then shows the parser what went through. At the end of `from`, closes
/// `to` for writing, and the other direction carries on until it ends too.
async fn pump<R, W>(mut from: R, mut to: W, mut parser: FrameParser) -> io::Result<()>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
let mut buf = vec![0u8; 16 * 1024];
loop {
let n = from.read(&mut buf).await?;
if n == 0 {
let _ = to.shutdown().await;
return Ok(());
}
to.write_all(&buf[..n]).await?;
to.flush().await?;
parser.feed(&buf[..n]);
}
}
struct Frame {
fin: bool,
opcode: u8,
mask: Option<[u8; 4]>,
remaining: u64,
offset: usize,
}
/// A data message being reassembled from its frames.
struct Assembly {
opcode: u8,
compressed: bool,
size: u64,
payload: Capped,
error: Option<String>,
}
/// One direction's inflate context, kept across messages unless the side that sends them
/// negotiated no context takeover.
struct Inflater {
decompress: Decompress,
no_context_takeover: bool,
/// Inflating failed, so the context is lost for the messages after.
failed: bool,
}
impl Inflater {
fn inflate(&mut self, mut input: &[u8], out: &mut Capped) -> Result<(), String> {
let mut buf = [0u8; 8192];
loop {
let (read_before, written_before) =
(self.decompress.total_in(), self.decompress.total_out());
let status = self
.decompress
.decompress(input, &mut buf, FlushDecompress::Sync)
.map_err(|e| e.to_string())?;
let read = (self.decompress.total_in() - read_before) as usize;
let written = (self.decompress.total_out() - written_before) as usize;
out.push(&buf[..written]);
input = &input[read..];
if status == Status::StreamEnd {
// A final block ends the deflate stream; the next message starts a new one.
self.decompress.reset(false);
}
if (input.is_empty() && written < buf.len()) || (read == 0 && written == 0) {
return Ok(());
}
}
}
}
/// Reads WebSocket frames (RFC 6455 5.2) from one direction's bytes, records each message and
/// control frame. Stops recording at a frame it cannot read, as the stream is then unknown.
struct FrameParser {
direction: MessageDirection,
flow: Arc<WsFlow>,
limit: usize,
header: Vec<u8>,
frame: Option<Frame>,
message: Option<Assembly>,
control: Vec<u8>,
inflater: Option<Inflater>,
broken: bool,
}
/// The length of the frame header that starts with `header`, once its first two bytes are in.
fn header_len(header: &[u8]) -> Option<usize> {
let second = *header.get(1)?;
let length = match second & 0x7f {
126 => 2,
127 => 8,
_ => 0,
};
let mask = if second & 0x80 != 0 { 4 } else { 0 };
Some(2 + length + mask)
}
impl FrameParser {
fn new(
direction: MessageDirection,
flow: Arc<WsFlow>,
limit: usize,
no_context_takeover: Option<bool>,
) -> Self {
Self {
direction,
flow,
limit,
header: Vec::with_capacity(14),
frame: None,
message: None,
control: Vec::new(),
inflater: no_context_takeover.map(|no_context_takeover| Inflater {
decompress: Decompress::new(false),
no_context_takeover,
failed: false,
}),
broken: false,
}
}
fn feed(&mut self, mut data: &[u8]) {
while !data.is_empty() && !self.broken {
let Some(frame) = self.frame.as_mut() else {
let need = header_len(&self.header).unwrap_or(2);
let take = (need - self.header.len()).min(data.len());
self.header.extend_from_slice(&data[..take]);
data = &data[take..];
if header_len(&self.header) == Some(self.header.len()) {
self.begin_frame();
}
continue;
};
let n = frame.remaining.min(data.len() as u64) as usize;
let mut chunk = data[..n].to_vec();
if let Some(mask) = frame.mask {
for (i, byte) in chunk.iter_mut().enumerate() {
*byte ^= mask[(frame.offset + i) % 4];
}
}
frame.offset += n;
frame.remaining -= n as u64;
let (fin, opcode, done) = (frame.fin, frame.opcode, frame.remaining == 0);
data = &data[n..];
self.payload(opcode, &chunk);
if done {
self.frame = None;
self.end_frame(fin, opcode);
}
}
}
fn begin_frame(&mut self) {
let header = std::mem::take(&mut self.header);
let (first, second) = (header[0], header[1]);
let fin = first & 0x80 != 0;
let rsv1 = first & 0x40 != 0;
let opcode = first & 0x0f;
let (length, rest) = match second & 0x7f {
126 => (u16::from_be_bytes([header[2], header[3]]) as u64, &header[4..]),
127 => (u64::from_be_bytes(header[2..10].try_into().unwrap()), &header[10..]),
n => (n as u64, &header[2..]),
};
let mask = (second & 0x80 != 0).then(|| [rest[0], rest[1], rest[2], rest[3]]);
self.header = header;
self.header.clear();
match opcode {
OP_CLOSE | OP_PING | OP_PONG => self.control.clear(),
OP_TEXT | OP_BINARY => {
self.message = Some(Assembly {
opcode,
compressed: rsv1 && self.inflater.is_some(),
size: 0,
payload: Capped::new(self.limit),
error: None,
});
}
OP_CONTINUATION => {}
// An opcode no extension here defines: what follows cannot be read.
_ => {
self.broken = true;
return;
}
}
self.frame = Some(Frame { fin, opcode, mask, remaining: length, offset: 0 });
if length == 0 {
self.frame = None;
self.end_frame(fin, opcode);
}
}
fn payload(&mut self, opcode: u8, chunk: &[u8]) {
if opcode >= OP_CLOSE {
// At most what a control frame may carry, however long a broken one claims to be.
let room = MAX_CONTROL_PAYLOAD.saturating_sub(self.control.len());
self.control.extend_from_slice(&chunk[..chunk.len().min(room)]);
return;
}
let Some(message) = self.message.as_mut() else {
return;
};
message.size += chunk.len() as u64;
match self.inflater.as_mut().filter(|_| message.compressed) {
Some(inflater) if !inflater.failed => {
if let Err(e) = inflater.inflate(chunk, &mut message.payload) {
inflater.failed = true;
message.error = Some(format!("Failed to inflate the message: {e}"));
}
}
Some(_) => {
message.payload.push(chunk);
message.error.get_or_insert_with(|| {
"Not inflated: an earlier message failed to, which lost the context".into()
});
}
None => message.payload.push(chunk),
}
}
fn end_frame(&mut self, fin: bool, opcode: u8) {
if opcode >= OP_CLOSE {
let payload = std::mem::take(&mut self.control);
let kind = match opcode {
OP_PING => MessageKind::WebSocketPing,
OP_PONG => MessageKind::WebSocketPong,
_ => {
let code =
(payload.len() >= 2).then(|| u16::from_be_bytes([payload[0], payload[1]]));
let reason = String::from_utf8_lossy(payload.get(2..).unwrap_or_default());
MessageKind::WebSocketClose { code, reason: reason.into_owned() }
}
};
let mut message = StreamMessage::new(self.direction, kind);
message.size = payload.len() as u64;
message.data = payload.into();
self.flow.emit(message);
if opcode == OP_CLOSE {
self.flow.closed(self.direction);
}
return;
}
if !fin {
return;
}
let Some(mut assembly) = self.message.take() else {
return;
};
if assembly.compressed
&& let Some(inflater) = self.inflater.as_mut()
{
if !inflater.failed
&& let Err(e) = inflater.inflate(&DEFLATE_TAIL, &mut assembly.payload)
{
inflater.failed = true;
assembly.error = Some(format!("Failed to inflate the message: {e}"));
}
if inflater.no_context_takeover && !inflater.failed {
inflater.decompress.reset(false);
}
}
let kind = if assembly.opcode == OP_TEXT {
MessageKind::WebSocketText
} else {
MessageKind::WebSocketBinary
};
let mut message = StreamMessage::new(self.direction, kind);
let (data, truncated) = assembly.payload.take();
message.size = assembly.size;
message.data = data;
message.truncated = truncated;
message.compressed = assembly.compressed;
message.error = assembly.error;
self.flow.emit(message);
}
}
+549 -21
View File
@@ -6,6 +6,8 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::Receiver;
use std::time::{Duration, SystemTime};
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD as BASE64;
use bytes::Bytes;
use futures_util::{SinkExt, StreamExt};
use http_body_util::combinators::BoxBody;
@@ -28,10 +30,26 @@ use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::handshake::derive_accept_key;
use tokio_tungstenite::tungstenite::protocol::Role;
use yaak_proxy::{CaStore, ProxyConfig, ProxyEvent, ProxyHandle, start_proxy};
use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame;
use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
use yaak_proxy::{
CaStore, MessageDirection, MessageKind, ProxyConfig, ProxyEvent, ProxyHandle, StreamKind,
StreamMessage, start_proxy,
};
const BIG_BODY: usize = 64 * 1024;
/// The longest a streaming test may run: one that hangs fails with a message instead of
/// holding CI until its job is cancelled.
const STREAM_TEST_TIMEOUT: Duration = Duration::from_secs(60);
async fn within<T>(test: impl Future<Output = T>) -> T {
match tokio::time::timeout(STREAM_TEST_TIMEOUT, test).await {
Ok(result) => result,
Err(_) => panic!("the test did not finish within {STREAM_TEST_TIMEOUT:?}"),
}
}
macro_rules! find {
($events:expr, $pat:pat => $out:expr) => {
$events
@@ -166,7 +184,7 @@ fn websocket_echo(mut req: Request<Incoming>) -> Response<TestBody> {
let io = TokioIo::new(on_upgrade.await.unwrap());
let mut ws = WebSocketStream::from_raw_socket(io, Role::Server, None).await;
while let Some(Ok(msg)) = ws.next().await {
if msg.is_text() {
if msg.is_text() || msg.is_binary() {
ws.send(msg).await.unwrap();
}
}
@@ -231,6 +249,10 @@ struct Harness {
impl Harness {
fn start(capture_body_limit: usize) -> Self {
Self::start_with(capture_body_limit, |_| {})
}
fn start_with(capture_body_limit: usize, configure: impl FnOnce(&mut ProxyConfig)) -> Self {
let dir = TempDir::new().unwrap();
let ca_dir = dir.path().join("ca");
let created = CaStore::load_or_create(&ca_dir).unwrap();
@@ -242,6 +264,7 @@ impl Harness {
let mut config = ProxyConfig::new(0, ca.clone());
config.capture_body_limit = capture_body_limit;
config.upstream_root_certs = vec![pki.ca_der.clone()];
configure(&mut config);
let mut proxy = start_proxy(config).unwrap();
let events = proxy.take_event_rx().unwrap();
Self { proxy, events, ca, pki, _dir: dir }
@@ -381,6 +404,7 @@ fn is_terminal(event: &ProxyEvent) -> bool {
ProxyEvent::ResponseBodyComplete { .. }
| ProxyEvent::Error { .. }
| ProxyEvent::Cancelled { .. }
| ProxyEvent::StreamEnd { .. }
)
}
@@ -395,7 +419,10 @@ fn flow_id(event: &ProxyEvent) -> u64 {
| ProxyEvent::ResponseBodyChunk { id, .. }
| ProxyEvent::ResponseBodyComplete { id, .. }
| ProxyEvent::Error { id, .. }
| ProxyEvent::Cancelled { id, .. } => *id,
| ProxyEvent::Cancelled { id, .. }
| ProxyEvent::StreamStart { id, .. }
| ProxyEvent::Message { id, .. }
| ProxyEvent::StreamEnd { id, .. } => *id,
}
}
@@ -677,33 +704,534 @@ async fn bodies_over_the_cap_are_truncated_but_delivered() {
#[tokio::test(flavor = "multi_thread")]
async fn websocket_upgrade_passes_through() {
let h = Harness::start(1024);
let upstream = spawn_upstream(Some(h.pki.server.clone())).await;
let port = upstream.port();
within(async {
let h = Harness::start(1024);
let upstream = spawn_upstream(Some(h.pki.server.clone())).await;
let port = upstream.port();
let tcp = h.connect_tunnel(port).await;
let tcp = h.connect_tunnel(port).await;
let tls = TlsConnector::from(h.tls_client_config())
.connect(ServerName::try_from("localhost").unwrap(), tcp)
.await
.unwrap();
let url = format!("wss://localhost:{port}/ws");
let (mut ws, _) = tokio_tungstenite::client_async(url.as_str(), tls).await.unwrap();
ws.send(Message::text("hello")).await.unwrap();
assert_eq!(ws.next().await.unwrap().unwrap().into_text().unwrap().as_str(), "hello");
ws.send(Message::binary(vec![0u8, 1, 2, 255])).await.unwrap();
assert_eq!(ws.next().await.unwrap().unwrap().into_data().as_ref(), &[0u8, 1, 2, 255]);
ws.send(Message::Ping(Bytes::from_static(b"p"))).await.unwrap();
assert!(
matches!(ws.next().await.unwrap().unwrap(), Message::Pong(p) if p.as_ref() == b"p")
);
ws.close(Some(CloseFrame { code: CloseCode::Normal, reason: "bye".into() })).await.unwrap();
while let Some(Ok(_)) = ws.next().await {}
let events = h.flow_events();
let req_url = find!(events, ProxyEvent::RequestStart { url, .. } => url.clone());
assert_eq!(req_url, format!("https://localhost:{port}/ws"));
let (status, version) = find!(events,
ProxyEvent::ResponseStart { status, http_version, .. } => (*status, http_version.clone()));
assert_eq!((status, version.as_str()), (101, "HTTP/1.1"));
let kind = find!(events, ProxyEvent::StreamStart { kind, .. } => *kind);
assert_eq!(kind, StreamKind::WebSocket);
// The flow ends with the close handshake, not with the 101.
assert!(!events.iter().any(|e| matches!(e, ProxyEvent::ResponseBodyComplete { .. })));
let error = find!(events, ProxyEvent::StreamEnd { error, .. } => error.clone());
assert_eq!(error, None);
for direction in [MessageDirection::Client, MessageDirection::Server] {
let sent = messages(&events, direction);
let expected_control = match direction {
MessageDirection::Client => MessageKind::WebSocketPing,
MessageDirection::Server => MessageKind::WebSocketPong,
};
let close = MessageKind::WebSocketClose { code: Some(1000), reason: "bye".into() };
let kinds: Vec<_> = sent.iter().map(|m| m.kind.clone()).collect();
assert_eq!(
kinds,
vec![
MessageKind::WebSocketText,
MessageKind::WebSocketBinary,
expected_control,
close
],
"{direction:?}"
);
assert_eq!(&sent[0].data[..], b"hello");
assert_eq!(&sent[1].data[..], &[0u8, 1, 2, 255]);
assert_eq!(&sent[2].data[..], b"p");
}
})
.await
}
/// The messages a flow recorded from one side, in order.
fn messages(events: &[ProxyEvent], direction: MessageDirection) -> Vec<StreamMessage> {
events
.iter()
.filter_map(|e| match e {
ProxyEvent::Message { message, .. } if message.direction == direction => {
Some(message.clone())
}
_ => None,
})
.collect()
}
/// A WebSocket frame as written on the wire, masked when `mask` is given.
fn ws_frame(fin: bool, rsv1: bool, opcode: u8, payload: &[u8], mask: Option<[u8; 4]>) -> Vec<u8> {
let mut out = vec![(u8::from(fin) << 7) | (u8::from(rsv1) << 6) | opcode];
let mask_bit = if mask.is_some() { 0x80 } else { 0 };
match payload.len() {
n if n < 126 => out.push(mask_bit | n as u8),
n if n <= u16::MAX as usize => {
out.push(mask_bit | 126);
out.extend((n as u16).to_be_bytes());
}
n => {
out.push(mask_bit | 127);
out.extend((n as u64).to_be_bytes());
}
}
match mask {
Some(mask) => {
out.extend(mask);
out.extend(payload.iter().enumerate().map(|(i, b)| b ^ mask[i % 4]));
}
None => out.extend(payload),
}
out
}
/// permessage-deflate's compression of one message, in a context kept across messages.
fn deflate_message(compress: &mut flate2::Compress, data: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(data.len() + 64);
compress.compress_vec(data, &mut out, flate2::FlushCompress::Sync).unwrap();
assert!(out.ends_with(&[0, 0, 0xff, 0xff]));
out.truncate(out.len() - 4);
out
}
/// A server that switches to WebSocket with permessage-deflate and then sends back every byte
/// it gets, as it got it, keeping a copy.
fn raw_websocket_echo(
mut req: Request<Incoming>,
received: Arc<std::sync::Mutex<Vec<u8>>>,
) -> 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 {
let mut io = TokioIo::new(on_upgrade.await.unwrap());
let mut buf = vec![0u8; 4096];
loop {
let n = io.read(&mut buf).await.unwrap_or(0);
if n == 0 {
break;
}
received.lock().unwrap().extend_from_slice(&buf[..n]);
io.write_all(&buf[..n]).await.unwrap();
}
});
Response::builder()
.status(StatusCode::SWITCHING_PROTOCOLS)
.header(UPGRADE, "websocket")
.header(CONNECTION, "Upgrade")
.header("sec-websocket-accept", accept)
.header("sec-websocket-extensions", "permessage-deflate")
.body(full(Bytes::new()))
.unwrap()
}
#[tokio::test(flavor = "multi_thread")]
async fn websocket_frames_pass_byte_for_byte() {
within(async {
let h = Harness::start(1024 * 1024);
let received = Arc::new(std::sync::Mutex::new(Vec::new()));
let handler = {
let received = received.clone();
move |req| {
let received = received.clone();
async move { Ok::<_, Infallible>(raw_websocket_echo(req, received)) }
}
};
let upstream =
spawn_upstream_with(Some(h.pki.server.clone()), Arc::new(AtomicUsize::new(0)), handler)
.await;
let mut tls = h.tls_tunnel(upstream.port(), &[b"http/1.1"]).await;
let handshake = format!(
"GET /raw HTTP/1.1\r\nHost: localhost:{}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\
Sec-WebSocket-Extensions: permessage-deflate\r\n\r\n",
upstream.port()
);
tls.write_all(handshake.as_bytes()).await.unwrap();
let mut head = Vec::new();
while !head.ends_with(b"\r\n\r\n") {
head.push(tls.read_u8().await.unwrap());
}
assert!(head.starts_with(b"HTTP/1.1 101"), "{}", String::from_utf8_lossy(&head));
let mask = Some([0x12, 0x34, 0x56, 0x78]);
let mut compress = flate2::Compress::new(flate2::Compression::default(), false);
let big: Vec<u8> = (0..200_000u32).map(|i| (i % 7) as u8).collect();
let mut wire = Vec::new();
// A text message in two fragments, with a ping between them.
wire.extend(ws_frame(false, false, 0x1, b"Hel", mask));
wire.extend(ws_frame(true, false, 0x9, b"x", mask));
wire.extend(ws_frame(true, false, 0x0, b"lo", mask));
// Two compressed messages sharing one deflate context, the second over 64 KB inflated.
let compressed = deflate_message(&mut compress, b"compressed text compressed text");
wire.extend(ws_frame(true, true, 0x1, &compressed, mask));
let compressed = deflate_message(&mut compress, &big);
wire.extend(ws_frame(true, true, 0x2, &compressed, mask));
let mut close = 4000u16.to_be_bytes().to_vec();
close.extend(b"done");
wire.extend(ws_frame(true, false, 0x8, &close, mask));
// Odd-sized writes, so frames and headers split across reads.
for piece in wire.chunks(7) {
tls.write_all(piece).await.unwrap();
}
// TLS may hold the last records back until flushed.
tls.flush().await.unwrap();
let mut echoed = vec![0u8; wire.len()];
tls.read_exact(&mut echoed).await.unwrap();
assert_eq!(echoed, wire, "the client gets back exactly what the server sent");
assert_eq!(*received.lock().unwrap(), wire, "the server gets exactly what the client sent");
drop(tls);
let events = h.flow_events();
// The server sent the client's masked frames back, so both sides recorded the same.
for direction in [MessageDirection::Client, MessageDirection::Server] {
let sent = messages(&events, direction);
let kinds: Vec<_> = sent.iter().map(|m| m.kind.clone()).collect();
assert_eq!(
kinds,
vec![
MessageKind::WebSocketPing,
MessageKind::WebSocketText,
MessageKind::WebSocketText,
MessageKind::WebSocketBinary,
MessageKind::WebSocketClose { code: Some(4000), reason: "done".into() },
],
"{direction:?}"
);
assert_eq!(&sent[1].data[..], b"Hello");
assert_eq!(sent[1].size, 5);
assert_eq!(&sent[2].data[..], b"compressed text compressed text");
assert!(sent[2].compressed && sent[2].error.is_none(), "{:?}", sent[2].error);
assert_eq!(&sent[3].data[..], &big[..], "inflated with the first message's context");
assert!(sent[3].size < big.len() as u64, "size on the wire, compressed");
}
let error = find!(events, ProxyEvent::StreamEnd { error, .. } => error.clone());
assert_eq!(error, None, "ended by the close frames");
})
.await
}
/// Round trips over a WebSocket through the proxy, in microseconds, sorted.
async fn websocket_round_trips(websocket_capture: bool, n: usize) -> Vec<u128> {
let h = Harness::start_with(1024 * 1024, |c| c.websocket_capture = websocket_capture);
let upstream = spawn_upstream(Some(h.pki.server.clone())).await;
let tcp = h.connect_tunnel(upstream.port()).await;
let tls = TlsConnector::from(h.tls_client_config())
.connect(ServerName::try_from("localhost").unwrap(), tcp)
.await
.unwrap();
let url = format!("wss://localhost:{port}/ws");
let url = format!("wss://localhost:{}/ws", upstream.port());
let (mut ws, _) = tokio_tungstenite::client_async(url.as_str(), tls).await.unwrap();
for text in ["ping", "pong"] {
ws.send(Message::text(text)).await.unwrap();
let reply = ws.next().await.unwrap().unwrap();
assert_eq!(reply.into_text().unwrap().as_str(), text);
let payload = "x".repeat(256);
let mut times = Vec::with_capacity(n);
for i in 0..n + 100 {
let start = std::time::Instant::now();
ws.send(Message::text(payload.clone())).await.unwrap();
ws.next().await.unwrap().unwrap();
// The first hundred warm up.
if i >= 100 {
times.push(start.elapsed().as_micros());
}
}
ws.close(None).await.unwrap();
times.sort_unstable();
times
}
let events = h.flow_events();
let req_url = find!(events, ProxyEvent::RequestStart { url, .. } => url.clone());
assert_eq!(req_url, format!("https://localhost:{port}/ws"));
let (status, version) = find!(events,
ProxyEvent::ResponseStart { status, http_version, .. } => (*status, http_version.clone()));
assert_eq!((status, version.as_str()), (101, "HTTP/1.1"));
let (size, error) = find!(events, ProxyEvent::ResponseBodyComplete { size, error, .. } => (*size, error.clone()));
assert_eq!((size, error), (0, None));
#[tokio::test(flavor = "multi_thread")]
async fn websocket_capture_adds_no_measurable_latency() {
within(async {
const N: usize = 2000;
let blind = websocket_round_trips(false, N).await;
let captured = websocket_round_trips(true, N).await;
let pct = |t: &[u128], p: usize| t[(t.len() * p / 100).min(t.len() - 1)];
eprintln!(
"WebSocket round trip through the proxy, {N} x 256 B text, microseconds: \
blind splice p50 {} p90 {} p99 {}; captured p50 {} p90 {} p99 {}",
pct(&blind, 50),
pct(&blind, 90),
pct(&blind, 99),
pct(&captured, 50),
pct(&captured, 90),
pct(&captured, 99),
);
// Generous, so a busy machine does not fail it: within a millisecond of the blind splice at
// the median.
assert!(pct(&captured, 50) < pct(&blind, 50) + 1000);
})
.await
}
/// An event stream whose chunks the test sends one at a time.
fn event_stream(chunks: tokio::sync::mpsc::UnboundedReceiver<&'static str>) -> Response<TestBody> {
let frames = futures_util::stream::unfold(chunks, |mut chunks| async move {
let chunk = chunks.recv().await?;
Some((Ok(Frame::data(Bytes::from(chunk))), chunks))
});
Response::builder()
.header(CONTENT_TYPE, "text/event-stream")
.body(BodyExt::boxed(StreamBody::new(frames)))
.unwrap()
}
/// The next event of a flow, skipping body chunks and headers.
fn next_stream_message(h: &Harness) -> StreamMessage {
loop {
match h.events.recv_timeout(Duration::from_secs(10)).expect("no message") {
ProxyEvent::Message { message, .. } => return message,
ProxyEvent::ResponseBodyComplete { .. } => panic!("the stream ended first"),
_ => {}
}
}
}
#[tokio::test(flavor = "multi_thread")]
async fn server_sent_events_are_recorded_as_they_arrive() {
within(async {
let h = Harness::start(1024);
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let rx = Arc::new(std::sync::Mutex::new(Some(rx)));
let handler = move |_req| {
let rx = rx.lock().unwrap().take().unwrap();
async move { Ok::<_, Infallible>(event_stream(rx)) }
};
let upstream =
spawn_upstream_with(Some(h.pki.server.clone()), Arc::new(AtomicUsize::new(0)), handler)
.await;
let url = format!("https://localhost:{}/events", upstream.port());
let mut resp = h.client().get(&url).send().await.unwrap();
tx.send("event: greet\nid: 1\ndata: hello\n").unwrap();
tx.send("data: world\n\n").unwrap();
// Recorded while the stream is still open.
let first = next_stream_message(&h);
assert_eq!(&first.data[..], b"hello\nworld");
assert_eq!(
first.kind,
MessageKind::SseEvent {
event: Some("greet".into()),
id: Some("1".into()),
retry: None,
raw: "event: greet\nid: 1\ndata: hello\ndata: world\n\n".into(),
}
);
tokio::time::sleep(Duration::from_millis(500)).await;
tx.send(": a comment\r\nretry: 1500\r\ndata: two\r\n\r\n").unwrap();
let second = next_stream_message(&h);
assert_eq!(&second.data[..], b"two");
assert!(matches!(second.kind, MessageKind::SseEvent { retry: Some(1500), .. }));
drop(tx);
let mut body = Vec::new();
while let Some(chunk) = resp.chunk().await.unwrap() {
body.extend_from_slice(&chunk);
}
assert!(body.ends_with(b"data: two\r\n\r\n"), "the client gets the stream unchanged");
let events = h.flow_events();
let error = find!(events, ProxyEvent::ResponseBodyComplete { error, .. } => error.clone());
assert_eq!(error, None);
})
.await
}
/// A gRPC message on the wire: its five-byte prefix, then `payload`.
fn grpc_frame(flag: u8, payload: &[u8]) -> Vec<u8> {
let mut out = vec![flag];
out.extend((payload.len() as u32).to_be_bytes());
out.extend(payload);
out
}
/// A protobuf message with one string field, number 1, as `HelloRequest { name }` encodes.
fn proto_string(value: &str) -> Vec<u8> {
let mut out = vec![0x0a, value.len() as u8];
out.extend(value.as_bytes());
out
}
/// A gRPC server for `pkg.Greeter`: `SayHello` answers once, `StreamHellos` three times, and
/// `Fail` at once with a NOT_FOUND status in its headers.
async fn greeter(req: Request<Incoming>) -> Result<Response<TestBody>, Infallible> {
let path = req.uri().path().to_string();
let web = req.headers()[CONTENT_TYPE].to_str().unwrap().starts_with("application/grpc-web");
// Asked for by a header of the test's own: clients built with gzip support ask for it on
// every request.
let gzip = req.headers().contains_key("x-test-gzip");
let request = req.into_body().collect().await.unwrap().to_bytes();
let request = if web { BASE64.decode(&request).unwrap() } else { request.to_vec() };
let name = String::from_utf8_lossy(&request[7..]).into_owned();
let replies: Vec<String> = match path.as_str() {
"/pkg.Greeter/SayHello" => vec![format!("Hello {name}")],
"/pkg.Greeter/StreamHellos" => (1..=3).map(|i| format!("Hello {name} #{i}")).collect(),
_ => {
return Ok(Response::builder()
.header(CONTENT_TYPE, "application/grpc")
.header("grpc-status", "5")
.header("grpc-message", "no%20such%20greeter")
.body(full(Bytes::new()))
.unwrap());
}
};
if web {
// gRPC-Web's text flavor: each frame base64 of its own, the trailers in the last.
let mut body = String::new();
for reply in &replies {
body.push_str(&BASE64.encode(grpc_frame(0, &proto_string(reply))));
}
body.push_str(&BASE64.encode(grpc_frame(0x80, b"grpc-status: 0\r\ngrpc-message: \r\n")));
let resp = Response::builder().header(CONTENT_TYPE, "application/grpc-web-text");
// Compressed for HTTP as a whole when the client accepts it, as a CDN does.
if gzip {
let mut encoder =
flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
std::io::Write::write_all(&mut encoder, body.as_bytes()).unwrap();
let body = encoder.finish().unwrap();
return Ok(resp.header("content-encoding", "gzip").body(full(body)).unwrap());
}
return Ok(resp.body(full(body)).unwrap());
}
let frames = futures_util::stream::iter(replies)
.then(|reply| async move {
tokio::time::sleep(Duration::from_millis(20)).await;
Ok(Frame::data(Bytes::from(grpc_frame(0, &proto_string(&reply)))))
})
.chain(futures_util::stream::once(async {
let mut trailers = hyper::HeaderMap::new();
trailers.insert("grpc-status", HeaderValue::from_static("0"));
Ok(Frame::trailers(trailers))
}));
Ok(Response::builder()
.header(CONTENT_TYPE, "application/grpc")
.body(BodyExt::boxed(StreamBody::new(frames)))
.unwrap())
}
async fn grpc_call(
sender: &mut http2::SendRequest<TestBody>,
port: u16,
method: &str,
name: &str,
) -> (Vec<u8>, Option<hyper::HeaderMap>, hyper::HeaderMap) {
let req = Request::builder()
.method("POST")
.uri(format!("https://localhost:{port}/pkg.Greeter/{method}"))
.header(CONTENT_TYPE, "application/grpc")
.header("te", "trailers")
.body(full(grpc_frame(0, &proto_string(name))))
.unwrap();
let resp = sender.send_request(req).await.unwrap();
let headers = resp.headers().clone();
let collected = resp.into_body().collect().await.unwrap();
let trailers = collected.trailers().cloned();
(collected.to_bytes().to_vec(), trailers, headers)
}
#[tokio::test(flavor = "multi_thread")]
async fn grpc_messages_and_status_are_recorded() {
within(async {
let h = Harness::start(1024);
let upstream =
spawn_upstream_with(Some(h.pki.server.clone()), Arc::new(AtomicUsize::new(0)), greeter)
.await;
let port = upstream.port();
let mut sender = h.h2_client(port).await;
let (body, trailers, _) = grpc_call(&mut sender, port, "SayHello", "Ada").await;
assert_eq!(body, grpc_frame(0, &proto_string("Hello Ada")));
assert_eq!(trailers.unwrap()["grpc-status"], "0", "trailers pass through");
let events = h.flow_events();
assert_eq!(find!(events, ProxyEvent::StreamStart { kind, .. } => *kind), StreamKind::Grpc);
let sent = messages(&events, MessageDirection::Client);
assert_eq!(sent.len(), 1);
assert_eq!((&sent[0].data[..], sent[0].size), (&proto_string("Ada")[..], 5));
let received = messages(&events, MessageDirection::Server);
assert_eq!(received.len(), 2);
assert_eq!(&received[0].data[..], &proto_string("Hello Ada")[..]);
assert_eq!(received[1].kind, MessageKind::GrpcTrailers { status: Some(0), message: None });
let (_, _, _) = grpc_call(&mut sender, port, "StreamHellos", "Bo").await;
let events = h.flow_events();
let received = messages(&events, MessageDirection::Server);
let replies: Vec<_> = received.iter().map(|m| m.data.clone()).collect();
assert_eq!(replies.len(), 4);
for (i, reply) in replies[..3].iter().enumerate() {
assert_eq!(&reply[..], &proto_string(&format!("Hello Bo #{}", i + 1))[..]);
}
let (body, _, headers) = grpc_call(&mut sender, port, "Fail", "Cy").await;
assert!(body.is_empty());
assert_eq!(headers["grpc-status"], "5");
let events = h.flow_events();
let received = messages(&events, MessageDirection::Server);
assert_eq!(
received[0].kind,
MessageKind::GrpcTrailers { status: Some(5), message: Some("no such greeter".into()) }
);
})
.await
}
#[tokio::test(flavor = "multi_thread")]
async fn grpc_web_text_is_recorded() {
within(async {
let h = Harness::start(1024);
let upstream =
spawn_upstream_with(Some(h.pki.server.clone()), Arc::new(AtomicUsize::new(0)), greeter)
.await;
let url = format!("https://localhost:{}/pkg.Greeter/StreamHellos", upstream.port());
// Plain, then with the whole response gzipped for HTTP.
for gzip in [false, true] {
let mut req = h
.client()
.post(&url)
.header(CONTENT_TYPE, "application/grpc-web-text")
.body(BASE64.encode(grpc_frame(0, &proto_string("Di"))));
if gzip {
req = req.header("x-test-gzip", "1");
}
let resp = req.send().await.unwrap();
assert_eq!(resp.status(), 200);
resp.bytes().await.unwrap();
let events = h.flow_events();
// As the server sent it: a client with gzip support may undo it before the test sees it.
let encoding =
headers(&events, false).into_iter().find(|(n, _)| n == "content-encoding");
assert_eq!(encoding.is_some(), gzip, "the response passed compressed");
let kind = find!(events, ProxyEvent::StreamStart { kind, .. } => *kind);
assert_eq!(kind, StreamKind::GrpcWeb);
let sent = messages(&events, MessageDirection::Client);
assert_eq!(&sent[0].data[..], &proto_string("Di")[..]);
let received = messages(&events, MessageDirection::Server);
assert_eq!(received.len(), 4, "gzip: {gzip}");
assert_eq!(&received[2].data[..], &proto_string("Hello Di #3")[..]);
assert_eq!(
received[3].kind,
MessageKind::GrpcTrailers { status: Some(0), message: Some(String::new()) }
);
}
})
.await
}
#[tokio::test(flavor = "multi_thread")]
+2 -1
View File
@@ -111,7 +111,8 @@
"client:tauri-before-build": "npm run bootstrap",
"client:tauri-before-dev": "node scripts/run-workspaces-dev.mjs apps/yaak-client",
"proxy:tauri-before-build": "npm run bootstrap",
"proxy:tauri-before-bundle": "run-s bootstrap proxy:vendor-helper",
"proxy:tauri-before-bundle": "run-s bootstrap proxy:vendor-helper proxy:sign-vendored",
"proxy:sign-vendored": "node scripts/sign-proxy-vendored.cjs",
"proxy:vendor-helper": "node scripts/vendor-proxy-helper.cjs",
"proxy:tauri-before-dev": "node scripts/run-workspaces-dev.mjs apps/yaak-proxy"
},
+36
View File
@@ -0,0 +1,36 @@
#!/usr/bin/env node
/**
* Signs the executables Yaak Proxy's bundle takes from Yaak's vendored tools with hardened
* runtime, as notarization requires of every executable in the app. Runs after vendoring, in
* the bundle's before-build step, with the identity `scripts/bundle-proxy.mjs` chose. Without
* one, or off macOS, the bundle is unsigned anyway and this does nothing.
*/
const { execFileSync } = require("child_process");
const { existsSync } = require("fs");
const path = require("path");
const rootDir = path.join(__dirname, "..");
const identity = process.env.APPLE_SIGNING_IDENTITY;
if (process.platform !== "darwin" || !identity) {
console.log("No signing identity, so the vendored protoc is left unsigned");
process.exit(0);
}
const clientDir = path.join(rootDir, "crates-tauri", "yaak-app-client");
const protoc = path.join(clientDir, "vendored", "protoc", "yaakprotoc");
const entitlements = path.join(clientDir, "macos", "entitlements.yaakprotoc.plist");
if (!existsSync(protoc)) {
console.error(`${protoc} is missing: run \`npm run vendor\` first`);
process.exit(1);
}
execFileSync(
"codesign",
["--force", "--options", "runtime", "--entitlements", entitlements, "--sign", identity, protoc],
{ stdio: "inherit" },
);
console.log(`Signed ${path.relative(rootDir, protoc)}`);