[fix] Retry connection timeouts to IMDSv2 (#130)

Author: jplock

Cargo.lock

@@ -1,4 +1,4 @@
-aws-lambda-powertools[tracer,parser]==3.4.0
+aws-lambda-powertools[tracer,parser]==3.4.1
 cryptography==43.0.3
 hpke==0.3.2
 pksuid==1.1.2

api/dependencies/requirements.txt

@@ -1,6 +1,6 @@
 [tool.black]
 line-length = 120
-target-version = ['py312']
+target-version = ['py313']
 include = '\.pyi?$'
 extend-exclude = '''
 (

api/pyproject.toml

@@ -1,3 +1,3 @@
 black==24.10.0
-aws-lambda-powertools[all,aws-sdk]==3.4.0
+aws-lambda-powertools[all,aws-sdk]==3.4.1
 boto3-stubs[dynamodb,kms]

api/requirements-dev.txt

@@ -1 +1 @@
-aws-lambda-powertools==3.4.0
+aws-lambda-powertools==3.4.1

canary/dependencies/requirements.txt

@@ -1,6 +1,6 @@
 [tool.black]
 line-length = 120
-target-version = ['py312']
+target-version = ['py313']
 include = '\.pyi?$'
 extend-exclude = '''
 (

canary/pyproject.toml

@@ -1,3 +1,3 @@
 black==24.10.0
-aws-lambda-powertools[all,aws-sdk]==3.4.0
+aws-lambda-powertools[all,aws-sdk]==3.4.1
 boto3-stubs[dynamodb,kms]

canary/requirements-dev.txt

@@ -14,13 +14,13 @@ path = "src/main.rs"
 
 [dependencies]
 anyhow = { version = "=1.0.95", default-features = false }
-aws-lc-rs = { version = "=1.12.0", default-features = false }
+aws-lc-rs = { version = "=1.12.2", default-features = false }
 byteorder = { version = "=1.5.0", default-features = false }
 cel-interpreter = { version = "=0.9.0", default-features = false, features = ["json", "chrono"] }
 chrono = { version = "=0.4.39", default-features = false, features = ["now"] }
-data-encoding = { version = "=2.6.0", default-features = false, features = ["alloc"] }
+data-encoding = { version = "=2.7.0", default-features = false, features = ["alloc"] }
 serde = { version = "=1.0.217", default-features = false, features = ["derive"] }
-serde_json = { version = "=1.0.134", default-features = false }
-rustls = { version = "=0.23.20", default-features = false, features = ["aws_lc_rs"] }
+serde_json = { version = "=1.0.137", default-features = false }
+rustls = { version = "=0.23.21", default-features = false, features = ["aws_lc_rs"] }
 vsock = { version = "=0.5.1", default-features = false }
 zeroize = { version = "=1.8.1", default-features = false, features = ["zeroize_derive"] }

enclave/Cargo.toml

@@ -15,16 +15,17 @@ path = "src/main.rs"
 
 [dependencies]
 anyhow = { version = "=1.0.95", default-features = false }
-aws-config = { version = "=1.5.12", default-features = false, features = ["rustls", "rt-tokio", "behavior-version-latest"] }
+aws-config = { version = "=1.5.15", default-features = false, features = ["rustls", "rt-tokio", "behavior-version-latest"] }
 aws-credential-types = { version = "=1.2.1", default-features = false }
+aws-smithy-runtime-api = { version = "=1.7.3", default-features = false, features = ["client"] }
 axum = { version = "=0.8.1", default-features = false, features = ["http1", "json", "tokio", "tracing"] }
 byteorder = { version = "=1.5.0", default-features = false }
-clap = { version = "=4.5.23", default-features = false, features = ["std", "derive", "env"] }
+clap = { version = "=4.5.27", default-features = false, features = ["std", "derive", "env"] }
 fastrand = { version = "=2.3.0", default-features = false }
 serde = { version = "=1.0.217", default-features = false, features = ["derive"] }
-serde_json = { version = "=1.0.134", default-features = false }
-thiserror = { version = "=2.0.9", default-features = false }
-tokio = { version = "=1.42.0", default-features = false, features = ["rt-multi-thread", "process", "tracing"] }
+serde_json = { version = "=1.0.137", default-features = false }
+thiserror = { version = "=2.0.11", default-features = false }
+tokio = { version = "=1.43.0", default-features = false, features = ["rt-multi-thread", "process", "tracing"] }
 tracing = { version = "=0.1.41", default-features = false, features = ["log"] }
 tracing-subscriber = { version = "=0.3.19", default-features = false, features = ["ansi", "env-filter", "fmt", "json"] }
 vsock = { version = "=0.5.1", default-features = false }

parent/Cargo.toml

@@ -1,9 +1,10 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: MIT-0
 
-use aws_config::imds::client::Client;
+use aws_config::imds::client::{Client, ImdsResponseRetryClassifier};
 use aws_config::imds::credentials::ImdsCredentialsProvider;
 use aws_credential_types::provider::ProvideCredentials;
+use aws_smithy_runtime_api::client::retries::classifiers::SharedRetryClassifier;
 
 use crate::constants;
 use crate::errors::AppError;
@@ -14,6 +15,9 @@ pub async fn load_credentials(profile: Option<String>) -> Result<Credential, App
         .endpoint("http://169.254.169.254:80") // hardcode IMDS IPv4 address to avoid checking for credentials on the file system
         .expect("valid URL")
         .token_ttl(constants::IMDS_TOKEN_TTL)
+        .retry_classifier(SharedRetryClassifier::new(
+            ImdsResponseRetryClassifier::default().with_retry_connect_timeouts(true),
+        ))
         .build();
 
     let imds = {