From 72ef9f07a04ce013f90f71f20c2a6ac2b4c05806 Mon Sep 17 00:00:00 2001 From: Richard Davison Date: Tue, 14 Jan 2025 15:12:26 +0100 Subject: [PATCH] Impl aws sigv4 --- Cargo.lock | 165 ++++++++++++++++++++++++++++++++++++++++-------- Cargo.toml | 3 + src/aws_auth.rs | 129 +++++++++++++++++++++++++++++++++++++ src/client.rs | 36 ++++++++--- src/db.rs | 1 + src/main.rs | 40 +++++++++++- 6 files changed, 336 insertions(+), 38 deletions(-) create mode 100644 src/aws_auth.rs diff --git a/Cargo.lock b/Cargo.lock index 1cbf0132..9fb2203d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "actix-codec" @@ -255,6 +255,21 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "android-tzdata" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anstream" version = "0.6.18" @@ -409,6 +424,20 @@ dependencies = [ "paste", ] +[[package]] +name = "aws-sign-v4" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a35b1c56648ef2a4eefb5a9e4152bf05310c48c459b9e661a8ea80517e0c2d7" +dependencies = [ + "chrono", + "hex", + "http 1.2.0", + "ring", + "sha256", + "url", +] + [[package]] name = "axum" version = "0.8.1" @@ -493,7 +522,7 @@ dependencies = [ "bitflags", "cexpr", "clang-sys", - "itertools 0.12.1", + "itertools 0.11.0", "lazy_static", "lazycell", "log", @@ -557,9 +586,9 @@ checksum = "3eeab4423108c5d7c744f4d234de88d18d636100093ae04caf4825134b9c3a32" [[package]] name = "borsh" -version = "1.5.3" +version = "1.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2506947f73ad44e344215ccd6403ac2ae18cd8e046e581a441bf8d199f257f03" +checksum = "9fb65153674e51d3a42c8f27b05b9508cea85edfaade8aa46bc8fc18cecdfef3" dependencies = [ "borsh-derive", "cfg_aliases", @@ -567,9 +596,9 @@ dependencies = [ [[package]] name = "borsh-derive" -version = "1.5.3" +version = "1.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2593a3b8b938bd68373196c9832f516be11fa487ef4ae745eb282e6a56a7244" +checksum = "a396e17ad94059c650db3d253bb6e25927f1eb462eede7e7a153bb6e75dce0a7" dependencies = [ "once_cell", "proc-macro-crate", @@ -723,6 +752,20 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chrono" +version = "0.4.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e36cc9d416881d2e24f9a963be5fb1cd90966419ac844274161d10488b3e825" +dependencies = [ + "android-tzdata", + "iana-time-zone", + "js-sys", + "num-traits", + "wasm-bindgen", + "windows-targets 0.52.6", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -1456,6 +1499,12 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + [[package]] name = "hickory-proto" version = "0.24.2" @@ -1646,6 +1695,29 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.61" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "235e081f3925a06703c2d0117ea8b91f042756fd6e7a6e5d901e8ca1a996b220" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "wasm-bindgen", + "windows-core 0.52.0", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "1.5.0" @@ -1852,9 +1924,9 @@ checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" [[package]] name = "itertools" -version = "0.12.1" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" dependencies = [ "either", ] @@ -1885,9 +1957,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.76" +version = "0.3.77" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6717b6b5b077764fb5966237269cb3c64edddde4b14ce42647430a78ced9e7b7" +checksum = "1cfaf33c695fc6e08064efbc1f72ec937429614f25eef83af942d0e227c3a28f" dependencies = [ "once_cell", "wasm-bindgen", @@ -2320,10 +2392,12 @@ dependencies = [ "assert_cmd", "average", "aws-lc-rs", + "aws-sign-v4", "axum", "base64", "byte-unit", "bytes", + "chrono", "clap", "float-cmp", "float-ord", @@ -2543,9 +2617,9 @@ dependencies = [ [[package]] name = "prettyplease" -version = "0.2.27" +version = "0.2.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "483f8c21f64f3ea09fe0f30f5d48c3e8eefe5dac9129f0075f76593b4c1da705" +checksum = "6924ced06e1f7dfe3fa48d57b9f74f55d8915f5036121bef647ef4b204895fac" dependencies = [ "proc-macro2", "syn 2.0.96", @@ -3149,6 +3223,30 @@ dependencies = [ "digest", ] +[[package]] +name = "sha2" +version = "0.10.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "793db75ad2bcafc3ffa7c68b215fee268f537982cd901d132f89c6343f3a3dc8" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sha256" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18278f6a914fa3070aa316493f7d2ddfb9ac86ebc06fa3b83bffda487e9065b0" +dependencies = [ + "async-trait", + "bytes", + "hex", + "sha2", + "tokio", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -3956,20 +4054,21 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" [[package]] name = "wasm-bindgen" -version = "0.2.99" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a474f6281d1d70c17ae7aa6a613c87fce69a127e2624002df63dcb39d6cf6396" +checksum = "1edc8929d7499fc4e8f0be2262a241556cfc54a0bea223790e71446f2aab1ef5" dependencies = [ "cfg-if", "once_cell", + "rustversion", "wasm-bindgen-macro", ] [[package]] name = "wasm-bindgen-backend" -version = "0.2.99" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5f89bb38646b4f81674e8f5c3fb81b562be1fd936d84320f3264486418519c79" +checksum = "2f0a0651a5c2bc21487bde11ee802ccaf4c51935d0d3d42a6101f98161700bc6" dependencies = [ "bumpalo", "log", @@ -3981,9 +4080,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.49" +version = "0.4.50" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38176d9b44ea84e9184eff0bc34cc167ed044f816accfe5922e54d84cf48eca2" +checksum = "555d470ec0bc3bb57890405e5d4322cc9ea83cebb085523ced7be4144dac1e61" dependencies = [ "cfg-if", "js-sys", @@ -3994,9 +4093,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.99" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cc6181fd9a7492eef6fef1f33961e3695e4579b9872a6f7c83aee556666d4fe" +checksum = "7fe63fc6d09ed3792bd0897b314f53de8e16568c2b3f7982f468c0bf9bd0b407" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -4004,9 +4103,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.99" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30d7a95b763d3c45903ed6c81f156801839e5ee968bb07e534c44df0fcd330c2" +checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" dependencies = [ "proc-macro2", "quote", @@ -4017,15 +4116,18 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.99" +version = "0.2.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "943aab3fdaaa029a6e0271b35ea10b72b943135afe9bffca82384098ad0e06a6" +checksum = "1a05d73b933a847d6cccdda8f838a22ff101ad9bf93e33684f39c1f5f0eece3d" +dependencies = [ + "unicode-ident", +] [[package]] name = "web-sys" -version = "0.3.76" +version = "0.3.77" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04dd7223427d52553d3702c004d3b2fe07c148165faa56313cb00211e31c12bc" +checksum = "33b6dd2ef9186f1f2072e409e99cd22a975331a6b3591b12c764e0e55c60d5d2" dependencies = [ "js-sys", "wasm-bindgen", @@ -4077,7 +4179,16 @@ version = "0.58.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dd04d41d93c4992d421894c18c8b43496aa748dd4c081bac0dc93eb0489272b6" dependencies = [ - "windows-core", + "windows-core 0.58.0", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-core" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" +dependencies = [ "windows-targets 0.52.6", ] diff --git a/Cargo.toml b/Cargo.toml index 707175e8..9b43de91 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,6 +45,9 @@ tokio = { version = "1.38.1", features = ["full"] } ratatui = { version = "0.29.0", default-features = false, features = [ "crossterm", ] } +aws-sign-v4 = "0.3" +chrono = "0.4" +bytes = "1" hyper = { version = "1.4", features = ["client", "http1", "http2"] } diff --git a/src/aws_auth.rs b/src/aws_auth.rs new file mode 100644 index 00000000..aa37c6b6 --- /dev/null +++ b/src/aws_auth.rs @@ -0,0 +1,129 @@ +use crate::client::ClientError; +use anyhow::Result; + +use bytes::Bytes; +use hyper::{ + header::{self, HeaderName}, + HeaderMap, +}; +use url::Url; + +pub struct AwsSignatureConfig { + pub access_key: String, + pub secret_key: String, + pub session_token: Option, + pub service: String, + pub region: String, +} + +// Initialize unsignable headers as a static constant +static UNSIGNABLE_HEADERS: [HeaderName; 8] = [ + header::ACCEPT, + header::ACCEPT_ENCODING, + header::USER_AGENT, + header::EXPECT, + header::RANGE, + header::CONNECTION, + HeaderName::from_static("presigned-expires"), + HeaderName::from_static("x-amzn-trace-id"), +]; + +impl AwsSignatureConfig { + pub fn sign_request( + &self, + method: &str, + headers: &mut HeaderMap, + url: &Url, + body: Option, + ) -> Result<(), ClientError> { + let datetime = chrono::Utc::now(); + + let header_amz_date = datetime + .format("%Y%m%dT%H%M%SZ") + .to_string() + .parse() + .unwrap(); + + if !headers.contains_key(header::HOST) { + let host = url + .host_str() + .ok_or_else(|| ClientError::SigV4Error("URL must contain a host"))?; + headers.insert( + header::HOST, + host.parse() + .map_err(|_| ClientError::SigV4Error("Invalid host header name"))?, + ); + } + headers.insert("x-amz-date", header_amz_date); + + if let Some(session_token) = &self.session_token { + headers.insert("x-amz-security-token", session_token.parse().unwrap()); + } + + headers.remove(header::AUTHORIZATION); + + //remove and store headers in a vec from unsignable_headers + let removed_headers: Vec<(header::HeaderName, header::HeaderValue)> = UNSIGNABLE_HEADERS + .iter() + .filter_map(|k| headers.remove(k).map(|v| (k.clone(), v))) + .collect(); + + let body = body.as_deref().unwrap_or_default(); + headers.insert( + header::CONTENT_LENGTH, + body.len().to_string().parse().unwrap(), + ); + + let aws_sign = aws_sign_v4::AwsSign::new( + method, + url.as_str(), + &datetime, + headers, + &self.region, + &self.access_key, + &self.secret_key, + &self.service, + body, + ); + + let signature = aws_sign.sign(); + + //insert headers + for (key, value) in removed_headers { + headers.insert(key, value); + } + + headers.insert( + header::AUTHORIZATION, + signature + .parse() + .map_err(|_| ClientError::SigV4Error("Invalid authorization header name"))?, + ); + + Ok(()) + } + + pub fn new( + access_key: &str, + secret_key: &str, + signing_params: &str, + session_token: Option, + ) -> Result { + let parts: Vec<&str> = signing_params + .strip_prefix("aws:amz:") + .unwrap_or_default() + .split(':') + .collect(); + if parts.len() != 2 { + anyhow::bail!("Invalid AWS signing params format. Expected aws:amz:region:service"); + } + + Ok(Self { + access_key: access_key.into(), + secret_key: secret_key.into(), + session_token, + region: parts[0].to_string(), + service: parts[1].to_string(), + }) + } +} diff --git a/src/client.rs b/src/client.rs index f5f0dfee..96c0a048 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,3 +1,4 @@ +use bytes::Bytes; use http_body_util::{BodyExt, Full}; use hyper::{http, Method}; use hyper_util::rt::{TokioExecutor, TokioIo}; @@ -18,13 +19,14 @@ use tokio::{ use url::{ParseError, Url}; use crate::{ + aws_auth::AwsSignatureConfig, pcg64si::Pcg64Si, url_generator::{UrlGenerator, UrlGeneratorError}, ConnectToEntry, }; -type SendRequestHttp1 = hyper::client::conn::http1::SendRequest>; -type SendRequestHttp2 = hyper::client::conn::http2::SendRequest>; +type SendRequestHttp1 = hyper::client::conn::http1::SendRequest>; +type SendRequestHttp2 = hyper::client::conn::http2::SendRequest>; #[derive(Debug, Clone, Copy)] pub struct ConnectionTime { @@ -161,6 +163,8 @@ pub enum ClientError { UrlGeneratorError(#[from] UrlGeneratorError), #[error(transparent)] UrlParseError(#[from] ParseError), + #[error("AWS SigV4 signature error: {0}")] + SigV4Error(&'static str), } pub struct Client { @@ -176,6 +180,7 @@ pub struct Client { pub disable_keepalive: bool, pub insecure: bool, pub proxy_url: Option, + pub aws_config: Option, #[cfg(unix)] pub unix_socket: Option, #[cfg(feature = "vsock")] @@ -540,7 +545,7 @@ impl Client { } #[inline] - fn request(&self, url: &Url) -> Result>, ClientError> { + fn request(&self, url: &Url) -> Result>, ClientError> { let use_proxy = self.proxy_url.is_some() && url.scheme() == "http"; let mut builder = http::Request::builder() @@ -556,15 +561,28 @@ impl Client { self.http_version }); - *builder - .headers_mut() - .ok_or(ClientError::GetHeaderFromBuilderError)? = self.headers.clone(); + let bytes = self.body.map(Bytes::from_static); - if let Some(body) = self.body { - Ok(builder.body(Full::new(body))?) + let body = if let Some(body) = &bytes { + Full::new(body.clone()) } else { - Ok(builder.body(Full::default())?) + Full::default() + }; + + let mut headers = self.headers.clone(); + + // Apply AWS SigV4 if configured + if let Some(aws_config) = &self.aws_config { + aws_config.sign_request(self.method.as_str(), &mut headers, url, bytes)? } + + *builder + .headers_mut() + .ok_or(ClientError::GetHeaderFromBuilderError)? = headers; + + let request = builder.body(body)?; + + Ok(request) } async fn work_http1( diff --git a/src/db.rs b/src/db.rs index 7b82ff01..20b041e5 100644 --- a/src/db.rs +++ b/src/db.rs @@ -88,6 +88,7 @@ mod test_db { disable_keepalive: false, insecure: false, proxy_url: None, + aws_config: None, #[cfg(unix)] unix_socket: None, #[cfg(feature = "vsock")] diff --git a/src/main.rs b/src/main.rs index 44a8a179..571ef6c2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ use anyhow::Context; +use aws_auth::AwsSignatureConfig; use clap::Parser; use crossterm::tty::IsTty; use hickory_resolver::config::{ResolverConfig, ResolverOpts}; @@ -24,6 +25,7 @@ use std::{ use url::Url; use url_generator::UrlGenerator; +mod aws_auth; mod client; mod db; mod histogram; @@ -151,8 +153,18 @@ Note: If qps is specified, burst will be ignored", body_path: Option, #[arg(help = "Content-Type.", short = 'T')] content_type: Option, - #[arg(help = "Basic authentication, username:password", short = 'a')] + #[arg( + help = "Basic authentication or AWS credentials, username:password", + short = 'a' + )] basic_auth: Option, + #[arg(help = "AWS session token", long = "aws-session")] + aws_session: Option, + #[arg( + help = "AWS SigV4 signing params (format: aws:amz:region:service)", + long = "aws-sigv4" + )] + aws_sigv4: Option, #[arg(help = "HTTP proxy", short = 'x')] proxy: Option, #[arg( @@ -302,9 +314,32 @@ impl FromStr for VsockAddr { } async fn run() -> anyhow::Result<()> { - let opts: Opts = Opts::parse(); + let mut opts: Opts = Opts::parse(); let work_mode = opts.work_mode(); + // Parse AWS credentials from basic auth if AWS signing is requested + let aws_config = if let Some(signing_params) = opts.aws_sigv4 { + if let Some(auth) = &opts.basic_auth { + let parts: Vec<&str> = auth.split(':').collect(); + if parts.len() != 2 { + anyhow::bail!("Invalid AWS credentials format. Expected access_key:secret_key"); + } + let access_key = parts[0]; + let secret_key = parts[1]; + let session_token = opts.aws_session.take(); + Some(AwsSignatureConfig::new( + access_key, + secret_key, + &signing_params, + session_token, + )?) + } else { + anyhow::bail!("AWS credentials (--auth) required when using --aws-sigv4"); + } + } else { + None + }; + let parse_http_version = |is_http2: bool, version: Option<&str>| match (is_http2, version) { (true, Some(_)) => anyhow::bail!("--http2 and --http-version are exclusive"), (true, None) => Ok(http::Version::HTTP_2), @@ -487,6 +522,7 @@ async fn run() -> anyhow::Result<()> { let resolver = hickory_resolver::AsyncResolver::tokio(config, resolver_opts); let client = Arc::new(client::Client { + aws_config, http_version, proxy_http_version, url_generator,