Add recursion detection middleware to the default middleware stack (#1003) (cbd61ab1) · Commits · Public Repositories / Smithy Rust

.pre-commit-hooks/license-header.sh

+1 −1

Original line number	Diff line number	Diff line
		@@ -6,4 +6,4 @@ set -e
		# SPDX-License-Identifier: Apache-2.0.
		#

		cd "$(git rev-parse --show-toplevel)/tools/sdk-lints" && cargo run -- check --license
		cd "$(git rev-parse --show-toplevel)/tools/sdk-lints" && cargo run -- check --license --changelog

CHANGELOG.next.toml

+18 −0

Original line number	Diff line number	Diff line
		@@ -39,6 +39,24 @@ meta = { "breaking" = false, "tada" = false, "bug" = true }
		references = ["smithy-rs#998", "aws-sdk-rust#359"]
		author = "rcoh"

		[[aws-sdk-rust]]
		message = "`aws_http::AwsErrorRetryPolicy` was moved to `aws_http::retry::AwsErrorRetryPolicy`."
		meta = { "breaking" = true, "tada" = false, "bug" = false }
		references = ["smithy-rs#1003"]
		author = "rcoh"

		[[smithy-rs]]
		message = "The signature of `aws_smithy_protocol_test::validate_headers` was made more flexible but may require adjusting invocations slightly."
		meta = { "breaking" = true, "tada" = false, "bug" = false }
		references = ["smithy-rs#1003"]
		author = "rcoh"

		[[aws-sdk-rust]]
		message = "Add recursion detection middleware to the default stack"
		meta = { "breaking" = false, "tada" = false, "bug" = false }
		references = ["smithy-rs#1003"]
		author = "rcoh"

		[[aws-sdk-rust]]
		message = "aws_types::Config is now `Clone`"
		meta = { "breaking" = false, "tada" = false, "bug" = false }

aws/rust-runtime/aws-http/Cargo.toml

+5 −0

Original line number	Diff line number	Diff line
		@@ -14,14 +14,19 @@ aws-types = { path = "../aws-types" }
		http = "0.2.3"
		lazy_static = "1"
		tracing = "0.1"
		percent-encoding = "2.1.0"

		[dev-dependencies]
		async-trait = "0.1.50"
		aws-smithy-async = { path = "../../../rust-runtime/aws-smithy-async", features = ["rt-tokio"] }
		aws-smithy-protocol-test = { path = "../../../rust-runtime/aws-smithy-protocol-test" }
		env_logger = "0.9"
		http = "0.2.3"
		tokio = { version = "1", features = ["macros", "rt", "rt-multi-thread", "test-util"] }
		tracing-subscriber = { version = "0.2.16", features = ["fmt"] }
		proptest = "1"
		serde = { version = "1", features = ["derive"]}
		serde_json = "1"

		[package.metadata.docs.rs]
		all-features = true

aws/rust-runtime/aws-http/src/lib.rs

+6 −251

Original line number	Diff line number	Diff line
		@@ -16,256 +16,11 @@
		/// Credentials middleware
		pub mod auth;

		/// User agent middleware
		pub mod user_agent;

		use aws_smithy_http::result::SdkError;
		use aws_smithy_http::retry::ClassifyResponse;
		use aws_smithy_types::retry::{ErrorKind, ProvideErrorKind, RetryKind};
		use std::time::Duration;

		/// A retry policy that models AWS error codes as outlined in the SEP
		///
		/// In order of priority:
		/// 1. The `x-amz-retry-after` header is checked
		/// 2. The modeled error retry mode is checked
		/// 3. The code is checked against a predetermined list of throttling errors & transient error codes
		/// 4. The status code is checked against a predetermined list of status codes
		#[non_exhaustive]
		#[derive(Clone, Debug)]
		pub struct AwsErrorRetryPolicy;

		const TRANSIENT_ERROR_STATUS_CODES: &[u16] = &[500, 502, 503, 504];
		const THROTTLING_ERRORS: &[&str] = &[
		"Throttling",
		"ThrottlingException",
		"ThrottledException",
		"RequestThrottledException",
		"TooManyRequestsException",
		"ProvisionedThroughputExceededException",
		"TransactionInProgressException",
		"RequestLimitExceeded",
		"BandwidthLimitExceeded",
		"LimitExceededException",
		"RequestThrottled",
		"SlowDown",
		"PriorRequestNotComplete",
		"EC2ThrottledException",
		];
		const TRANSIENT_ERRORS: &[&str] = &["RequestTimeout", "RequestTimeoutException"];

		impl AwsErrorRetryPolicy {
		/// Create an `AwsErrorRetryPolicy` with the default set of known error & status codes
		pub fn new() -> Self {
		AwsErrorRetryPolicy
		}
		}

		impl Default for AwsErrorRetryPolicy {
		fn default() -> Self {
		Self::new()
		}
		}

		impl<T, E> ClassifyResponse<T, SdkError<E>> for AwsErrorRetryPolicy
		where
		E: ProvideErrorKind,
		{
		fn classify(&self, err: Result<&T, &SdkError<E>>) -> RetryKind {
		let (err, response) = match err {
		Ok(_) => return RetryKind::NotRetryable,
		Err(SdkError::ServiceError { err, raw }) => (err, raw),
		Err(SdkError::DispatchFailure(err)) => {
		return if err.is_timeout() \|\| err.is_io() {
		RetryKind::Error(ErrorKind::TransientError)
		} else if let Some(ek) = err.is_other() {
		RetryKind::Error(ek)
		} else {
		RetryKind::NotRetryable
		}
		}
		Err(_) => return RetryKind::NotRetryable,
		};
		if let Some(retry_after_delay) = response
		.http()
		.headers()
		.get("x-amz-retry-after")
		.and_then(\|header\| header.to_str().ok())
		.and_then(\|header\| header.parse::<u64>().ok())
		{
		return RetryKind::Explicit(Duration::from_millis(retry_after_delay));
		}
		if let Some(kind) = err.retryable_error_kind() {
		return RetryKind::Error(kind);
		};
		if let Some(code) = err.code() {
		if THROTTLING_ERRORS.contains(&code) {
		return RetryKind::Error(ErrorKind::ThrottlingError);
		}
		if TRANSIENT_ERRORS.contains(&code) {
		return RetryKind::Error(ErrorKind::TransientError);
		}
		};
		if TRANSIENT_ERROR_STATUS_CODES.contains(&response.http().status().as_u16()) {
		return RetryKind::Error(ErrorKind::TransientError);
		};
		// TODO(https://github.com/awslabs/smithy-rs/issues/966): IDPCommuncation error needs to be retried
		RetryKind::NotRetryable
		}
		}

		#[cfg(test)]
		mod test {
		use crate::AwsErrorRetryPolicy;
		use aws_smithy_http::body::SdkBody;
		use aws_smithy_http::operation;
		use aws_smithy_http::result::{SdkError, SdkSuccess};
		use aws_smithy_http::retry::ClassifyResponse;
		use aws_smithy_types::retry::{ErrorKind, ProvideErrorKind, RetryKind};
		use std::time::Duration;

		struct UnmodeledError;

		struct CodedError {
		code: &'static str,
		}

		impl ProvideErrorKind for UnmodeledError {
		fn retryable_error_kind(&self) -> Option<ErrorKind> {
		None
		}

		fn code(&self) -> Option<&str> {
		None
		}
		}

		impl ProvideErrorKind for CodedError {
		fn retryable_error_kind(&self) -> Option<ErrorKind> {
		None
		}

		fn code(&self) -> Option<&str> {
		Some(self.code)
		}
		}
		/// Recursion Detection middleware
		pub mod recursion_detection;

		fn make_err<E>(
		err: E,
		raw: http::Response<&'static str>,
		) -> Result<SdkSuccess<()>, SdkError<E>> {
		Err(SdkError::ServiceError {
		err,
		raw: operation::Response::new(raw.map(\|b\| SdkBody::from(b))),
		})
		}
		/// AWS-specific retry logic
		pub mod retry;

		#[test]
		fn not_an_error() {
		let policy = AwsErrorRetryPolicy::new();
		let test_response = http::Response::new("OK");
		assert_eq!(
		policy.classify(make_err(UnmodeledError, test_response).as_ref()),
		RetryKind::NotRetryable
		);
		}

		#[test]
		fn classify_by_response_status() {
		let policy = AwsErrorRetryPolicy::new();
		let test_resp = http::Response::builder()
		.status(500)
		.body("error!")
		.unwrap();
		assert_eq!(
		policy.classify(make_err(UnmodeledError, test_resp).as_ref()),
		RetryKind::Error(ErrorKind::TransientError)
		);
		}

		#[test]
		fn classify_by_response_status_not_retryable() {
		let policy = AwsErrorRetryPolicy::new();
		let test_resp = http::Response::builder()
		.status(408)
		.body("error!")
		.unwrap();
		assert_eq!(
		policy.classify(make_err(UnmodeledError, test_resp).as_ref()),
		RetryKind::NotRetryable
		);
		}

		#[test]
		fn classify_by_error_code() {
		let test_response = http::Response::new("OK");
		let policy = AwsErrorRetryPolicy::new();

		assert_eq!(
		policy.classify(make_err(CodedError { code: "Throttling" }, test_response).as_ref()),
		RetryKind::Error(ErrorKind::ThrottlingError)
		);

		let test_response = http::Response::new("OK");
		assert_eq!(
		policy.classify(
		make_err(
		CodedError {
		code: "RequestTimeout"
		},
		test_response
		)
		.as_ref()
		),
		RetryKind::Error(ErrorKind::TransientError)
		)
		}

		#[test]
		fn classify_generic() {
		let err = aws_smithy_types::Error::builder().code("SlowDown").build();
		let test_response = http::Response::new("OK");
		let policy = AwsErrorRetryPolicy::new();
		assert_eq!(
		policy.classify(make_err(err, test_response).as_ref()),
		RetryKind::Error(ErrorKind::ThrottlingError)
		);
		}

		#[test]
		fn classify_by_error_kind() {
		struct ModeledRetries;
		let test_response = http::Response::new("OK");
		impl ProvideErrorKind for ModeledRetries {
		fn retryable_error_kind(&self) -> Option<ErrorKind> {
		Some(ErrorKind::ClientError)
		}

		fn code(&self) -> Option<&str> {
		// code should not be called when `error_kind` is provided
		unimplemented!()
		}
		}

		let policy = AwsErrorRetryPolicy::new();

		assert_eq!(
		policy.classify(make_err(ModeledRetries, test_response).as_ref()),
		RetryKind::Error(ErrorKind::ClientError)
		);
		}

		#[test]
		fn test_retry_after_header() {
		let policy = AwsErrorRetryPolicy::new();
		let test_response = http::Response::builder()
		.header("x-amz-retry-after", "5000")
		.body("retry later")
		.unwrap();

		assert_eq!(
		policy.classify(make_err(UnmodeledError, test_response).as_ref()),
		RetryKind::Explicit(Duration::from_millis(5000))
		);
		}
		}
		/// User agent middleware
		pub mod user_agent;

aws/rust-runtime/aws-http/src/recursion_detection.rs

0 → 100644

+167 −0

Original line number	Diff line number	Diff line
		/*
		* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
		* SPDX-License-Identifier: Apache-2.0.
		*/

		use crate::recursion_detection::env::TRACE_ID;
		use aws_smithy_http::middleware::MapRequest;
		use aws_smithy_http::operation::Request;
		use aws_types::os_shim_internal::Env;
		use http::HeaderValue;
		use percent_encoding::{percent_encode, CONTROLS};
		use std::borrow::Cow;

		/// Recursion Detection Middleware
		///
		/// This middleware inspects the value of the `AWS_LAMBDA_FUNCTION_NAME` and `_X_AMZ_TRACE_ID` environment
		/// variables to detect if the request is being invoked in a lambda function. If it is, the `X-Amzn-Trace-Id` header
		/// will be set. This enables downstream services to prevent accidentally infinitely recursive invocations spawned
		/// from lambda.
		#[non_exhaustive]
		#[derive(Default, Debug, Clone)]
		pub struct RecursionDetectionStage {
		env: Env,
		}

		impl RecursionDetectionStage {
		/// Creates a new `RecursionDetectionStage`
		pub fn new() -> Self {
		Self::default()
		}
		}

		impl MapRequest for RecursionDetectionStage {
		type Error = std::convert::Infallible;

		fn apply(&self, request: Request) -> Result<Request, Self::Error> {
		request.augment(\|mut req, _conf\| {
		augument_request(&mut req, &self.env);
		Ok(req)
		})
		}
		}

		const TRACE_ID_HEADER: &str = "x-amzn-trace-id";

		mod env {
		pub(super) const LAMBDA_FUNCTION_NAME: &str = "AWS_LAMBDA_FUNCTION_NAME";
		pub(super) const TRACE_ID: &str = "_X_AMZ_TRACE_ID";
		}

		/// Set the trace id header from the request
		fn augument_request<B>(req: &mut http::Request<B>, env: &Env) {
		if req.headers().contains_key(TRACE_ID_HEADER) {
		return;
		}
		if let (Ok(_function_name), Ok(trace_id)) =
		(env.get(env::LAMBDA_FUNCTION_NAME), env.get(TRACE_ID))
		{
		req.headers_mut()
		.insert(TRACE_ID_HEADER, encode_header(trace_id.as_bytes()));
		}
		}

		/// Encodes a byte slice as a header.
		///
		/// ASCII control characters are percent encoded which ensures that all byte sequences are valid headers
		fn encode_header(value: &[u8]) -> HeaderValue {
		let value: Cow<'_, str> = percent_encode(value, CONTROLS).into();
		HeaderValue::from_bytes(value.as_bytes()).expect("header is encoded, header must be valid")
		}

		#[cfg(test)]
		mod test {
		use crate::recursion_detection::{encode_header, RecursionDetectionStage};
		use aws_smithy_http::body::SdkBody;
		use aws_smithy_http::middleware::MapRequest;
		use aws_smithy_http::operation;
		use aws_smithy_protocol_test::{assert_ok, validate_headers};
		use aws_types::os_shim_internal::Env;
		use http::HeaderValue;
		use proptest::{prelude::*, proptest};
		use serde::Deserialize;
		use std::collections::HashMap;

		proptest! {
		#[test]
		fn header_encoding_never_panics(s in any::<Vec<u8>>()) {
		encode_header(&s);
		}
		}

		#[test]
		fn every_char() {
		let buff = (0..=255).collect::<Vec<u8>>();
		assert_eq!(
		encode_header(&buff),
		HeaderValue::from_static(
		r##"%00%01%02%03%04%05%06%07%08%09%0A%0B%0C%0D%0E%0F%10%11%12%13%14%15%16%17%18%19%1A%1B%1C%1D%1E%1F !"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\]^_`abcdefghijklmnopqrstuvwxyz{\|}~%7F%80%81%82%83%84%85%86%87%88%89%8A%8B%8C%8D%8E%8F%90%91%92%93%94%95%96%97%98%99%9A%9B%9C%9D%9E%9F%A0%A1%A2%A3%A4%A5%A6%A7%A8%A9%AA%AB%AC%AD%AE%AF%B0%B1%B2%B3%B4%B5%B6%B7%B8%B9%BA%BB%BC%BD%BE%BF%C0%C1%C2%C3%C4%C5%C6%C7%C8%C9%CA%CB%CC%CD%CE%CF%D0%D1%D2%D3%D4%D5%D6%D7%D8%D9%DA%DB%DC%DD%DE%DF%E0%E1%E2%E3%E4%E5%E6%E7%E8%E9%EA%EB%EC%ED%EE%EF%F0%F1%F2%F3%F4%F5%F6%F7%F8%F9%FA%FB%FC%FD%FE%FF"##
		)
		);
		}

		#[test]
		fn run_tests() {
		let test_cases: Vec<TestCase> =
		serde_json::from_str(include_str!("../test-data/recursion-detection.json"))
		.expect("invalid test case");
		for test_case in test_cases {
		check(test_case)
		}
		}

		#[derive(Deserialize)]
		#[serde(rename_all = "camelCase")]
		struct TestCase {
		env: HashMap<String, String>,
		request_headers_before: Vec<String>,
		request_headers_after: Vec<String>,
		}

		impl TestCase {
		fn env(&self) -> Env {
		Env::from(self.env.clone())
		}

		/// Headers on the input request
		fn request_headers_before(&self) -> impl Iterator<Item = (&str, &str)> {
		Self::split_headers(&self.request_headers_before)
		}

		/// Headers on the output request
		fn request_headers_after(&self) -> impl Iterator<Item = (&str, &str)> {
		Self::split_headers(&self.request_headers_after)
		}

		/// Split text headers on `: `
		fn split_headers(headers: &[String]) -> impl Iterator<Item = (&str, &str)> {
		headers
		.iter()
		.map(\|header\| header.split_once(": ").expect("header must contain :"))
		}
		}

		fn check(test_case: TestCase) {
		let env = test_case.env();
		let mut req = http::Request::builder();
		for (k, v) in test_case.request_headers_before() {
		req = req.header(k, v);
		}
		let req = req.body(SdkBody::empty()).expect("must be valid");
		let req = operation::Request::new(req);
		let augmented_req = RecursionDetectionStage { env }
		.apply(req)
		.expect("stage must succeed");
		for k in augmented_req.http().headers().keys() {
		assert_eq!(
		augmented_req.http().headers().get_all(k).iter().count(),
		1,
		"No duplicated headers"
		)
		}
		assert_ok(validate_headers(
		&augmented_req.http(),
		test_case.request_headers_after(),
		))
		}
		}