Codex Handbook
aws-auth/src/lib.rs 261 lines
mod config;mod signing;use std::time::SystemTime;use aws_credential_types::provider::ProvideCredentials;use aws_credential_types::provider::SharedCredentialsProvider;use bytes::Bytes;use http::HeaderMap;use http::Method;use thiserror::Error;/// AWS auth configuration used to resolve credentials and sign requests.#[derive(Debug, Clone, PartialEq, Eq)]pub struct AwsAuthConfig {    pub profile: Option<String>,    pub region: Option<String>,    pub service: String,}/// Generic HTTP request shape consumed by SigV4 signing.#[derive(Debug, Clone, PartialEq, Eq)]pub struct AwsRequestToSign {    pub method: Method,    pub url: String,    pub headers: HeaderMap,    pub body: Bytes,}/// Signed request parts returned to the caller.#[derive(Debug, Clone, PartialEq, Eq)]pub struct AwsSignedRequest {    pub url: String,    pub headers: HeaderMap,}/// Errors returned by credential loading or SigV4 signing.#[derive(Debug, Error)]pub enum AwsAuthError {    #[error("AWS service name must not be empty")]    EmptyService,    #[error("AWS SDK config did not resolve a credentials provider")]    MissingCredentialsProvider,    #[error("AWS SDK config did not resolve a region")]    MissingRegion,    #[error("failed to load AWS credentials: {0}")]    Credentials(#[from] aws_credential_types::provider::error::CredentialsError),    #[error("request URL is not a valid URI: {0}")]    InvalidUri(#[source] http::uri::InvalidUri),    #[error("failed to construct HTTP request for signing: {0}")]    BuildHttpRequest(#[source] http::Error),    #[error("request contains a non-UTF8 header value: {0}")]    InvalidHeaderValue(#[source] http::header::ToStrError),    #[error("failed to build signable request: {0}")]    SigningRequest(#[source] aws_sigv4::http_request::SigningError),    #[error("failed to build SigV4 signing params: {0}")]    SigningParams(String),    #[error("SigV4 signing failed: {0}")]    SigningFailure(#[source] aws_sigv4::http_request::SigningError),}/// Loaded AWS auth context that can sign outbound HTTP requests.#[derive(Clone)]pub struct AwsAuthContext {    credentials_provider: SharedCredentialsProvider,    region: String,    service: String,}impl std::fmt::Debug for AwsAuthContext {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.debug_struct("AwsAuthContext")            .field("region", &self.region)            .field("service", &self.service)            .finish_non_exhaustive()    }}impl AwsAuthContext {    pub async fn load(config: AwsAuthConfig) -> Result<Self, AwsAuthError> {        let sdk_config = config::load_sdk_config(&config).await?;        let credentials_provider = config::credentials_provider(&sdk_config)?;        let region = config::resolved_region(&sdk_config)?;        Ok(Self {            credentials_provider,            region,            service: config.service.trim().to_string(),        })    }    pub fn region(&self) -> &str {        &self.region    }    pub fn service(&self) -> &str {        &self.service    }    pub async fn sign(&self, request: AwsRequestToSign) -> Result<AwsSignedRequest, AwsAuthError> {        self.sign_at(request, SystemTime::now()).await    }    async fn sign_at(        &self,        request: AwsRequestToSign,        time: SystemTime,    ) -> Result<AwsSignedRequest, AwsAuthError> {        let credentials = self.credentials_provider.provide_credentials().await?;        signing::sign_request(&credentials, &self.region, &self.service, request, time)    }}impl AwsAuthError {    /// Returns whether retrying the outbound request can reasonably recover from this auth error.    pub fn is_retryable(&self) -> bool {        match self {            AwsAuthError::Credentials(error) => matches!(                error,                aws_credential_types::provider::error::CredentialsError::ProviderTimedOut(_)                    | aws_credential_types::provider::error::CredentialsError::ProviderError(_)            ),            AwsAuthError::EmptyService            | AwsAuthError::MissingCredentialsProvider            | AwsAuthError::MissingRegion            | AwsAuthError::InvalidUri(_)            | AwsAuthError::BuildHttpRequest(_)            | AwsAuthError::InvalidHeaderValue(_)            | AwsAuthError::SigningRequest(_)            | AwsAuthError::SigningParams(_)            | AwsAuthError::SigningFailure(_) => false,        }    }}#[cfg(test)]mod tests {    use std::time::Duration;    use std::time::UNIX_EPOCH;    use aws_credential_types::Credentials;    use aws_credential_types::provider::error::CredentialsError;    use pretty_assertions::assert_eq;    use super::*;    fn test_context(session_token: Option<&str>) -> AwsAuthContext {        AwsAuthContext {            credentials_provider: SharedCredentialsProvider::new(Credentials::new(                "AKIDEXAMPLE",                "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",                session_token.map(str::to_string),                /*expires_after*/ None,                "unit-test",            )),            region: "us-east-1".to_string(),            service: "bedrock".to_string(),        }    }    fn test_request() -> AwsRequestToSign {        let mut headers = HeaderMap::new();        headers.insert(            http::header::CONTENT_TYPE,            http::HeaderValue::from_static("application/json"),        );        headers.insert("x-test-header", http::HeaderValue::from_static("present"));        AwsRequestToSign {            method: Method::POST,            url: "https://bedrock-runtime.us-east-1.amazonaws.com/v1/responses".to_string(),            headers,            body: Bytes::from_static(br#"{"model":"openai.gpt-oss-120b-1:0"}"#),        }    }    #[tokio::test]    async fn sign_adds_sigv4_headers_and_preserves_existing_headers() {        let signed = test_context(/*session_token*/ None)            .sign_at(                test_request(),                UNIX_EPOCH + Duration::from_secs(1_700_000_000),            )            .await            .expect("request should sign");        assert_eq!(            signing::header_value(&signed.headers, http::header::CONTENT_TYPE.as_str()),            Some("application/json".to_string())        );        assert_eq!(            signing::header_value(&signed.headers, "x-test-header"),            Some("present".to_string())        );        assert_eq!(            signed.url,            "https://bedrock-runtime.us-east-1.amazonaws.com/v1/responses"        );        assert!(            signing::header_value(&signed.headers, http::header::AUTHORIZATION.as_str())                .is_some_and(|value| value.starts_with("AWS4-HMAC-SHA256 "))        );        assert!(signing::header_value(&signed.headers, "x-amz-date").is_some());    }    #[test]    fn credentials_provider_failures_are_retryable() {        assert!(            AwsAuthError::Credentials(CredentialsError::provider_error("temporarily unavailable"))                .is_retryable()        );        assert!(            AwsAuthError::Credentials(CredentialsError::provider_timed_out(Duration::from_secs(1)))                .is_retryable()        );    }    #[test]    fn deterministic_aws_auth_errors_are_not_retryable() {        assert!(!AwsAuthError::EmptyService.is_retryable());        assert!(            !AwsAuthError::Credentials(CredentialsError::not_loaded_no_source()).is_retryable()        );        assert!(            !AwsAuthError::Credentials(CredentialsError::invalid_configuration("bad profile"))                .is_retryable()        );        assert!(            !AwsAuthError::Credentials(CredentialsError::unhandled("unexpected response"))                .is_retryable()        );    }    #[tokio::test]    async fn sign_includes_session_token_when_credentials_have_one() {        let signed = test_context(Some("session-token"))            .sign_at(                test_request(),                UNIX_EPOCH + Duration::from_secs(1_700_000_000),            )            .await            .expect("request should sign");        assert_eq!(            signing::header_value(&signed.headers, "x-amz-security-token"),            Some("session-token".to_string())        );    }    #[tokio::test]    async fn load_rejects_empty_service_name() {        let err = AwsAuthContext::load(AwsAuthConfig {            profile: None,            region: None,            service: "   ".to_string(),        })        .await        .expect_err("empty service should be rejected");        assert_eq!(err.to_string(), "AWS service name must not be empty");    }}