use anyhow::Result;
use async_trait::async_trait;
use reqwest::{
header::{HeaderMap, HeaderName, HeaderValue},
redirect::Policy,
Client, Response, StatusCode,
};
#[cfg(any(feature = "rustls-tls", feature = "native-tls"))]
use reqwest::{Certificate, Identity};
use serde_json::Value;
use std::fmt;
#[cfg(any(feature = "rustls-tls", feature = "native-tls"))]
use std::fs::read_to_string;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use url::Host;
pub const DEFAULT_PROVIDER_TIMEOUT_SECS: u64 = 600;
pub const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 30;
pub type RequestBuilderDecorator =
Arc<dyn Fn(reqwest::RequestBuilder) -> Result<reqwest::RequestBuilder> + Send + Sync>;
pub struct ApiClient {
client: Client,
host: String,
auth: AuthMethod,
default_headers: HeaderMap,
default_query: Vec<(String, String)>,
timeout: Duration,
tls_config: Option<TlsConfig>,
request_builder: Option<RequestBuilderDecorator>,
transport_policy: TransportPolicy,
}
#[derive(Clone)]
enum TransportPolicy {
Default,
HttpsOnly,
LoopbackHttp,
SameOrigin(url::Origin),
}
pub enum AuthMethod {
NoAuth,
BearerToken(String),
ApiKey { header_name: String, key: String },
Custom(Box<dyn AuthProvider>),
}
#[derive(Debug, Clone)]
pub struct TlsCertKeyPair {
pub cert_path: PathBuf,
pub key_path: PathBuf,
}
#[derive(Debug, Clone)]
pub struct TlsConfig {
pub client_identity: Option<TlsCertKeyPair>,
pub ca_cert_path: Option<PathBuf>,
}
impl TlsConfig {
pub fn new() -> Self {
Self {
client_identity: None,
ca_cert_path: None,
}
}
pub fn with_client_cert_and_key(mut self, cert_path: PathBuf, key_path: PathBuf) -> Self {
self.client_identity = Some(TlsCertKeyPair {
cert_path,
key_path,
});
self
}
pub fn with_ca_cert(mut self, path: PathBuf) -> Self {
self.ca_cert_path = Some(path);
self
}
pub fn is_configured(&self) -> bool {
self.client_identity.is_some() || self.ca_cert_path.is_some()
}
#[cfg(any(feature = "rustls-tls", feature = "native-tls"))]
fn load_identity(&self) -> Result<Option<Identity>> {
if let Some(cert_key_pair) = &self.client_identity {
let cert_pem = read_to_string(&cert_key_pair.cert_path)
.map_err(|e| anyhow::anyhow!("Failed to read client certificate: {}", e))?;
let key_pem = read_to_string(&cert_key_pair.key_path)
.map_err(|e| anyhow::anyhow!("Failed to read client private key: {}", e))?;
#[cfg(not(feature = "native-tls"))]
let identity = {
let combined_pem = format!("{}\n{}", cert_pem, key_pem);
Identity::from_pem(combined_pem.as_bytes()).map_err(|e| {
anyhow::anyhow!("Failed to create identity from cert and key: {}", e)
})?
};
#[cfg(feature = "native-tls")]
let identity = {
let pkcs8_key_pem = convert_key_to_pkcs8_pem(&key_pem)?;
Identity::from_pkcs8_pem(cert_pem.as_bytes(), pkcs8_key_pem.as_bytes()).map_err(
|e| anyhow::anyhow!("Failed to create identity from cert and key: {}", e),
)?
};
Ok(Some(identity))
} else {
Ok(None)
}
}
#[cfg(any(feature = "rustls-tls", feature = "native-tls"))]
fn load_ca_certificates(&self) -> Result<Vec<Certificate>> {
match &self.ca_cert_path {
Some(ca_path) => {
let ca_pem = read_to_string(ca_path)
.map_err(|e| anyhow::anyhow!("Failed to read CA certificate: {}", e))?;
let certs = Certificate::from_pem_bundle(ca_pem.as_bytes())
.map_err(|e| anyhow::anyhow!("Failed to parse CA certificate bundle: {}", e))?;
Ok(certs)
}
None => Ok(Vec::new()),
}
}
}
impl Default for TlsConfig {
fn default() -> Self {
Self::new()
}
}
#[cfg(feature = "native-tls")]
const RSA_ALGORITHM_ID: pkcs8::AlgorithmIdentifierRef<'static> = pkcs8::AlgorithmIdentifierRef {
oid: pkcs8::ObjectIdentifier::new_unwrap("1.2.840.113549.1.1.1"),
parameters: Some(pkcs8::der::asn1::AnyRef::NULL),
};
#[cfg(feature = "native-tls")]
fn convert_key_to_pkcs8_pem(key_pem_str: &str) -> Result<String> {
use pkcs8::der::{asn1::OctetStringRef, Decode, Encode};
let parsed =
pem::parse(key_pem_str).map_err(|e| anyhow::anyhow!("Failed to parse PEM key: {}", e))?;
match parsed.tag() {
"PRIVATE KEY" => Ok(key_pem_str.to_string()),
"RSA PRIVATE KEY" => {
let private_key = OctetStringRef::new(parsed.contents())?;
let info = pkcs8::PrivateKeyInfoRef::new(RSA_ALGORITHM_ID, private_key);
let der_bytes = info
.to_der()
.map_err(|e| anyhow::anyhow!("Failed to encode PKCS#8: {}", e))?;
Ok(pem::encode(&pem::Pem::new("PRIVATE KEY", der_bytes)))
}
"EC PRIVATE KEY" => {
let ec_key = sec1::EcPrivateKey::from_der(parsed.contents())
.map_err(|e| anyhow::anyhow!("Failed to parse EC key: {}", e))?;
let curve_oid = ec_key
.parameters
.and_then(|p| p.named_curve())
.ok_or_else(|| {
anyhow::anyhow!(
"EC key missing curve parameters. Convert to PKCS#8: \
openssl pkey -in key.pem -out key-pkcs8.pem"
)
})?;
let algorithm = pkcs8::AlgorithmIdentifierRef {
oid: sec1::ALGORITHM_OID,
parameters: Some((&curve_oid).into()),
};
let private_key = OctetStringRef::new(parsed.contents())?;
let info = pkcs8::PrivateKeyInfoRef::new(algorithm, private_key);
let der_bytes = info
.to_der()
.map_err(|e| anyhow::anyhow!("Failed to encode PKCS#8: {}", e))?;
Ok(pem::encode(&pem::Pem::new("PRIVATE KEY", der_bytes)))
}
tag => Err(anyhow::anyhow!(
"Unsupported key format '{}'. Expected PKCS#8, PKCS#1, or SEC1. \
Convert with: openssl pkey -in key.pem -out key-pkcs8.pem",
tag
)),
}
}
#[async_trait]
pub trait AuthProvider: Send + Sync {
async fn get_auth_header(&self) -> Result<(String, String)>;
async fn refresh_credentials(&self) -> Result<()> {
anyhow::bail!("credential refresh not supported")
}
}
pub struct ApiResponse {
pub status: StatusCode,
pub payload: Option<Value>,
}
impl fmt::Debug for AuthMethod {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
AuthMethod::NoAuth => f.debug_tuple("NoAuth").finish(),
AuthMethod::BearerToken(_) => f.debug_tuple("BearerToken").field(&"[hidden]").finish(),
AuthMethod::ApiKey { header_name, .. } => f
.debug_struct("ApiKey")
.field("header_name", header_name)
.field("key", &"[hidden]")
.finish(),
AuthMethod::Custom(_) => f.debug_tuple("Custom").field(&"[provider]").finish(),
}
}
}
impl ApiResponse {
pub async fn from_response(response: Response) -> Result<Self> {
let status = response.status();
let payload = response.json().await.ok();
Ok(Self { status, payload })
}
}
pub struct ApiRequestBuilder<'a> {
client: &'a ApiClient,
path: &'a str,
headers: HeaderMap,
streaming: bool,
}
impl ApiClient {
pub fn new_with_tls(
host: String,
auth: AuthMethod,
tls_config: Option<TlsConfig>,
) -> Result<Self> {
Self::with_timeout_and_tls(
host,
auth,
Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS),
tls_config,
)
}
pub fn with_timeout_and_tls(
host: String,
auth: AuthMethod,
timeout: Duration,
tls_config: Option<TlsConfig>,
) -> Result<Self> {
let mut client_builder = Self::client_builder(timeout);
if let Some(ref config) = tls_config {
client_builder = Self::configure_tls(client_builder, config)?;
}
let client = client_builder.build()?;
Ok(Self {
client,
host,
auth,
default_headers: HeaderMap::new(),
default_query: Vec::new(),
timeout,
tls_config,
request_builder: None,
transport_policy: TransportPolicy::Default,
})
}
pub fn host(&self) -> &str {
&self.host
}
pub fn timeout(&self) -> Duration {
self.timeout
}
pub fn default_header(&self, name: &str) -> Option<&str> {
self.default_headers
.get(name)
.and_then(|value| value.to_str().ok())
}
fn client_builder(timeout: Duration) -> reqwest::ClientBuilder {
Client::builder()
.connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS))
.read_timeout(timeout)
}
pub fn http_client(tls_config: Option<&TlsConfig>) -> Result<Client> {
let mut client_builder =
Self::client_builder(Duration::from_secs(DEFAULT_PROVIDER_TIMEOUT_SECS));
if let Some(tls_config) = tls_config {
client_builder = Self::configure_tls(client_builder, tls_config)?;
}
Ok(client_builder.build()?)
}
fn rebuild_client(&mut self) -> Result<()> {
let mut client_builder =
Self::client_builder(self.timeout).default_headers(self.default_headers.clone());
client_builder = Self::configure_transport(client_builder, &self.transport_policy);
if let Some(ref tls_config) = self.tls_config {
client_builder = Self::configure_tls(client_builder, tls_config)?;
}
self.client = client_builder.build()?;
Ok(())
}
fn configure_transport(
client_builder: reqwest::ClientBuilder,
transport_policy: &TransportPolicy,
) -> reqwest::ClientBuilder {
match transport_policy {
TransportPolicy::Default => client_builder,
TransportPolicy::HttpsOnly => client_builder.https_only(true),
TransportPolicy::LoopbackHttp => {
client_builder
.no_proxy()
.redirect(Policy::custom(|attempt| {
if attempt.previous().len() > 10 {
return attempt.error("too many redirects");
}
let url = attempt.url();
let is_loopback_http = url.scheme() == "http"
&& match url.host() {
Some(Host::Domain(host)) => host.eq_ignore_ascii_case("localhost"),
Some(Host::Ipv4(address)) => address.is_loopback(),
Some(Host::Ipv6(address)) => address.is_loopback(),
None => false,
};
if url.scheme() == "https" || is_loopback_http {
attempt.follow()
} else {
attempt.error("redirect violates the loopback transport policy")
}
}))
}
TransportPolicy::SameOrigin(origin) => {
let origin = origin.clone();
client_builder.redirect(Policy::custom(move |attempt| {
if attempt.previous().len() >= 10 {
return attempt.error("too many redirects");
}
if attempt.url().origin() == origin {
attempt.follow()
} else {
attempt.error("redirect crosses the authenticated request origin")
}
}))
}
}
}
#[cfg(any(feature = "rustls-tls", feature = "native-tls"))]
fn configure_tls(
mut client_builder: reqwest::ClientBuilder,
tls_config: &TlsConfig,
) -> Result<reqwest::ClientBuilder> {
if tls_config.is_configured() {
if let Some(identity) = tls_config.load_identity()? {
client_builder = client_builder.identity(identity);
}
let ca_certs = tls_config.load_ca_certificates()?;
for ca_cert in ca_certs {
client_builder = client_builder.add_root_certificate(ca_cert);
}
}
Ok(client_builder)
}
#[cfg(not(any(feature = "rustls-tls", feature = "native-tls")))]
fn configure_tls(
client_builder: reqwest::ClientBuilder,
tls_config: &TlsConfig,
) -> Result<reqwest::ClientBuilder> {
if tls_config.is_configured() {
return Err(anyhow::anyhow!(
"Custom TLS configuration requires the `rustls-tls` or `native-tls` feature"
));
}
Ok(client_builder)
}
pub fn with_headers(mut self, headers: HeaderMap) -> Result<Self> {
self.default_headers = headers;
self.rebuild_client()?;
Ok(self)
}
pub fn with_query(mut self, params: Vec<(String, String)>) -> Self {
self.default_query = params;
self
}
pub fn with_header(mut self, key: &str, value: &str) -> Result<Self> {
let header_name = HeaderName::from_bytes(key.as_bytes())?;
let header_value = HeaderValue::from_str(value)?;
self.default_headers.insert(header_name, header_value);
self.rebuild_client()?;
Ok(self)
}
pub fn with_request_builder(mut self, request_builder: RequestBuilderDecorator) -> Self {
self.request_builder = Some(request_builder);
self
}
pub fn with_auth(mut self, auth: AuthMethod) -> Self {
self.auth = auth;
self
}
pub fn with_https_only(mut self) -> Result<Self> {
self.transport_policy = TransportPolicy::HttpsOnly;
self.rebuild_client()?;
Ok(self)
}
pub fn with_loopback_http_only(mut self) -> Result<Self> {
self.transport_policy = TransportPolicy::LoopbackHttp;
self.rebuild_client()?;
Ok(self)
}
pub fn with_same_origin_redirects(mut self) -> Result<Self> {
let origin = url::Url::parse(&self.host)
.map_err(|error| anyhow::anyhow!("Invalid base URL: {}", error))?
.origin();
self.transport_policy = TransportPolicy::SameOrigin(origin);
self.rebuild_client()?;
Ok(self)
}
pub fn request<'a>(&'a self, path: &'a str) -> ApiRequestBuilder<'a> {
ApiRequestBuilder {
client: self,
path,
headers: HeaderMap::new(),
streaming: false,
}
}
pub async fn refresh_credentials(&self) -> Result<()> {
match &self.auth {
AuthMethod::Custom(provider) => provider.refresh_credentials().await,
_ => anyhow::bail!("credential refresh not supported"),
}
}
pub async fn api_post(&self, path: &str, payload: &Value) -> Result<ApiResponse> {
self.request(path).api_post(payload).await
}
pub async fn response_post(&self, path: &str, payload: &Value) -> Result<Response> {
self.request(path).response_post(payload).await
}
pub async fn api_get(&self, path: &str) -> Result<ApiResponse> {
self.request(path).api_get().await
}
pub async fn response_get(&self, path: &str) -> Result<Response> {
self.request(path).response_get().await
}
fn build_url(&self, path: &str) -> Result<url::Url> {
use url::Url;
let mut base_url =
Url::parse(&self.host).map_err(|e| anyhow::anyhow!("Invalid base URL: {}", e))?;
let base_path = base_url.path();
if !base_path.is_empty() && base_path != "/" && !base_path.ends_with('/') {
base_url.set_path(&format!("{}/", base_path));
}
let mut url = base_url
.join(path)
.map_err(|e| anyhow::anyhow!("Failed to construct URL: {}", e))?;
for (key, value) in &self.default_query {
url.query_pairs_mut().append_pair(key, value);
}
Ok(url)
}
}
impl<'a> ApiRequestBuilder<'a> {
pub fn header(mut self, key: &str, value: &str) -> Result<Self> {
let header_name = HeaderName::from_bytes(key.as_bytes())?;
let header_value = HeaderValue::from_str(value)?;
self.headers.insert(header_name, header_value);
Ok(self)
}
#[allow(dead_code)]
pub fn headers(mut self, headers: HeaderMap) -> Self {
self.headers.extend(headers);
self
}
pub fn model_headers(self, model_config: &crate::model::ModelConfig) -> Result<Self> {
match &model_config.request_headers {
Some(headers) => headers
.iter()
.try_fold(self, |builder, (key, value)| builder.header(key, value)),
None => Ok(self),
}
}
pub fn streaming(mut self, streaming: bool) -> Self {
self.streaming = streaming;
self
}
pub async fn api_post(self, payload: &Value) -> Result<ApiResponse> {
let response = self.response_post(payload).await?;
ApiResponse::from_response(response).await
}
async fn send_bounded(&self, request: reqwest::RequestBuilder) -> Result<Response> {
if self.streaming {
Ok(crate::http_status::send_bounded(request, self.client.timeout).await?)
} else {
Ok(request.send().await?)
}
}
pub async fn response_post(self, payload: &Value) -> Result<Response> {
let request = self.send_request(|url, client| client.post(url)).await?;
self.send_bounded(request.json(payload)).await
}
pub async fn multipart_post(self, form: reqwest::multipart::Form) -> Result<Response> {
let request = self.send_request(|url, client| client.post(url)).await?;
Ok(request.multipart(form).send().await?)
}
pub async fn api_get(self) -> Result<ApiResponse> {
let response = self.response_get().await?;
ApiResponse::from_response(response).await
}
pub async fn response_get(self) -> Result<Response> {
let request = self.send_request(|url, client| client.get(url)).await?;
self.send_bounded(request).await
}
async fn send_request<F>(&self, request_builder: F) -> Result<reqwest::RequestBuilder>
where
F: FnOnce(url::Url, &Client) -> reqwest::RequestBuilder,
{
let url = self.client.build_url(self.path)?;
let headers = self.headers.clone();
let mut request = request_builder(url, &self.client.client);
request = request.headers(headers);
if !self.streaming {
request = request.timeout(self.client.timeout);
}
if let Some(decorator) = &self.client.request_builder {
request = decorator(request)?;
}
request = match &self.client.auth {
AuthMethod::NoAuth => request,
AuthMethod::BearerToken(token) => {
request.header("Authorization", format!("Bearer {}", token))
}
AuthMethod::ApiKey { header_name, key } => request.header(header_name.as_str(), key),
AuthMethod::Custom(provider) => {
let (header_name, header_value) = provider.get_auth_header().await?;
request.header(header_name, header_value)
}
};
Ok(request)
}
}
impl fmt::Debug for ApiClient {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ApiClient")
.field("host", &self.host)
.field("auth", &"[auth method]")
.field("timeout", &self.timeout)
.field("default_headers", &self.default_headers)
.finish_non_exhaustive()
}
}
#[cfg(all(test, feature = "native-tls"))]
mod native_tls_tests {
use super::convert_key_to_pkcs8_pem;
const PKCS8_RSA_KEY: &str = "\
-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCzmc7SPaB07+gK
p/VETtmayzNASEaaGIEKO5eOrlrKfIMdmSyB38Ia4dCyxb8WKAROeBtaev8AXfKh
5sPRPi42SNF77d171QWKU4WKXRWWpO2UWwzrXc/BaCNr0Ey9S4ej3v0DmdPkW+pY
PZ5lzHUT6wTutGwJsceCN+/lgPmazaZvfnTdLlRChLyPFjwGkGFwmu7YZACdbIqd
CW5oIs+weAK2L9+HZaLi0TWcnJrnONs00GHTm2DwLkl8ulKRn6+XygWZLpY4owE6
/x4v4CrTKStvfJlPiErcR/t13vZ+kZq3XxcgIY5V0HEz+AT57/MvvGomQXyMmQnN
bQrmdbvRAgMBAAECggEABd6knqSmuDeogeOmHYzFqhxSwRqzmwDqw7i5ne9HFpfN
6bNaO+biF0KS0Wrq0Tiu1lkwzKrZJ2wKnF24P+AMPtgvp4EZRF0sjI4o/sIGzPj4
LklxAUa8Pd4PTUxeG7Aqn1GP+5SFEzEpVieyY/GUEHn1e+8v2jvCEViOabXmwndA
tGHScHwMA0qW1smCOLj12mfFJJjJjlj7QveJFSJ94wo8XjdZvzLDonQoAf4vN99f
rf40zgAx8VQk1UmOPYgc7mqN3JerjWX3rG820UWilLnAA35rrVxjzeV9yfAbgpVC
qHPaOaN4tBVTb3z9w94wK/uDjd4Px26Qq0Yipi4xGQKBgQD7RXE6hNlx1hW7pdWw
hrxdH9ofptBzJfQwYnO0cq6KFF6+QQYzJcMAZAKEytDzfjkH1iHURtBFA1ndKBSI
B1G8iR7Jie4szDghJ+zS9QHHhXLRprSqyamRefnLe/614YlaP7i9wq/sHg8KKm5b
g+S49EfuKctAenMx8sDPDG7Y2QKBgQC2+xOZTE5uPPitOyR8eMUF1ye9zeONzyJ2
l2rjEZuIOWd12wsrgGVTRXoecpqFNIxcmK2mnBtrWCJ/A5xFXjQjIZ1huJ/oNm+t
bI/GtlM5BknfEP6m8NVtMx5ubzqDIlsgRpXCWco3ytaKZJtdnJhZAOmBogAKUiPr
NL2afWbfuQKBgA722MJybPzBkjPFsY88xvUI05W0+o8RAJTtGBT0C37/rZDJEJR1
OELKsfe0mHMX3k+gKg2ZVgf2JiDspSRgwzZmOCYqh7u1QdZ/qTP7EWsPgIIJ2pKd
RfL6/6xlqwqr7uEoEFmEwbNfKughFMdweGunaK0/YfJqGHguC9uI0wUpAoGBAKDM
U5TLscC+Y+oXpyCbIMjZIltxqx7bf/Wnao7Q0lUL3Rd8vnwkAOcEjyRiodedLhvR
MAjR9maGtQnZKmLrPfYBfER78komTE2isVZ2svvRwuj8Dky2J1gnK/7wUAMdFedv
H/wC2+nbnl1PvBivnFHas1jv/AkV1erEFjrFKLpBAoGABQ29M3xsVv7++/ivl0JM
lNzUZwWbG4GI9Yx/zBMJpjiPi9Bkbe2Iu7EgLuiaTnktLSzm9hwVEIxd6a6X/wRA
cGr4L/FDExxjHQPzYpeGQ5fA9HlDWZzQ7Ou+nZ/O+OwXuvnmjXTXyB0djVJHOmLU
gOfZqLKYsL5zn6pefvXu0dU=
-----END PRIVATE KEY-----";
const PKCS1_RSA_KEY: &str = "\
-----BEGIN RSA PRIVATE KEY-----
MIIEowIBAAKCAQEAs5nO0j2gdO/oCqf1RE7ZmsszQEhGmhiBCjuXjq5aynyDHZks
gd/CGuHQssW/FigETngbWnr/AF3yoebD0T4uNkjRe+3de9UFilOFil0VlqTtlFsM
613PwWgja9BMvUuHo979A5nT5FvqWD2eZcx1E+sE7rRsCbHHgjfv5YD5ms2mb350
3S5UQoS8jxY8BpBhcJru2GQAnWyKnQluaCLPsHgCti/fh2Wi4tE1nJya5zjbNNBh
05tg8C5JfLpSkZ+vl8oFmS6WOKMBOv8eL+Aq0ykrb3yZT4hK3Ef7dd72fpGat18X
ICGOVdBxM/gE+e/zL7xqJkF8jJkJzW0K5nW70QIDAQABAoIBAAXepJ6kprg3qIHj
ph2MxaocUsEas5sA6sO4uZ3vRxaXzemzWjvm4hdCktFq6tE4rtZZMMyq2SdsCpxd
uD/gDD7YL6eBGURdLIyOKP7CBsz4+C5JcQFGvD3eD01MXhuwKp9Rj/uUhRMxKVYn
smPxlBB59XvvL9o7whFYjmm15sJ3QLRh0nB8DANKltbJgji49dpnxSSYyY5Y+0L3
iRUifeMKPF43Wb8yw6J0KAH+LzffX63+NM4AMfFUJNVJjj2IHO5qjdyXq41l96xv
NtFFopS5wAN+a61cY83lfcnwG4KVQqhz2jmjeLQVU298/cPeMCv7g43eD8dukKtG
IqYuMRkCgYEA+0VxOoTZcdYVu6XVsIa8XR/aH6bQcyX0MGJztHKuihRevkEGMyXD
AGQChMrQ8345B9Yh1EbQRQNZ3SgUiAdRvIkeyYnuLMw4ISfs0vUBx4Vy0aa0qsmp
kXn5y3v+teGJWj+4vcKv7B4PCipuW4PkuPRH7inLQHpzMfLAzwxu2NkCgYEAtvsT
mUxObjz4rTskfHjFBdcnvc3jjc8idpdq4xGbiDlnddsLK4BlU0V6HnKahTSMXJit
ppwba1gifwOcRV40IyGdYbif6DZvrWyPxrZTOQZJ3xD+pvDVbTMebm86gyJbIEaV
wlnKN8rWimSbXZyYWQDpgaIAClIj6zS9mn1m37kCgYAO9tjCcmz8wZIzxbGPPMb1
CNOVtPqPEQCU7RgU9At+/62QyRCUdThCyrH3tJhzF95PoCoNmVYH9iYg7KUkYMM2
ZjgmKoe7tUHWf6kz+xFrD4CCCdqSnUXy+v+sZasKq+7hKBBZhMGzXyroIRTHcHhr
p2itP2Hyahh4LgvbiNMFKQKBgQCgzFOUy7HAvmPqF6cgmyDI2SJbcase23/1p2qO
0NJVC90XfL58JADnBI8kYqHXnS4b0TAI0fZmhrUJ2Spi6z32AXxEe/JKJkxNorFW
drL70cLo/A5MtidYJyv+8FADHRXnbx/8Atvp255dT7wYr5xR2rNY7/wJFdXqxBY6
xSi6QQKBgAUNvTN8bFb+/vv4r5dCTJTc1GcFmxuBiPWMf8wTCaY4j4vQZG3tiLux
IC7omk55LS0s5vYcFRCMXemul/8EQHBq+C/xQxMcYx0D82KXhkOXwPR5Q1mc0Ozr
vp2fzvjsF7r55o1018gdHY1SRzpi1IDn2aiymLC+c5+qXn717tHV
-----END RSA PRIVATE KEY-----";
const SEC1_EC_KEY: &str = "\
-----BEGIN EC PRIVATE KEY-----
MHcCAQEEIIQyP9q8dmPTyzxgCAF22fS5w6GJP7lVrkq3gdG216TtoAoGCCqGSM49
AwEHoUQDQgAEAxC2ZMkrbwEImb3S2BD8qA9rquFbozW6d7VrhVk1tOqnIuM+vXkz
ShGoCNbfNS+COlPMRAujyDlATZcLs9p4tA==
-----END EC PRIVATE KEY-----";
#[test]
fn test_pkcs8_key_passthrough() {
let result = convert_key_to_pkcs8_pem(PKCS8_RSA_KEY).unwrap();
assert_eq!(result, PKCS8_RSA_KEY.to_string());
}
#[test]
fn test_pkcs1_rsa_key_conversion() {
let result = convert_key_to_pkcs8_pem(PKCS1_RSA_KEY).unwrap();
assert!(result.contains("-----BEGIN PRIVATE KEY-----"));
assert!(result.contains("-----END PRIVATE KEY-----"));
assert!(!result.contains("RSA PRIVATE KEY"));
let re_parsed = pem::parse(&result).unwrap();
assert_eq!(re_parsed.tag(), "PRIVATE KEY");
use pkcs8::der::Decode;
pkcs8::PrivateKeyInfoRef::from_der(re_parsed.contents()).unwrap();
}
#[test]
fn test_sec1_ec_key_conversion() {
let result = convert_key_to_pkcs8_pem(SEC1_EC_KEY).unwrap();
assert!(result.contains("-----BEGIN PRIVATE KEY-----"));
assert!(result.contains("-----END PRIVATE KEY-----"));
assert!(!result.contains("EC PRIVATE KEY"));
let re_parsed = pem::parse(&result).unwrap();
assert_eq!(re_parsed.tag(), "PRIVATE KEY");
use pkcs8::der::Decode;
pkcs8::PrivateKeyInfoRef::from_der(re_parsed.contents()).unwrap();
}
#[test]
fn test_unsupported_key_format() {
let bad_pem =
"-----BEGIN DSA PRIVATE KEY-----\nMIIBuwIBAAKB\n-----END DSA PRIVATE KEY-----";
let result = convert_key_to_pkcs8_pem(bad_pem);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("Unsupported key format"),);
}
#[test]
fn test_invalid_pem() {
let result = convert_key_to_pkcs8_pem("not a pem");
assert!(result.is_err());
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::SocketAddr;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
async fn spawn_chunked_server(gap_ms: u64, chunks: usize) -> SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
loop {
let Ok((mut sock, _)) = listener.accept().await else {
break;
};
tokio::spawn(async move {
let mut buf = [0u8; 8192];
let _ = sock.read(&mut buf).await;
if sock
.write_all(
b"HTTP/1.1 200 OK\r\n\
content-type: text/event-stream\r\n\
transfer-encoding: chunked\r\n\r\n",
)
.await
.is_err()
{
return;
}
for i in 0..chunks {
if i > 0 {
tokio::time::sleep(Duration::from_millis(gap_ms)).await;
}
let data = format!("data: {}\n\n", i);
let chunk = format!("{:x}\r\n{}\r\n", data.len(), data);
if sock.write_all(chunk.as_bytes()).await.is_err() {
return;
}
let _ = sock.flush().await;
}
let _ = sock.write_all(b"0\r\n\r\n").await;
});
}
});
addr
}
fn client_with_timeout(addr: SocketAddr, timeout_ms: u64) -> ApiClient {
let mut client = ApiClient::with_timeout_and_tls(
format!("http://{}", addr),
AuthMethod::NoAuth,
Duration::from_millis(timeout_ms),
None,
)
.unwrap();
client.client = Client::builder()
.no_proxy()
.connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS))
.read_timeout(client.timeout)
.build()
.unwrap();
client
}
async fn drain_counting_data_lines(mut response: Response) -> Result<usize, reqwest::Error> {
let mut body = Vec::new();
while let Some(chunk) = response.chunk().await? {
body.extend_from_slice(&chunk);
}
Ok(String::from_utf8_lossy(&body).matches("data:").count())
}
#[tokio::test]
async fn https_only_rejects_http_after_client_rebuild() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let client = ApiClient::new_with_tls(
format!("http://{addr}"),
AuthMethod::BearerToken("secret".to_string()),
None,
)
.unwrap()
.with_https_only()
.unwrap()
.with_header("x-test", "value")
.unwrap();
assert!(client.response_get("models").await.is_err());
assert!(
tokio::time::timeout(Duration::from_millis(100), listener.accept())
.await
.is_err()
);
}
#[tokio::test]
async fn loopback_transport_rejects_remote_http_redirect() {
for status in ["307 Temporary Redirect", "308 Permanent Redirect"] {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = [0u8; 4096];
let _ = socket.read(&mut request).await;
let response = format!(
"HTTP/1.1 {status}\r\nLocation: http://192.0.2.1/capture\r\nContent-Length: 0\r\n\r\n"
);
socket.write_all(response.as_bytes()).await.unwrap();
});
let client = ApiClient::new_with_tls(
format!("http://{addr}"),
AuthMethod::BearerToken("secret".to_string()),
None,
)
.unwrap()
.with_loopback_http_only()
.unwrap()
.with_header("x-test", "value")
.unwrap();
let error = client
.response_post("chat", &serde_json::json!({ "secret": "prompt" }))
.await
.unwrap_err();
assert!(
format!("{error:#}").contains("redirect violates the loopback transport policy"),
"unexpected redirect error: {error:#}"
);
}
}
#[tokio::test]
async fn streaming_request_survives_beyond_total_timeout() {
let addr = spawn_chunked_server(50, 12).await;
let client = client_with_timeout(addr, 400);
let response = client
.request("v1/messages")
.streaming(true)
.response_post(&serde_json::json!({}))
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let count = drain_counting_data_lines(response).await.unwrap();
assert_eq!(count, 12);
}
#[tokio::test]
async fn streaming_request_fails_when_stream_stalls() {
let addr = spawn_chunked_server(5_000, 2).await;
let client = client_with_timeout(addr, 400);
let response = client
.request("v1/messages")
.streaming(true)
.response_post(&serde_json::json!({}))
.await
.unwrap();
let err = drain_counting_data_lines(response)
.await
.expect_err("stalled stream should time out, not complete");
assert!(err.is_timeout(), "expected a timeout error, got: {err}");
}
#[tokio::test]
async fn non_streaming_request_enforces_total_deadline() {
let addr = spawn_chunked_server(50, 12).await;
let client = client_with_timeout(addr, 400);
let response = client
.request("v1/messages")
.response_post(&serde_json::json!({}))
.await
.unwrap();
let err = drain_counting_data_lines(response)
.await
.expect_err("total deadline should cut off the response body");
assert!(err.is_timeout(), "expected a timeout error, got: {err}");
}
#[tokio::test]
async fn streaming_request_times_out_before_response_headers() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
loop {
let Ok((mut sock, _)) = listener.accept().await else {
break;
};
tokio::spawn(async move {
let mut buf = [0u8; 8192];
while sock.read(&mut buf).await.is_ok_and(|n| n > 0) {}
});
}
});
let client = client_with_timeout(addr, 400);
let started = std::time::Instant::now();
let err = client
.request("v1/messages")
.streaming(true)
.response_post(&serde_json::json!({}))
.await
.expect_err("the phase before the response body must stay bounded");
assert!(
started.elapsed() < Duration::from_secs(5),
"should fail near the configured timeout, took {:?}",
started.elapsed()
);
assert!(matches!(
crate::errors::ProviderError::from(err),
crate::errors::ProviderError::NetworkError(message)
if message.starts_with("Request timed out")
));
}
#[tokio::test]
async fn streaming_error_body_shares_send_deadline() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut buf = [0u8; 8192];
let _ = socket.read(&mut buf).await;
tokio::time::sleep(Duration::from_millis(300)).await;
socket
.write_all(b"HTTP/1.1 500 Internal Server Error\r\ncontent-length: 1\r\n\r\n")
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
let _ = socket.write_all(b"x").await;
});
let client = client_with_timeout(addr, 400);
let started = std::time::Instant::now();
let response = client
.request("v1/messages")
.streaming(true)
.response_post(&serde_json::json!({}))
.await
.unwrap();
crate::http_status::handle_status(response)
.await
.unwrap_err();
assert!(
started.elapsed() < Duration::from_millis(550),
"send and error body used separate deadlines: {:?}",
started.elapsed()
);
}
#[test]
fn test_model_headers_applied_and_override_static_headers() {
let runtime = tokio::runtime::Runtime::new().unwrap();
runtime.block_on(async {
let client = ApiClient::new_with_tls(
"http://localhost:8080".to_string(),
AuthMethod::NoAuth,
None,
)
.unwrap()
.with_header("x-static", "static-value")
.unwrap()
.with_header("queue_threshold", "1000")
.unwrap();
let model_config = crate::model::ModelConfig::new("test-model").with_request_headers(
Some(std::collections::HashMap::from([
("queue_threshold".to_string(), "500".to_string()),
("Idempotency-Key".to_string(), "abc-123".to_string()),
])),
);
let request = client
.request("/test")
.model_headers(&model_config)
.unwrap()
.send_request(|url, client| client.get(url))
.await
.unwrap();
let headers = request.build().unwrap().headers().clone();
let get = |name: &str| {
headers
.get(name)
.and_then(|value| value.to_str().ok())
.map(str::to_string)
};
assert_eq!(get("queue_threshold"), Some("500".to_string()));
assert_eq!(get("Idempotency-Key"), Some("abc-123".to_string()));
});
}
#[test]
fn test_model_headers_rejects_invalid_header_name() {
let client = ApiClient::new_with_tls(
"http://localhost:8080".to_string(),
AuthMethod::NoAuth,
None,
)
.unwrap();
let model_config = crate::model::ModelConfig::new("test-model").with_request_headers(Some(
std::collections::HashMap::from([("bad header name".to_string(), "value".to_string())]),
));
assert!(client
.request("/test")
.model_headers(&model_config)
.is_err());
}
#[test]
fn test_request_builder_decorator() {
let runtime = tokio::runtime::Runtime::new().unwrap();
runtime.block_on(async {
let client = ApiClient::new_with_tls(
"http://localhost:8080".to_string(),
AuthMethod::BearerToken("test-token".to_string()),
None,
)
.unwrap()
.with_request_builder(Arc::new(|request| {
Ok(request.header("test-my-session-id", "test-session_id-456"))
}));
let request = client
.request("/test")
.send_request(|url, client| client.get(url))
.await
.unwrap();
let headers = request.build().unwrap().headers().clone();
let actual = headers
.get("test-my-session-id")
.and_then(|value| value.to_str().ok());
assert_eq!(actual, Some("test-session_id-456"));
});
}
}