Skip to content

Commit 4e547e6

Browse files
authored
Use DefaultCredentialsChain AWS authentication in remote_storage (#8440)
PR #8299 has switched the storage scrubber to use `DefaultCredentialsChain`. Now we do this for `remote_storage`, as it allows us to use `remote_storage` from inside kubernetes. Most of the diff is due to `GenericRemoteStorage::from_config` becoming `async fn`.
1 parent 3d582b2 commit 4e547e6

File tree

23 files changed

+220
-157
lines changed

23 files changed

+220
-157
lines changed

libs/remote_storage/src/lib.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -443,7 +443,7 @@ impl<Other: RemoteStorage> GenericRemoteStorage<Arc<Other>> {
443443
}
444444

445445
impl GenericRemoteStorage {
446-
pub fn from_config(storage_config: &RemoteStorageConfig) -> anyhow::Result<Self> {
446+
pub async fn from_config(storage_config: &RemoteStorageConfig) -> anyhow::Result<Self> {
447447
let timeout = storage_config.timeout;
448448
Ok(match &storage_config.storage {
449449
RemoteStorageKind::LocalFs { local_path: path } => {
@@ -458,7 +458,7 @@ impl GenericRemoteStorage {
458458
std::env::var("AWS_ACCESS_KEY_ID").unwrap_or_else(|_| "<none>".into());
459459
info!("Using s3 bucket '{}' in region '{}' as a remote storage, prefix in bucket: '{:?}', bucket endpoint: '{:?}', profile: {profile}, access_key_id: {access_key_id}",
460460
s3_config.bucket_name, s3_config.bucket_region, s3_config.prefix_in_bucket, s3_config.endpoint);
461-
Self::AwsS3(Arc::new(S3Bucket::new(s3_config, timeout)?))
461+
Self::AwsS3(Arc::new(S3Bucket::new(s3_config, timeout).await?))
462462
}
463463
RemoteStorageKind::AzureContainer(azure_config) => {
464464
let storage_account = azure_config

libs/remote_storage/src/s3_bucket.rs

Lines changed: 24 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,10 @@ use std::{
1616

1717
use anyhow::{anyhow, Context as _};
1818
use aws_config::{
19-
environment::credentials::EnvironmentVariableCredentialsProvider,
20-
imds::credentials::ImdsCredentialsProvider,
21-
meta::credentials::CredentialsProviderChain,
22-
profile::ProfileFileCredentialsProvider,
23-
provider_config::ProviderConfig,
19+
default_provider::credentials::DefaultCredentialsChain,
2420
retry::{RetryConfigBuilder, RetryMode},
25-
web_identity_token::WebIdentityTokenCredentialsProvider,
2621
BehaviorVersion,
2722
};
28-
use aws_credential_types::provider::SharedCredentialsProvider;
2923
use aws_sdk_s3::{
3024
config::{AsyncSleep, IdentityCache, Region, SharedAsyncSleep},
3125
error::SdkError,
@@ -76,40 +70,27 @@ struct GetObjectRequest {
7670
}
7771
impl S3Bucket {
7872
/// Creates the S3 storage, errors if incorrect AWS S3 configuration provided.
79-
pub fn new(remote_storage_config: &S3Config, timeout: Duration) -> anyhow::Result<Self> {
73+
pub async fn new(remote_storage_config: &S3Config, timeout: Duration) -> anyhow::Result<Self> {
8074
tracing::debug!(
8175
"Creating s3 remote storage for S3 bucket {}",
8276
remote_storage_config.bucket_name
8377
);
8478

85-
let region = Some(Region::new(remote_storage_config.bucket_region.clone()));
86-
87-
let provider_conf = ProviderConfig::without_region().with_region(region.clone());
88-
89-
let credentials_provider = {
90-
// uses "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"
91-
CredentialsProviderChain::first_try(
92-
"env",
93-
EnvironmentVariableCredentialsProvider::new(),
94-
)
95-
// uses "AWS_PROFILE" / `aws sso login --profile <profile>`
96-
.or_else(
97-
"profile-sso",
98-
ProfileFileCredentialsProvider::builder()
99-
.configure(&provider_conf)
100-
.build(),
101-
)
102-
// uses "AWS_WEB_IDENTITY_TOKEN_FILE", "AWS_ROLE_ARN", "AWS_ROLE_SESSION_NAME"
103-
// needed to access remote extensions bucket
104-
.or_else(
105-
"token",
106-
WebIdentityTokenCredentialsProvider::builder()
107-
.configure(&provider_conf)
108-
.build(),
109-
)
110-
// uses imds v2
111-
.or_else("imds", ImdsCredentialsProvider::builder().build())
112-
};
79+
let region = Region::new(remote_storage_config.bucket_region.clone());
80+
let region_opt = Some(region.clone());
81+
82+
// https://docs.aws.amazon.com/sdkref/latest/guide/standardized-credentials.html
83+
// https://docs.rs/aws-config/latest/aws_config/default_provider/credentials/struct.DefaultCredentialsChain.html
84+
// Incomplete list of auth methods used by this:
85+
// * "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"
86+
// * "AWS_PROFILE" / `aws sso login --profile <profile>`
87+
// * "AWS_WEB_IDENTITY_TOKEN_FILE", "AWS_ROLE_ARN", "AWS_ROLE_SESSION_NAME"
88+
// * http (ECS/EKS) container credentials
89+
// * imds v2
90+
let credentials_provider = DefaultCredentialsChain::builder()
91+
.region(region)
92+
.build()
93+
.await;
11394

11495
// AWS SDK requires us to specify how the RetryConfig should sleep when it wants to back off
11596
let sleep_impl: Arc<dyn AsyncSleep> = Arc::new(TokioSleep::new());
@@ -118,9 +99,9 @@ impl S3Bucket {
11899
#[allow(deprecated)] /* TODO: https://github.com/neondatabase/neon/issues/7665 */
119100
BehaviorVersion::v2023_11_09(),
120101
)
121-
.region(region)
102+
.region(region_opt)
122103
.identity_cache(IdentityCache::lazy().build())
123-
.credentials_provider(SharedCredentialsProvider::new(credentials_provider))
104+
.credentials_provider(credentials_provider)
124105
.sleep_impl(SharedAsyncSleep::from(sleep_impl));
125106

126107
let sdk_config: aws_config::SdkConfig = std::thread::scope(|s| {
@@ -1041,8 +1022,8 @@ mod tests {
10411022

10421023
use crate::{RemotePath, S3Bucket, S3Config};
10431024

1044-
#[test]
1045-
fn relative_path() {
1025+
#[tokio::test]
1026+
async fn relative_path() {
10461027
let all_paths = ["", "some/path", "some/path/"];
10471028
let all_paths: Vec<RemotePath> = all_paths
10481029
.iter()
@@ -1085,8 +1066,9 @@ mod tests {
10851066
max_keys_per_list_response: Some(5),
10861067
upload_storage_class: None,
10871068
};
1088-
let storage =
1089-
S3Bucket::new(&config, std::time::Duration::ZERO).expect("remote storage init");
1069+
let storage = S3Bucket::new(&config, std::time::Duration::ZERO)
1070+
.await
1071+
.expect("remote storage init");
10901072
for (test_path_idx, test_path) in all_paths.iter().enumerate() {
10911073
let result = storage.relative_path_to_s3_object(test_path);
10921074
let expected = expected_outputs[prefix_idx][test_path_idx];

libs/remote_storage/tests/test_real_azure.rs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ struct EnabledAzure {
3131
impl EnabledAzure {
3232
async fn setup(max_keys_in_list_response: Option<i32>) -> Self {
3333
let client = create_azure_client(max_keys_in_list_response)
34+
.await
3435
.context("Azure client creation")
3536
.expect("Azure client creation failed");
3637

@@ -187,7 +188,7 @@ impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs {
187188
}
188189
}
189190

190-
fn create_azure_client(
191+
async fn create_azure_client(
191192
max_keys_per_list_response: Option<i32>,
192193
) -> anyhow::Result<Arc<GenericRemoteStorage>> {
193194
use rand::Rng;
@@ -221,6 +222,8 @@ fn create_azure_client(
221222
timeout: Duration::from_secs(120),
222223
};
223224
Ok(Arc::new(
224-
GenericRemoteStorage::from_config(&remote_storage_config).context("remote storage init")?,
225+
GenericRemoteStorage::from_config(&remote_storage_config)
226+
.await
227+
.context("remote storage init")?,
225228
))
226229
}

libs/remote_storage/tests/test_real_s3.rs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,7 @@ struct EnabledS3 {
197197
impl EnabledS3 {
198198
async fn setup(max_keys_in_list_response: Option<i32>) -> Self {
199199
let client = create_s3_client(max_keys_in_list_response)
200+
.await
200201
.context("S3 client creation")
201202
.expect("S3 client creation failed");
202203

@@ -352,7 +353,7 @@ impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs {
352353
}
353354
}
354355

355-
fn create_s3_client(
356+
async fn create_s3_client(
356357
max_keys_per_list_response: Option<i32>,
357358
) -> anyhow::Result<Arc<GenericRemoteStorage>> {
358359
use rand::Rng;
@@ -385,7 +386,9 @@ fn create_s3_client(
385386
timeout: RemoteStorageConfig::DEFAULT_TIMEOUT,
386387
};
387388
Ok(Arc::new(
388-
GenericRemoteStorage::from_config(&remote_storage_config).context("remote storage init")?,
389+
GenericRemoteStorage::from_config(&remote_storage_config)
390+
.await
391+
.context("remote storage init")?,
389392
))
390393
}
391394

pageserver/ctl/src/main.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ async fn main() -> anyhow::Result<()> {
179179
.get("remote_storage")
180180
.expect("need remote_storage");
181181
let config = RemoteStorageConfig::from_toml(toml_item)?;
182-
let storage = remote_storage::GenericRemoteStorage::from_config(&config);
182+
let storage = remote_storage::GenericRemoteStorage::from_config(&config).await;
183183
let cancel = CancellationToken::new();
184184
storage
185185
.unwrap()

pageserver/src/bin/pageserver.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -385,7 +385,7 @@ fn start_pageserver(
385385
let shutdown_pageserver = tokio_util::sync::CancellationToken::new();
386386

387387
// Set up remote storage client
388-
let remote_storage = create_remote_storage_client(conf)?;
388+
let remote_storage = BACKGROUND_RUNTIME.block_on(create_remote_storage_client(conf))?;
389389

390390
// Set up deletion queue
391391
let (deletion_queue, deletion_workers) = DeletionQueue::new(
@@ -701,7 +701,7 @@ fn start_pageserver(
701701
}
702702
}
703703

704-
fn create_remote_storage_client(
704+
async fn create_remote_storage_client(
705705
conf: &'static PageServerConf,
706706
) -> anyhow::Result<GenericRemoteStorage> {
707707
let config = if let Some(config) = &conf.remote_storage_config {
@@ -711,7 +711,7 @@ fn create_remote_storage_client(
711711
};
712712

713713
// Create the client
714-
let mut remote_storage = GenericRemoteStorage::from_config(config)?;
714+
let mut remote_storage = GenericRemoteStorage::from_config(config).await?;
715715

716716
// If `test_remote_failures` is non-zero, wrap the client with a
717717
// wrapper that simulates failures.

pageserver/src/consumption_metrics.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ pub async fn collect_metrics(
9696
.expect("Failed to create http client with timeout");
9797

9898
let bucket_client = if let Some(bucket_config) = metric_collection_bucket {
99-
match GenericRemoteStorage::from_config(bucket_config) {
99+
match GenericRemoteStorage::from_config(bucket_config).await {
100100
Ok(client) => Some(client),
101101
Err(e) => {
102102
// Non-fatal error: if we were given an invalid config, we will proceed

pageserver/src/deletion_queue.rs

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -828,9 +828,9 @@ mod test {
828828
}
829829
}
830830

831-
fn setup(test_name: &str) -> anyhow::Result<TestSetup> {
831+
async fn setup(test_name: &str) -> anyhow::Result<TestSetup> {
832832
let test_name = Box::leak(Box::new(format!("deletion_queue__{test_name}")));
833-
let harness = TenantHarness::create(test_name)?;
833+
let harness = TenantHarness::create(test_name).await?;
834834

835835
// We do not load() the harness: we only need its config and remote_storage
836836

@@ -844,7 +844,9 @@ mod test {
844844
},
845845
timeout: RemoteStorageConfig::DEFAULT_TIMEOUT,
846846
};
847-
let storage = GenericRemoteStorage::from_config(&storage_config).unwrap();
847+
let storage = GenericRemoteStorage::from_config(&storage_config)
848+
.await
849+
.unwrap();
848850

849851
let mock_control_plane = MockControlPlane::new();
850852

@@ -922,7 +924,9 @@ mod test {
922924
#[tokio::test]
923925
async fn deletion_queue_smoke() -> anyhow::Result<()> {
924926
// Basic test that the deletion queue processes the deletions we pass into it
925-
let ctx = setup("deletion_queue_smoke").expect("Failed test setup");
927+
let ctx = setup("deletion_queue_smoke")
928+
.await
929+
.expect("Failed test setup");
926930
let client = ctx.deletion_queue.new_client();
927931
client.recover(HashMap::new())?;
928932

@@ -992,7 +996,9 @@ mod test {
992996

993997
#[tokio::test]
994998
async fn deletion_queue_validation() -> anyhow::Result<()> {
995-
let ctx = setup("deletion_queue_validation").expect("Failed test setup");
999+
let ctx = setup("deletion_queue_validation")
1000+
.await
1001+
.expect("Failed test setup");
9961002
let client = ctx.deletion_queue.new_client();
9971003
client.recover(HashMap::new())?;
9981004

@@ -1051,7 +1057,9 @@ mod test {
10511057
#[tokio::test]
10521058
async fn deletion_queue_recovery() -> anyhow::Result<()> {
10531059
// Basic test that the deletion queue processes the deletions we pass into it
1054-
let mut ctx = setup("deletion_queue_recovery").expect("Failed test setup");
1060+
let mut ctx = setup("deletion_queue_recovery")
1061+
.await
1062+
.expect("Failed test setup");
10551063
let client = ctx.deletion_queue.new_client();
10561064
client.recover(HashMap::new())?;
10571065

pageserver/src/pgdatadir_mapping.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2031,7 +2031,7 @@ mod tests {
20312031
#[tokio::test]
20322032
async fn aux_files_round_trip() -> anyhow::Result<()> {
20332033
let name = "aux_files_round_trip";
2034-
let harness = TenantHarness::create(name)?;
2034+
let harness = TenantHarness::create(name).await?;
20352035

20362036
pub const TIMELINE_ID: TimelineId =
20372037
TimelineId::from_array(hex!("11223344556677881122334455667788"));

0 commit comments

Comments
 (0)