bitrouter 1.0.0-alpha.34

//! Raw BitRouter Cloud HTTP API client.

use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use reqwest::Url;
use reqwest::header::{ACCEPT, AUTHORIZATION, HeaderMap, HeaderValue};
use reqwest::{Method, StatusCode, Version};

use crate::cloud::account::manager::CredentialManager;

/// A raw HTTP request to a relative BitRouter Cloud endpoint.
pub struct ApiRequest {
    method: Method,
    endpoint: String,
    headers: HeaderMap,
    body: Option<Vec<u8>>,
}

impl std::fmt::Debug for ApiRequest {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let header_names = self.headers.keys().collect::<Vec<_>>();
        let endpoint = redacted_endpoint(&self.endpoint);
        let body = self
            .body
            .as_ref()
            .map(|bytes| format!("<redacted> ({} bytes)", bytes.len()));
        f.debug_struct("ApiRequest")
            .field("method", &self.method)
            .field("endpoint", &endpoint)
            .field("header_names", &header_names)
            .field("body", &body)
            .finish()
    }
}

impl ApiRequest {
    /// Create a request with no custom headers or body.
    pub fn new(method: Method, endpoint: impl Into<String>) -> Self {
        Self {
            method,
            endpoint: endpoint.into(),
            headers: HeaderMap::new(),
            body: None,
        }
    }

    /// Replace the request's custom headers.
    pub fn with_headers(mut self, headers: HeaderMap) -> Self {
        self.headers = headers;
        self
    }

    /// Attach an exact byte body to the request.
    pub fn with_body(mut self, body: impl Into<Vec<u8>>) -> Self {
        self.body = Some(body.into());
        self
    }

    /// Return the configured HTTP method.
    pub fn method(&self) -> &Method {
        &self.method
    }

    /// Return the relative endpoint supplied by the caller.
    pub fn endpoint(&self) -> &str {
        &self.endpoint
    }

    /// Return the custom request headers.
    pub fn headers(&self) -> &HeaderMap {
        &self.headers
    }
}

/// A raw HTTP response whose body remains available as a reqwest stream.
pub struct ApiResponse {
    response: reqwest::Response,
}

impl std::fmt::Debug for ApiResponse {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let url = redacted_endpoint(self.response.url().as_str());
        let header_names = self.response.headers().keys().collect::<Vec<_>>();
        f.debug_struct("ApiResponse")
            .field("status", &self.response.status())
            .field("version", &self.response.version())
            .field("url", &url)
            .field("header_names", &header_names)
            .finish_non_exhaustive()
    }
}

impl ApiResponse {
    /// Return the response status without consuming the body.
    pub fn status(&self) -> StatusCode {
        self.response.status()
    }

    /// Return the HTTP protocol version.
    pub fn version(&self) -> Version {
        self.response.version()
    }

    /// Return the response headers.
    pub fn headers(&self) -> &HeaderMap {
        self.response.headers()
    }

    /// Return the final response URL. Redirect following is disabled, so this
    /// is always the confined request URL.
    pub fn url(&self) -> &Url {
        self.response.url()
    }

    /// Consume the wrapper and return the streaming reqwest response.
    pub fn into_response(self) -> reqwest::Response {
        self.response
    }
}

/// Authenticated client for arbitrary relative BitRouter Cloud endpoints.
pub struct CloudApiClient {
    base_url: Url,
    http: reqwest::Client,
    manager: Arc<CredentialManager>,
}

impl std::fmt::Debug for CloudApiClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CloudApiClient")
            .field("base_url", &self.base_url)
            .field("credentials", &"<redacted>")
            .finish_non_exhaustive()
    }
}

impl CloudApiClient {
    /// Build a client from the current hosted account credential.
    pub async fn from_manager(manager: Arc<CredentialManager>) -> Result<Self> {
        let credential = manager
            .current()
            .await
            .context("reading BitRouter Cloud credentials")?
            .with_context(|| {
                format!(
                    "no stored credentials — run `{} cloud login` first",
                    bitrouter_sdk::invocation::name()
                )
            })?;
        let base_url = Url::parse(credential.base_url()).with_context(|| {
            format!(
                "parsing the BitRouter Cloud URL stored in {}",
                manager.path().display()
            )
        })?;
        let http = reqwest::Client::builder()
            .user_agent(concat!("bitrouter/", env!("CARGO_PKG_VERSION")))
            .connect_timeout(Duration::from_secs(30))
            .redirect(reqwest::redirect::Policy::none())
            .build()
            .context("building BitRouter Cloud API client")?;
        Ok(Self {
            base_url,
            http,
            manager,
        })
    }

    /// Return the stored login URL used as the request origin.
    pub fn base_url(&self) -> &Url {
        &self.base_url
    }

    /// Resolve a relative endpoint against the stored login origin.
    pub fn endpoint_url(&self, endpoint: &str) -> Result<Url> {
        resolve_endpoint(&self.base_url, endpoint)
    }

    /// Send one authenticated request without interpreting its status or
    /// buffering its response body.
    pub async fn execute(&self, request: ApiRequest) -> Result<ApiResponse> {
        let url = self.endpoint_url(&request.endpoint)?;
        let mut headers = request.headers;
        if !headers.contains_key(AUTHORIZATION) {
            let bearer = self.current_bearer().await?;
            let value = HeaderValue::from_str(&format!("Bearer {bearer}"))
                .context("stored credential cannot be represented as an HTTP header")?;
            headers.insert(AUTHORIZATION, value);
        }
        if !headers.contains_key(ACCEPT) {
            headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
        }
        let mut builder = self.http.request(request.method, url).headers(headers);
        if let Some(body) = request.body {
            builder = builder.body(body);
        }
        let response = builder
            .send()
            .await
            .map_err(reqwest::Error::without_url)
            .context("sending BitRouter Cloud request")?;
        Ok(ApiResponse { response })
    }

    async fn current_bearer(&self) -> Result<String> {
        self.manager
            .session()
            .resolve_bearer(None, Some(self.base_url.as_str()))
            .await
            .map(|credential| credential.secret().to_owned())
            .map_err(|error| anyhow::anyhow!(error))
    }
}

fn resolve_endpoint(base_url: &Url, endpoint: &str) -> Result<Url> {
    if endpoint.is_empty() {
        anyhow::bail!("API endpoint cannot be empty");
    }
    if endpoint.starts_with("//") || endpoint.starts_with("\\\\") || Url::parse(endpoint).is_ok() {
        anyhow::bail!("API endpoint must be a relative path on the logged-in origin");
    }
    let mut origin = base_url.clone();
    origin.set_path("/");
    origin.set_query(None);
    origin.set_fragment(None);
    let resolved = origin.join(endpoint).context("resolving API endpoint")?;
    if resolved.origin() != origin.origin() {
        anyhow::bail!("API endpoint resolves outside the logged-in origin");
    }
    if resolved.fragment().is_some() {
        anyhow::bail!("API endpoint fragments are not supported");
    }
    Ok(resolved)
}

fn redacted_endpoint(endpoint: &str) -> String {
    match endpoint.find(['?', '#']) {
        Some(index) => format!("{}<redacted>", &endpoint[..=index]),
        None => endpoint.to_owned(),
    }
}

#[cfg(test)]
mod tests {
    use bitrouter_ai::providers::hosted::credentials::StoredCredential;
    use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderValue};
    use wiremock::matchers::{header, method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    use super::*;
    async fn api_client(base_url: &str) -> anyhow::Result<(tempfile::TempDir, CloudApiClient)> {
        let directory = tempfile::tempdir()?;
        let manager = Arc::new(CredentialManager::with_client(
            directory.path().join("account-credentials.json"),
            reqwest::Client::new(),
        ));
        manager
            .save(StoredCredential::api_key(
                "brk_AAAAAAAAAAAAAAAA.secret".to_owned(),
                base_url.to_owned(),
            ))
            .await?;
        let client = CloudApiClient::from_manager(manager).await?;
        Ok((directory, client))
    }

    #[test]
    fn resolves_only_relative_paths_on_the_login_origin() -> anyhow::Result<()> {
        let base = reqwest::Url::parse("https://api.bitrouter.ai/oauth")?;
        assert_eq!(
            resolve_endpoint(&base, "/v1/models?owned=true")?.as_str(),
            "https://api.bitrouter.ai/v1/models?owned=true"
        );
        assert_eq!(
            resolve_endpoint(&base, "v1/models")?.as_str(),
            "https://api.bitrouter.ai/v1/models"
        );
        for endpoint in [
            "https://evil.example/v1/models",
            "//evil.example/v1/models",
            "\\\\evil.example/v1/models",
            "/v1/models#fragment",
            "mailto:security@example.com",
        ] {
            assert!(
                resolve_endpoint(&base, endpoint).is_err(),
                "{endpoint} must be rejected"
            );
        }
        Ok(())
    }

    #[test]
    fn api_request_debug_redacts_headers_and_body() {
        let mut headers = HeaderMap::new();
        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_static("Bearer request-secret"),
        );
        let request = ApiRequest::new(reqwest::Method::POST, "/v1/responses?api_key=query-secret")
            .with_headers(headers)
            .with_body(br#"{"input":"private prompt"}"#.to_vec());

        let rendered = format!("{request:?}");

        assert!(!rendered.contains("request-secret"));
        assert!(!rendered.contains("query-secret"));
        assert!(!rendered.contains("private prompt"));
        assert!(rendered.contains("<redacted>"));
    }

    #[tokio::test]
    async fn sends_stored_api_key_and_preserves_response_stream() -> anyhow::Result<()> {
        let server = MockServer::start().await;
        let (_directory, client) = api_client(&server.uri()).await?;
        Mock::given(method("GET"))
            .and(path("/v1/models"))
            .and(header(
                "user-agent",
                concat!("bitrouter/", env!("CARGO_PKG_VERSION")),
            ))
            .and(header(
                "authorization",
                "Bearer brk_AAAAAAAAAAAAAAAA.secret",
            ))
            .respond_with(
                ResponseTemplate::new(200)
                    .insert_header("x-api-key", "response-secret")
                    .set_body_raw("{\"data\":[]}", "application/json"),
            )
            .mount(&server)
            .await;

        let response = client
            .execute(ApiRequest::new(
                reqwest::Method::GET,
                "/v1/models?api_key=query-secret",
            ))
            .await?;

        assert_eq!(response.status(), reqwest::StatusCode::OK);
        let rendered = format!("{response:?}");
        assert!(!rendered.contains("query-secret"));
        assert!(!rendered.contains("response-secret"));
        assert_eq!(response.into_response().text().await?, "{\"data\":[]}");
        let requests = match server.received_requests().await {
            Some(requests) => requests,
            None => anyhow::bail!("wiremock did not record the request"),
        };
        assert_eq!(requests.len(), 1);
        Ok(())
    }

    #[tokio::test]
    async fn does_not_follow_cross_origin_redirects() -> anyhow::Result<()> {
        let destination = MockServer::start().await;
        let origin = MockServer::start().await;
        let (_directory, client) = api_client(&origin.uri()).await?;
        Mock::given(method("GET"))
            .and(path("/v1/models"))
            .respond_with(
                ResponseTemplate::new(302)
                    .insert_header("location", format!("{}/stolen", destination.uri())),
            )
            .mount(&origin)
            .await;

        let response = client
            .execute(ApiRequest::new(reqwest::Method::GET, "/v1/models"))
            .await?;

        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
        let destination_requests = match destination.received_requests().await {
            Some(requests) => requests,
            None => anyhow::bail!("wiremock did not record destination requests"),
        };
        assert!(destination_requests.is_empty());
        let origin_requests = match origin.received_requests().await {
            Some(requests) => requests,
            None => anyhow::bail!("wiremock did not record origin requests"),
        };
        assert_eq!(origin_requests.len(), 1);
        Ok(())
    }

    #[tokio::test]
    async fn raw_authorization_bypasses_near_expiry_oauth_resolution() -> anyhow::Result<()> {
        let server = MockServer::start().await;
        let directory = tempfile::tempdir()?;
        let manager = Arc::new(CredentialManager::with_client(
            directory.path().join("account-credentials.json"),
            reqwest::Client::new(),
        ));
        manager
            .save(StoredCredential::from(
                bitrouter_ai::providers::hosted::credentials::Credentials {
                    access_token: "near-expiry-access".to_owned(),
                    refresh_token: Some("refresh-token".to_owned()),
                    expires_at: chrono::Utc::now() + chrono::Duration::seconds(10),
                    refresh_token_expires_at: None,
                    token_type: "Bearer".to_owned(),
                    scope: "inference:invoke".to_owned(),
                    client_id: "bitrouter-cli".to_owned(),
                    authorization_server: server.uri(),
                    namespace_id: Some("ns-test".to_owned()),
                    subject: None,
                },
            ))
            .await?;
        let client = CloudApiClient::from_manager(manager).await?;
        Mock::given(method("GET"))
            .and(path("/v1/models"))
            .and(header("authorization", "Bearer user-override"))
            .respond_with(ResponseTemplate::new(200))
            .mount(&server)
            .await;
        let mut headers = HeaderMap::new();
        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_static("Bearer user-override"),
        );
        headers.append("x-test", HeaderValue::from_static("one"));
        headers.append("x-test", HeaderValue::from_static("two"));

        client
            .execute(ApiRequest::new(reqwest::Method::GET, "/v1/models").with_headers(headers))
            .await?;

        let requests = match server.received_requests().await {
            Some(requests) => requests,
            None => anyhow::bail!("wiremock did not record the request"),
        };
        let request = requests
            .first()
            .context("wiremock did not receive a request")?;
        assert_eq!(request.headers[AUTHORIZATION], "Bearer user-override");
        let values = request
            .headers
            .get_all("x-test")
            .iter()
            .map(HeaderValue::to_str)
            .collect::<std::result::Result<Vec<_>, _>>()?;
        assert_eq!(values, ["one", "two"]);
        let metadata_requests = requests
            .iter()
            .filter(|request| request.url.path() == "/.well-known/oauth-authorization-server")
            .count();
        let refresh_requests = requests
            .iter()
            .filter(|request| request.url.path() == "/oauth/token")
            .count();
        assert_eq!(metadata_requests, 0);
        assert_eq!(refresh_requests, 0);
        Ok(())
    }
}