From 6330e78aca76f39870e5bdc947fb04d52e0ae7ca Mon Sep 17 00:00:00 2001 From: justcodebruh Date: Tue, 17 Mar 2026 13:55:30 -0400 Subject: [PATCH] Add custom CA bundle support to BraintrustClientBuilder --- Cargo.toml | 2 +- src/http.rs | 127 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/logger.rs | 32 +++++++++++-- 4 files changed, 157 insertions(+), 5 deletions(-) create mode 100644 src/http.rs diff --git a/Cargo.toml b/Cargo.toml index bfb938a..9ae31df 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,7 +22,7 @@ futures = "0.3" chrono = { version = "0.4", features = ["serde"] } crossbeam = "0.8" indexmap = "2" -reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "stream", "rustls-tls"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "stream", "rustls-tls-native-roots"] } serde = { version = "1", features = ["derive"] } serde_json = "1" serde_repr = "0.1" diff --git a/src/http.rs b/src/http.rs new file mode 100644 index 0000000..6b2cb85 --- /dev/null +++ b/src/http.rs @@ -0,0 +1,127 @@ +use std::path::Path; + +use reqwest::Client; + +use crate::error::{BraintrustError, Result}; + +pub(crate) fn build_http_client( + timeout: std::time::Duration, + ca_bundle: Option<&Path>, +) -> Result { + let mut builder = Client::builder().timeout(timeout); + + if let Some(ca_bundle) = ca_bundle { + let pem = std::fs::read(ca_bundle).map_err(|err| { + BraintrustError::InvalidConfig(format!( + "failed to read CA bundle {}: {}", + ca_bundle.display(), + err + )) + })?; + let certs = reqwest::Certificate::from_pem_bundle(&pem).map_err(|err| { + BraintrustError::InvalidConfig(format!( + "failed to parse PEM certificates from CA bundle {}: {}", + ca_bundle.display(), + err + )) + })?; + if certs.is_empty() { + return Err(BraintrustError::InvalidConfig(format!( + "CA bundle {} did not contain any PEM certificates", + ca_bundle.display() + ))); + } + for cert in certs { + builder = builder.add_root_certificate(cert); + } + } + + builder + .build() + .map_err(|err| BraintrustError::InvalidConfig(err.to_string())) +} + +#[cfg(test)] +pub(crate) mod tests { + use std::path::PathBuf; + use std::time::Duration; + + use super::*; + + pub(crate) const VALID_TEST_CERT_PEM: &str = r#"-----BEGIN CERTIFICATE----- +MIICpDCCAYwCCQDtlc4RX+IuODANBgkqhkiG9w0BAQsFADAUMRIwEAYDVQQDDAls +b2NhbGhvc3QwHhcNMjYwMzE3MTY1MzAyWhcNMjYwMzE4MTY1MzAyWjAUMRIwEAYD +VQQDDAlsb2NhbGhvc3QwggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQDX +K/y/7AlhzPBkIbiEiCt/l1Qfa99h8FdblOe8BCeJkpoW4Fw10mZnWBgX6peZMF7j +p4rjtIJTWkfl8eNoTPdOkYfmi6B3AwAZzl7VQCMgE0gCFkIrgXrkeLqP+q231UxE +wKgilRG3DWFfELZCQeFtq0jSBcnyWybw+o9SgajaQ+SJg7lbgT6o+8AwHQ54HBo+ +VVZJ2CybZvmijQXGiVCMpZ34nxJVW/i6AbsFwp+CLMHOFjrpLuZpv61EnZaGsqsF +RG/VPiNca769Dr8YG4RtPRBKvyDMnUqEDkGwYXrhAVxvI3kKlQq3MHppGCSsjnVl +oqhWm//sE7znMJtuzIf7AgMBAAEwDQYJKoZIhvcNAQELBQADggEBAD1zS7eOkfU2 +IzxjW7MAJce5JrAcRWWe3L2ORx+y+PS4uI0ms1FM4AopZ2FxXdbSSXLf5bqC2f2i +qy+8YbVdZacFtFLmnZicCXP86Na5JUYxZERDyqKN4GFwSrfELwLsuv9TWpir+p/H +3XxQ/8/eJdTHOunNtl4BVUefjGp9PVNb6NFvLDkSkNN37KcjNpB9jPVK970uZ5lb +kOx6ulbMXpNH73h5rwzgs6FbVbcAavPJKYGr170rDRxidpfRz3ex+RBQvcfFQeRx +NP64Q8OosOHraKRn7bvST7bXvGFZUp06aIFrlwdmSQPXU/6o4zYNmkR4RVv4VvQ7 +cb0bfZ7fHHs= +-----END CERTIFICATE----- +"#; + + pub(crate) fn write_temp_bundle(contents: &str) -> PathBuf { + let path = std::env::temp_dir().join(format!( + "braintrust-sdk-rust-ca-bundle-{}.pem", + uuid::Uuid::new_v4() + )); + std::fs::write(&path, contents).expect("write temp bundle"); + path + } + + #[test] + fn rejects_missing_bundle_path() { + let path = std::env::temp_dir().join(format!( + "braintrust-sdk-rust-missing-{}.pem", + uuid::Uuid::new_v4() + )); + + let err = build_http_client(Duration::from_secs(1), Some(&path)).unwrap_err(); + + assert!(matches!(err, BraintrustError::InvalidConfig(_))); + assert!(err.to_string().contains("failed to read CA bundle")); + } + + #[test] + fn rejects_empty_bundle() { + let path = write_temp_bundle(""); + + let err = build_http_client(Duration::from_secs(1), Some(&path)).unwrap_err(); + std::fs::remove_file(&path).expect("remove temp bundle"); + + assert!(matches!(err, BraintrustError::InvalidConfig(_))); + assert!(err + .to_string() + .contains("did not contain any PEM certificates")); + } + + #[test] + fn rejects_malformed_bundle() { + let path = write_temp_bundle( + "-----BEGIN CERTIFICATE-----\nnot-base64\n-----END CERTIFICATE-----\n", + ); + + let err = build_http_client(Duration::from_secs(1), Some(&path)).unwrap_err(); + std::fs::remove_file(&path).expect("remove temp bundle"); + + assert!(matches!(err, BraintrustError::InvalidConfig(_))); + assert!(err.to_string().contains("failed to parse PEM certificates")); + } + + #[test] + fn accepts_valid_bundle() { + let path = write_temp_bundle(VALID_TEST_CERT_PEM); + + let client = build_http_client(Duration::from_secs(1), Some(&path)); + std::fs::remove_file(&path).expect("remove temp bundle"); + + assert!(client.is_ok()); + } +} diff --git a/src/lib.rs b/src/lib.rs index 186a9d3..f3ee380 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ mod error; pub mod eval; mod experiments; mod extractors; +mod http; mod json_merge; mod log_queue; mod logger; diff --git a/src/logger.rs b/src/logger.rs index 0491256..ef3fee2 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -18,6 +19,7 @@ use crate::experiments::api::{ ExperimentRegisterResponse, ExperimentRegistrar, }; use crate::experiments::{BaseExperimentInfo, ExperimentBuilder}; +use crate::http::build_http_client; use crate::log_queue::{LogQueue, LogQueueConfig}; use crate::span::SpanSubmitter; use crate::types::{ParentSpanInfo, SpanPayload}; @@ -187,6 +189,7 @@ pub struct BraintrustClientBuilder { api_url: Option, org_name: Option, default_project: Option, + ca_bundle: Option, queue_size: usize, blocking_login: bool, skip_login: bool, @@ -214,6 +217,7 @@ impl BraintrustClientBuilder { api_url: std::env::var("BRAINTRUST_API_URL").ok(), org_name: std::env::var("BRAINTRUST_ORG_NAME").ok(), default_project: std::env::var("BRAINTRUST_DEFAULT_PROJECT").ok(), + ca_bundle: None, queue_size: DEFAULT_QUEUE_SIZE, blocking_login: false, skip_login: false, @@ -275,6 +279,12 @@ impl BraintrustClientBuilder { self } + /// Set a custom PEM CA bundle used for HTTPS requests made by the SDK client. + pub fn ca_bundle(mut self, path: impl Into) -> Self { + self.ca_bundle = Some(path.into()); + self + } + /// Set the internal queue size for buffering log events. pub fn queue_size(mut self, size: usize) -> Self { self.queue_size = size; @@ -313,10 +323,7 @@ impl BraintrustClientBuilder { let api_url = Url::parse(&api_url_str) .map_err(|e| BraintrustError::InvalidConfig(format!("invalid api_url: {}", e)))?; - let http_client = reqwest::Client::builder() - .timeout(REQUEST_TIMEOUT) - .build() - .map_err(|e| BraintrustError::InvalidConfig(e.to_string()))?; + let http_client = build_http_client(REQUEST_TIMEOUT, self.ca_bundle.as_deref())?; // Create login state (initially empty, populated by login) let login_state = LoginState::new(); @@ -1265,6 +1272,23 @@ mod tests { assert!(matches!(err, BraintrustError::InvalidConfig(_))); } + #[tokio::test] + async fn builder_accepts_custom_ca_bundle_when_login_is_skipped() { + let path = crate::http::tests::write_temp_bundle(crate::http::tests::VALID_TEST_CERT_PEM); + + let client = BraintrustClient::builder() + .app_url("https://example.com") + .api_url("https://example.com") + .ca_bundle(&path) + .skip_login(true) + .build() + .await; + + std::fs::remove_file(&path).expect("remove temp bundle"); + + assert!(client.is_ok()); + } + #[tokio::test] async fn project_registration_is_cached() { let server = MockServer::start().await;