diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 71ed05c0adf..e998fe1abbb 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3385,6 +3385,8 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "bytes", + "futures-util", "litellm-auth-types", "litellm-http", "moka", @@ -4041,6 +4043,7 @@ dependencies = [ "litellm-llms-types", "litellm-secrets", "litellm-tracing", + "reqwest", "rstest", "serde_json", "tokio", diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 65b411ffaef..1a7e2cd6782 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] litellm-auth-types.workspace = true litellm-http.workspace = true +reqwest.workspace = true moka = { workspace = true, features = ["sync"] } serde.workspace = true @@ -26,5 +27,6 @@ aws-smithy-runtime-api = "1.13.0" [dev-dependencies] rstest.workspace = true litellm-http = { workspace = true, features = ["test-support"] } -reqwest.workspace = true +bytes.workspace = true +futures-util.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/auth-aws/src/aws.rs b/litellm-rust/crates/auth-aws/src/aws.rs index eaa033d631f..d83b01c24ad 100644 --- a/litellm-rust/crates/auth-aws/src/aws.rs +++ b/litellm-rust/crates/auth-aws/src/aws.rs @@ -1,4 +1,3 @@ -use std::collections::BTreeMap; use std::time::Duration; use std::time::{SystemTime, UNIX_EPOCH}; @@ -12,6 +11,7 @@ use aws_sigv4::http_request::{ }; use aws_sigv4::sign::v4; use aws_smithy_runtime_api::client::identity::Identity; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; use litellm_auth_types::AwsParams; @@ -435,12 +435,12 @@ fn default_session_name() -> String { /// client header cannot change the canonical request and invalidate the /// signature. Signing everything instead makes the request 403 on a header the /// caller supplied, on a deployment that works on the Python path. -pub fn aws_signature_headers(headers: &BTreeMap) -> BTreeMap { +pub fn aws_signature_headers(headers: &HeaderMap) -> HeaderMap { headers .iter() .filter(|(name, _)| { - let name = name.to_ascii_lowercase(); - AWS_SIGNED_HEADER_NAMES.contains(&name.as_str()) + let name = name.as_str(); + AWS_SIGNED_HEADER_NAMES.contains(&name) || name.starts_with("x-amz-") || name.starts_with("x-amzn-") }) @@ -456,45 +456,63 @@ pub fn is_sigv4_computed_header(name: &str) -> bool { SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str()) } -pub fn sign_post( +#[derive(Clone, Copy, Debug)] +pub struct SigningScope<'a> { + pub region: &'a str, + pub service: &'a str, + pub credentials: &'a Credentials, + pub time: SystemTime, +} + +pub fn sign_headers( + method: &str, url: &str, body: &[u8], - headers: &BTreeMap, - region: &str, - service: &str, - credentials: &Credentials, - signing_time: SystemTime, -) -> Result, Error> { - let identity: Identity = credentials.clone().into(); + headers: &HeaderMap, + scope: SigningScope<'_>, +) -> Result { + let identity: Identity = scope.credentials.clone().into(); let params = v4::SigningParams::builder() .identity(&identity) - .region(region) - .name(service) - .time(signing_time) + .region(scope.region) + .name(scope.service) + .time(scope.time) .settings(SigningSettings::default()) .build() .map(SigningParams::from) .map_err(|error| Error::AwsSigningParameters(error.to_string()))?; let header_refs = headers .iter() - .map(|(name, value)| (name.as_str(), value.as_str())); - let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) - .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; + .map(|(name, value)| { + value + .to_str() + .map(|value| (name.as_str(), value)) + .map_err(|_| Error::AwsSignableRequest(format!("{name} is not visible ASCII"))) + }) + .collect::, _>>()?; + let request = SignableRequest::new( + method, + url, + header_refs.into_iter(), + SignableBody::Bytes(body), + ) + .map_err(|error| Error::AwsSignableRequest(error.to_string()))?; let (instructions, _) = sign(request, ¶ms) .map_err(|error| Error::AwsSigning(error.to_string()))? .into_parts(); - Ok(instructions + instructions .headers() .map(|(name, value)| { - let normalized_name = match name { - "authorization" => "Authorization", - "x-amz-date" => "X-Amz-Date", - "x-amz-security-token" => "X-Amz-Security-Token", - _ => name, - }; - (normalized_name.to_string(), value.to_string()) + HeaderName::from_bytes(name.as_bytes()) + .map_err(|error| error.to_string()) + .and_then(|name| { + HeaderValue::from_str(value) + .map(|value| (name, value)) + .map_err(|error| error.to_string()) + }) }) - .collect()) + .collect::>() + .map_err(Error::AwsSigning) } /// Model-id and region parsing shared by every Bedrock route. @@ -639,6 +657,7 @@ pub fn host_supplied_credentials(params: &AwsParams) -> Option { mod tests { use super::*; use crate::constants::{AWS_REGION, BEDROCK_SERVICE}; + use reqwest::header::CONTENT_TYPE; fn no_env(_: &str) -> Option { None @@ -731,12 +750,12 @@ mod tests { ); } - fn parity_inputs() -> (String, Vec, BTreeMap) { + fn parity_inputs() -> (String, Vec, HeaderMap) { ( "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" .to_string(), br#"{"input":"hello"}"#.to_vec(), - BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]), + HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]), ) } @@ -961,9 +980,9 @@ mod tests { // cannot change the canonical request. Signing it instead makes the // request 403 the moment anything on the wire rewrites or drops it. let (url, body, mut headers) = parity_inputs(); - headers.insert("x-request-id".to_string(), "abc-123".to_string()); - headers.insert("Accept-Encoding".to_string(), "gzip".to_string()); - headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string()); + headers.insert("x-request-id", HeaderValue::from_static("abc-123")); + headers.insert("Accept-Encoding", HeaderValue::from_static("gzip")); + headers.insert("x-amzn-trace-id", HeaderValue::from_static("Root=1-abc")); let signable = aws_signature_headers(&headers); assert!(!signable.contains_key("x-request-id")); @@ -979,18 +998,21 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &signable, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - SystemTime::UNIX_EPOCH, + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: SystemTime::UNIX_EPOCH, + }, ) .expect("signs"); - let authorization = signed - .get("Authorization") + let authorization = signed["Authorization"] + .to_str() .expect("carries an authorization header"); assert!( !authorization.contains("x-request-id"), @@ -1012,29 +1034,24 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &headers, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + }, ) .expect("golden signature"); + assert_eq!(signed["X-Amz-Date"], "20240102T030405Z"); + assert_eq!(signed["X-Amz-Security-Token"], "session-token"); assert_eq!( - signed.get("X-Amz-Date").map(String::as_str), - Some("20240102T030405Z") - ); - assert_eq!( - signed.get("X-Amz-Security-Token").map(String::as_str), - Some("session-token") - ); - assert_eq!( - signed.get("Authorization").map(String::as_str), - Some( - "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" - ) + signed["Authorization"], + "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" ); } @@ -1048,14 +1065,17 @@ mod tests { None, "test", ); - let signed = sign_post( + let signed = sign_headers( + "POST", &url, &body, &headers, - "us-east-1", - BEDROCK_SERVICE, - &credentials, - UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + SigningScope { + region: "us-east-1", + service: BEDROCK_SERVICE, + credentials: &credentials, + time: UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + }, ) .expect("signature"); assert!(!signed.contains_key("X-Amz-Security-Token")); @@ -1068,7 +1088,7 @@ mod tests { let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?; let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec(); let headers = - BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]); + HeaderMap::from_iter([(CONTENT_TYPE, HeaderValue::from_static("application/json"))]); let credentials = AwsAuthService::default() .resolve_credentials( AwsAuthConfig { @@ -1087,23 +1107,25 @@ mod tests { let url = format!( "https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke" ); - let signed_headers = sign_post( + let signed_headers = sign_headers( + "POST", &url, &body, &headers, - region, - BEDROCK_SERVICE, - &credentials, - SystemTime::now(), + SigningScope { + region, + service: BEDROCK_SERVICE, + credentials: &credentials, + time: SystemTime::now(), + }, )?; - let mut request = client.post(&url).body(body.clone()); - for (name, value) in &headers { - request = request.header(name, value); - } - for (name, value) in signed_headers { - request = request.header(name, value); - } - let response = request.send().await?; + let response = client + .post(&url) + .headers(headers.clone()) + .headers(signed_headers) + .body(body.clone()) + .send() + .await?; let status = response.status(); let response_body = response.text().await?; let snippet: String = response_body.chars().take(240).collect(); diff --git a/litellm-rust/crates/auth-aws/src/signer.rs b/litellm-rust/crates/auth-aws/src/signer.rs index 3868fdd939b..6c2cb4fd162 100644 --- a/litellm-rust/crates/auth-aws/src/signer.rs +++ b/litellm-rust/crates/auth-aws/src/signer.rs @@ -1,11 +1,13 @@ -use std::{collections::BTreeMap, time::SystemTime}; +use std::time::SystemTime; + +use aws_credential_types::Credentials; +use reqwest::{Request, header::HeaderMap}; use crate::{ - AwsAuthService, AwsCredentialSource, Error, aws_signature_headers, is_sigv4_computed_header, - sign_post, + AwsAuthService, AwsCredentialSource, Error, SigningScope, aws_signature_headers, + is_sigv4_computed_header, sign_headers, }; -use aws_credential_types::Credentials; -use litellm_http::outbound::{RequestSigner, UnsignedRequest}; +use litellm_http::outbound::RequestSigner; #[derive(Clone, Debug)] pub struct SigV4Signer { @@ -45,41 +47,48 @@ impl SigV4Signer { } impl RequestSigner for SigV4Signer { - fn sign( - &self, - request: UnsignedRequest<'_>, - ) -> Result, litellm_http::Error> { - if let Some((name, _)) = request - .headers - .iter() - .find(|(name, _)| is_sigv4_computed_header(name)) + fn sign(&self, request: &Request) -> Result { + if let Some(name) = request + .headers() + .keys() + .find(|name| is_sigv4_computed_header(name.as_str())) { - return Err(litellm_http::Error::ComputedHeader(name.clone())); + return Err(litellm_http::Error::ComputedHeader(name.to_string())); } - let headers: BTreeMap = request.headers.iter().cloned().collect(); - sign_post( - request.url, - request.body, - &aws_signature_headers(&headers), - &self.region, - self.service, - &self.credentials, - (self.clock)(), + let body = match request.body() { + None => &[], + Some(body) => body.as_bytes().ok_or_else(|| { + litellm_http::Error::Signature("a streaming body cannot be signed".into()) + })?, + }; + sign_headers( + request.method().as_str(), + request.url().as_str(), + body, + &aws_signature_headers(request.headers()), + SigningScope { + region: &self.region, + service: self.service, + credentials: &self.credentials, + time: (self.clock)(), + }, ) - .map(|signature| signature.into_iter().collect()) .map_err(|error| litellm_http::Error::Signature(error.to_string())) } } #[cfg(test)] mod tests { + use reqwest::{Body, Method, Url}; + use serde_json::{Value, json as json_value}; use std::time::{Duration, UNIX_EPOCH}; - use litellm_http::outbound::OutboundRequest; - use serde_json::{Value, json}; + use litellm_http::outbound::{json, signed}; use super::*; + const URL: &str = "https://textract.us-east-1.amazonaws.com/"; + fn fixed_clock() -> SystemTime { UNIX_EPOCH + Duration::from_secs(1_700_000_000) } @@ -93,75 +102,83 @@ mod tests { .with_clock(fixed_clock) } + fn target() -> (String, String) { + ("X-Amz-Target".into(), "Textract.DetectDocumentText".into()) + } + fn authorization(body: &Value, service: &'static str) -> String { - OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], - body, - None, + signed( + json(URL, &[target()], body, None).unwrap(), &signer(service), ) .unwrap() - .header("Authorization") - .unwrap() - .to_string() + .headers()["authorization"] + .to_str() + .unwrap() + .to_string() } #[test] fn the_signature_verifies_against_the_bytes_that_are_sent() { - let sent = OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())], - &json!({"Document": {"Bytes": "aGk="}}), - None, + let sent = signed( + json( + URL, + &[target()], + &json_value!({"Document": {"Bytes": "aGk="}}), + None, + ) + .unwrap(), &signer("textract"), ) .unwrap(); - let unsigned: BTreeMap = sent + let unsigned: HeaderMap = sent .headers() .iter() - .filter(|(name, _)| !is_sigv4_computed_header(name)) - .cloned() + .filter(|(name, _)| !is_sigv4_computed_header(name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) .collect(); - let recomputed = sign_post( - sent.url(), - sent.body(), + let recomputed = sign_headers( + "POST", + sent.url().as_str(), + sent.body().unwrap().as_bytes().unwrap(), &aws_signature_headers(&unsigned), - "us-east-1", - "textract", - &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), - fixed_clock(), + SigningScope { + region: "us-east-1", + service: "textract", + credentials: &Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"), + time: fixed_clock(), + }, ) .unwrap(); - assert_eq!( - sent.header("Authorization"), - Some(recomputed["Authorization"].as_str()) - ); + assert_eq!(sent.headers()["authorization"], recomputed["authorization"]); } #[test] fn the_signature_depends_on_the_body_and_the_service() { - let original = authorization(&json!({"text": "card 4111"}), "textract"); + let original = authorization(&json_value!({"text": "card 4111"}), "textract"); assert_ne!( original, - authorization(&json!({"text": "card [REDACTED]"}), "textract") + authorization(&json_value!({"text": "card [REDACTED]"}), "textract") ); assert_ne!( original, - authorization(&json!({"text": "card 4111"}), "bedrock") + authorization(&json_value!({"text": "card 4111"}), "bedrock") ); assert!(original.contains("/us-east-1/textract/aws4_request")); } #[test] fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() { - let error = OutboundRequest::signed_json( - "https://textract.us-east-1.amazonaws.com/".into(), - vec![("authorization".into(), "Bearer caller".into())], - &json!({}), - None, + let error = signed( + json( + URL, + &[("authorization".into(), "Bearer caller".into())], + &json_value!({}), + None, + ) + .unwrap(), &signer("textract"), ) .unwrap_err(); @@ -171,4 +188,17 @@ mod tests { litellm_http::Error::ComputedHeader("authorization".into()) ); } + + #[test] + fn a_streaming_body_cannot_be_signed() { + let mut request = Request::new(Method::POST, Url::parse(URL).unwrap()); + *request.body_mut() = Some(Body::wrap_stream(futures_util::stream::once(async { + Ok::<_, std::io::Error>(bytes::Bytes::from_static(b"{}")) + }))); + + assert!(matches!( + signed(request, &signer("textract")), + Err(litellm_http::Error::Signature(_)) + )); + } } diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index eafb4d2976b..ae14756baa3 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -24,6 +24,8 @@ pub enum Error { Client(String), #[error("request body could not be serialized: {0}")] RequestBody(String), + #[error("invalid request: {0}")] + Request(String), #[error("request forwards a header the signer computes: {0}")] ComputedHeader(String), #[error("request signing failed: {0}")] diff --git a/litellm-rust/crates/http/src/outbound.rs b/litellm-rust/crates/http/src/outbound.rs index c2cfb00d79b..09b5d0179d0 100644 --- a/litellm-rust/crates/http/src/outbound.rs +++ b/litellm-rust/crates/http/src/outbound.rs @@ -1,131 +1,68 @@ -//! The request a route hands to the transport. The body is serialized once, -//! when the request is built, and a [`RequestSigner`] sees those exact bytes. +//! Building the request a route hands to the transport. The body is serialized +//! once, into the [`reqwest::Request`] that is sent, and a [`RequestSigner`] +//! reads those exact bytes back out of it. //! //! Host hooks may rewrite the wire request (redaction, guardrails) and a //! signature such as AWS SigV4 covers the body, so a route builds this after -//! its hooks ran and cannot change or re-serialize it afterwards. +//! its hooks ran and signs last. use std::time::Duration; +use reqwest::{ + Request, + header::{CONTENT_TYPE, HeaderMap}, +}; use serde::Serialize; -use crate::{ - Error, - request::{HeaderPolicy, has_header, with_headers}, -}; - -#[derive(Clone, Copy, Debug)] -pub struct UnsignedRequest<'a> { - pub url: &'a str, - pub headers: &'a [(String, String)], - pub body: &'a [u8], -} +use crate::Error; /// Returns the headers to add to the request; it never sees a mutable request. pub trait RequestSigner: Send + Sync { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error>; + fn sign(&self, request: &Request) -> Result; } -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct OutboundRequest { - url: String, - headers: Vec<(String, String)>, - body: Vec, +pub fn json( + url: &str, + headers: &[(String, String)], + body: &impl Serialize, timeout: Option, +) -> Result { + let body = serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?; + let with_headers = headers + .iter() + .fold(http::Request::post(url), |builder, (name, value)| { + builder.header(name, value) + }); + let has_content_type = with_headers + .headers_ref() + .is_some_and(|headers| headers.contains_key(CONTENT_TYPE)); + let builder = if has_content_type { + with_headers + } else { + with_headers.header(CONTENT_TYPE, "application/json") + }; + let mut request = builder + .body(body) + .map_err(|error| Error::Request(error.to_string())) + .and_then(|request| { + Request::try_from(request) + .map_err(|error| Error::Request(error.without_url().to_string())) + })?; + *request.timeout_mut() = timeout; + Ok(request) } -impl OutboundRequest { - pub fn json( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - ) -> Result { - Self::build(url, headers, body, timeout, None) - } - - pub fn signed_json( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - signer: &dyn RequestSigner, - ) -> Result { - Self::build(url, headers, body, timeout, Some(signer)) - } - - fn build( - url: String, - headers: Vec<(String, String)>, - body: &impl Serialize, - timeout: Option, - signer: Option<&dyn RequestSigner>, - ) -> Result { - let body = - serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?; - let content_type = (!has_header(&headers, "content-type")) - .then(|| ("content-type".to_string(), "application/json".to_string())); - let unsigned: Vec<(String, String)> = headers.into_iter().chain(content_type).collect(); - let signature = signer - .map(|signer| { - signer.sign(UnsignedRequest { - url: &url, - headers: &unsigned, - body: &body, - }) - }) - .transpose()? - .unwrap_or_default(); - Ok(Self { - url, - headers: unsigned.into_iter().chain(signature).collect(), - body, - timeout, - }) - } - - pub fn url(&self) -> &str { - &self.url - } - - pub fn headers(&self) -> &[(String, String)] { - &self.headers - } - - pub fn header(&self, name: &str) -> Option<&str> { - self.headers - .iter() - .find(|(key, _)| key.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.as_str()) - } - - pub fn body(&self) -> &[u8] { - &self.body - } - - pub fn timeout(&self) -> Option { - self.timeout - } - - pub async fn send(self, client: &crate::Client) -> Result { - let builder = with_headers( - client.post(&self.url).body(self.body), - &self.headers, - HeaderPolicy::All, - ); - match self.timeout { - Some(timeout) => builder.timeout(timeout), - None => builder, - } - .send() - .await - } +pub fn signed(mut request: Request, signer: &dyn RequestSigner) -> Result { + let signature = signer.sign(&request)?; + request.headers_mut().extend(signature); + Ok(request) } #[cfg(test)] mod tests { use std::sync::Mutex; + use reqwest::header::HeaderValue; use serde_json::json; use super::*; @@ -134,61 +71,73 @@ mod tests { struct Recording(Mutex>); impl RequestSigner for Recording { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { - *self.0.lock().unwrap() = request.body.to_vec(); - Ok(vec![("authorization".into(), "signed".into())]) + fn sign(&self, request: &Request) -> Result { + *self.0.lock().unwrap() = request.body().unwrap().as_bytes().unwrap().to_vec(); + Ok(HeaderMap::from_iter([( + reqwest::header::AUTHORIZATION, + HeaderValue::from_static("signed"), + )])) } } #[test] fn the_signer_sees_exactly_the_bytes_that_are_sent() { let signer = Recording::default(); - let request = OutboundRequest::signed_json( - "https://provider.test/".into(), - vec![("x-caller".into(), "kept".into())], - &json!({"b": 1, "a": [true, null]}), - None, + let request = signed( + json( + "https://provider.test/", + &[("x-caller".into(), "kept".into())], + &json!({"b": 1, "a": [true, null]}), + Some(Duration::from_secs(7)), + ) + .unwrap(), &signer, ) .unwrap(); - assert_eq!(request.body(), signer.0.lock().unwrap().as_slice()); - assert_eq!(request.header("authorization"), Some("signed")); - assert_eq!(request.header("x-caller"), Some("kept")); + assert_eq!( + request.body().unwrap().as_bytes().unwrap(), + signer.0.lock().unwrap().as_slice() + ); + assert_eq!(request.headers()["authorization"], "signed"); + assert_eq!(request.headers()["x-caller"], "kept"); + assert_eq!(request.timeout(), Some(&Duration::from_secs(7))); } #[test] fn the_content_type_is_part_of_what_the_signer_sees() { struct RequiresContentType; impl RequestSigner for RequiresContentType { - fn sign(&self, request: UnsignedRequest<'_>) -> Result, Error> { - has_header(request.headers, "content-type") - .then(Vec::new) + fn sign(&self, request: &Request) -> Result { + request + .headers() + .contains_key(CONTENT_TYPE) + .then(HeaderMap::new) .ok_or_else(|| Error::Signature("content-type was not signed".into())) } } - let defaulted = OutboundRequest::signed_json( - "u".into(), - Vec::new(), - &json!({}), - None, + let defaulted = signed( + json("https://u.test/", &[], &json!({}), None).unwrap(), &RequiresContentType, ) .unwrap(); - assert_eq!(defaulted.header("content-type"), Some("application/json")); + assert_eq!(defaulted.headers()[CONTENT_TYPE], "application/json"); - let provider = OutboundRequest::signed_json( - "u".into(), - vec![("Content-Type".into(), "application/x-amz-json-1.1".into())], - &json!({}), - None, + let provider = signed( + json( + "https://u.test/", + &[("Content-Type".into(), "application/x-amz-json-1.1".into())], + &json!({}), + None, + ) + .unwrap(), &RequiresContentType, ) .unwrap(); assert_eq!( - provider.header("content-type"), - Some("application/x-amz-json-1.1") + provider.headers()[CONTENT_TYPE], + "application/x-amz-json-1.1" ); assert_eq!(provider.headers().len(), 1); } @@ -197,14 +146,33 @@ mod tests { fn a_signer_failure_produces_no_request() { struct Refuses; impl RequestSigner for Refuses { - fn sign(&self, _request: UnsignedRequest<'_>) -> Result, Error> { + fn sign(&self, _request: &Request) -> Result { Err(Error::ComputedHeader("authorization".into())) } } - assert_eq!( - OutboundRequest::signed_json("u".into(), Vec::new(), &json!({}), None, &Refuses), - Err(Error::ComputedHeader("authorization".into())) - ); + let error = signed( + json("https://u.test/", &[], &json!({}), None).unwrap(), + &Refuses, + ) + .unwrap_err(); + assert_eq!(error, Error::ComputedHeader("authorization".into())); + } + + #[test] + fn a_request_that_cannot_be_sent_fails_when_it_is_built() { + assert!(matches!( + json("not a url", &[], &json!({}), None), + Err(Error::Request(_)) + )); + assert!(matches!( + json( + "https://u.test/", + &[("x-bad".into(), "\u{7f}".into())], + &json!({}), + None + ), + Err(Error::Request(_)) + )); } } diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index cfebe8c6b01..5a3dd19ade2 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -53,6 +53,18 @@ pub async fn execute_http_request( client.execute(request).await } +pub fn header_pairs(headers: &reqwest::header::HeaderMap) -> Vec<(String, String)> { + headers + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.to_string(), value.to_string())) + }) + .collect() +} + pub fn truncate_error_body(body: &str) -> String { if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS { return body.to_string(); diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index 130786e788d..fa3fad067a9 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -6,7 +6,6 @@ license.workspace = true repository.workspace = true [dependencies] -futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-cache-response.workspace = true litellm-core-utils.workspace = true @@ -16,7 +15,10 @@ litellm-inference.workspace = true litellm-llms.workspace = true litellm-llms-types.workspace = true litellm-secrets.workspace = true + serde_json = { workspace = true, features = ["preserve_order"] } +futures-util.workspace = true +reqwest.workspace = true tracing.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index c52ff19ac08..15a45747997 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -5,7 +5,6 @@ use litellm_auth::AuthServices; use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest}; use litellm_http::{ Client, - outbound::OutboundRequest, request::{truncate_error_body, with_default_headers}, }; use litellm_llms::base_llm::{ @@ -96,7 +95,7 @@ pub(super) async fn execute( headers: wire.headers, signer: authenticated.signer, }, - wire.url, + &wire.url, &wire.body, call.timeout, )?; @@ -170,10 +169,10 @@ pub(super) fn as_response_error(err: Error) -> Error { pub(super) fn outbound_request( authenticated: Authenticated, - url: String, + url: &str, body: &Value, timeout: Option, -) -> Result { +) -> Result { litellm_inference::outbound::outbound_request( authenticated, url, diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 8b6154ecbee..292995d842c 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -288,7 +288,7 @@ async fn send( ) -> Result { let request = outbound_request( authenticated, - url.to_string(), + url, body, Some(timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS))), )?; @@ -859,12 +859,12 @@ mod tests { headers: call.wire.headers, signer: call.signer, }, - call.wire.url, + &call.wire.url, &call.wire.body, None, ) .unwrap(); - let authorization = signed.header("authorization").unwrap_or_default(); + let authorization = signed.headers()["authorization"].to_str().unwrap(); assert!( authorization.contains(&format!("/us-west-2/{service}/aws4_request")), "{authorization}" diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs index 5336d67895d..c4096767188 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/aws_textract.rs @@ -1,9 +1,12 @@ -use std::{collections::BTreeMap, time::SystemTime}; +use std::time::SystemTime; -use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; +use litellm_auth_aws::{Credentials, SigningScope, aws_signature_headers, sign_headers}; use rstest::rstest; use time::{PrimitiveDateTime, format_description}; -use wiremock::Request; +use wiremock::{ + Request, + http::{HeaderMap, HeaderName}, +}; use super::*; @@ -42,21 +45,26 @@ fn expected_authorization(url: &str, sent: &Request) -> String { .unwrap() .assume_utc() .into(); - let headers: BTreeMap = ["content-type", "x-amz-target"] + let headers: HeaderMap = ["content-type", "x-amz-target"] .into_iter() - .map(|name| (name.to_string(), sent.header(name).unwrap().to_string())) + .map(|name| (HeaderName::from_static(name), sent.headers[name].clone())) .collect(); - sign_post( + sign_headers( + "POST", url, &sent.body, &aws_signature_headers(&headers), - "eu-west-1", - "textract", - &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), - signed_at, + SigningScope { + region: "eu-west-1", + service: "textract", + credentials: &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + time: signed_at, + }, ) - .unwrap()["Authorization"] - .clone() + .unwrap()["authorization"] + .to_str() + .unwrap() + .to_string() } /// The recorded URL names wiremock's host, not the address the client signed for. diff --git a/litellm-rust/crates/inference-responses/src/handler.rs b/litellm-rust/crates/inference-responses/src/handler.rs index 02182d23ae3..cbbc8be139d 100644 --- a/litellm-rust/crates/inference-responses/src/handler.rs +++ b/litellm-rust/crates/inference-responses/src/handler.rs @@ -79,7 +79,7 @@ pub(super) async fn execute( headers: wire.headers, signer: authenticated.signer, }, - wire.url, + &wire.url, &wire.body, Some(call.timeout.unwrap_or(Duration::from_secs(600))), )?; diff --git a/litellm-rust/crates/inference-transcription/src/handler.rs b/litellm-rust/crates/inference-transcription/src/handler.rs index 5a7238c5e36..6541ae8263e 100644 --- a/litellm-rust/crates/inference-transcription/src/handler.rs +++ b/litellm-rust/crates/inference-transcription/src/handler.rs @@ -48,7 +48,7 @@ pub(super) async fn execute( let authenticated = resolve_auth(auth, environment, &env_lookup).await?; let outbound = litellm_inference::outbound::outbound_request( authenticated, - url, + &url, &body, Some( call.timeout diff --git a/litellm-rust/crates/inference/src/outbound.rs b/litellm-rust/crates/inference/src/outbound.rs index 41fe3aa62b8..9d9b1e606c2 100644 --- a/litellm-rust/crates/inference/src/outbound.rs +++ b/litellm-rust/crates/inference/src/outbound.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_http::outbound::OutboundRequest; +use litellm_http::outbound::{json, signed}; use litellm_llms::base_llm::auth::Authenticated; use serde_json::Value; @@ -11,10 +11,10 @@ use serde_json::Value; fields(status) )] pub async fn send( - request: OutboundRequest, + request: reqwest::Request, client: &litellm_http::Client, ) -> Result { - request.send(client).await.inspect(|response| { + client.execute(request).await.inspect(|response| { tracing::Span::current().record("status", response.status().as_u16()); }) } @@ -23,13 +23,14 @@ pub async fn send( /// bytes that are sent. pub fn outbound_request( authenticated: Authenticated, - url: String, + url: &str, body: &Value, timeout: Option, -) -> Result { +) -> Result { let Authenticated { headers, signer } = authenticated; + let request = json(url, &headers, body, timeout)?; match signer { - None => OutboundRequest::json(url, headers, body, timeout), - Some(signer) => OutboundRequest::signed_json(url, headers, body, timeout, &signer), + None => Ok(request), + Some(signer) => signed(request, &signer), } } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 21ad4ea8025..a9ddcb87086 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -7,7 +7,8 @@ use litellm_host::interceptors::WireRequest; use litellm_http::{ Client, ClientVariant, HttpClientConfig, HttpClientPool, media::{MediaFetcher, UrlPolicy}, - outbound::{OutboundRequest, RequestSigner}, + outbound::{RequestSigner, json, signed}, + request::header_pairs, transport, }; use serde::{Serialize, de::DeserializeOwned}; @@ -128,22 +129,14 @@ pub async fn ocr( ) -> Result { let http = config.prepare_request(request, client, hooks).await?; let url = http.url().to_string(); - let headers = http.headers().to_vec(); - let response = http - .send(client.provider_http()) + let headers = header_pairs(http.headers()); + let response = client + .provider_http() + .execute(http) .await .map_err(transport_error)?; if !response.status().is_success() { - let headers = response - .headers() - .iter() - .filter_map(|(name, value)| { - value - .to_str() - .ok() - .map(|value| (name.to_string(), value.to_string())) - }) - .collect(); + let headers = header_pairs(response.headers()); return match read_response_bytes(response, request.connection.max_response_bytes).await { Err(Error::Transport(transport::Error::Http { status, body })) => { Err(config.get_error_class(body, status, headers)) @@ -225,7 +218,7 @@ pub async fn transform_request_body( body: B, signer: Option<&dyn RequestSigner>, hooks: &dyn CallHooks, -) -> Result { +) -> Result { let composed = litellm_core_utils::call_arguments::compose_body( &request.optional_params, &body, @@ -241,17 +234,16 @@ pub async fn transform_request_body( }); } config.validate_request_body(&changed.body)?; - let timeout = Some(request.connection.timeout); + let built = json( + url, + &changed.headers, + &changed.body, + Some(request.connection.timeout), + )?; Ok(match signer { - Some(signer) => OutboundRequest::signed_json( - url.into(), - changed.headers, - &changed.body, - timeout, - signer, - ), - None => OutboundRequest::json(url.into(), changed.headers, &changed.body, timeout), - }?) + Some(signer) => signed(built, signer)?, + None => built, + }) } fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest { @@ -267,10 +259,10 @@ pub fn build_http_request( url: String, headers: Vec<(String, String)>, body: &impl Serialize, -) -> Result { - Ok(OutboundRequest::json( - url, - headers, +) -> Result { + Ok(json( + &url, + &headers, body, Some(request.connection.timeout), )?) diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index 2da0e1abbfc..73e7b64638d 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,15 +1,15 @@ -use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; -use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; - -use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; -use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; -use litellm_http::outbound::{OutboundRequest, RequestSigner}; -use litellm_secrets::source::Secrets; use serde::{ Serialize, de::{DeserializeOwned, IntoDeserializer}, }; use serde_json::{Map, Value}; +use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; + +use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; +use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; +use litellm_http::outbound::RequestSigner; +use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; +use litellm_secrets::source::Secrets; use crate::base_llm::ocr::{ error::Error, @@ -391,7 +391,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> impl Future> + Send { + ) -> impl Future> + Send { async move { let params = self.map_ocr_params(&request.optional_params, &request.model)?; let environment = self.validate_environment(request, client).await?; diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 9b991b75e5b..fc73ce7c014 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -1,11 +1,5 @@ use std::collections::BTreeMap; -use litellm_core_utils::{ - call_arguments::{CallArguments, compose_body}, - params::OpaqueParams, - url_utils::ApiUrl, -}; -use litellm_http::outbound::OutboundRequest; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::{Map, Value, json}; @@ -18,6 +12,11 @@ use crate::base_llm::ocr::{ decode_and_normalize_response, }, }; +use litellm_core_utils::{ + call_arguments::{CallArguments, compose_body}, + params::OpaqueParams, + url_utils::ApiUrl, +}; use litellm_llms_types::formats::ocr::{ LiteLLMOcrResponse, OcrDocument, OcrPage, OcrResponseFormat, OcrUsageInfo, }; @@ -173,7 +172,7 @@ impl BaseOcrConfig for ReductoParseV3Config { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -262,7 +261,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { request: &PreparedOcrRequest, client: &OcrClient, hooks: &dyn CallHooks, - ) -> Result { + ) -> Result { prepare_upload_request(self, request, client, hooks).await } } @@ -275,7 +274,7 @@ async fn prepare_upload_request, -) -> Result { +) -> Result { let params = config.map_ocr_params(&request.optional_params, &request.model)?; let headers = config.validate_environment(request, client).await?; let url = config.get_complete_url(request, ¶ms, &headers)?;