feat: allow passing a session token to AWS stores

pull/24376/head
Marco Neumann 2021-05-27 11:53:12 +02:00
parent d5df7c381a
commit 1cbe4031d3
4 changed files with 51 additions and 9 deletions

View File

@ -34,6 +34,9 @@
# AWS_DEFAULT_REGION=us-east-2
# INFLUXDB_IOX_BUCKET=bucket-name
#
# Optionally:
# AWS_SESSION_TOKEN=token
#
# If using an s3 compatible storage
# AWS_ENDPOINT = http://localhost:9000
#

View File

@ -296,6 +296,7 @@ impl AmazonS3 {
region: impl Into<String>,
bucket_name: impl Into<String>,
endpoint: Option<impl Into<String>>,
session_token: Option<impl Into<String>>,
) -> Result<Self> {
let region = region.into();
let region: rusoto_core::Region = match endpoint {
@ -309,14 +310,23 @@ impl AmazonS3 {
let http_client = rusoto_core::request::HttpClient::new()
.expect("Current implementation of rusoto_core has no way for this to fail");
let client = match (access_key_id, secret_access_key) {
(Some(access_key_id), Some(secret_access_key)) => {
let client = match (access_key_id, secret_access_key, session_token) {
(Some(access_key_id), Some(secret_access_key), Some(session_token)) => {
let credentials_provider = StaticProvider::new(
access_key_id.into(),
secret_access_key.into(),
Some(session_token.into()),
None,
);
rusoto_s3::S3Client::new_with(http_client, credentials_provider, region)
}
(Some(access_key_id), Some(secret_access_key), None) => {
let credentials_provider =
StaticProvider::new_minimal(access_key_id.into(), secret_access_key.into());
rusoto_s3::S3Client::new_with(http_client, credentials_provider, region)
}
(None, Some(_)) => return Err(Error::MissingAccessKey),
(Some(_), None) => return Err(Error::MissingSecretAccessKey),
(None, Some(_), _) => return Err(Error::MissingAccessKey),
(Some(_), None, _) => return Err(Error::MissingSecretAccessKey),
_ => {
let credentials_provider = InstanceMetadataProvider::new();
rusoto_s3::S3Client::new_with(http_client, credentials_provider, region)
@ -451,6 +461,7 @@ mod tests {
region: String,
bucket: String,
endpoint: Option<String>,
token: Option<String>,
}
// Helper macro to skip tests if TEST_INTEGRATION and the AWS environment variables are not set.
@ -502,6 +513,7 @@ mod tests {
bucket: env::var("INFLUXDB_IOX_BUCKET")
.expect("already checked INFLUXDB_IOX_BUCKET"),
endpoint: env::var("AWS_ENDPOINT").ok(),
token: env::var("AWS_SESSION_TOKEN").ok(),
}
}
}};
@ -533,6 +545,7 @@ mod tests {
config.region,
config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);
@ -554,6 +567,7 @@ mod tests {
config.region,
&config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);
@ -585,6 +599,7 @@ mod tests {
config.region,
&config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);
@ -627,6 +642,7 @@ mod tests {
config.region,
&config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);
@ -664,6 +680,7 @@ mod tests {
config.region,
&config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);
@ -711,6 +728,7 @@ mod tests {
config.region,
&config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);
@ -756,6 +774,7 @@ mod tests {
config.region,
config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);
@ -781,6 +800,7 @@ mod tests {
config.region,
&config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);
@ -818,6 +838,7 @@ mod tests {
config.region,
&config.bucket,
config.endpoint,
config.token,
)
.expect("Valid S3 config"),
);

View File

@ -361,7 +361,7 @@ Possible values (case insensitive):
)]
pub aws_default_region: String,
/// When using Amazon s3 compatibility storage service, set this to the
/// When using Amazon S3 compatibility storage service, set this to the
/// endpoint.
///
/// Must also set `--object-store=s3`, `--bucket`. Can also set `--aws-default-region`
@ -372,6 +372,16 @@ Possible values (case insensitive):
#[structopt(long = "--aws-endpoint", env = "AWS_ENDPOINT")]
pub aws_endpoint: Option<String>,
/// When using Amazon S3 as an object store, set this to the session token. This is handy when using a federated
/// login / SSO and you fetch credentials via the UI.
///
/// Is it assumed that the session is valid as long as the IOx server is running.
///
/// Prefer the environment variable over the command line flag in shared
/// environments.
#[structopt(long = "--aws-session-token", env = "AWS_SESSION_TOKEN")]
pub aws_session_token: Option<String>,
/// When using Google Cloud Storage as the object store, set this to the
/// path to the JSON file that contains the Google credentials.
///

View File

@ -336,14 +336,22 @@ impl TryFrom<&Config> for ObjectStore {
config.aws_secret_access_key.as_ref(),
config.aws_default_region.as_str(),
config.aws_endpoint.as_ref(),
config.aws_session_token.as_ref(),
) {
(Some(bucket), key_id, secret_key, region, endpoint) => {
(Some(bucket), key_id, secret_key, region, endpoint, session_token) => {
Ok(Self::new_amazon_s3(
AmazonS3::new(key_id, secret_key, region, bucket, endpoint)
.context(InvalidS3Config)?,
AmazonS3::new(
key_id,
secret_key,
region,
bucket,
endpoint,
session_token,
)
.context(InvalidS3Config)?,
))
}
(bucket, _, _, _, _) => {
(bucket, _, _, _, _, _) => {
let mut missing_args = vec![];
if bucket.is_none() {