klieo-ops 3.5.0

//! Governed outbound HTTP client for tool egress.
//!
//! `governed_client(governor, hosts)` returns a `ClientWithMiddleware` that
//! acquires an egress permit from the `Governor` before every outbound
//! request. Requests to hosts not in the allowlist are rejected pre-flight
//! with a middleware error containing "not in allowlist". Requests to
//! allowlisted hosts that exceed the per-host rate cap surface as a
//! middleware error wrapping the governor's reason.

use super::trait_::{Governor, GovernorError};
use crate::error::OpsError;
use async_trait::async_trait;
use http::Extensions;
use reqwest_middleware::{ClientBuilder, ClientWithMiddleware, Middleware, Next};
use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;

/// Cap on the whole request (connect + headers + body). Without it a stuck
/// upstream pins an egress permit indefinitely and starves the governor.
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
/// Cap on the TCP/TLS connect phase alone, so an unreachable host fails fast.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);

/// Build a `ClientWithMiddleware` that governs every outbound request.
///
/// Each request acquires an egress permit from `governor` before being
/// dispatched. Hosts not present in `hosts` are rejected pre-flight.
///
/// # Errors
///
/// Returns `OpsError::Unavailable` if the underlying `reqwest::Client`
/// cannot be constructed.
///
/// # Example
///
/// ```ignore
/// use std::sync::Arc;
/// let client = klieo_ops::governor::governed_client(
///     governor.clone(),
///     &["payout-api.internal"],
/// )?;
/// let resp = client.get("https://payout-api.internal/v1/health").send().await?;
/// ```
pub fn governed_client(
    governor: Arc<dyn Governor>,
    hosts: &[&str],
) -> Result<ClientWithMiddleware, OpsError> {
    let base = reqwest::Client::builder()
        .timeout(REQUEST_TIMEOUT)
        .connect_timeout(CONNECT_TIMEOUT)
        .build()
        .map_err(|e| OpsError::Unavailable {
            component: format!("reqwest::Client::builder: {e}"),
        })?;
    let allowlist: HashSet<String> = hosts.iter().map(|h| h.to_ascii_lowercase()).collect();
    let middleware = GovernorMiddleware {
        governor,
        allowlist,
    };
    Ok(ClientBuilder::new(base).with(middleware).build())
}

struct GovernorMiddleware {
    governor: Arc<dyn Governor>,
    allowlist: HashSet<String>,
}

#[async_trait]
impl Middleware for GovernorMiddleware {
    async fn handle(
        &self,
        req: reqwest::Request,
        extensions: &mut Extensions,
        next: Next<'_>,
    ) -> reqwest_middleware::Result<reqwest::Response> {
        let host = req
            .url()
            .host_str()
            .ok_or_else(|| {
                reqwest_middleware::Error::Middleware(anyhow::anyhow!(
                    "governed_client: URL has no host: {}",
                    req.url()
                ))
            })?
            .to_string();

        if !self.allowlist.contains(&host) {
            return Err(reqwest_middleware::Error::Middleware(anyhow::anyhow!(
                "governed_client: host `{host}` not in allowlist"
            )));
        }

        let _permit = self
            .governor
            .acquire_egress(host.clone())
            .await
            .map_err(governor_err_to_mw)?;

        next.run(req, extensions).await
    }
}

fn governor_err_to_mw(e: GovernorError) -> reqwest_middleware::Error {
    reqwest_middleware::Error::Middleware(anyhow::anyhow!(
        "governed_client: governor refused egress: {e}"
    ))
}