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