From 1129f7aa939458a31f6f3eedd8746ee89e9559d7 Mon Sep 17 00:00:00 2001 From: Gregory Schier Date: Fri, 6 Nov 2026 09:53:56 -0700 Subject: [PATCH] Capture WebSocket, SSE and gRPC (#814) Co-authored-by: Claude Opus 5.5 --- Cargo.lock | 71 ++ apps/yaak-proxy/harness/fixtures.ts | 6 + apps/yaak-proxy/lib/flowViews.test.ts | 6 + crates-proxy/yaak-proxy-lib/Cargo.toml | 11 +- .../yaak-proxy-lib/bindings/gen_models.ts | 137 ++- .../yaak-proxy-lib/bindings/gen_rpc.ts | 92 +- .../migrations/00000006_flow_messages.sql | 44 + crates-proxy/yaak-proxy-lib/src/client.rs | 2 +- crates-proxy/yaak-proxy-lib/src/control.rs | 92 +- crates-proxy/yaak-proxy-lib/src/db.rs | 19 +- crates-proxy/yaak-proxy-lib/src/flows.rs | 24 +- crates-proxy/yaak-proxy-lib/src/grpc.rs | 595 +++++++++++++ crates-proxy/yaak-proxy-lib/src/har.rs | 78 +- crates-proxy/yaak-proxy-lib/src/lib.rs | 826 +++++++++++++++++- crates-proxy/yaak-proxy-lib/src/messages.rs | 368 ++++++++ crates-proxy/yaak-proxy-lib/src/models.rs | 77 ++ crates-proxy/yaak-proxy-lib/tests/streams.rs | 437 +++++++++ crates-tauri/yaak-app-proxy/src/lib.rs | 46 +- .../yaak-app-proxy/tauri.bundle.conf.json | 4 + crates/yaak-grpc/src/lib.rs | 8 + crates/yaak-proxy/Cargo.toml | 4 +- crates/yaak-proxy/src/body.rs | 66 +- crates/yaak-proxy/src/decode.rs | 378 ++++++++ crates/yaak-proxy/src/grpc.rs | 346 ++++++++ crates/yaak-proxy/src/lib.rs | 34 +- crates/yaak-proxy/src/messages.rs | 113 +++ crates/yaak-proxy/src/request.rs | 110 ++- crates/yaak-proxy/src/sse.rs | 306 +++++++ crates/yaak-proxy/src/ws.rs | 419 +++++++++ crates/yaak-proxy/tests/proxy.rs | 570 +++++++++++- package.json | 3 +- scripts/sign-proxy-vendored.cjs | 36 + 32 files changed, 5228 insertions(+), 100 deletions(-) create mode 100644 crates-proxy/yaak-proxy-lib/migrations/00000006_flow_messages.sql create mode 100644 crates-proxy/yaak-proxy-lib/src/grpc.rs create mode 100644 crates-proxy/yaak-proxy-lib/src/messages.rs create mode 100644 crates-proxy/yaak-proxy-lib/tests/streams.rs create mode 100644 crates/yaak-proxy/src/decode.rs create mode 100644 crates/yaak-proxy/src/grpc.rs create mode 100644 crates/yaak-proxy/src/messages.rs create mode 100644 crates/yaak-proxy/src/sse.rs create mode 100644 crates/yaak-proxy/src/ws.rs create mode 100644 scripts/sign-proxy-vendored.cjs diff --git a/Cargo.lock b/Cargo.lock index 9ff5f085a..005fe1520 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", diff --git a/apps/yaak-proxy/harness/fixtures.ts b/apps/yaak-proxy/harness/fixtures.ts index bc6966b45..57086aefc 100644 --- a/apps/yaak-proxy/harness/fixtures.ts +++ b/apps/yaak-proxy/harness/fixtures.ts @@ -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 }); diff --git a/apps/yaak-proxy/lib/flowViews.test.ts b/apps/yaak-proxy/lib/flowViews.test.ts index 0902ca2bf..23c6bc288 100644 --- a/apps/yaak-proxy/lib/flowViews.test.ts +++ b/apps/yaak-proxy/lib/flowViews.test.ts @@ -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 { diff --git a/crates-proxy/yaak-proxy-lib/Cargo.toml b/crates-proxy/yaak-proxy-lib/Cargo.toml index 05845b2c7..e4a62cf6e 100644 --- a/crates-proxy/yaak-proxy-lib/Cargo.toml +++ b/crates-proxy/yaak-proxy-lib/Cargo.toml @@ -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"] } diff --git a/crates-proxy/yaak-proxy-lib/bindings/gen_models.ts b/crates-proxy/yaak-proxy-lib/bindings/gen_models.ts index 153dd80f9..8acd262a0 100644 --- a/crates-proxy/yaak-proxy-lib/bindings/gen_models.ts +++ b/crates-proxy/yaak-proxy-lib/bindings/gen_models.ts @@ -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, +/** + * The services the descriptors define, by full name. + */ +services: Array, }; + +/** + * 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"; diff --git a/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts b/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts index 629597ae7..d80119161 100644 --- a/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts +++ b/crates-proxy/yaak-proxy-lib/bindings/gen_rpc.ts @@ -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, }; + 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, */ 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; export type GetProxyStateResponse = { state: ProxyState, @@ -214,6 +262,28 @@ total: number, */ matching: number | null, }; +export type ListGrpcSchemasRequest = Record; + +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, +/** + * Messages the flow has in all. + */ +total: number, }; + export type ListWorkspacesRequest = Record; export type ListWorkspacesResponse = { @@ -222,6 +292,11 @@ export type ListWorkspacesResponse = { */ yaakFound: boolean, workspaces: Array, }; +/** + * 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], 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, /** diff --git a/crates-proxy/yaak-proxy-lib/migrations/00000006_flow_messages.sql b/crates-proxy/yaak-proxy-lib/migrations/00000006_flow_messages.sql new file mode 100644 index 000000000..7852d5f90 --- /dev/null +++ b/crates-proxy/yaak-proxy-lib/migrations/00000006_flow_messages.sql @@ -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 `.msg.`, 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 +); diff --git a/crates-proxy/yaak-proxy-lib/src/client.rs b/crates-proxy/yaak-proxy-lib/src/client.rs index 914fa4c3a..099c94a82 100644 --- a/crates-proxy/yaak-proxy-lib/src/client.rs +++ b/crates-proxy/yaak-proxy-lib/src/client.rs @@ -230,6 +230,6 @@ impl OfflineStore { pub async fn har(&self, filter: &FlowFilter) -> Result { 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 } } diff --git a/crates-proxy/yaak-proxy-lib/src/control.rs b/crates-proxy/yaak-proxy-lib/src/control.rs index fca098329..9355a43b2 100644 --- a/crates-proxy/yaak-proxy-lib/src/control.rs +++ b/crates-proxy/yaak-proxy-lib/src/control.rs @@ -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, + /// Why a gRPC message could not be decoded. + #[serde(skip_serializing_if = "Option::is_none")] + pub decode_error: Option, +} + +#[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) -> Result { 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) -> Result { 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 Result<(Option, Option), 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 { for (key, value) in url::form_urlencoded::parse(query.as_bytes()) { if key == "bodyLimit" { diff --git a/crates-proxy/yaak-proxy-lib/src/db.rs b/crates-proxy/yaak-proxy-lib/src/db.rs index dcbb88a59..012aba937 100644 --- a/crates-proxy/yaak-proxy-lib/src/db.rs +++ b/crates-proxy/yaak-proxy-lib/src/db.rs @@ -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:?}"); } diff --git a/crates-proxy/yaak-proxy-lib/src/flows.rs b/crates-proxy/yaak-proxy-lib/src/flows.rs index 0ae72f4c5..46ea4db86 100644 --- a/crates-proxy/yaak-proxy-lib/src/flows.rs +++ b/crates-proxy/yaak-proxy-lib/src/flows.rs @@ -32,7 +32,8 @@ pub struct FlowFilter { /// The method, any case. #[ts(optional)] pub method: Option, - /// 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, /// 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) { + /// `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) { let mut sql = String::from("1"); let mut params: Vec = 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, 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 } fn count_matching(conn: &rusqlite::Connection, filter: &FlowFilter) -> Result { - 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, _>>()?; 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 = 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" diff --git a/crates-proxy/yaak-proxy-lib/src/grpc.rs b/crates-proxy/yaak-proxy-lib/src/grpc.rs new file mode 100644 index 000000000..1d0087cb0 --- /dev/null +++ b/crates-proxy/yaak-proxy-lib/src/grpc.rs @@ -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, + /// The services the descriptors define, by full name. + pub services: Vec, +} + +/// 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, + /// The schema decoded with. None when schemaless. + pub schema: Option, + /// Why it decoded schemaless, or what to know about the decode. + pub note: Option, +} + +/// 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>>; + +/// The decoded descriptor pools, by schema id, and the hosts reflection was asked of. +#[derive(Clone, Default)] +pub struct GrpcState { + protoc: Arc>>, + pools: Arc>>, + /// By `host:port`. + lookups: Arc>>, +} + +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 { + 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) -> Result { + 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 { + 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, 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::, _>>().map_err(|e| e.to_string()) + }) +} + +fn get(db: &ProxyQueryManager, id: &str) -> Result { + 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 { + let id = generate_prefixed_id("gs"); + let services: Vec = 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 { + 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 = 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 { + 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::().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 { + 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 { + 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 { + if let Some(pool) = self.pools.lock().ok().and_then(|p| p.get(id).cloned()) { + return Ok(pool); + } + let bytes: Vec = 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> { + let host = authority.rsplit_once(':').map_or(authority, |(h, _)| h); + let mut schemas: Vec = + 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 { + 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 { + 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) -> DecodedGrpcMessage { + let mut notes: Vec = 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> { + 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 { + 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"); + } +} diff --git a/crates-proxy/yaak-proxy-lib/src/har.rs b/crates-proxy/yaak-proxy-lib/src/har.rs index fc08df008..49409a2aa 100644 --- a/crates-proxy/yaak-proxy-lib/src/har.rs +++ b/crates-proxy/yaak-proxy-lib/src/har.rs @@ -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, #[serde(rename = "_error", skip_serializing_if = "Option::is_none")] pub error: Option, + /// 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>, +} + +#[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 { +/// 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 { 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, 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) -> 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())); diff --git a/crates-proxy/yaak-proxy-lib/src/lib.rs b/crates-proxy/yaak-proxy-lib/src/lib.rs index f27f7277c..ff169f997 100644 --- a/crates-proxy/yaak-proxy-lib/src/lib.rs +++ b/crates-proxy/yaak-proxy-lib/src/lib.rs @@ -7,9 +7,11 @@ pub mod control; pub mod convert; pub mod db; pub mod flows; +pub mod grpc; pub mod har; pub mod helper; pub mod macos; +pub mod messages; pub mod models; pub mod save; pub mod settings; @@ -21,9 +23,11 @@ use crate::client::ControlClient; use crate::control::{ControlFile, ControlInfo, ControlServer, Owner}; use crate::convert::DroppedHeader; use crate::db::ProxyQueryManager; +use crate::grpc::{DecodedGrpcMessage, GrpcSchema, GrpcState}; use crate::helper::HelperState; use crate::macos::{SystemProxyBackend, SystemProxyStatus}; -use crate::models::{Flow, FlowBody, FlowIden, ModelPayload, ProxyHeader, find_header}; +use crate::messages::FlowMessage; +use crate::models::{Flow, FlowBody, FlowIden, ModelPayload, ProxyHeader, StreamKind, find_header}; use crate::save::{ClientDb, SaveFlowRequest, SaveFlowResponse, WorkspaceSummary}; use crate::settings::{HttpVersion, ProxySettings}; use base64::Engine; @@ -40,6 +44,7 @@ use std::sync::{Arc, Mutex, OnceLock}; use std::time::{Duration, SystemTime}; use ts_rs::TS; use yaak_database::{ModelChangeEvent, UpdateSource, UpsertModelInfo}; +pub use yaak_grpc::GrpcConfig; use yaak_models::blob_manager::{BlobManager, BodyChunkWriter}; use yaak_proxy::{CaStore, ProxyConfig, ProxyEvent, ProxyHandle}; use yaak_rpc::{RpcError, RpcEventEmitter, define_rpc}; @@ -70,6 +75,8 @@ pub struct ProxyCtx { pub db: ProxyQueryManager, pub blobs: BlobManager, pub events: RpcEventEmitter, + /// gRPC schemas as decoded, and the protoc that compiles attached protos. + pub grpc: GrpcState, } impl ProxyCtx { @@ -94,9 +101,17 @@ impl ProxyCtx { db, blobs, events, + grpc: GrpcState::default(), }) } + /// The protoc that compiles `.proto` files attached for gRPC decoding, as the app bundles + /// it. Without one, a `protoc` on `PATH` is used. + pub fn with_protoc(self, config: yaak_grpc::GrpcConfig) -> Self { + self.grpc.set_protoc(config); + self + } + /// The process this runs in, which decides how the system proxy gets admin rights. pub fn with_owner(mut self, owner: Owner) -> Self { self.owner = owner; @@ -167,6 +182,7 @@ impl ProxyCtx { } let settings = self.settings()?; + self.end_stale_streams(); let ca = CaStore::load_or_create(&self.ca_dir()) .map_err(|e| format!("Failed to load proxy CA: {e}"))?; let mut config = ProxyConfig::new(port.unwrap_or(settings.port), Arc::new(ca)); @@ -223,6 +239,22 @@ impl ProxyCtx { Ok(()) } + /// Marks streams an engine left open, as one that crashed does, as ended. Only the engine's + /// owner, about to start it, may: until then another process's streams may be open. + fn end_stale_streams(&self) { + let ended = self.db.with_tx(|c| { + c.conn().resolve().execute( + "UPDATE flows SET error = ?1 WHERE stream_kind IS NOT NULL \ + AND elapsed_ms IS NULL AND error IS NULL AND NOT cancelled", + [STOPPED_MID_STREAM], + )?; + Ok::<_, yaak_database::Error>(()) + }); + if let Err(e) = ended { + warn!("Failed to end streams left open: {e}"); + } + } + /// Stops the engine and turns off a system proxy that points at it. pub fn stop(&self) -> Result<(), String> { self.system_proxy_on.store(false, Ordering::SeqCst); @@ -947,8 +979,146 @@ pub struct HelperRegisterRequest {} #[ts(export, export_to = "gen_rpc.ts")] pub struct HelperUnregisterRequest {} +#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +#[serde(rename_all = "camelCase")] +pub struct ListMessagesRequest { + pub flow_id: String, + /// Only messages after this `seq`, for following a stream as it grows. + #[ts(optional, type = "number")] + pub after_seq: Option, + /// At most this many, the first after `afterSeq`. + #[ts(optional, type = "number")] + pub limit: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +#[serde(rename_all = "camelCase")] +pub struct ListMessagesResponse { + /// In order, by `seq`. + pub messages: Vec, + /// Messages the flow has in all. + #[ts(type = "number")] + pub total: i64, +} + +#[derive(Debug, Clone, Deserialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +#[serde(rename_all = "camelCase")] +pub struct GetMessagePayloadRequest { + pub flow_id: String, + #[ts(type = "number")] + pub seq: i64, +} + +#[derive(Debug, Clone, Serialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +#[serde(rename_all = "camelCase")] +pub struct GetMessagePayloadResponse { + /// The payload as stored (unmasked, inflated, decompressed), in base64. + pub data: String, + /// It is cut at the capture limit. + pub truncated: bool, +} + +#[derive(Debug, Clone, Deserialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +#[serde(rename_all = "camelCase")] +pub struct DecodeGrpcMessageRequest { + pub flow_id: String, + #[ts(type = "number")] + pub seq: i64, +} + +#[derive(Deserialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +pub struct ListGrpcSchemasRequest {} + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +#[serde(rename_all = "camelCase")] +pub struct AddGrpcSchemaRequest { + /// `api.example.com`, `localhost:50051`, `*.example.com` or `*`. + pub host_pattern: String, + /// `.proto` files, and directories they import from. + pub proto_paths: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +pub struct RemoveGrpcSchemaRequest { + pub id: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +pub struct ReflectGrpcHostRequest { + /// `host` or `host:port`, over TLS on 443 unless a port is given, or either after + /// `http://` to ask without TLS. + pub host: String, +} + +/// A message of a streaming flow was recorded. +#[derive(Debug, Clone, Serialize, TS)] +#[ts(export, export_to = "gen_rpc.ts")] +pub struct MessagePayload { + pub message: FlowMessage, +} + // -- Handlers -- +fn list_messages( + ctx: &ProxyCtx, + req: ListMessagesRequest, +) -> Result { + find_flow(ctx, &req.flow_id)?; + let messages = + messages::list(&ctx.db, &req.flow_id, req.after_seq, req.limit).map_err(rpc_error)?; + let total = messages::count(&ctx.db, &req.flow_id).map_err(rpc_error)?; + Ok(ListMessagesResponse { messages, total }) +} + +fn get_message_payload( + ctx: &ProxyCtx, + req: GetMessagePayloadRequest, +) -> Result { + let message = messages::get(&ctx.db, &req.flow_id, req.seq).map_err(rpc_error)?; + let data = messages::payload(&ctx.db, &ctx.blobs, &message).map_err(rpc_error)?; + Ok(GetMessagePayloadResponse { data: BASE64.encode(data), truncated: message.truncated }) +} + +async fn decode_grpc_message( + ctx: ProxyCtx, + req: DecodeGrpcMessageRequest, +) -> Result { + let flow = find_flow(&ctx, &req.flow_id)?; + grpc::decode(&ctx.db, &ctx.blobs, &ctx.grpc, &flow, req.seq).await.map_err(rpc_error) +} + +fn list_grpc_schemas( + ctx: &ProxyCtx, + _req: ListGrpcSchemasRequest, +) -> Result, RpcError> { + grpc::list(&ctx.db).map_err(rpc_error) +} + +async fn add_grpc_schema(ctx: ProxyCtx, req: AddGrpcSchemaRequest) -> Result { + grpc::add(&ctx.db, &ctx.grpc, &req.host_pattern, &req.proto_paths).await.map_err(rpc_error) +} + +fn remove_grpc_schema(ctx: &ProxyCtx, req: RemoveGrpcSchemaRequest) -> Result { + grpc::remove(&ctx.db, &ctx.grpc, &req.id).map_err(rpc_error) +} + +async fn reflect_grpc_host( + ctx: ProxyCtx, + req: ReflectGrpcHostRequest, +) -> Result { + let (host, plaintext) = grpc::reflection_host(&req.host); + grpc::reflect(&ctx.db, &ctx.grpc, &host, plaintext).await.map_err(rpc_error) +} + fn execute_action(ctx: &ProxyCtx, invocation: ActionInvocation) -> Result { match invocation { ActionInvocation::Global { action } => match action { @@ -1086,7 +1256,7 @@ async fn export_har(ctx: ProxyCtx, req: ExportHarRequest) -> Result, response_body: Option, + /// Written to the database already, as a streaming flow is when its stream opens. + persisted: bool, + /// Cleared while it streamed: nothing more of it is kept. + cleared: bool, + next_seq: i64, + /// Changed since it was last written, by a message or a status. + dirty: bool, + last_write: Option, +} + +impl InFlight { + /// Something shown of the flow changed: one written while it streams is written again on + /// the stream flow interval, so it does not wait for a message to show it. + fn changed(&mut self) { + if self.persisted { + self.dirty = true; + } + } +} + +/// A message waiting to be written, with its payload. +struct PendingMessage { + id: u64, + message: messages::FlowMessage, + data: bytes::Bytes, } /// Where the event loop stores flows, and whether it does. @@ -1289,10 +1499,108 @@ struct FlowSink { } fn run_event_loop(rx: std::sync::mpsc::Receiver, sink: FlowSink) { - let FlowSink { db, blobs, events, attribution, recording } = sink; + use std::sync::mpsc::{RecvTimeoutError, TryRecvError}; let mut in_flight: HashMap = HashMap::new(); + let mut pending: Vec = Vec::new(); + let mut last_due_check = std::time::Instant::now(); - while let Ok(event) = rx.recv() { + loop { + // Messages are written once no more are waiting, or a batch is full; flows changed by + // them once their interval is up. + let event = if !pending.is_empty() { + match rx.try_recv() { + Ok(event) => event, + Err(TryRecvError::Empty) => { + sink.write_messages(&mut in_flight, &mut pending); + sink.write_due_flows(&mut in_flight); + continue; + } + Err(TryRecvError::Disconnected) => break, + } + } else if let Some(due) = next_flow_write(&in_flight) { + let wait = due.saturating_duration_since(std::time::Instant::now()); + match rx.recv_timeout(wait) { + Ok(event) => event, + Err(RecvTimeoutError::Timeout) => { + sink.write_due_flows(&mut in_flight); + continue; + } + Err(RecvTimeoutError::Disconnected) => break, + } + } else { + match rx.recv() { + Ok(event) => event, + Err(_) => break, + } + }; + if let ProxyEvent::Message { id, message } = event { + record_message(&mut in_flight, &mut pending, id, &message); + if pending.len() >= MESSAGE_BATCH { + sink.write_messages(&mut in_flight, &mut pending); + } + } else { + // Before anything else of the flows, so a flow's messages are in before it ends. + sink.write_messages(&mut in_flight, &mut pending); + sink.handle(&mut in_flight, event); + } + // Traffic that never lets up never empties the queue, so flows due to be written are + // looked for here too, every so often, their messages written first. + if last_due_check.elapsed() >= DUE_CHECK_INTERVAL { + last_due_check = std::time::Instant::now(); + sink.write_messages(&mut in_flight, &mut pending); + sink.write_due_flows(&mut in_flight); + } + } + + sink.write_messages(&mut in_flight, &mut pending); + // The engine stopped: a stream still open ends with it. Other flows in flight are dropped, + // as they always have been. + for (_, mut f) in in_flight.drain() { + if f.persisted { + f.flow.error = Some(STOPPED_MID_STREAM.to_string()); + sink.finish_flow(f); + } + } +} + +/// When the next changed streaming flow is due to be written. +fn next_flow_write(in_flight: &HashMap) -> Option { + in_flight + .values() + .filter(|f| f.dirty) + .map(|f| { + f.last_write.map_or_else(std::time::Instant::now, |t| t + STREAM_FLOW_WRITE_INTERVAL) + }) + .min() +} + +fn record_message( + in_flight: &mut HashMap, + pending: &mut Vec, + id: u64, + message: &yaak_proxy::StreamMessage, +) { + let Some(f) = in_flight.get_mut(&id).filter(|f| !f.cleared) else { + return; + }; + if let yaak_proxy::MessageKind::GrpcTrailers { status, message } = &message.kind { + f.flow.grpc_status = *status; + f.flow.grpc_message = message.clone(); + } + let seq = f.next_seq; + f.next_seq += 1; + f.flow.message_count += 1; + f.dirty = true; + pending.push(PendingMessage { + id, + message: messages::FlowMessage::from_engine(&f.flow.id, seq, message), + data: message.data.clone(), + }); +} + +impl FlowSink { + fn handle(&self, in_flight: &mut HashMap, event: ProxyEvent) { + let FlowSink { blobs, attribution, recording, .. } = self; match event { ProxyEvent::RequestStart { id, @@ -1306,7 +1614,7 @@ fn run_event_loop(rx: std::sync::mpsc::Receiver, sink: FlowSink) { // Paused when it started: it passes through, and none of its events are kept. // Judged by its start rather than now, as the loop can run behind. if !recording.lock().map(|r| r.was_on_at(started_at)).unwrap_or(true) { - continue; + return; } // Begun as the connection was accepted; this covers one accepted before the // engine's port was known. @@ -1329,9 +1637,38 @@ fn run_event_loop(rx: std::sync::mpsc::Receiver, sink: FlowSink) { connection, request_body: Some(BodyChunkWriter::new(blobs.clone(), request_body)), response_body: Some(BodyChunkWriter::new(blobs.clone(), response_body)), + persisted: false, + cleared: false, + next_seq: 0, + dirty: false, + last_write: None, }, ); } + ProxyEvent::StreamStart { id, kind } => { + if let Some(f) = in_flight.get_mut(&id) { + let kind = StreamKind::from(kind); + f.flow.stream_kind = Some(kind); + if kind.is_grpc() + && let Ok(url) = url::Url::parse(&f.flow.url) + && let Some((service, method)) = yaak_proxy::service_and_method(url.path()) + { + f.flow.grpc_service = Some(service.to_string()); + f.flow.grpc_method = Some(method.to_string()); + } + // Shown while it streams, rather than only once it ends. + self.write_stream_flow(f); + } + } + // Recorded by the loop, which batches them. + ProxyEvent::Message { .. } => {} + ProxyEvent::StreamEnd { id, elapsed_ms, error } => { + if let Some(mut f) = in_flight.remove(&id) { + f.flow.elapsed_ms = Some(clamp_ms(elapsed_ms)); + f.flow.error = error; + self.finish_flow(f); + } + } ProxyEvent::RequestHeader { id, name, value } => { if let Some(f) = in_flight.get_mut(&id) { f.flow.request_headers.push(ProxyHeader { name, value }); @@ -1344,6 +1681,7 @@ fn run_event_loop(rx: std::sync::mpsc::Receiver, sink: FlowSink) { f.flow.request_size += len as i64; f.flow.request_truncated |= data.len() < len; write_body_chunk(&mut f.request_body, &data); + f.changed(); } } ProxyEvent::RequestBodyComplete { id, size, truncated } => { @@ -1351,6 +1689,7 @@ fn run_event_loop(rx: std::sync::mpsc::Receiver, sink: FlowSink) { f.flow.request_size = size as i64; f.flow.request_truncated = truncated; finish_body(f.request_body.take()); + f.changed(); } } ProxyEvent::ResponseStart { id, status, http_version, remote_addr, elapsed_ms } => { @@ -1359,16 +1698,22 @@ fn run_event_loop(rx: std::sync::mpsc::Receiver, sink: FlowSink) { f.flow.upstream_http_version = Some(http_version); f.flow.remote_addr = remote_addr.map(|a| a.to_string()); f.flow.elapsed_headers_ms = Some(clamp_ms(elapsed_ms)); + f.changed(); } } ProxyEvent::ResponseHeader { id, name, value } => { if let Some(f) = in_flight.get_mut(&id) { f.flow.response_headers.push(ProxyHeader { name, value }); + f.changed(); } } - ProxyEvent::ResponseBodyChunk { id, data, .. } => { + // Counted per chunk too, for a stream written while it is open. + ProxyEvent::ResponseBodyChunk { id, len, data } => { if let Some(f) = in_flight.get_mut(&id) { + f.flow.response_size += len as i64; + f.flow.response_truncated |= data.len() < len; write_body_chunk(&mut f.response_body, &data); + f.changed(); } } ProxyEvent::ResponseBodyComplete { @@ -1385,25 +1730,204 @@ fn run_event_loop(rx: std::sync::mpsc::Receiver, sink: FlowSink) { f.flow.elapsed_ms = Some(clamp_ms(elapsed_ms)); f.flow.error = error; f.flow.cancelled = cancelled; - finish_flow(&db, &events, &attribution, f); + self.finish_flow(f); } } ProxyEvent::Cancelled { id, elapsed_ms } => { if let Some(mut f) = in_flight.remove(&id) { f.flow.elapsed_ms = Some(clamp_ms(elapsed_ms)); f.flow.cancelled = true; - finish_flow(&db, &events, &attribution, f); + self.finish_flow(f); } } ProxyEvent::Error { id, error, elapsed_ms } => { if let Some(mut f) = in_flight.remove(&id) { f.flow.elapsed_ms = Some(clamp_ms(elapsed_ms)); f.flow.error = Some(error); - finish_flow(&db, &events, &attribution, f); + self.finish_flow(f); } } } } + + /// Writes `pending` in one transaction, each payload over [`messages::INLINE_LIMIT`] first + /// to `blobs.sqlite`, and tells the app of each. A flow cleared meanwhile keeps nothing more. + fn write_messages( + &self, + in_flight: &mut HashMap, + pending: &mut Vec, + ) { + if pending.is_empty() { + return; + } + let batch = std::mem::take(pending); + let mut rows = Vec::with_capacity(batch.len()); + for PendingMessage { id, mut message, data } in batch { + let inline = if data.len() > messages::INLINE_LIMIT { + let mut writer = BodyChunkWriter::new(self.blobs.clone(), message.blob_id()); + if let Err(e) = writer.write(&data).and_then(|_| writer.finish()) { + warn!("Failed to store a proxy message payload: {e}"); + message.stored_size = 0; + message + .error + .get_or_insert_with(|| format!("Failed to store the payload: {e}")); + } + None + } else { + Some(data) + }; + rows.push((id, message, inline)); + } + + let mut orphaned = Vec::new(); + let mut newly_cleared = Vec::new(); + let written = self.db.with_tx(|ctx| { + let conn = ctx.conn().resolve(); + let mut checked = std::collections::HashSet::new(); + for (id, message, inline) in &rows { + let Some(f) = in_flight.get_mut(id) else { + continue; + }; + // Cleared since it was written: what comes after is not kept either. + if f.persisted && checked.insert(*id) && !flow_exists(conn, &f.flow.id) { + f.cleared = true; + newly_cleared.push(*id); + } + if f.cleared { + orphaned.push(message.blob_id()); + continue; + } + messages::insert(conn, message, inline.as_deref())?; + // Before the commit, as flows are, so a clear's event comes after it. + self.events.emit("message_write", &MessagePayload { message: message.clone() }); + } + Ok::<_, yaak_database::Error>(()) + }); + if let Err(e) = written { + warn!("Failed to write proxy messages: {e}"); + } + if !orphaned.is_empty() { + let _ = self.blobs.with_tx(|b| orphaned.iter().try_for_each(|id| b.delete_chunks(id))); + } + for id in newly_cleared { + if let Some(f) = in_flight.get_mut(&id) { + self.stop_cleared(f); + } + } + } + + /// A stream cleared while open keeps nothing more: its bodies stop being written, and what + /// they wrote since the clear is deleted, as the stream may run long before it ends. + fn stop_cleared(&self, f: &mut InFlight) { + f.cleared = true; + f.request_body = None; + f.response_body = None; + self.forget_cleared(&f.flow.id); + } + + /// Writes the streaming flows whose messages or status changed, at most once per + /// [`STREAM_FLOW_WRITE_INTERVAL`] each. + fn write_due_flows(&self, in_flight: &mut HashMap) { + let now = std::time::Instant::now(); + for f in in_flight.values_mut() { + let due = f.last_write.is_none_or(|t| now >= t + STREAM_FLOW_WRITE_INTERVAL); + if f.dirty && due { + self.write_stream_flow(f); + } + } + } + + /// Writes a streaming flow as it is so far, for the app to show while it streams. Unless it + /// was cleared since it was last written: then nothing more of it is kept. + fn write_stream_flow(&self, f: &mut InFlight) { + f.dirty = false; + f.last_write = Some(std::time::Instant::now()); + if f.cleared { + return; + } + let mut flow = f.flow.clone(); + flow.content_type = find_header(&flow.response_headers, "content-type").map(String::from); + // Not waited for here; the flow's end waits for it. + if let Some(app) = self.attribution.get(f.connection, Duration::ZERO) { + flow.client_pid = Some(app.pid); + flow.client_name = Some(app.name); + flow.client_bundle_id = app.bundle_id; + flow.client_path = Some(app.path); + } + let persisted = f.persisted; + let written = self.db.with_tx(|ctx| self.write_flow(ctx, flow, persisted)); + match written { + Ok(true) => f.persisted = true, + Ok(false) => self.stop_cleared(f), + Err(e) => warn!("Failed to write proxy flow: {e}"), + } + } + + fn finish_flow(&self, f: InFlight) { + finish_body(f.request_body); + finish_body(f.response_body); + let id = f.flow.id.clone(); + if f.cleared { + self.forget_cleared(&id); + return; + } + + let mut flow = f.flow; + flow.content_type = find_header(&flow.response_headers, "content-type").map(String::from); + if let Some(app) = self.attribution.get(f.connection, ATTRIBUTION_WAIT) { + flow.client_pid = Some(app.pid); + flow.client_name = Some(app.name); + flow.client_bundle_id = app.bundle_id; + flow.client_path = Some(app.path); + } + let written = self.db.with_tx(|ctx| self.write_flow(ctx, flow, f.persisted)); + match written { + Ok(true) => {} + Ok(false) => self.forget_cleared(&id), + Err(e) => warn!("Failed to write proxy flow: {e}"), + } + } + + /// Writes `flow` and tells the app. One written before, as a stream is while open, keeps + /// what was recorded on it since, the request it was saved as; one cleared since stays + /// cleared, and this returns false. + fn write_flow( + &self, + ctx: &yaak_database::DbContext, + mut flow: Flow, + persisted: bool, + ) -> Result { + if persisted { + let Some(stored) = ctx.find_optional::(FlowIden::Id, &flow.id) else { + return Ok(false); + }; + flow.saved_request_id = stored.saved_request_id; + } + let (saved, created) = ctx.upsert(&flow, &UpdateSource::Background)?; + // Sent while this holds the writer, before the commit, so a clear's event, sent after + // its own commit, always comes after it. + self.events.emit( + "model_write", + &ModelPayload { model: saved, change: ModelChangeEvent::Upsert { created } }, + ); + Ok(true) + } + + /// Deletes what a flow cleared while it streamed wrote after the clear. + fn forget_cleared(&self, flow_id: &str) { + let deleted = self.blobs.with_tx(|b| { + b.delete_chunks(&Flow::body_id(flow_id, FlowBody::Request))?; + b.delete_chunks(&Flow::body_id(flow_id, FlowBody::Response))?; + b.delete_chunks_like(&format!("{flow_id}.msg.%")) + }); + if let Err(e) = deleted { + warn!("Failed to delete a cleared flow's bodies: {e}"); + } + } +} + +fn flow_exists(conn: &rusqlite::Connection, flow_id: &str) -> bool { + conn.query_row("SELECT 1 FROM flows WHERE id = ?1", [flow_id], |_| Ok(())).is_ok() } fn clamp_ms(ms: u64) -> i32 { @@ -1431,38 +1955,6 @@ fn finish_body(writer: Option) { } } -fn finish_flow( - db: &ProxyQueryManager, - events: &RpcEventEmitter, - attribution: &Attribution, - f: InFlight, -) { - finish_body(f.request_body); - finish_body(f.response_body); - - let mut flow = f.flow; - flow.content_type = find_header(&flow.response_headers, "content-type").map(String::from); - if let Some(app) = attribution.get(f.connection, ATTRIBUTION_WAIT) { - flow.client_pid = Some(app.pid); - flow.client_name = Some(app.name); - flow.client_bundle_id = app.bundle_id; - flow.client_path = Some(app.path); - } - let written = db.with_tx(|ctx| { - let (saved, created) = ctx.upsert(&flow, &UpdateSource::Background)?; - // Sent while this holds the writer, before the commit, so a clear's event, sent after - // its own commit, always comes after it. - events.emit( - "model_write", - &ModelPayload { model: saved, change: ModelChangeEvent::Upsert { created } }, - ); - Ok::<_, yaak_database::Error>(()) - }); - if let Err(e) = written { - warn!("Failed to write proxy flow: {e}"); - } -} - // -- Router + Schema -- define_rpc! { @@ -1470,17 +1962,23 @@ define_rpc! { commands { execute_action(ActionInvocation) -> bool, format_body(FormatBodyRequest) -> String, + get_message_payload(GetMessagePayloadRequest) -> GetMessagePayloadResponse, list_actions(ListActionsRequest) -> ListActionsResponse, list_flows(flows::FlowFilter) -> ListFlowsResponse, + list_grpc_schemas(ListGrpcSchemasRequest) -> Vec, + list_messages(ListMessagesRequest) -> ListMessagesResponse, + remove_grpc_schema(RemoveGrpcSchemaRequest) -> bool, set_recording(SetRecordingRequest) -> bool, set_system_proxy_at_launch(SetSystemProxyAtLaunchRequest) -> bool, set_http_version(SetHttpVersionRequest) -> HttpVersion, } async_commands { + add_grpc_schema(AddGrpcSchemaRequest) -> GrpcSchema, app_icon(AppIconRequest) -> Option, ca_status(CaStatusRequest) -> CaStatus, ca_trust(CaTrustRequest) -> CaStatus, clear_flows(ClearFlowsRequest) -> usize, + decode_grpc_message(DecodeGrpcMessageRequest) -> DecodedGrpcMessage, export_har(ExportHarRequest) -> usize, filter_body(FilterBodyRequest) -> FilterBodyResponse, get_flow_body(GetFlowBodyRequest) -> GetFlowBodyResponse, @@ -1489,6 +1987,7 @@ define_rpc! { 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, @@ -1497,6 +1996,7 @@ define_rpc! { } events { flows_cleared(FlowsClearedPayload), + message_write(MessagePayload), model_write(ModelPayload), proxy_state_changed(ProxyStatePayload), } @@ -1545,6 +2045,250 @@ mod tests { assert!(log.was_on_at(SystemTime::now())); } + /// An event loop writing to `ctx`'s databases, fed by the returned sender. + fn event_loop( + ctx: &ProxyCtx, + ) -> (std::sync::mpsc::Sender, std::thread::JoinHandle<()>) { + let (tx, rx) = std::sync::mpsc::channel(); + let sink = FlowSink { + db: ctx.db.clone(), + blobs: ctx.blobs.clone(), + events: ctx.events.clone(), + attribution: Attribution::new(0), + recording: ctx.recording.clone(), + }; + (tx, std::thread::spawn(move || run_event_loop(rx, sink))) + } + + fn open_websocket(tx: &std::sync::mpsc::Sender, id: u64, url: &str) { + tx.send(ProxyEvent::RequestStart { + id, + started_at: SystemTime::now(), + method: "GET".into(), + url: url.into(), + http_version: "HTTP/1.1".into(), + connection: id, + client_addr: "127.0.0.1:1".parse().unwrap(), + }) + .unwrap(); + tx.send(ProxyEvent::ResponseStart { + id, + status: 101, + http_version: "HTTP/1.1".into(), + remote_addr: None, + elapsed_ms: 1, + }) + .unwrap(); + tx.send(ProxyEvent::StreamStart { id, kind: yaak_proxy::StreamKind::WebSocket }).unwrap(); + } + + fn ws_message( + id: u64, + direction: yaak_proxy::MessageDirection, + kind: yaak_proxy::MessageKind, + data: &[u8], + ) -> ProxyEvent { + ProxyEvent::Message { + id, + message: yaak_proxy::StreamMessage { + direction, + kind, + at: SystemTime::now(), + size: data.len() as u64, + data: bytes::Bytes::copy_from_slice(data), + truncated: false, + compressed: false, + error: None, + }, + } + } + + fn wait_for_flow(ctx: &ProxyCtx, check: impl Fn(&Flow) -> bool) -> Flow { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + loop { + let flows = ctx.db.with_conn(|c| c.find_all::().unwrap()); + if let Some(flow) = flows.into_iter().find(&check) { + return flow; + } + assert!(std::time::Instant::now() < deadline, "no flow as expected"); + std::thread::sleep(Duration::from_millis(10)); + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn websocket_messages_are_stored_and_exported() { + use yaak_proxy::{MessageDirection as D, MessageKind as K}; + let dir = tempfile::tempdir().unwrap(); + let (events, event_rx) = RpcEventEmitter::new(); + let ctx = ProxyCtx::new(dir.path(), events).unwrap(); + let (tx, event_loop) = event_loop(&ctx); + open_websocket(&tx, 1, "https://example.com/socket"); + let big: Vec = (0..10_000u32).map(|i| i as u8).collect(); + tx.send(ws_message(1, D::Client, K::WebSocketText, b"hi")).unwrap(); + tx.send(ws_message(1, D::Server, K::WebSocketBinary, &big)).unwrap(); + let close = K::WebSocketClose { code: Some(1000), reason: "bye".into() }; + tx.send(ws_message(1, D::Client, close, b"\x03\xe8bye")).unwrap(); + + // Written while open, its count following its messages. + let open = wait_for_flow(&ctx, |f| f.message_count == 3); + assert_eq!((open.stream_kind, open.elapsed_ms), (Some(StreamKind::Websocket), None)); + tx.send(ProxyEvent::StreamEnd { id: 1, elapsed_ms: 50, error: None }).unwrap(); + drop(tx); + event_loop.join().unwrap(); + + let flow = flows::list(&ctx.db, &Default::default()).unwrap().remove(0); + assert_eq!((flow.elapsed_ms, flow.error.as_deref()), (Some(50), None)); + let stored = messages::list(&ctx.db, &flow.id, None, None).unwrap(); + let kinds: Vec<_> = stored.iter().map(|m| (m.seq, m.kind, m.direction)).collect(); + assert_eq!( + kinds, + vec![ + (0, messages::MessageKind::WebsocketText, messages::MessageDirection::Client), + (1, messages::MessageKind::WebsocketBinary, messages::MessageDirection::Server), + (2, messages::MessageKind::WebsocketClose, messages::MessageDirection::Client), + ] + ); + assert_eq!(stored[0].preview.as_deref(), Some("hi")); + assert_eq!(stored[1].preview, None); + assert_eq!( + (stored[2].meta.close_code, stored[2].meta.close_reason.as_deref()), + (Some(1000), Some("bye")) + ); + // The large payload is in the blob store, the small ones in their rows. + assert_eq!(messages::payload(&ctx.db, &ctx.blobs, &stored[1]).unwrap(), big); + assert!(ctx.blobs.connect().body_exists(&stored[1].blob_id()).unwrap()); + assert!(!ctx.blobs.connect().body_exists(&stored[0].blob_id()).unwrap()); + assert_eq!(messages::payload(&ctx.db, &ctx.blobs, &stored[0]).unwrap(), b"hi"); + + let har = + har::export(&ctx.db, &ctx.blobs, std::slice::from_ref(&flow), 1 << 20).await.unwrap(); + let har = serde_json::to_value(&har).unwrap(); + let sent = &har["log"]["entries"][0]["_webSocketMessages"]; + assert_eq!(sent.as_array().unwrap().len(), 2, "data messages only, as devtools exports"); + assert_eq!((sent[0]["type"].as_str(), sent[0]["opcode"].as_u64()), (Some("send"), Some(1))); + assert_eq!(sent[0]["data"], "hi"); + assert_eq!( + (sent[1]["type"].as_str(), sent[1]["opcode"].as_u64()), + (Some("receive"), Some(2)) + ); + assert_eq!(sent[1]["data"], BASE64.encode(&big)); + + let names: Vec<_> = event_rx.try_iter().map(|e| e.event).collect(); + assert_eq!(names.iter().filter(|n| **n == "message_write").count(), 3, "{names:?}"); + + ctx.clear_flows().unwrap(); + assert!(!ctx.blobs.connect().body_exists(&stored[1].blob_id()).unwrap()); + } + + #[test] + fn a_stream_cleared_while_open_stays_cleared() { + use yaak_proxy::{MessageDirection as D, MessageKind as K}; + let dir = tempfile::tempdir().unwrap(); + let (events, _event_rx) = RpcEventEmitter::new(); + let ctx = ProxyCtx::new(dir.path(), events).unwrap(); + let (tx, event_loop) = event_loop(&ctx); + open_websocket(&tx, 1, "https://example.com/socket"); + tx.send(ws_message(1, D::Client, K::WebSocketText, b"before")).unwrap(); + let flow = wait_for_flow(&ctx, |f| f.message_count == 1); + ctx.clear_flows().unwrap(); + + tx.send(ws_message(1, D::Server, K::WebSocketBinary, &[7; 8192])).unwrap(); + // A body chunk past a blob chunk's size would be written at once, were it kept. + let data = bytes::Bytes::from(vec![b'x'; 2 << 20]); + tx.send(ProxyEvent::ResponseBodyChunk { id: 1, len: data.len(), data }).unwrap(); + std::thread::sleep(Duration::from_millis(500)); + let blobs = ctx.blobs.connect().list_body_ids().unwrap(); + assert!(blobs.iter().all(|b| !b.starts_with(&flow.id)), "while open: {blobs:?}"); + + tx.send(ProxyEvent::StreamEnd { id: 1, elapsed_ms: 9, error: None }).unwrap(); + drop(tx); + event_loop.join().unwrap(); + assert_eq!(flows::count(&ctx.db).unwrap(), 0, "not brought back by its end"); + assert_eq!(messages::count(&ctx.db, &flow.id).unwrap(), 0); + let blobs = ctx.blobs.connect().list_body_ids().unwrap(); + assert!(blobs.iter().all(|b| !b.starts_with(&flow.id)), "{blobs:?}"); + } + + #[test] + fn an_open_stream_is_written_while_events_never_let_up() { + use yaak_proxy::{MessageDirection as D, MessageKind as K}; + let dir = tempfile::tempdir().unwrap(); + let (events, _event_rx) = RpcEventEmitter::new(); + let ctx = ProxyCtx::new(dir.path(), events).unwrap(); + let (tx, event_loop) = event_loop(&ctx); + open_websocket(&tx, 1, "https://example.com/socket"); + wait_for_flow(&ctx, |f| f.stream_kind.is_some()); + + // A backlog of messages, each followed by another event, which costs the loop a + // transaction for the message before it, so its queue does not empty for a while. + const SENT: i64 = 10_000; + for _ in 0..SENT { + tx.send(ws_message(1, D::Client, K::WebSocketText, b"tick")).unwrap(); + let data = bytes::Bytes::from_static(b"x"); + tx.send(ProxyEvent::ResponseBodyChunk { id: 1, len: 1, data }).unwrap(); + } + let counted = wait_for_flow(&ctx, |f| f.message_count > 0); + assert!(counted.message_count < SENT, "written before the backlog was through"); + drop(tx); + event_loop.join().unwrap(); + } + + #[test] + fn an_open_call_shows_its_response_before_any_message() { + let dir = tempfile::tempdir().unwrap(); + let (events, _event_rx) = RpcEventEmitter::new(); + let ctx = ProxyCtx::new(dir.path(), events).unwrap(); + let (tx, event_loop) = event_loop(&ctx); + tx.send(ProxyEvent::RequestStart { + id: 1, + started_at: SystemTime::now(), + method: "POST".into(), + url: "https://example.com/pkg.Greeter/Watch".into(), + http_version: "HTTP/2".into(), + connection: 1, + client_addr: "127.0.0.1:1".parse().unwrap(), + }) + .unwrap(); + tx.send(ProxyEvent::StreamStart { id: 1, kind: yaak_proxy::StreamKind::Grpc }).unwrap(); + wait_for_flow(&ctx, |f| f.stream_kind.is_some()); + tx.send(ProxyEvent::ResponseStart { + id: 1, + status: 200, + http_version: "HTTP/2".into(), + remote_addr: None, + elapsed_ms: 3, + }) + .unwrap(); + tx.send(ProxyEvent::ResponseHeader { + id: 1, + name: "content-type".into(), + value: "application/grpc".into(), + }) + .unwrap(); + // No message yet, and the server may send none for a long while. + let open = wait_for_flow(&ctx, |f| f.status == Some(200)); + assert_eq!( + (open.elapsed_ms, open.content_type.as_deref()), + (None, Some("application/grpc")) + ); + drop(tx); + event_loop.join().unwrap(); + } + + #[test] + fn streams_still_open_end_when_the_engine_stops() { + let dir = tempfile::tempdir().unwrap(); + let (events, _event_rx) = RpcEventEmitter::new(); + let ctx = ProxyCtx::new(dir.path(), events).unwrap(); + let (tx, event_loop) = event_loop(&ctx); + open_websocket(&tx, 1, "https://example.com/socket"); + wait_for_flow(&ctx, |f| f.stream_kind.is_some()); + drop(tx); + event_loop.join().unwrap(); + let flow = flows::list(&ctx.db, &Default::default()).unwrap().remove(0); + assert_eq!(flow.error.as_deref(), Some(STOPPED_MID_STREAM)); + } + /// A flow whose response is `body` under `content_type`, stored in `ctx` fn flow_with_response(ctx: &ProxyCtx, id: &str, content_type: &str, body: &[u8]) { let flow = Flow { diff --git a/crates-proxy/yaak-proxy-lib/src/messages.rs b/crates-proxy/yaak-proxy-lib/src/messages.rs new file mode 100644 index 000000000..190f7715d --- /dev/null +++ b/crates-proxy/yaak-proxy-lib/src/messages.rs @@ -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 { + KINDS.into_iter().find(|k| k.as_str() == s) + } + + /// The WebSocket opcode, for HAR's `_webSocketMessages`. + pub fn websocket_opcode(self) -> Option { + 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, + /// A WebSocket close frame's reason. + #[serde(skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub close_reason: Option, + /// An SSE event's `event` field. + #[serde(skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub event: Option, + /// An SSE event's `id` field. + #[serde(skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub event_id: Option, + /// An SSE event's `retry` field, in milliseconds. + #[serde(skip_serializing_if = "Option::is_none")] + #[ts(optional, type = "number")] + pub retry: Option, + /// 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, + /// gRPC trailers' `grpc-status`. + #[serde(skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub grpc_status: Option, + /// gRPC trailers' `grpc-message`, decoded. + #[serde(skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub grpc_message: Option, +} + +/// 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, + pub meta: MessageMeta, + /// The payload's start as text, when it is text. + pub preview: Option, +} + +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::::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 { + 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 { + 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, + limit: Option, +) -> Result, 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::, _>>().map_err(|e| e.to_string()) + }) +} + +/// Messages recorded for a flow. +pub fn count(db: &ProxyQueryManager, flow_id: &str) -> Result { + 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, 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 { + 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, String> { + let inline: Option> = 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()) +} diff --git a/crates-proxy/yaak-proxy-lib/src/models.rs b/crates-proxy/yaak-proxy-lib/src/models.rs index 15d35747e..606323b38 100644 --- a/crates-proxy/yaak-proxy-lib/src/models.rs +++ b/crates-proxy/yaak-proxy-lib/src/models.rs @@ -66,6 +66,64 @@ pub struct Flow { pub client_bundle_id: Option, /// Its app bundle's path, or its executable's outside an app. pub client_path: Option, + /// 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, + /// 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, + /// A gRPC call's method, from its path. + pub grpc_method: Option, + /// A gRPC call's status, from its trailers. Not 0 is a failed call. + pub grpc_status: Option, + /// A gRPC call's `grpc-message`, decoded. + pub grpc_message: Option, +} + +/// 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::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 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::(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")?, }) } } diff --git a/crates-proxy/yaak-proxy-lib/tests/streams.rs b/crates-proxy/yaak-proxy-lib/tests/streams.rs new file mode 100644 index 000000000..42a13462c --- /dev/null +++ b/crates-proxy/yaak-proxy-lib/tests/streams.rs @@ -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; +type TestBody = UnsyncBoxBody; + +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(test: impl Future) -> 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) -> TestBody { + Full::new(data.into()).map_err(|never| match never {}).boxed_unsync() +} + +fn grpc_frame(flag: u8, payload: &[u8]) -> Vec { + 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) -> Vec { + 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> + Send>>; +/// A gRPC reflection service, its type erased. +type Reflection = + std::sync::Arc) -> ReflectionFuture + Send + Sync>; + +fn reflection_service(service: S) -> Reflection +where + S: tonic::codegen::Service< + Request, + Response = Response, + 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, + reflection: Option, +) -> Result, 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) -> 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) { + 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 { + 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::().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 +} diff --git a/crates-tauri/yaak-app-proxy/src/lib.rs b/crates-tauri/yaak-app-proxy/src/lib.rs index ac694c4b2..5856a3c8e 100644 --- a/crates-tauri/yaak-app-proxy/src/lib.rs +++ b/crates-tauri/yaak-app-proxy/src/lib.rs @@ -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(win: &WebviewWindow) { }); } +/// 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(app: &tauri::AppHandle) -> Option { + 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>, @@ -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)); + } +} diff --git a/crates-tauri/yaak-app-proxy/tauri.bundle.conf.json b/crates-tauri/yaak-app-proxy/tauri.bundle.conf.json index b5feca0f4..bc2dfbde3 100644 --- a/crates-tauri/yaak-app-proxy/tauri.bundle.conf.json +++ b/crates-tauri/yaak-app-proxy/tauri.bundle.conf.json @@ -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, diff --git a/crates/yaak-grpc/src/lib.rs b/crates/yaak-grpc/src/lib.rs index 8911108a4..7624ae7bc 100644 --- a/crates/yaak-grpc/src/lib.rs +++ b/crates/yaak-grpc/src/lib.rs @@ -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 { + serialize_dynamic_message_json(msg) +} + pub fn serialize_options() -> SerializeOptions { SerializeOptions::new().skip_default_fields(false) } diff --git a/crates/yaak-proxy/Cargo.toml b/crates/yaak-proxy/Cargo.toml index aa055f251..736231482 100644 --- a/crates/yaak-proxy/Cargo.toml +++ b/crates/yaak-proxy/Cargo.toml @@ -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" diff --git a/crates/yaak-proxy/src/body.rs b/crates/yaak-proxy/src/body.rs index cc814525f..c13e54a54 100644 --- a/crates/yaak-proxy/src/body.rs +++ b/crates/yaak-proxy/src/body.rs @@ -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), + Grpc(GrpcParser), +} + +impl BodyParser { + fn feed(&mut self, data: &[u8]) -> Vec { + match self { + Self::Sse(p) => p.feed(data), + Self::Grpc(p) => p.feed(data), + } + } + + fn trailers(&mut self, trailers: &hyper::HeaderMap) -> Option { + 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 { + 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>, + parser: Option, } 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) -> Self { + self.parser = parser; + self + } + + fn emit_messages(&self, messages: impl IntoIterator) { + 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); } } diff --git a/crates/yaak-proxy/src/decode.rs b/crates/yaak-proxy/src/decode.rs new file mode 100644 index 000000000..3488f857a --- /dev/null +++ b/crates/yaak-proxy/src/decode.rs @@ -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), + /// 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>, + /// 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, + /// 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, 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, 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 = + 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}"); + } +} diff --git a/crates/yaak-proxy/src/grpc.rs b/crates/yaak-proxy/src/grpc.rs new file mode 100644 index 000000000..d7b053ddb --- /dev/null +++ b/crates/yaak-proxy/src/grpc.rs @@ -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, + /// gRPC-Web's base64 flavor, with base64 not yet a whole quantum. + base64: Option>, + prefix: Vec, + /// 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 { + 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 { + 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 { + 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) { + 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 { + 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 { + 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()); + } +} diff --git a/crates/yaak-proxy/src/lib.rs b/crates/yaak-proxy/src/lib.rs index f7401932e..81f70685d 100644 --- a/crates/yaak-proxy/src/lib.rs +++ b/crates/yaak-proxy/src/lib.rs @@ -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, + }, } /// A running engine. Dropping it stops the listener and joins its thread. @@ -225,6 +249,7 @@ pub(crate) struct Engine { handshake_failures: Mutex>, /// 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 { passthrough: Mutex::new(HashSet::new()), handshake_failures: Mutex::new(HashMap::new()), http2: AtomicBool::new(config.http2), + websocket_capture: config.websocket_capture, via: via_marker(), }); diff --git a/crates/yaak-proxy/src/messages.rs b/crates/yaak-proxy/src/messages.rs new file mode 100644 index 000000000..bd21f0471 --- /dev/null +++ b/crates/yaak-proxy/src/messages.rs @@ -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, + reason: String, + }, + SseEvent { + event: Option, + id: Option, + retry: Option, + /// 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, + message: Option, + }, +} + +/// 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, +} + +/// 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) + } +} diff --git a/crates/yaak-proxy/src/request.rs b/crates/yaak-proxy/src/request.rs index 6b76ab3ff..bdc7d5ed0 100644 --- a/crates/yaak-proxy/src/request.rs +++ b/crates/yaak-proxy/src/request.rs @@ -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; @@ -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, + /// The request body's `Content-Encoding`, which gRPC-Web may use. + content_encoding: Option, +} + /// 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, + grpc: Option, } impl Flow { @@ -83,12 +102,13 @@ impl Flow { req: &Request, 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, - remote_addr: Option, - ) -> Response { + /// How the response's body is read for messages, if it streams them, after recording that + /// it does. + fn response_parser(&self, resp: &Response) -> Option { + 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, remote_addr: Option) { 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, + remote_addr: Option, + ) -> Response { + 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, + remote_addr: Option, + ) -> (Response, Arc) { + 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 { 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}"); diff --git a/crates/yaak-proxy/src/sse.rs b/crates/yaak-proxy/src/sse.rs new file mode 100644 index 000000000..e94fdaeef --- /dev/null +++ b/crates/yaak-proxy/src/sse.rs @@ -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, + id: Option, + retry: Option, + 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>, + 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 { + // 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 { + 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 { + // 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 { + 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 { + 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 { + 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"); + } +} diff --git a/crates/yaak-proxy/src/ws.rs b/crates/yaak-proxy/src/ws.rs new file mode 100644 index 000000000..dc5753908 --- /dev/null +++ b/crates/yaak-proxy/src/ws.rs @@ -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 { + 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, + /// 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) -> 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) { + 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, + deflate: Option, +) -> 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(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, +} + +/// 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, + limit: usize, + header: Vec, + frame: Option, + message: Option, + control: Vec, + inflater: Option, + 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 { + 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, + limit: usize, + no_context_takeover: Option, + ) -> 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); + } +} diff --git a/crates/yaak-proxy/tests/proxy.rs b/crates/yaak-proxy/tests/proxy.rs index 9e911d16f..322bf35dc 100644 --- a/crates/yaak-proxy/tests/proxy.rs +++ b/crates/yaak-proxy/tests/proxy.rs @@ -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(test: impl Future) -> 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) -> Response { 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 { + 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 { + 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 { + 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, + received: Arc>>, +) -> Response { + 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 = (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 { + 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 { + 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 { + 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 { + 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) -> Result, 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 = 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, + port: u16, + method: &str, + name: &str, +) -> (Vec, Option, 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")] diff --git a/package.json b/package.json index 3e2022546..9a83a9673 100644 --- a/package.json +++ b/package.json @@ -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" }, diff --git a/scripts/sign-proxy-vendored.cjs b/scripts/sign-proxy-vendored.cjs new file mode 100644 index 000000000..1735a1941 --- /dev/null +++ b/scripts/sign-proxy-vendored.cjs @@ -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)}`);