diff --git a/Cargo.lock b/Cargo.lock index 320c097..b4131eb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2679,6 +2679,9 @@ dependencies = [ "bytes", "fs4", "hex", + "http", + "http-body", + "http-body-util", "nuncio-cal", "nuncio-cli", "nuncio-contacts", @@ -2696,6 +2699,7 @@ dependencies = [ "tokio", "tokio-stream", "tonic", + "tower", "tracing", "tracing-appender", "tracing-subscriber", diff --git a/Cargo.toml b/Cargo.toml index b045994..b926cbd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -80,6 +80,10 @@ flate2 = "1.0" tar = "0.4" zip = { version = "8.6", default-features = false, features = ["deflate"] } tonic = "0.14" +tower = "0.5" +http = "1" +http-body = "1" +http-body-util = "0.1" tonic-prost = "0.14" prost = "0.14" prost-types = "0.14" diff --git a/crates/nunciod/Cargo.toml b/crates/nunciod/Cargo.toml index fd79f30..f47c835 100644 --- a/crates/nunciod/Cargo.toml +++ b/crates/nunciod/Cargo.toml @@ -30,6 +30,9 @@ tracing-subscriber.workspace = true tracing-appender.workspace = true tempfile.workspace = true tonic.workspace = true +tower.workspace = true +http.workspace = true +http-body.workspace = true async-trait.workspace = true subtle.workspace = true hex.workspace = true @@ -41,4 +44,5 @@ fs4.workspace = true nuncio-cli = { path = "../nuncio-cli" } sqlx.workspace = true wiremock.workspace = true +http-body-util.workspace = true tokio = { workspace = true, features = ["test-util"] } diff --git a/crates/nunciod/src/grpc.rs b/crates/nunciod/src/grpc.rs index 8fdf6cd..6562824 100644 --- a/crates/nunciod/src/grpc.rs +++ b/crates/nunciod/src/grpc.rs @@ -2458,25 +2458,41 @@ impl BearerAuthInterceptor { impl tonic::service::Interceptor for BearerAuthInterceptor { fn call(&mut self, request: Request<()>) -> Result, Status> { + // Only ever log the peer and a fixed reason string. The token value, + // the raw `authorization` header, and anything derived from either + // (length, prefix, substring) are NEVER logged -- doing so would leak + // the bearer credential into the logs. This runs inside the per-RPC + // span (see `rpc_trace`), so these events carry the request_id. + let peer = request + .remote_addr() + .map(|addr| addr.to_string()) + .unwrap_or_else(|| "unknown".to_string()); + + let reject = |reason: &'static str| { + tracing::warn!(peer = %peer, reason, "gRPC auth rejected"); + Status::unauthenticated(reason) + }; + let header = request .metadata() .get("authorization") - .ok_or_else(|| Status::unauthenticated("missing authorization metadata"))?; + .ok_or_else(|| reject("missing authorization metadata"))?; let header_str = header .to_str() - .map_err(|_| Status::unauthenticated("authorization metadata is not valid ASCII"))?; - let provided = header_str.strip_prefix("Bearer ").ok_or_else(|| { - Status::unauthenticated("expected 'Bearer ' authorization scheme") - })?; + .map_err(|_| reject("authorization metadata is not valid ASCII"))?; + let provided = header_str + .strip_prefix("Bearer ") + .ok_or_else(|| reject("expected 'Bearer ' authorization scheme"))?; let matches: bool = provided .as_bytes() .ct_eq(self.expected_token.as_bytes()) .into(); if matches { + tracing::debug!(peer = %peer, "gRPC auth accepted"); Ok(request) } else { - Err(Status::unauthenticated("invalid bearer token")) + Err(reject("invalid bearer token")) } } } @@ -2861,6 +2877,12 @@ pub async fn serve_on_listener_with_overrides_and_shutdown( ); Server::builder() + // Wrap every RPC in a correlation span (request_id + method + peer) + // before the resource-limit tuning below, so the span brackets the + // whole request -- routing, the per-service `BearerAuthInterceptor`, + // and the handler -- and its request_id is inherited by every event + // emitted while the request is handled. + .layer(crate::rpc_trace::RpcTraceLayer) .concurrency_limit_per_connection(GRPC_CONCURRENCY_LIMIT_PER_CONNECTION) .timeout(GRPC_REQUEST_TIMEOUT) .max_concurrent_streams(GRPC_MAX_CONCURRENT_STREAMS) @@ -2888,6 +2910,104 @@ mod tests { use nuncio_proto::v1::system_client::SystemClient; use tonic::Code; + use crate::test_tracing::with_recorder; + use tonic::service::Interceptor; + use tracing::Level; + + /// The bearer token used by the auth-logging tests. A test asserts this + /// exact value never appears in any captured log field. + const TEST_BEARER_TOKEN: &str = "s3cr3t-bearer-token-value"; + + fn bearer_request(header: Option<&str>) -> Request<()> { + let mut request = Request::new(()); + if let Some(value) = header { + request.metadata_mut().insert( + "authorization", + value.parse().expect("valid ASCII metadata value"), + ); + } + request + } + + #[test] + fn auth_accept_logs_debug_without_leaking_the_token() { + let (recorder, result) = with_recorder(|| { + let mut interceptor = BearerAuthInterceptor::new(TEST_BEARER_TOKEN); + interceptor.call(bearer_request(Some(&format!("Bearer {TEST_BEARER_TOKEN}")))) + }); + assert!(result.is_ok(), "a valid bearer token is accepted"); + + let events = recorder.events(); + let accept = events + .iter() + .find(|e| e.message() == "gRPC auth accepted") + .expect("an accept must be logged"); + assert_eq!(accept.level, Level::DEBUG, "accept is logged at DEBUG"); + + assert!( + recorder + .all_field_values() + .iter() + .all(|v| !v.contains(TEST_BEARER_TOKEN)), + "the bearer token must never appear in any log field" + ); + } + + #[test] + fn auth_reject_bad_token_logs_warn_without_leaking_the_token() { + let presented = "attacker-guess-token"; + let (recorder, result) = with_recorder(|| { + let mut interceptor = BearerAuthInterceptor::new(TEST_BEARER_TOKEN); + interceptor.call(bearer_request(Some(&format!("Bearer {presented}")))) + }); + let status = result.expect_err("a mismatched token is rejected"); + assert_eq!(status.code(), Code::Unauthenticated); + + let events = recorder.events(); + let reject = events + .iter() + .find(|e| e.message() == "gRPC auth rejected") + .expect("a reject must be logged"); + assert_eq!(reject.level, Level::WARN, "reject is logged at WARN"); + assert_eq!( + reject.fields.get("reason").map(String::as_str), + Some("invalid bearer token"), + "the reject must carry a fixed reason string" + ); + + // Neither the real token nor the presented guess may appear anywhere. + assert!( + recorder + .all_field_values() + .iter() + .all(|v| !v.contains(TEST_BEARER_TOKEN) && !v.contains(presented)), + "no token material may appear in any log field" + ); + } + + #[test] + fn auth_reject_missing_header_logs_warn_reason() { + let (recorder, result) = with_recorder(|| { + let mut interceptor = BearerAuthInterceptor::new(TEST_BEARER_TOKEN); + interceptor.call(bearer_request(None)) + }); + assert!( + result.is_err(), + "a missing authorization header is rejected" + ); + + let events = recorder.events(); + let reject = events + .iter() + .find(|e| e.message() == "gRPC auth rejected") + .expect("a reject must be logged"); + assert_eq!(reject.level, Level::WARN); + assert_eq!( + reject.fields.get("reason").map(String::as_str), + Some("missing authorization metadata") + ); + } + /// Proves a stored second-granularity instant survives the proto /// boundary exactly: `map_email_to_proto`'s `received_at` decodes back /// to the identical unix-seconds value via `timestamp_to_unix_secs`. diff --git a/crates/nunciod/src/lib.rs b/crates/nunciod/src/lib.rs index b37803f..a87fa7e 100644 --- a/crates/nunciod/src/lib.rs +++ b/crates/nunciod/src/lib.rs @@ -11,9 +11,12 @@ pub mod logging; pub mod orchestrator; pub mod outbox; pub mod pagination; +pub mod rpc_trace; pub mod send; pub mod sync; pub mod sync_dispatcher; +#[cfg(test)] +pub(crate) mod test_tracing; pub use orchestrator::SelfHealingSyncOrchestrator; /// Default PERSISTENT database path for `nunciod`: `~/.nuncio/nuncio.db` diff --git a/crates/nunciod/src/rpc_trace.rs b/crates/nunciod/src/rpc_trace.rs new file mode 100644 index 0000000..0931742 --- /dev/null +++ b/crates/nunciod/src/rpc_trace.rs @@ -0,0 +1,401 @@ +//! Per-RPC tracing span and request-id correlation for the gRPC transport. +//! +//! This is a [`tower::Layer`] mounted on the tonic `Server` so that every +//! inbound RPC is wrapped in a single root span carrying: +//! +//! * a generated `request_id` -- a lightweight per-process identifier +//! (process-start tag + monotonic counter) that lets an operator follow one +//! request across every log line it produced, without pulling in a heavier +//! UUID dependency; and +//! * the gRPC method path and the peer socket address. +//! +//! The inner service (routing + the per-service `BearerAuthInterceptor` + the +//! handler future) is entered inside that span, so *every* event emitted while +//! the request is handled -- including the auth accept/reject logs and any +//! event a handler emits -- inherits the `request_id` automatically. The span +//! opens with a DEBUG "received" event and closes, once the response body has +//! been fully drained, with an INFO event carrying the terminal gRPC status +//! code and the elapsed wall-clock time. +//! +//! ## Why the response body is wrapped +//! +//! A gRPC call's final status lives in the HTTP/2 *trailers* (`grpc-status`), +//! which are only produced once the response body finishes streaming (an +//! immediate error is instead a trailers-only response with `grpc-status` in +//! the leading headers). The layer therefore wraps the response body in +//! [`TracedBody`], which observes the terminal status as it flows past and +//! emits the completion event when the stream ends (or when the body is +//! dropped, e.g. on client disconnect), so a long-lived server-streaming RPC +//! is not logged as "complete" until it actually ends. + +use std::future::Future; +use std::pin::Pin; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::OnceLock; +use std::task::{Context, Poll}; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; + +use http::{HeaderMap, Request, Response}; +use http_body::{Body, Frame, SizeHint}; +use tonic::body::Body as TonicBody; +use tonic::transport::server::TcpConnectInfo; +use tonic::Code; +use tower::{Layer, Service}; +use tracing::{Instrument, Span}; + +/// HTTP/2 trailer (or leading header, for a trailers-only error response) +/// carrying the final gRPC status code. +const GRPC_STATUS_HEADER: &str = "grpc-status"; + +/// Generates a correlation id unique within this daemon process. +/// +/// The id is `-`: the process tag is the daemon's start +/// time in unix nanoseconds (captured once), and the counter increments per +/// call. This is deliberately not a UUID -- it needs only to be unique within +/// a single running daemon so its logs can be correlated, which a per-process +/// tag plus a monotonic counter satisfies without a new dependency. +fn next_request_id() -> String { + static COUNTER: AtomicU64 = AtomicU64::new(0); + static PROCESS_TAG: OnceLock = OnceLock::new(); + + let tag = *PROCESS_TAG.get_or_init(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos() as u64) + .unwrap_or(0) + }); + let n = COUNTER.fetch_add(1, Ordering::Relaxed); + format!("{tag:016x}-{n:08x}") +} + +/// Reads the `grpc-status` code from a header/trailer map, if present. +fn grpc_status_code(headers: &HeaderMap) -> Option { + headers + .get(GRPC_STATUS_HEADER) + .map(|value| Code::from_bytes(value.as_bytes())) +} + +/// [`tower::Layer`] that wraps each RPC in a correlation span. Mounted on the +/// tonic `Server` via `Server::builder().layer(RpcTraceLayer)`. +#[derive(Clone, Copy, Debug, Default)] +pub struct RpcTraceLayer; + +impl Layer for RpcTraceLayer { + type Service = RpcTrace; + + fn layer(&self, inner: S) -> Self::Service { + RpcTrace { inner } + } +} + +/// The [`RpcTraceLayer`]'s service: opens the span, runs the inner service +/// (routing, auth interceptor, handler) inside it, and wraps the response body +/// so the terminal status can be observed at stream end. +#[derive(Clone, Debug)] +pub struct RpcTrace { + inner: S, +} + +impl Service> for RpcTrace +where + S: Service, Response = Response> + Clone + Send + 'static, + S::Future: Send + 'static, + S::Error: Send + 'static, +{ + type Response = Response>; + type Error = S::Error; + type Future = + Pin> + Send + 'static>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: Request) -> Self::Future { + // `poll_ready` was called on `self.inner`, so it -- not the fresh + // clone -- is the readied service that must handle this request. Swap + // the readied service out and leave the clone behind for the next + // call (the standard tower clone-and-replace pattern that avoids + // calling an un-readied clone). + let clone = self.inner.clone(); + let mut inner = std::mem::replace(&mut self.inner, clone); + + let method = req.uri().path().to_string(); + let peer = req + .extensions() + .get::() + .and_then(TcpConnectInfo::remote_addr) + .map(|addr| addr.to_string()) + .unwrap_or_else(|| "unknown".to_string()); + let request_id = next_request_id(); + + let span = tracing::info_span!( + "grpc.request", + request_id = %request_id, + rpc.method = %method, + peer = %peer, + ); + let start = Instant::now(); + + Box::pin(async move { + // Enter the span for the synchronous routing + interceptor call so + // the auth accept/reject logs inherit the request_id, and emit the + // open event under it. + let call = span.in_scope(|| { + tracing::debug!("gRPC request received"); + inner.call(req) + }); + // Instrument the handler future so events emitted across its await + // points also inherit the span. + let response = call.instrument(span.clone()).await?; + + // A trailers-only error response carries `grpc-status` in the + // leading headers; a normal response carries it in the trailers, + // observed by `TracedBody` as the body drains. + let initial_status = grpc_status_code(response.headers()); + let (parts, body) = response.into_parts(); + let traced = TracedBody::new(body, span, start, initial_status); + Ok(Response::from_parts(parts, traced)) + }) + } +} + +/// Response body wrapper that observes the terminal `grpc-status` and logs the +/// request's completion (status + elapsed) exactly once, when the body stream +/// ends or is dropped. +/// +/// All fields are `Unpin`, so the pin projection in [`Body::poll_frame`] is a +/// safe `Pin::new` over the inner body rather than an `unsafe` projection. +#[derive(Debug)] +pub struct TracedBody { + inner: B, + span: Span, + start: Instant, + status: Option, + logged: bool, +} + +impl TracedBody { + fn new(inner: B, span: Span, start: Instant, initial_status: Option) -> Self { + Self { + inner, + span, + start, + status: initial_status, + logged: false, + } + } + + /// Emits the terminal completion event once. Idempotent: subsequent calls + /// (e.g. a `Drop` after the stream already ended) are no-ops. + fn log_completion(&mut self) { + if self.logged { + return; + } + self.logged = true; + let elapsed_ms = self.start.elapsed().as_millis() as u64; + let status = self.status.unwrap_or(Code::Unknown); + self.span.in_scope(|| { + tracing::info!( + grpc.status = ?status, + elapsed_ms, + "gRPC request completed" + ); + }); + } +} + +impl Body for TracedBody +where + B: Body + Unpin, +{ + type Data = B::Data; + type Error = B::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + // Safe because `Self: Unpin` (all fields are `Unpin`). + let this = self.get_mut(); + let polled = Pin::new(&mut this.inner).poll_frame(cx); + match &polled { + Poll::Ready(Some(Ok(frame))) => { + if let Some(trailers) = frame.trailers_ref() { + if let Some(code) = grpc_status_code(trailers) { + this.status = Some(code); + } + } + } + Poll::Ready(Some(Err(_))) | Poll::Ready(None) => this.log_completion(), + Poll::Pending => {} + } + polled + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + + fn size_hint(&self) -> SizeHint { + self.inner.size_hint() + } +} + +impl Drop for TracedBody { + fn drop(&mut self) { + // Covers stream ends that do not surface a final frame here (client + // disconnect, an outer layer dropping the body): the request is still + // logged as completed exactly once. + self.log_completion(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_tracing::with_recorder; + use http_body_util::BodyExt; + use std::collections::HashSet; + use std::convert::Infallible; + use tower::ServiceExt; + use tracing::Level; + + /// Minimal inner service that returns a response with `grpc-status` set in + /// the leading headers (a trailers-only response, tonic's shape for an + /// immediate status), so the layer's terminal-status extraction can be + /// exercised without a live server. + #[derive(Clone)] + struct StatusService { + code: i32, + } + + impl Service> for StatusService { + type Response = Response; + type Error = Infallible; + type Future = + Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _req: Request) -> Self::Future { + let code = self.code; + Box::pin(async move { + // A handler-emitted event: it must inherit the span so it + // carries the request_id set by the layer. + tracing::info!("handler ran"); + let mut response = Response::new(TonicBody::empty()); + response + .headers_mut() + .insert(GRPC_STATUS_HEADER, code.to_string().parse().unwrap()); + Ok(response) + }) + } + } + + fn request() -> Request { + Request::builder() + .uri("/nuncio.v1.Mail/Sync") + .body(TonicBody::empty()) + .unwrap() + } + + #[test] + fn rpc_produces_span_with_request_id_and_logs_terminal_status() { + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + + let (recorder, ()) = with_recorder(|| { + runtime.block_on(async { + let mut svc = RpcTraceLayer.layer(StatusService { code: 5 }); + let response = svc.ready().await.unwrap().call(request()).await.unwrap(); + // Drain the body so the completion event fires at stream end. + response.into_body().collect().await.unwrap(); + }); + }); + + let spans = recorder.spans(); + let rpc_span = spans + .iter() + .find(|s| s.name == "grpc.request") + .expect("a grpc.request span must be created"); + let request_id = rpc_span + .fields + .get("request_id") + .expect("the span must carry a request_id field"); + assert!(!request_id.is_empty(), "request_id must be populated"); + assert_eq!( + rpc_span.fields.get("rpc.method").map(String::as_str), + Some("/nuncio.v1.Mail/Sync"), + "the span must carry the gRPC method path" + ); + + let events = recorder.events(); + let completion = events + .iter() + .find(|e| e.message() == "gRPC request completed") + .expect("the terminal completion event must be logged"); + assert_eq!( + completion.level, + Level::INFO, + "the completion event is logged at INFO" + ); + assert_eq!( + completion.fields.get("grpc.status").map(String::as_str), + Some("NotFound"), + "the completion event must carry the terminal gRPC status code" + ); + assert!( + completion.fields.contains_key("elapsed_ms"), + "the completion event must carry the elapsed time" + ); + + // The handler ran inside the span, so its own event was captured too + // (proving handler events inherit the correlation context). + assert!( + events.iter().any(|e| e.message() == "handler ran"), + "a handler-emitted event must be observed under the request span" + ); + } + + #[test] + fn ok_status_is_recorded_at_completion() { + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let (recorder, ()) = with_recorder(|| { + runtime.block_on(async { + let mut svc = RpcTraceLayer.layer(StatusService { code: 0 }); + let response = svc.ready().await.unwrap().call(request()).await.unwrap(); + response.into_body().collect().await.unwrap(); + }); + }); + + let events = recorder.events(); + let completion = events + .iter() + .find(|e| e.message() == "gRPC request completed") + .expect("completion event"); + assert_eq!( + completion.fields.get("grpc.status").map(String::as_str), + Some("Ok") + ); + } + + #[test] + fn grpc_status_code_parses_known_codes_and_ignores_absent() { + let mut headers = HeaderMap::new(); + assert!(grpc_status_code(&headers).is_none()); + headers.insert(GRPC_STATUS_HEADER, "3".parse().unwrap()); + assert_eq!(grpc_status_code(&headers), Some(Code::InvalidArgument)); + } + + #[test] + fn request_ids_are_unique_within_the_process() { + let ids: HashSet = (0..1000).map(|_| next_request_id()).collect(); + assert_eq!(ids.len(), 1000, "request ids must not collide"); + } +} diff --git a/crates/nunciod/src/sync_dispatcher.rs b/crates/nunciod/src/sync_dispatcher.rs index 2278c64..8e46819 100644 --- a/crates/nunciod/src/sync_dispatcher.rs +++ b/crates/nunciod/src/sync_dispatcher.rs @@ -48,6 +48,7 @@ use std::sync::{Arc, Mutex, PoisonError}; use std::time::Duration; use thiserror::Error; use tokio::sync::{watch, Semaphore}; +use tracing::Instrument; /// Environment variable bounding how many account syncs may run at once. pub const CONCURRENCY_ENV_VAR: &str = "NUNCIO_SYNC_CONCURRENCY"; @@ -285,18 +286,30 @@ impl SyncDispatcher { let inner = self.inner.clone(); let id = account_id.to_string(); - tokio::spawn(async move { - let outcome = inner.run_guarded(&id).await; - // Remove the in-flight entry before publishing so a request - // arriving after completion starts a fresh sync rather than - // joining an already-finished one. - inner - .inflight - .lock() - .unwrap_or_else(PoisonError::into_inner) - .remove(&id); - let _ = tx.send(Some(outcome)); - }); + // The actual sync runs in a detached task, which does NOT + // inherit the caller's span automatically. Capture the current + // span (the RPC's correlation span when a `Mail/Sync` triggered + // this) and instrument the task with it, so the sync's logs are + // correlatable to the request that kicked it off. A request + // that COALESCES onto an already-running sync joins the owner's + // task and is therefore correlated to the owner's request_id, + // not its own. + let span = tracing::Span::current(); + tokio::spawn( + async move { + let outcome = inner.run_guarded(&id).await; + // Remove the in-flight entry before publishing so a + // request arriving after completion starts a fresh sync + // rather than joining an already-finished one. + inner + .inflight + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(&id); + let _ = tx.send(Some(outcome)); + } + .instrument(span), + ); rx } }; @@ -319,12 +332,19 @@ impl SyncDispatcher { /// production syncer additionally reports them on the event bus). pub async fn sync_all(&self, account_ids: Vec) -> usize { let mut set = tokio::task::JoinSet::new(); + // Propagate the caller's span (the RPC correlation span for an + // all-accounts `Mail/Sync`) into each per-account task so the fan-out + // stays correlated to the triggering request. + let span = tracing::Span::current(); for account_id in account_ids { let dispatcher = self.clone(); - set.spawn(async move { - let result = dispatcher.sync_account(&account_id).await; - (account_id, result) - }); + set.spawn( + async move { + let result = dispatcher.sync_account(&account_id).await; + (account_id, result) + } + .instrument(span.clone()), + ); } let mut total = 0usize; diff --git a/crates/nunciod/src/test_tracing.rs b/crates/nunciod/src/test_tracing.rs new file mode 100644 index 0000000..9c68c9b --- /dev/null +++ b/crates/nunciod/src/test_tracing.rs @@ -0,0 +1,133 @@ +//! Shared in-test `tracing` capture helper. +//! +//! Installs a thread-local subscriber that records every span (with its +//! fields) and every event (with its level, fields, and message) into an +//! in-memory buffer, so tests can assert on what the code logged -- including +//! the negative assertion that a secret never appears in any recorded field. +//! Deterministic and fully offline: no global subscriber, no I/O. + +#![cfg(test)] + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use tracing::field::{Field, Visit}; +use tracing::span::Attributes; +use tracing::{Event, Id, Level, Subscriber}; +use tracing_subscriber::layer::Context as LayerContext; +use tracing_subscriber::prelude::*; +use tracing_subscriber::registry::LookupSpan; + +/// A captured span: its name and the string-rendered fields it opened with. +#[derive(Debug, Clone)] +pub(crate) struct RecordedSpan { + pub name: String, + pub fields: HashMap, +} + +/// A captured event: its level and its string-rendered fields (including the +/// `message`). +#[derive(Debug, Clone)] +pub(crate) struct RecordedEvent { + pub level: Level, + pub fields: HashMap, +} + +impl RecordedEvent { + /// The event's `message` field, or the empty string if it had none. + pub fn message(&self) -> &str { + self.fields.get("message").map(String::as_str).unwrap_or("") + } +} + +/// In-memory sink of everything the subscriber captured during a test. +#[derive(Debug, Default)] +pub(crate) struct Recorder { + spans: Mutex>, + events: Mutex>, +} + +impl Recorder { + /// A snapshot of every captured span. + pub fn spans(&self) -> Vec { + self.spans.lock().unwrap_or_else(|e| e.into_inner()).clone() + } + + /// A snapshot of every captured event. + pub fn events(&self) -> Vec { + self.events + .lock() + .unwrap_or_else(|e| e.into_inner()) + .clone() + } + + /// Every field value recorded across all spans and events, so a test can + /// assert that a secret never appears anywhere in the telemetry. + pub fn all_field_values(&self) -> Vec { + let mut values = Vec::new(); + for span in self.spans().iter() { + values.extend(span.fields.values().cloned()); + } + for event in self.events().iter() { + values.extend(event.fields.values().cloned()); + } + values + } +} + +struct FieldVisitor(HashMap); + +impl Visit for FieldVisitor { + fn record_str(&mut self, field: &Field, value: &str) { + self.0.insert(field.name().to_string(), value.to_string()); + } + fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { + self.0 + .entry(field.name().to_string()) + .or_insert_with(|| format!("{value:?}")); + } +} + +struct CaptureLayer(Arc); + +impl tracing_subscriber::Layer for CaptureLayer +where + S: Subscriber + for<'a> LookupSpan<'a>, +{ + fn on_new_span(&self, attrs: &Attributes<'_>, _id: &Id, _ctx: LayerContext<'_, S>) { + let mut visitor = FieldVisitor(HashMap::new()); + attrs.record(&mut visitor); + self.0 + .spans + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(RecordedSpan { + name: attrs.metadata().name().to_string(), + fields: visitor.0, + }); + } + + fn on_event(&self, event: &Event<'_>, _ctx: LayerContext<'_, S>) { + let mut visitor = FieldVisitor(HashMap::new()); + event.record(&mut visitor); + self.0 + .events + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push(RecordedEvent { + level: *event.metadata().level(), + fields: visitor.0, + }); + } +} + +/// Runs `f` with a thread-local capturing subscriber installed and returns the +/// recorder alongside `f`'s result. Because the subscriber is thread-local, +/// any async work driven inside `f` must run on the calling thread (e.g. a +/// current-thread runtime's `block_on`) for its telemetry to be captured. +pub(crate) fn with_recorder(f: impl FnOnce() -> T) -> (Arc, T) { + let recorder = Arc::new(Recorder::default()); + let subscriber = tracing_subscriber::registry().with(CaptureLayer(recorder.clone())); + let value = tracing::subscriber::with_default(subscriber, f); + (recorder, value) +}