Custom certs

Cargo.toml

@@ -12,13 +12,13 @@ repository = "https://github.com/braintrustdata/bt"
 actix-web = "4.11.0"
 anyhow = "1.0.89"
 backoff = { version = "0.4.0", features = ["tokio"] }
-braintrust-sdk-rust = { git = "https://github.com/braintrustdata/braintrust-sdk-rust", rev = "c8c7c7a8d9189164584adef71449a26d4ae8be2b" }
+braintrust-sdk-rust = { git = "https://github.com/braintrustdata/braintrust-sdk-rust", rev = "43ba73edbf5220b57090e049feb094b60a92fcd4" }
 clap = { version = "4.5.20", features = ["derive", "env"] }
 crossterm = "0.28.1"
 futures-util = "0.3.31"
 indicatif = "0.17.8"
 ratatui = "0.29.0"
-reqwest = { version = "0.12.7", default-features = false, features = ["json", "rustls-tls"] }
+reqwest = { version = "0.12.7", default-features = false, features = ["json", "rustls-tls-native-roots"] }
 serde = { version = "1.0.210", features = ["derive"] }
 serde_json = "1.0.128"
 sha2 = "0.10.8"
@@ -33,7 +33,7 @@ urlencoding = "2"
 lingua = { git = "https://github.com/braintrustdata/lingua", rev = "3c79f2c427d12d3e2fe104910ef3c0768ad83770" }
 comfy-table = "7.2.2"
 base64 = "0.22"
-oauth2 = { version = "4.4", default-features = false, features = ["reqwest", "rustls-tls"] }
+oauth2 = { version = "4.4", default-features = false }
 getrandom = "0.3"
 chrono = { version = "0.4.40", features = ["clock"] }
 dirs = "5"

src/args.rs

@@ -1,4 +1,4 @@
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 use clap::Args;
 
@@ -66,6 +66,15 @@ pub struct BaseArgs {
     )]
     pub app_url: Option<String>,
 
+    /// Path to a PEM-encoded CA bundle used for HTTPS requests.
+    #[arg(
+        long = "ca-cert",
+        env = "BRAINTRUST_CA_CERT",
+        hide_env_values = true,
+        global = true
+    )]
+    pub ca_cert: Option<PathBuf>,
+
     /// Path to a .env file to load before running commands.
     #[arg(
         long,
@@ -84,3 +93,9 @@ pub struct CLIArgs<T: Args> {
     #[command(flatten)]
     pub args: T,
 }
+
+impl BaseArgs {
+    pub fn ca_cert(&self) -> Option<&Path> {
+        self.ca_cert.as_deref()
+    }
+}

src/auth.rs

@@ -12,20 +12,17 @@ use braintrust_sdk_rust::{BraintrustClient, LoginState};
 use clap::{Args, Subcommand};
 use crossterm::event::{self, Event, KeyCode, KeyEventKind};
 use dialoguer::{Confirm, Input, Password};
-use oauth2::basic::{BasicClient, BasicTokenType};
-use oauth2::reqwest::async_http_client;
+use oauth2::basic::BasicClient;
 use oauth2::{
-    AuthUrl, AuthorizationCode, ClientId, CsrfToken, EmptyExtraTokenFields, PkceCodeChallenge,
-    PkceCodeVerifier, RedirectUrl, RefreshToken, Scope, StandardTokenResponse, TokenResponse,
-    TokenUrl,
+    AuthUrl, ClientId, CsrfToken, PkceCodeChallenge, PkceCodeVerifier, RedirectUrl, Scope, TokenUrl,
 };
-use reqwest::Client;
 use serde::{Deserialize, Serialize};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpListener;
 
 use crate::{
     args::{BaseArgs, DEFAULT_API_URL, DEFAULT_APP_URL},
+    http::{build_http_client, build_http_client_from_builder},
     ui,
 };
 
@@ -169,7 +166,11 @@ pub async fn list_available_orgs(base: &BaseArgs) -> Result<Vec<AvailableOrg>> {
         .unwrap_or_else(|| DEFAULT_APP_URL.to_string());
     let api_key = match resolved.api_key {
         Some(api_key) => api_key,
-        None => login(base).await?.login.api_key,
+        None => login(base)
+            .await?
+            .login
+            .api_key()
+            .context("login state missing API key")?,
     };
 
     let mut orgs = fetch_login_orgs(&api_key, &app_url).await?;
@@ -330,7 +331,7 @@ pub async fn login_read_only(base: &BaseArgs) -> Result<LoginContext> {
     }
 
     let ctx = fast_login(base).await?;
-    if ctx.login.org_name.trim().is_empty() {
+    if ctx.login.org_name().unwrap_or_default().trim().is_empty() {
         login(base).await
     } else {
         Ok(ctx)
@@ -357,13 +358,17 @@ pub async fn fast_login(base: &BaseArgs) -> Result<LoginContext> {
         .clone()
         .unwrap_or_else(|| DEFAULT_APP_URL.to_string());
 
+    let login = LoginState::new();
+    login.set(
+        api_key,
+        String::new(),
+        org_name,
+        api_url.clone(),
+        app_url.clone(),
+    );
+
     Ok(LoginContext {
-        login: LoginState {
-            api_key,
-            org_id: String::new(),
-            org_name,
-            api_url: Some(api_url.clone()),
-        },
+        login,
         api_url,
         app_url,
     })
@@ -398,27 +403,32 @@ pub async fn login(base: &BaseArgs) -> Result<LoginContext> {
     if let Some(project) = &project {
         builder = builder.default_project(project);
     }
-
     let login = match builder.build().await {
         Ok(client) => client.wait_for_login().await?,
         Err(err) if auth.is_oauth => {
             let org_name = auth
                 .org_name
                 .clone()
                 .ok_or_else(|| anyhow::anyhow!("oauth profile is missing org_name: {err}"))?;
-            LoginState {
-                api_key: api_key.clone(),
-                org_id: String::new(),
+            let login = LoginState::new();
+            login.set(
+                api_key.clone(),
+                String::new(),
                 org_name,
-                api_url: auth.api_url.clone(),
-            }
+                auth.api_url
+                    .clone()
+                    .unwrap_or_else(|| DEFAULT_API_URL.to_string()),
+                auth.app_url
+                    .clone()
+                    .unwrap_or_else(|| DEFAULT_APP_URL.to_string()),
+            );
+            login
         }
         Err(err) => return Err(err.into()),
     };
 
     let api_url = login
-        .api_url
-        .clone()
+        .api_url()
         .or(auth.api_url.clone())
         .unwrap_or_else(|| DEFAULT_API_URL.to_string());
 
@@ -554,7 +564,6 @@ pub async fn resolved_auth_env(base: &BaseArgs) -> Result<Vec<(String, String)>>
     if let Some(org_name) = auth.org_name {
         envs.push(("BRAINTRUST_ORG_NAME".to_string(), org_name));
     }
-
     Ok(envs)
 }
 
@@ -1540,9 +1549,7 @@ fn print_saved_profiles(store: &AuthStore, json: bool) -> Result<()> {
 
 async fn fetch_login_orgs(api_key: &str, app_url: &str) -> Result<Vec<LoginOrgInfo>> {
     let login_url = format!("{}/api/apikey/login", app_url.trim_end_matches('/'));
-    let client = Client::builder()
-        .timeout(crate::http::DEFAULT_HTTP_TIMEOUT)
-        .build()
+    let client = build_http_client(crate::http::DEFAULT_HTTP_TIMEOUT)
         .context("failed to initialize HTTP client")?;
     let response = client
         .post(&login_url)
@@ -1937,41 +1944,72 @@ async fn exchange_oauth_authorization_code(
     code: &str,
     code_verifier: PkceCodeVerifier,
 ) -> Result<OAuthTokenResponse> {
-    let oauth_client = build_oauth_client(api_url, client_id, Some(redirect_uri))?;
-    let token_response = oauth_client
-        .exchange_code(AuthorizationCode::new(code.to_string()))
-        .set_pkce_verifier(code_verifier)
-        .request_async(async_http_client)
-        .await
-        .with_context(|| {
-            format!(
-                "failed to call oauth token endpoint {}/oauth/token",
-                api_url.trim_end_matches('/')
-            )
-        })?;
-    Ok(to_oauth_token_response(token_response))
+    let http_client = build_http_client_from_builder(
+        reqwest::Client::builder()
+            .timeout(crate::http::DEFAULT_HTTP_TIMEOUT)
+            .redirect(reqwest::redirect::Policy::none()),
+    )
+    .context("failed to initialize oauth HTTP client")?;
+    request_oauth_token(
+        &http_client,
+        api_url,
+        &[
+            ("grant_type", "authorization_code"),
+            ("client_id", client_id),
+            ("code", code),
+            ("redirect_uri", redirect_uri),
+            ("code_verifier", code_verifier.secret()),
+        ],
+    )
+    .await
 }
 
 async fn refresh_oauth_access_token(
     api_url: &str,
     refresh_token: &str,
     client_id: &str,
 ) -> Result<OAuthTokenResponse> {
-    let oauth_client = build_oauth_client(api_url, client_id, None)?;
-    let token_response = oauth_client
-        .exchange_refresh_token(&RefreshToken::new(refresh_token.to_string()))
-        .request_async(async_http_client)
-        .await
-        .with_context(|| {
-            format!(
-                "failed to call oauth token endpoint {}/oauth/token",
-                api_url.trim_end_matches('/')
-            )
-        })?;
-    Ok(to_oauth_token_response(token_response))
+    let http_client = build_http_client_from_builder(
+        reqwest::Client::builder()
+            .timeout(crate::http::DEFAULT_HTTP_TIMEOUT)
+            .redirect(reqwest::redirect::Policy::none()),
+    )
+    .context("failed to initialize oauth HTTP client")?;
+    request_oauth_token(
+        &http_client,
+        api_url,
+        &[
+            ("grant_type", "refresh_token"),
+            ("client_id", client_id),
+            ("refresh_token", refresh_token),
+        ],
+    )
+    .await
 }
 
-type OAuth2StdTokenResponse = StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>;
+async fn request_oauth_token(
+    http_client: &reqwest::Client,
+    api_url: &str,
+    params: &[(&str, &str)],
+) -> Result<OAuthTokenResponse> {
+    let token_url = format!("{}/oauth/token", api_url.trim_end_matches('/'));
+    let response = http_client
+        .post(&token_url)
+        .form(params)
+        .send()
+        .await
+        .with_context(|| format!("failed to call oauth token endpoint {token_url}"))?;
+    if !response.status().is_success() {
+        let status = response.status();
+        let body = response.text().await.unwrap_or_default();
+        bail!("oauth token request failed ({status}): {body}");
+    }
+
+    response
+        .json()
+        .await
+        .context("failed to parse oauth token response")
+}
 
 fn build_oauth_client(
     api_url: &str,
@@ -1998,16 +2036,6 @@ fn build_oauth_client(
     }
 }
 
-fn to_oauth_token_response(tokens: OAuth2StdTokenResponse) -> OAuthTokenResponse {
-    OAuthTokenResponse {
-        access_token: tokens.access_token().secret().to_string(),
-        refresh_token: tokens
-            .refresh_token()
-            .map(|token| token.secret().to_string()),
-        expires_in: tokens.expires_in().map(|duration| duration.as_secs()),
-    }
-}
-
 fn prompt_api_key() -> Result<String> {
     if !ui::is_interactive() {
         bail!("--api-key is required in non-interactive mode");
@@ -2652,6 +2680,7 @@ mod tests {
             no_input: false,
             api_url: None,
             app_url: None,
+            ca_cert: None,
             env_file: None,
         }
     }
@@ -3250,8 +3279,8 @@ mod tests {
             .await
             .expect("fast path should succeed");
 
-        assert_eq!(ctx.login.org_name, "acme");
-        assert_eq!(ctx.login.org_id, "");
+        assert_eq!(ctx.login.org_name().as_deref(), Some("acme"));
+        assert_eq!(ctx.login.org_id().as_deref(), Some(""));
         assert_eq!(ctx.api_url, "not-a-valid-url");
     }
 
@@ -3298,8 +3327,8 @@ mod tests {
             .await
             .expect("fast path should succeed with cfg org");
 
-        assert_eq!(ctx.login.api_key, "acme-secret");
-        assert_eq!(ctx.login.org_name, "acme-org");
+        assert_eq!(ctx.login.api_key().as_deref(), Some("acme-secret"));
+        assert_eq!(ctx.login.org_name().as_deref(), Some("acme-org"));
         assert_eq!(ctx.api_url, "https://api.acme.example");
         assert_eq!(ctx.app_url, "https://www.acme.example");
     }
@@ -3316,7 +3345,7 @@ mod tests {
             .await
             .expect("fast path should succeed");
 
-        assert_eq!(ctx.login.org_name, "acme");
+        assert_eq!(ctx.login.org_name().as_deref(), Some("acme"));
         assert_eq!(ctx.api_url, DEFAULT_API_URL);
         assert_eq!(ctx.app_url, DEFAULT_APP_URL);
     }

src/eval.rs

@@ -404,9 +404,7 @@ pub async fn run(base: BaseArgs, args: EvalArgs) -> Result<()> {
             allowed_org_name: args.dev_org_name.clone(),
             allowed_origins: collect_allowed_dev_origins(&args.dev_allowed_origin, &app_url),
             app_url,
-            http_client: Client::builder()
-                .timeout(crate::http::DEFAULT_HTTP_TIMEOUT)
-                .build()
+            http_client: crate::http::build_http_client(crate::http::DEFAULT_HTTP_TIMEOUT)
                 .context("failed to create dev server HTTP client")?,
         };
         return run_dev_server(state).await;

src/functions/mod.rs

@@ -468,7 +468,7 @@ pub(crate) async fn resolve_auth_context(base: &BaseArgs) -> Result<AuthContext>
     Ok(AuthContext {
         client,
         app_url: ctx.app_url,
-        org_id: ctx.login.org_id,
+        org_id: ctx.login.org_id().unwrap_or_default(),
     })
 }
 

src/functions/push.rs

@@ -3443,6 +3443,7 @@ mod tests {
             no_input: false,
             api_url: None,
             app_url: None,
+            ca_cert: None,
             env_file: None,
         }
     }