Skip to content

Commit

Permalink
fix(remote_storage): use cached credentials (#3128)
Browse files Browse the repository at this point in the history
IMDSv2 has limits, and if we query it on every s3 interaction we are
going to go over those limits. Changes the s3_bucket client
configuration to use:
- ChainCredentialsProvider to handle env variables or imds usage
- LazyCachingCredentialsProvider to actually cache any credentials

Related: awslabs/aws-sdk-rust#629
Possibly related: #3118
  • Loading branch information
koivunej committed Dec 16, 2022
1 parent d8f6d6f commit 73ea0a0
Showing 1 changed file with 18 additions and 29 deletions.
47 changes: 18 additions & 29 deletions libs/remote_storage/src/s3_bucket.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,13 @@
//! allowing multiple api users to independently work with the same S3 bucket, if
//! their bucket prefixes are both specified and different.
use std::env::var;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use aws_config::{
environment::credentials::EnvironmentVariableCredentialsProvider, imds,
imds::credentials::ImdsCredentialsProvider, meta::credentials::provide_credentials_fn,
environment::credentials::EnvironmentVariableCredentialsProvider,
imds::credentials::ImdsCredentialsProvider,
meta::credentials::{CredentialsProviderChain, LazyCachingCredentialsProvider},
};
use aws_sdk_s3::{
config::Config,
Expand All @@ -20,7 +19,6 @@ use aws_sdk_s3::{
Client, Endpoint, Region,
};
use aws_smithy_http::body::SdkBody;
use aws_types::credentials::{CredentialsError, ProvideCredentials};
use hyper::Body;
use tokio::{io, sync::Semaphore};
use tokio_util::io::ReaderStream;
Expand All @@ -31,8 +29,6 @@ use crate::{
Download, DownloadError, RemotePath, RemoteStorage, S3Config, REMOTE_STORAGE_PREFIX_SEPARATOR,
};

const DEFAULT_IMDS_TIMEOUT: Duration = Duration::from_secs(10);

pub(super) mod metrics {
use metrics::{register_int_counter_vec, IntCounterVec};
use once_cell::sync::Lazy;
Expand Down Expand Up @@ -122,30 +118,23 @@ impl S3Bucket {
"Creating s3 remote storage for S3 bucket {}",
aws_config.bucket_name
);

let credentials_provider = {
// uses "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"
let env_creds = EnvironmentVariableCredentialsProvider::new();
// uses imds v2
let imds = ImdsCredentialsProvider::builder().build();

// finally add caching.
// this might change in future, see https://github.com/awslabs/aws-sdk-rust/issues/629
LazyCachingCredentialsProvider::builder()
.load(CredentialsProviderChain::first_try("env", env_creds).or_else("imds", imds))
.build()
};

let mut config_builder = Config::builder()
.region(Region::new(aws_config.bucket_region.clone()))
.credentials_provider(provide_credentials_fn(|| async {
match var("AWS_ACCESS_KEY_ID").is_ok() && var("AWS_SECRET_ACCESS_KEY").is_ok() {
true => {
EnvironmentVariableCredentialsProvider::new()
.provide_credentials()
.await
}
false => {
let imds_client = imds::Client::builder()
.connect_timeout(DEFAULT_IMDS_TIMEOUT)
.read_timeout(DEFAULT_IMDS_TIMEOUT)
.build()
.await
.map_err(CredentialsError::unhandled)?;
ImdsCredentialsProvider::builder()
.imds_client(imds_client)
.build()
.provide_credentials()
.await
}
}
}));
.credentials_provider(credentials_provider);

if let Some(custom_endpoint) = aws_config.endpoint.clone() {
let endpoint = Endpoint::immutable(
Expand Down

0 comments on commit 73ea0a0

Please sign in to comment.