You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by tu...@apache.org on 2023/06/02 14:40:34 UTC
[arrow-rs] branch master updated: Fix support for ECS IAM credentials (#4310)
This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new fdf37a08e Fix support for ECS IAM credentials (#4310)
fdf37a08e is described below
commit fdf37a08e961117857933d016960b37d304a6910
Author: Raphael Taylor-Davies <17...@users.noreply.github.com>
AuthorDate: Fri Jun 2 15:40:28 2023 +0100
Fix support for ECS IAM credentials (#4310)
---
object_store/src/aws/credential.rs | 45 ++++++++++++++++++++++++++++++++-
object_store/src/aws/mod.rs | 51 ++++++++++++++++++++++++++------------
2 files changed, 79 insertions(+), 17 deletions(-)
diff --git a/object_store/src/aws/credential.rs b/object_store/src/aws/credential.rs
index 909dde072..be0ffa578 100644
--- a/object_store/src/aws/credential.rs
+++ b/object_store/src/aws/credential.rs
@@ -20,7 +20,7 @@ use crate::client::retry::RetryExt;
use crate::client::token::{TemporaryToken, TokenCache};
use crate::client::TokenProvider;
use crate::util::hmac_sha256;
-use crate::{Result, RetryConfig};
+use crate::{CredentialProvider, Result, RetryConfig};
use async_trait::async_trait;
use bytes::Buf;
use chrono::{DateTime, Utc};
@@ -542,6 +542,49 @@ async fn web_identity(
})
}
+/// Credentials sourced from a task IAM role
+///
+/// <https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html>
+#[derive(Debug)]
+pub struct TaskCredentialProvider {
+ pub url: String,
+ pub retry: RetryConfig,
+ pub client: Client,
+ pub cache: TokenCache<Arc<AwsCredential>>,
+}
+
+#[async_trait]
+impl CredentialProvider for TaskCredentialProvider {
+ type Credential = AwsCredential;
+
+ async fn get_credential(&self) -> Result<Arc<AwsCredential>> {
+ self.cache
+ .get_or_insert_with(|| task_credential(&self.client, &self.retry, &self.url))
+ .await
+ .map_err(|source| crate::Error::Generic {
+ store: STORE,
+ source,
+ })
+ }
+}
+
+/// <https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html>
+async fn task_credential(
+ client: &Client,
+ retry: &RetryConfig,
+ url: &str,
+) -> Result<TemporaryToken<Arc<AwsCredential>>, StdError> {
+ let creds: InstanceCredentials =
+ client.get(url).send_retry(retry).await?.json().await?;
+
+ let now = Utc::now();
+ let ttl = (creds.expiration - now).to_std().unwrap_or_default();
+ Ok(TemporaryToken {
+ token: Arc::new(creds.into()),
+ expiry: Some(Instant::now() + ttl),
+ })
+}
+
#[cfg(test)]
mod tests {
use super::*;
diff --git a/object_store/src/aws/mod.rs b/object_store/src/aws/mod.rs
index 8de4b7c6a..8a486f986 100644
--- a/object_store/src/aws/mod.rs
+++ b/object_store/src/aws/mod.rs
@@ -46,7 +46,9 @@ use url::Url;
pub use crate::aws::checksum::Checksum;
use crate::aws::client::{S3Client, S3Config};
-use crate::aws::credential::{InstanceCredentialProvider, WebIdentityProvider};
+use crate::aws::credential::{
+ InstanceCredentialProvider, TaskCredentialProvider, WebIdentityProvider,
+};
use crate::client::get::GetClientExt;
use crate::client::list::ListClientExt;
use crate::client::{
@@ -87,9 +89,6 @@ pub use credential::{AwsAuthorizer, AwsCredential};
/// Default metadata endpoint
static DEFAULT_METADATA_ENDPOINT: &str = "http://169.254.169.254";
-/// ECS metadata endpoint
-static ECS_METADATA_ENDPOINT: &str = "http://169.254.170.2";
-
/// A specialized `Error` for object store-related errors
#[derive(Debug, Snafu)]
#[allow(missing_docs)]
@@ -399,6 +398,8 @@ pub struct AmazonS3Builder {
checksum_algorithm: Option<ConfigValue<Checksum>>,
/// Metadata endpoint, see <https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html>
metadata_endpoint: Option<String>,
+ /// Container credentials URL, see <https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html>
+ container_credentials_relative_uri: Option<String>,
/// Client options
client_options: ClientOptions,
/// Credentials
@@ -529,6 +530,11 @@ pub enum AmazonS3ConfigKey {
/// - `metadata_endpoint`
MetadataEndpoint,
+ /// Set the container credentials relative URI
+ ///
+ /// <https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html>
+ ContainerCredentialsRelativeUri,
+
/// Client options
Client(ClientConfigKey),
}
@@ -548,6 +554,9 @@ impl AsRef<str> for AmazonS3ConfigKey {
Self::MetadataEndpoint => "aws_metadata_endpoint",
Self::UnsignedPayload => "aws_unsigned_payload",
Self::Checksum => "aws_checksum_algorithm",
+ Self::ContainerCredentialsRelativeUri => {
+ "aws_container_credentials_relative_uri"
+ }
Self::Client(opt) => opt.as_ref(),
}
}
@@ -578,6 +587,9 @@ impl FromStr for AmazonS3ConfigKey {
"aws_metadata_endpoint" | "metadata_endpoint" => Ok(Self::MetadataEndpoint),
"aws_unsigned_payload" | "unsigned_payload" => Ok(Self::UnsignedPayload),
"aws_checksum_algorithm" | "checksum_algorithm" => Ok(Self::Checksum),
+ "aws_container_credentials_relative_uri" => {
+ Ok(Self::ContainerCredentialsRelativeUri)
+ }
// Backwards compatibility
"aws_allow_http" => Ok(Self::Client(ClientConfigKey::AllowHttp)),
_ => match s.parse() {
@@ -625,15 +637,6 @@ impl AmazonS3Builder {
}
}
- // This env var is set in ECS
- // https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html
- if let Ok(metadata_relative_uri) =
- std::env::var("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")
- {
- builder.metadata_endpoint =
- Some(format!("{ECS_METADATA_ENDPOINT}{metadata_relative_uri}"));
- }
-
builder
}
@@ -691,6 +694,9 @@ impl AmazonS3Builder {
AmazonS3ConfigKey::Checksum => {
self.checksum_algorithm = Some(ConfigValue::Deferred(value.into()))
}
+ AmazonS3ConfigKey::ContainerCredentialsRelativeUri => {
+ self.container_credentials_relative_uri = Some(value.into())
+ }
AmazonS3ConfigKey::Client(key) => {
self.client_options = self.client_options.with_config(key, value)
}
@@ -758,6 +764,9 @@ impl AmazonS3Builder {
self.checksum_algorithm.as_ref().map(ToString::to_string)
}
AmazonS3ConfigKey::Client(key) => self.client_options.get_config_value(key),
+ AmazonS3ConfigKey::ContainerCredentialsRelativeUri => {
+ self.container_credentials_relative_uri.clone()
+ }
}
}
@@ -999,6 +1008,15 @@ impl AmazonS3Builder {
client,
self.retry_config.clone(),
)) as _
+ } else if let Some(uri) = self.container_credentials_relative_uri {
+ info!("Using Task credential provider");
+ Arc::new(TaskCredentialProvider {
+ url: format!("http://169.254.170.2{uri}"),
+ retry: self.retry_config.clone(),
+ // The instance metadata endpoint is access over HTTP
+ client: self.client_options.clone().with_allow_http(true).client()?,
+ cache: Default::default(),
+ }) as _
} else {
info!("Using Instance credential provider");
@@ -1199,9 +1217,10 @@ mod tests {
assert_eq!(builder.endpoint.unwrap(), aws_endpoint);
assert_eq!(builder.token.unwrap(), aws_session_token);
- let metadata_uri =
- format!("{ECS_METADATA_ENDPOINT}{container_creds_relative_uri}");
- assert_eq!(builder.metadata_endpoint.unwrap(), metadata_uri);
+ assert_eq!(
+ builder.container_credentials_relative_uri.unwrap(),
+ container_creds_relative_uri
+ );
assert_eq!(
builder.checksum_algorithm.unwrap().get().unwrap(),
Checksum::SHA256