Skip to content

Commit

Permalink
refactor: Extract a method to set headers for signing
Browse files Browse the repository at this point in the history
When signing a URL by adding query parameters, the headers are needed as
part of the signed string.

When signing a request, the headers need to be set on the request as
well.
  • Loading branch information
carols10cents committed Sep 28, 2023
1 parent af87fd3 commit 9b5b417
Showing 1 changed file with 47 additions and 39 deletions.
86 changes: 47 additions & 39 deletions object_store/src/aws/credential.rs
Original file line number Diff line number Diff line change
Expand Up @@ -124,19 +124,7 @@ impl<'a> AwsAuthorizer<'a> {
///
/// [AWS SigV4]: https://docs.aws.amazon.com/IAM/latest/UserGuide/create-signed-request.html
pub fn authorize(&self, request: &mut Request, pre_calculated_digest: Option<&[u8]>) {
if let Some(ref token) = self.credential.token {
let token_val = HeaderValue::from_str(token).unwrap();
request.headers_mut().insert(TOKEN_HEADER, token_val);
}

let host = &request.url()[url::Position::BeforeHost..url::Position::AfterPort];
let host_val = HeaderValue::from_str(host).unwrap();
request.headers_mut().insert("host", host_val);

let date = self.date.unwrap_or_else(Utc::now);
let date_str = date.format("%Y%m%dT%H%M%SZ").to_string();
let date_val = HeaderValue::from_str(&date_str).unwrap();
request.headers_mut().insert(DATE_HEADER, date_val);

let digest = match self.sign_payload {
false => UNSIGNED_PAYLOAD.to_string(),
Expand All @@ -152,8 +140,13 @@ impl<'a> AwsAuthorizer<'a> {
},
};

let header_digest = HeaderValue::from_str(&digest).unwrap();
request.headers_mut().insert(HASH_HEADER, header_digest);
// This allocation is needed because the borrow checker can't see that the mutable borrow of request's headers
// is disjoint from immutably borrowing the request's URL if `set_headers` takes a `&Url` instead...
let host = &request.url()[url::Position::BeforeHost..url::Position::AfterPort]
.to_string();

self.set_headers(request.headers_mut(), &host, date, &digest);
let (signed_headers, canonical_headers) = canonicalize_headers(request.headers());

// Each path segment must be URI-encoded twice (except for Amazon S3 which only gets URI-encoded once).
// see https://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html
Expand All @@ -163,7 +156,6 @@ impl<'a> AwsAuthorizer<'a> {
.to_string(),
};

let (signed_headers, canonical_headers) = canonicalize_headers(request.headers());
let canonical_query = canonicalize_query(request.url());

let scope = self.scope(date);
Expand Down Expand Up @@ -202,28 +194,15 @@ impl<'a> AwsAuthorizer<'a> {

let digest = UNSIGNED_PAYLOAD;

let canonical_uri = match self.service {
"s3" => url.path().to_string(),
_ => utf8_percent_encode(url.path(), &STRICT_PATH_ENCODE_SET)
.to_string(),
};
let host = &url[url::Position::BeforeHost..url::Position::AfterPort].to_string();

let mut headers = HeaderMap::new();
self.set_headers(&mut headers, &host, date, digest);

if let Some(ref token) = self.credential.token {
let token_val = HeaderValue::from_str(token).unwrap();
headers.insert(TOKEN_HEADER, token_val);
}
let host = &url[url::Position::BeforeHost..url::Position::AfterPort];
let host_val = HeaderValue::from_str(host).unwrap();
headers.insert("host", host_val);

let date_str = date.format("%Y%m%dT%H%M%SZ").to_string();
let date_val = HeaderValue::from_str(&date_str).unwrap();
headers.insert(DATE_HEADER, date_val);

let header_digest = HeaderValue::from_str(digest).unwrap();
headers.insert(HASH_HEADER, header_digest);
let canonical_uri = match self.service {
"s3" => url.path().to_string(),
_ => utf8_percent_encode(url.path(), &STRICT_PATH_ENCODE_SET).to_string(),
};

let (signed_headers, canonical_headers) = canonicalize_headers(&headers);
let canonical_query = canonicalize_query(url);
Expand All @@ -241,17 +220,43 @@ impl<'a> AwsAuthorizer<'a> {
digest,
);

let signature = self.credential
let signature =
self.credential
.sign(&string_to_sign, date, self.region, self.service);

url.query_pairs_mut()
.append_pair("X-Amz-Algorithm", ALGORITHM)
.append_pair("X-Amz-Credential", &format!("{}/{}", self.credential.key_id, scope))
.append_pair("X-Amz-Date", &date_str)
.append_pair(
"X-Amz-Credential",
&format!("{}/{}", self.credential.key_id, scope),
)
.append_pair("X-Amz-Date", &date.format("%Y%m%dT%H%M%SZ").to_string())
.append_pair("X-Amz-SignedHeaders", "host;x-amz-date")
.append_pair("X-Amz-Signature", &signature);
}

fn set_headers(
&self,
headers: &mut HeaderMap,
host: &str,
date: DateTime<Utc>,
digest: &str,
) {
if let Some(ref token) = self.credential.token {
let token_val = HeaderValue::from_str(token).unwrap();
headers.insert(TOKEN_HEADER, token_val);
}
let host_val = HeaderValue::from_str(host).unwrap();
headers.insert("host", host_val);

let date_str = date.format("%Y%m%dT%H%M%SZ").to_string();
let date_val = HeaderValue::from_str(&date_str).unwrap();
headers.insert(DATE_HEADER, date_val);

let header_digest = HeaderValue::from_str(digest).unwrap();
headers.insert(HASH_HEADER, header_digest);
}

fn string_to_sign(
&self,
date: DateTime<Utc>,
Expand Down Expand Up @@ -759,11 +764,14 @@ mod tests {

let mut url = Url::parse("https://ec2.amazon.com/").unwrap();
authorizer.sign(Method::GET, &mut url);
assert_eq!(url.query().unwrap(), "X-Amz-Algorithm=AWS4-HMAC-SHA256&\
assert_eq!(
url.query().unwrap(),
"X-Amz-Algorithm=AWS4-HMAC-SHA256&\
X-Amz-Credential=AKIAIOSFODNN7EXAMPLE%2F20220806%2Fus-east-1%2Fec2%2Faws4_request&\
X-Amz-Date=20220806T180134Z&\
X-Amz-SignedHeaders=host%3Bx-amz-date&\
X-Amz-Signature=653c3d8ea261fd826207df58bc2bb69fbb5003e9eb3c0ef06e4a51f2a81d8699");
X-Amz-Signature=653c3d8ea261fd826207df58bc2bb69fbb5003e9eb3c0ef06e4a51f2a81d8699"
);
}

#[test]
Expand Down

0 comments on commit 9b5b417

Please sign in to comment.