From 1ad8e9d2706fb8e2651a3ad58aa8aa714d569c8e Mon Sep 17 00:00:00 2001 From: Plucky Date: Fri, 31 Oct 2025 10:41:49 +0800 Subject: [PATCH 1/2] feat: add http_client request logging --- rig-core/src/http_client/mod.rs | 109 +++++++++++++++++++++++++++++-- rig-core/src/http_client/util.rs | 41 ++++++++++++ rig-core/src/prelude.rs | 2 + 3 files changed, 148 insertions(+), 4 deletions(-) create mode 100644 rig-core/src/http_client/util.rs diff --git a/rig-core/src/http_client/mod.rs b/rig-core/src/http_client/mod.rs index 4d8a04ea0..a9aa66f09 100644 --- a/rig-core/src/http_client/mod.rs +++ b/rig-core/src/http_client/mod.rs @@ -1,11 +1,14 @@ -use crate::http_client::sse::BoxedStream; +use crate::http_client::{sse::BoxedStream, util::Escape}; use bytes::Bytes; -use http::StatusCode; pub use http::{HeaderMap, HeaderValue, Method, Request, Response, Uri, request::Builder}; +use http::{StatusCode, request::Parts}; use reqwest::{Body, multipart::Form}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use tracing::Level; pub mod retry; pub mod sse; +pub mod util; use std::pin::Pin; @@ -75,6 +78,46 @@ pub fn with_bearer_auth(req: Builder, auth: &str) -> Result { Ok(req.header("Authorization", auth_header)) } +static LOG_HTTP_BODY_MAX: AtomicUsize = AtomicUsize::new(8 * 1024); + +/// Extension trait for client builders to configure HTTP logging. +/// +/// # Example +/// ``` +/// use rig::prelude::*; +/// use rig::providers::openai; +/// +/// let client = openai::Client::builder("api-key") +/// .max_log_body_preview(8192) +/// .build(); +/// ``` +pub trait HttpLogConfigExt { + /// Set the maximum number of bytes to preview from the body when logging in the `TRACE` level. + /// Defaults to 8192 bytes. Set to 0 to disable body preview logging. + /// + /// This method can be called on any client builder to configure HTTP logging before building the client. + fn max_log_body_preview(self, max_preview_bytes: usize) -> Self; +} + +impl HttpLogConfigExt for T +where + T: Sized, +{ + fn max_log_body_preview(self, max_preview_bytes: usize) -> Self { + LOG_HTTP_BODY_MAX.store(max_preview_bytes, Ordering::Relaxed); + self + } +} + +/// Set the maximum number of bytes to preview from the body when logging in the `TRACE` level. Defaults to 8192 bytes. Set to 0 to disable body preview logging. +pub fn set_max_log_body_preview(max_preview_bytes: usize) { + LOG_HTTP_BODY_MAX.store(max_preview_bytes, Ordering::Relaxed); +} + +fn body_preview_len() -> usize { + LOG_HTTP_BODY_MAX.load(Ordering::Relaxed) +} + /// A helper trait to make generic requests (both regular and SSE) possible. pub trait HttpClientExt: WasmCompatSend + WasmCompatSync { /// Send a HTTP request, get a response back (as bytes). Response must be able to be turned back into Bytes. @@ -116,10 +159,14 @@ impl HttpClientExt for reqwest::Client { U: From + WasmCompatSend, { let (parts, body) = req.into_parts(); + + let body_bytes: Bytes = body.into(); + log_request(&parts, &body_bytes); + let req = self .request(parts.method, parts.uri.to_string()) .headers(parts.headers) - .body(body.into()); + .body(body_bytes); async move { let response = req.send().await.map_err(instance_error)?; @@ -159,6 +206,9 @@ impl HttpClientExt for reqwest::Client { U: WasmCompatSend + 'static, { let (parts, body) = req.into_parts(); + + log_headers(&parts); + let req = self .request(parts.method, parts.uri.to_string()) .headers(parts.headers) @@ -202,10 +252,13 @@ impl HttpClientExt for reqwest::Client { { let (parts, body) = req.into_parts(); + let body_bytes: Bytes = body.into(); + log_request(&parts, &body_bytes); + let req = self .request(parts.method, parts.uri.to_string()) .headers(parts.headers) - .body(body.into()) + .body(body_bytes) .build() .map_err(|x| Error::Instance(x.into())) .unwrap(); @@ -246,3 +299,51 @@ impl HttpClientExt for reqwest::Client { } } } + +fn log_request(parts: &Parts, body: &Bytes) { + if tracing::enabled!(Level::TRACE) { + // Redact sensitive headers (e.g., Authorization) for logging + let mut redacted_headers = parts.headers.clone(); + redacted_headers.remove("Authorization"); + let preview_len = body_preview_len(); + + if preview_len > 0 { + let shown = std::cmp::min(preview_len, body.len()); + let preview = Escape::new(&body[..shown]); + tracing::trace!( + target: "rig::http", + method = %parts.method, + uri = %parts.uri, + body_len = body.len(), + body_preview_len = shown, + headers = ?redacted_headers, + body_preview = ?preview, + "sending HTTP request" + ); + } else { + tracing::trace!( + target: "rig::http", + method = %parts.method, + uri = %parts.uri, + body_len = body.len(), + headers = ?redacted_headers, + "sending HTTP request" + ); + } + } +} + +fn log_headers(parts: &Parts) { + if tracing::enabled!(Level::TRACE) { + // Redact sensitive headers (e.g., Authorization) for logging + let mut redacted_headers = parts.headers.clone(); + redacted_headers.remove("Authorization"); + tracing::trace!( + target: "rig::http", + method = %parts.method, + uri = %parts.uri, + headers = ?redacted_headers, + "sending HTTP request" + ); + } +} diff --git a/rig-core/src/http_client/util.rs b/rig-core/src/http_client/util.rs new file mode 100644 index 000000000..0be02b9f9 --- /dev/null +++ b/rig-core/src/http_client/util.rs @@ -0,0 +1,41 @@ +use std::fmt; +use std::str; + +/// A helper struct to escape bytes for logging. +pub(crate) struct Escape<'a>(&'a [u8]); + +impl<'a> Escape<'a> { + pub(crate) fn new(bytes: &'a [u8]) -> Self { + Escape(bytes) + } +} + +impl fmt::Debug for Escape<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // For valid UTF-8 strings, output directly for better readability + if let Ok(s) = str::from_utf8(self.0) { + return write!(f, "{}", s); + } + write!(f, "{}", self) + } +} + +impl fmt::Display for Escape<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for &c in self.0 { + match c { + b'\n' => write!(f, "\\n")?, + b'\r' => write!(f, "\\r")?, + b'\t' => write!(f, "\\t")?, + b'\\' => write!(f, "\\\\")?, + b'"' => write!(f, "\\\"")?, + b'\0' => write!(f, "\\0")?, + // ASCII printable (0x20-0x7e, excluding space which is 0x20) + c if (0x20..0x7f).contains(&c) => write!(f, "{}", c as char)?, + // Non-printable bytes + c => write!(f, "\\x{c:02x}")?, + } + } + Ok(()) + } +} diff --git a/rig-core/src/prelude.rs b/rig-core/src/prelude.rs index 11b83e674..102082254 100644 --- a/rig-core/src/prelude.rs +++ b/rig-core/src/prelude.rs @@ -16,3 +16,5 @@ pub use crate::client::image_generation::ImageGenerationClient; pub use crate::client::audio_generation::AudioGenerationClient; pub use crate::client::{VerifyClient, VerifyError}; + +pub use crate::http_client::HttpLogConfigExt; From 7f7c2577508e338caea11ec71d10d4025777bab6 Mon Sep 17 00:00:00 2001 From: Plucky Date: Wed, 12 Nov 2025 15:54:55 +0800 Subject: [PATCH 2/2] feat: introduce Escape struct for byte logging and refactor HTTP logging settings --- .../src/http_client/{util.rs => escape.rs} | 0 rig-core/src/http_client/mod.rs | 187 +++++++++++------- rig-core/src/prelude.rs | 2 - 3 files changed, 118 insertions(+), 71 deletions(-) rename rig-core/src/http_client/{util.rs => escape.rs} (100%) diff --git a/rig-core/src/http_client/util.rs b/rig-core/src/http_client/escape.rs similarity index 100% rename from rig-core/src/http_client/util.rs rename to rig-core/src/http_client/escape.rs diff --git a/rig-core/src/http_client/mod.rs b/rig-core/src/http_client/mod.rs index a9aa66f09..7bc05ec13 100644 --- a/rig-core/src/http_client/mod.rs +++ b/rig-core/src/http_client/mod.rs @@ -1,14 +1,16 @@ -use crate::http_client::{sse::BoxedStream, util::Escape}; +use crate::http_client::{escape::Escape, sse::BoxedStream}; use bytes::Bytes; pub use http::{HeaderMap, HeaderValue, Method, Request, Response, Uri, request::Builder}; use http::{StatusCode, request::Parts}; use reqwest::{Body, multipart::Form}; -use std::sync::atomic::{AtomicUsize, Ordering}; -use tracing::Level; +use std::sync::{ + Arc, OnceLock, + atomic::{AtomicBool, AtomicUsize, Ordering}, +}; +mod escape; pub mod retry; pub mod sse; -pub mod util; use std::pin::Pin; @@ -78,44 +80,71 @@ pub fn with_bearer_auth(req: Builder, auth: &str) -> Result { Ok(req.header("Authorization", auth_header)) } -static LOG_HTTP_BODY_MAX: AtomicUsize = AtomicUsize::new(8 * 1024); - -/// Extension trait for client builders to configure HTTP logging. -/// -/// # Example -/// ``` -/// use rig::prelude::*; -/// use rig::providers::openai; -/// -/// let client = openai::Client::builder("api-key") -/// .max_log_body_preview(8192) -/// .build(); -/// ``` -pub trait HttpLogConfigExt { - /// Set the maximum number of bytes to preview from the body when logging in the `TRACE` level. - /// Defaults to 8192 bytes. Set to 0 to disable body preview logging. - /// - /// This method can be called on any client builder to configure HTTP logging before building the client. - fn max_log_body_preview(self, max_preview_bytes: usize) -> Self; +#[derive(Clone, Debug)] +pub struct HttpLogSettings { + max_body_preview: Arc, + headers_enabled: Arc, } -impl HttpLogConfigExt for T -where - T: Sized, -{ - fn max_log_body_preview(self, max_preview_bytes: usize) -> Self { - LOG_HTTP_BODY_MAX.store(max_preview_bytes, Ordering::Relaxed); - self +impl Default for HttpLogSettings { + fn default() -> Self { + Self { + max_body_preview: Arc::new(AtomicUsize::new(8 * 1024)), + headers_enabled: Arc::new(AtomicBool::new(true)), + } + } +} + +impl HttpLogSettings { + pub fn new(max_preview_bytes: usize) -> Self { + Self { + max_body_preview: Arc::new(AtomicUsize::new(max_preview_bytes)), + headers_enabled: Arc::new(AtomicBool::new(true)), + } + } + + pub fn max_body_preview(&self) -> usize { + self.max_body_preview.load(Ordering::Relaxed) + } + + pub fn set_max_body_preview(&self, max_preview_bytes: usize) { + self.max_body_preview + .store(max_preview_bytes, Ordering::Relaxed); + } + + pub fn log_headers_enabled(&self) -> bool { + self.headers_enabled.load(Ordering::Relaxed) + } + + pub fn set_log_headers_enabled(&self, enabled: bool) { + self.headers_enabled.store(enabled, Ordering::Relaxed); } } -/// Set the maximum number of bytes to preview from the body when logging in the `TRACE` level. Defaults to 8192 bytes. Set to 0 to disable body preview logging. +fn http_log_settings() -> &'static HttpLogSettings { + static SETTINGS: OnceLock = OnceLock::new(); + SETTINGS.get_or_init(HttpLogSettings::default) +} + +/// Set the maximum number of bytes to preview from the body when logging at the `TRACE` level. +/// Defaults to 8192 bytes. Set to 0 to disable body preview logging. pub fn set_max_log_body_preview(max_preview_bytes: usize) { - LOG_HTTP_BODY_MAX.store(max_preview_bytes, Ordering::Relaxed); + http_log_settings().set_max_body_preview(max_preview_bytes); } -fn body_preview_len() -> usize { - LOG_HTTP_BODY_MAX.load(Ordering::Relaxed) +/// Get the current maximum number of bytes previewed from the body when logging at the `TRACE` level. +fn max_log_body_preview() -> usize { + http_log_settings().max_body_preview() +} + +/// Enable or disable header logging when tracing HTTP requests/responses. +pub fn set_log_headers_enabled(enabled: bool) { + http_log_settings().set_log_headers_enabled(enabled); +} + +/// Returns whether header logging is currently enabled. +fn log_headers_enabled() -> bool { + http_log_settings().log_headers_enabled() } /// A helper trait to make generic requests (both regular and SSE) possible. @@ -301,49 +330,69 @@ impl HttpClientExt for reqwest::Client { } fn log_request(parts: &Parts, body: &Bytes) { - if tracing::enabled!(Level::TRACE) { - // Redact sensitive headers (e.g., Authorization) for logging - let mut redacted_headers = parts.headers.clone(); - redacted_headers.remove("Authorization"); - let preview_len = body_preview_len(); - - if preview_len > 0 { - let shown = std::cmp::min(preview_len, body.len()); - let preview = Escape::new(&body[..shown]); - tracing::trace!( - target: "rig::http", - method = %parts.method, - uri = %parts.uri, - body_len = body.len(), - body_preview_len = shown, - headers = ?redacted_headers, - body_preview = ?preview, - "sending HTTP request" - ); - } else { - tracing::trace!( - target: "rig::http", - method = %parts.method, - uri = %parts.uri, - body_len = body.len(), - headers = ?redacted_headers, - "sending HTTP request" - ); - } + if !log_headers_enabled() { + return; } -} -fn log_headers(parts: &Parts) { - if tracing::enabled!(Level::TRACE) { - // Redact sensitive headers (e.g., Authorization) for logging - let mut redacted_headers = parts.headers.clone(); - redacted_headers.remove("Authorization"); + let redacted_headers = redact_sensitive_headers(&parts.headers); + let preview_len = max_log_body_preview(); + + if preview_len > 0 { + let shown = std::cmp::min(preview_len, body.len()); + let preview = Escape::new(&body[..shown]); + tracing::trace!( + target: "rig::http", + method = %parts.method, + uri = %parts.uri, + body_len = body.len(), + body_preview_len = shown, + headers = ?redacted_headers, + body_preview = ?preview, + "sending HTTP request" + ); + } else { tracing::trace!( target: "rig::http", method = %parts.method, uri = %parts.uri, + body_len = body.len(), headers = ?redacted_headers, "sending HTTP request" ); } } + +fn log_headers(parts: &Parts) { + if !log_headers_enabled() { + return; + } + + let redacted_headers = redact_sensitive_headers(&parts.headers); + tracing::trace!( + target: "rig::http", + method = %parts.method, + uri = %parts.uri, + headers = ?redacted_headers, + "sending HTTP request" + ); +} + +/// Redact sensitive headers (e.g., Authorization) for logging +fn redact_sensitive_headers(headers: &HeaderMap) -> HeaderMap { + let is_sensitive_header = |name: &str| { + let trimmed = name.trim(); + match trimmed.len() { + 13 => trimmed.eq_ignore_ascii_case("authorization"), + len if len > 4 && trimmed[len - 4..].eq_ignore_ascii_case("-key") => true, + _ => false, + } + }; + let mut filtered = HeaderMap::with_capacity(headers.len()); + for (name, value) in headers.iter() { + // avoid the closure allocation, inline check, avoid clones if possible + if !is_sensitive_header(name.as_str()) { + filtered.append(name, value.clone()); + } + } + filtered +} diff --git a/rig-core/src/prelude.rs b/rig-core/src/prelude.rs index 102082254..11b83e674 100644 --- a/rig-core/src/prelude.rs +++ b/rig-core/src/prelude.rs @@ -16,5 +16,3 @@ pub use crate::client::image_generation::ImageGenerationClient; pub use crate::client::audio_generation::AudioGenerationClient; pub use crate::client::{VerifyClient, VerifyError}; - -pub use crate::http_client::HttpLogConfigExt;