car-inference 0.61.2

//! Minimal AWS Signature Version 4 signer for the Bedrock Converse protocol.
//!
//! Hand-rolled over `hmac`/`sha2` — deliberately NOT the ~200-crate AWS SDK.
//! Covers the header-based signing `bedrock-runtime` needs: a POST with a JSON
//! body, region/service scope, and an optional session token. Validated against
//! AWS's published SigV4 `get-vanilla` test vector (see tests).

use hmac::{Hmac, Mac};
use sha2::{Digest, Sha256};

type HmacSha256 = Hmac<Sha256>;

/// AWS credentials for signing. `session_token` is set for temporary
/// (STS / role) credentials and added as `x-amz-security-token`.
pub struct AwsCredentials {
    pub access_key_id: String,
    pub secret_access_key: String,
    pub session_token: Option<String>,
}

impl AwsCredentials {
    /// Resolve the standard AWS credentials env-first, then the OS keychain
    /// (`car_secrets::resolve_env_or_keychain`), so a native user can store them
    /// with `car keys` instead of exporting env vars. Returns `None` when the
    /// required key pair is absent, so callers can surface a clear "configure AWS
    /// credentials" error instead of signing with empties.
    pub fn from_env() -> Option<Self> {
        let access_key_id =
            car_secrets::resolve_env_or_keychain("AWS_ACCESS_KEY_ID").filter(|s| !s.is_empty())?;
        let secret_access_key = car_secrets::resolve_env_or_keychain("AWS_SECRET_ACCESS_KEY")
            .filter(|s| !s.is_empty())?;
        let session_token =
            car_secrets::resolve_env_or_keychain("AWS_SESSION_TOKEN").filter(|s| !s.is_empty());
        Some(Self {
            access_key_id,
            secret_access_key,
            session_token,
        })
    }
}

fn hmac(key: &[u8], data: &[u8]) -> Vec<u8> {
    let mut mac = HmacSha256::new_from_slice(key).expect("HMAC accepts any key length");
    mac.update(data);
    mac.finalize().into_bytes().to_vec()
}

fn sha256_hex(data: &[u8]) -> String {
    let mut h = Sha256::new();
    h.update(data);
    hex::encode(h.finalize())
}

/// Percent-encode one path segment per RFC 3986 (SigV4 canonical-URI rules):
/// the unreserved set `A-Z a-z 0-9 - _ . ~` passes through; everything else
/// (notably `:`) becomes `%XX`. Bedrock model IDs contain `:` (e.g.
/// `anthropic.claude-...-v2:0`), so the path segment MUST be encoded
/// identically in the signed canonical request and the actual request URL —
/// otherwise the signature won't match.
pub fn uri_encode_segment(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for b in s.bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                out.push(b as char)
            }
            _ => out.push_str(&format!("%{:02X}", b)),
        }
    }
    out
}

/// The canonical URI SigV4 signs. Every service except S3 URI-encodes each path
/// segment a second time (`%3A` in the sent URL becomes `%253A`); S3 signs the
/// path exactly as sent. `/` separators are kept.
fn canonical_uri(service: &str, request_path: &str) -> String {
    if service == "s3" {
        return request_path.to_string();
    }
    request_path
        .split('/')
        .map(uri_encode_segment)
        .collect::<Vec<_>>()
        .join("/")
}

/// Compute the SigV4 header set for a request. `canonical_path` is the
/// already-percent-encoded path exactly as it appears in the request URL
/// (e.g. `/model/<encoded id>/converse`); for every service but S3 the signer
/// encodes it a second time for the canonical request, so callers pass the
/// path they SEND, never the doubly-encoded form.
/// `base_headers` must include `host`. Returns the full header set to SEND:
/// the base headers plus `x-amz-date`, `authorization`, and (when a session
/// token is present) `x-amz-security-token`. `amz_date` is `YYYYMMDDTHHMMSSZ`;
/// `date_stamp` is `YYYYMMDD` (the date portion of the same instant).
#[allow(clippy::too_many_arguments)]
pub fn signed_headers(
    creds: &AwsCredentials,
    region: &str,
    service: &str,
    method: &str,
    canonical_path: &str,
    canonical_query: &str,
    base_headers: &[(String, String)],
    body: &[u8],
    amz_date: &str,
    date_stamp: &str,
) -> Vec<(String, String)> {
    // Header set that is both signed and sent.
    let mut headers: Vec<(String, String)> = base_headers.to_vec();
    headers.push(("x-amz-date".into(), amz_date.to_string()));
    if let Some(tok) = &creds.session_token {
        headers.push(("x-amz-security-token".into(), tok.clone()));
    }

    // Canonical headers: lowercase name, trimmed value, sorted by name, each
    // rendered "name:value\n". signed_headers: the same names, ';'-joined.
    let mut canon: Vec<(String, String)> = headers
        .iter()
        .map(|(k, v)| (k.to_lowercase(), v.trim().to_string()))
        .collect();
    canon.sort_by(|a, b| a.0.cmp(&b.0));
    let canonical_headers: String = canon.iter().map(|(k, v)| format!("{k}:{v}\n")).collect();
    let signed_header_names = canon
        .iter()
        .map(|(k, _)| k.clone())
        .collect::<Vec<_>>()
        .join(";");

    let payload_hash = sha256_hex(body);
    let canonical_path = canonical_uri(service, canonical_path);

    // Note the blank line: `canonical_headers` already ends in '\n', then the
    // spec adds another before `signed_header_names`.
    let canonical_request = format!(
        "{method}\n{canonical_path}\n{canonical_query}\n{canonical_headers}\n{signed_header_names}\n{payload_hash}"
    );

    let scope = format!("{date_stamp}/{region}/{service}/aws4_request");
    let string_to_sign = format!(
        "AWS4-HMAC-SHA256\n{amz_date}\n{scope}\n{}",
        sha256_hex(canonical_request.as_bytes())
    );

    let k_date = hmac(
        format!("AWS4{}", creds.secret_access_key).as_bytes(),
        date_stamp.as_bytes(),
    );
    let k_region = hmac(&k_date, region.as_bytes());
    let k_service = hmac(&k_region, service.as_bytes());
    let k_signing = hmac(&k_service, b"aws4_request");
    let signature = hex::encode(hmac(&k_signing, string_to_sign.as_bytes()));

    let authorization = format!(
        "AWS4-HMAC-SHA256 Credential={}/{}, SignedHeaders={}, Signature={}",
        creds.access_key_id, scope, signed_header_names, signature
    );
    headers.push(("authorization".into(), authorization));
    headers
}

#[cfg(test)]
mod tests {
    use super::*;

    /// AWS's published SigV4 `get-vanilla` test vector. If this matches, the
    /// canonical-request / string-to-sign / signing-key / signature chain is
    /// correct — independent of Bedrock.
    #[test]
    fn matches_aws_get_vanilla_vector() {
        let creds = AwsCredentials {
            access_key_id: "AKIDEXAMPLE".into(),
            secret_access_key: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY".into(),
            session_token: None,
        };
        let base = vec![("host".to_string(), "example.amazonaws.com".to_string())];
        let headers = signed_headers(
            &creds,
            "us-east-1",
            "service",
            "GET",
            "/",
            "",
            &base,
            b"",
            "20150830T123600Z",
            "20150830",
        );
        let auth = headers
            .iter()
            .find(|(k, _)| k == "authorization")
            .map(|(_, v)| v.as_str())
            .unwrap();
        assert_eq!(
            auth,
            "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20150830/us-east-1/service/aws4_request, \
             SignedHeaders=host;x-amz-date, \
             Signature=5fa00fa31553b73ebf1942676e86291e8372ff2a2260956d9b8aae1d763fbf31"
        );
    }

    /// Known answer for a Bedrock Converse POST whose model id contains ':'.
    /// The expected Authorization header was NOT produced by this signer. It is
    /// the one AWS CLI 2.31.8 (botocore) signed, offline against a refused
    /// loopback endpoint with dummy credentials. Regenerate with:
    ///   AWS_ACCESS_KEY_ID=AKIDEXAMPLE \
    ///   AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY \
    ///   aws --debug --region us-east-1 --endpoint-url http://127.0.0.1:9 \
    ///     bedrock-runtime converse --model-id 'anthropic.claude-3-5-sonnet-20240620-v1:0' \
    ///     --messages '[{"role":"user","content":[{"text":"hi"}]}]'
    /// and read the `CanonicalRequest:` / `Authorization` debug lines (the date
    /// below is the one that run printed). botocore signed the path
    /// `.../v1%253A0/converse` while sending `.../v1%3A0/converse`.
    #[test]
    fn matches_aws_cli_bedrock_converse_with_colon_model_id() {
        let creds = AwsCredentials {
            access_key_id: "AKIDEXAMPLE".into(),
            secret_access_key: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY".into(),
            session_token: None,
        };
        let base = vec![
            ("content-type".to_string(), "application/json".to_string()),
            ("host".to_string(), "127.0.0.1:9".to_string()),
        ];
        // The exact bytes the CLI sent (Python json.dumps spacing, 61 bytes).
        let body = br#"{"messages": [{"role": "user", "content": [{"text": "hi"}]}]}"#;
        let path = format!(
            "/model/{}/converse",
            uri_encode_segment("anthropic.claude-3-5-sonnet-20240620-v1:0")
        );
        assert_eq!(
            path,
            "/model/anthropic.claude-3-5-sonnet-20240620-v1%3A0/converse"
        );
        let headers = signed_headers(
            &creds,
            "us-east-1",
            "bedrock",
            "POST",
            &path,
            "",
            &base,
            body,
            "20261018T133135Z",
            "20261018",
        );
        let auth = headers
            .iter()
            .find(|(k, _)| k == "authorization")
            .map(|(_, v)| v.as_str())
            .unwrap();
        assert_eq!(
            auth,
            "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20261018/us-east-1/bedrock/aws4_request, \
             SignedHeaders=content-type;host;x-amz-date, \
             Signature=a4acf1fa09b45d7dd9662bf742e5eda1bd21dbf0aac843a26b1c33e8cd566e73"
        );
    }

    #[test]
    fn canonical_uri_double_encodes_except_for_s3() {
        assert_eq!(
            canonical_uri("bedrock", "/model/a%3Ab/converse"),
            "/model/a%253Ab/converse"
        );
        assert_eq!(canonical_uri("bedrock", "/"), "/");
        assert_eq!(canonical_uri("s3", "/bucket/a%3Ab"), "/bucket/a%3Ab");
    }

    #[test]
    fn uri_encode_handles_bedrock_model_ids() {
        // The `:` in a Bedrock model id must encode to %3A.
        assert_eq!(
            uri_encode_segment("anthropic.claude-3-5-sonnet-20241022-v2:0"),
            "anthropic.claude-3-5-sonnet-20241022-v2%3A0"
        );
    }

    #[test]
    fn session_token_is_added_to_signed_headers() {
        let creds = AwsCredentials {
            access_key_id: "AKID".into(),
            secret_access_key: "secret".into(),
            session_token: Some("token123".into()),
        };
        let base = vec![("host".to_string(), "h".to_string())];
        let headers = signed_headers(
            &creds,
            "us-east-1",
            "bedrock",
            "POST",
            "/",
            "",
            &base,
            b"{}",
            "20150830T123600Z",
            "20150830",
        );
        assert!(headers
            .iter()
            .any(|(k, v)| k == "x-amz-security-token" && v == "token123"));
        // ...and it's covered by SignedHeaders.
        let auth = headers.iter().find(|(k, _)| k == "authorization").unwrap();
        assert!(auth.1.contains("x-amz-security-token"));
    }
}