From 15a387903ca46c2a5ce6b9c4c77fd6ad320f22d9 Mon Sep 17 00:00:00 2001 From: John Chrostek Date: Wed, 12 Nov 2025 09:55:34 -0500 Subject: [PATCH] [SVLS-7704] add support for SSM Parameter API key --- bottlecap/src/config/env.rs | 7 + bottlecap/src/config/mod.rs | 2 + bottlecap/src/config/yaml.rs | 1 + bottlecap/src/secrets/decrypt.rs | 303 +++++++++++++++++++++++-------- 4 files changed, 242 insertions(+), 71 deletions(-) diff --git a/bottlecap/src/config/env.rs b/bottlecap/src/config/env.rs index f0b6c075a..3b1a0bcc9 100644 --- a/bottlecap/src/config/env.rs +++ b/bottlecap/src/config/env.rs @@ -357,6 +357,11 @@ pub struct EnvConfig { /// The AWS KMS API key to use for the Datadog Agent. #[serde(deserialize_with = "deserialize_optional_string")] pub kms_api_key: Option, + /// @env `DD_API_KEY_SSM_ARN` + /// + /// The AWS Systems Manager Parameter Store parameter ARN containing the Datadog API key. + #[serde(deserialize_with = "deserialize_optional_string")] + pub api_key_ssm_arn: Option, /// @env `DD_SERVERLESS_LOGS_ENABLED` /// /// Enable logs for AWS Lambda. Default is `true`. @@ -609,6 +614,7 @@ fn merge_config(config: &mut Config, env_config: &EnvConfig) { // AWS Lambda merge_string!(config, env_config, api_key_secret_arn); merge_string!(config, env_config, kms_api_key); + merge_string!(config, env_config, api_key_ssm_arn); merge_option_to_value!(config, env_config, serverless_logs_enabled); merge_option_to_value!(config, env_config, serverless_flush_strategy); merge_option_to_value!(config, env_config, enhanced_metrics); @@ -952,6 +958,7 @@ mod tests { api_key_secret_arn: "arn:aws:secretsmanager:region:account:secret:datadog-api-key" .to_string(), kms_api_key: "test-kms-key".to_string(), + api_key_ssm_arn: String::default(), serverless_logs_enabled: false, serverless_flush_strategy: FlushStrategy::Periodically(PeriodicStrategy { interval: 60000, diff --git a/bottlecap/src/config/mod.rs b/bottlecap/src/config/mod.rs index 369e6ee73..863661848 100644 --- a/bottlecap/src/config/mod.rs +++ b/bottlecap/src/config/mod.rs @@ -337,6 +337,7 @@ pub struct Config { // AWS Lambda pub api_key_secret_arn: String, pub kms_api_key: String, + pub api_key_ssm_arn: String, pub serverless_logs_enabled: bool, pub serverless_flush_strategy: FlushStrategy, pub enhanced_metrics: bool, @@ -440,6 +441,7 @@ impl Default for Config { // AWS Lambda api_key_secret_arn: String::default(), kms_api_key: String::default(), + api_key_ssm_arn: String::default(), serverless_logs_enabled: true, serverless_flush_strategy: FlushStrategy::Default, enhanced_metrics: true, diff --git a/bottlecap/src/config/yaml.rs b/bottlecap/src/config/yaml.rs index c800455f7..022d00664 100644 --- a/bottlecap/src/config/yaml.rs +++ b/bottlecap/src/config/yaml.rs @@ -972,6 +972,7 @@ api_security_sample_delay: 60 # Seconds api_key_secret_arn: "arn:aws:secretsmanager:region:account:secret:datadog-api-key" .to_string(), kms_api_key: "test-kms-key".to_string(), + api_key_ssm_arn: String::default(), serverless_logs_enabled: false, serverless_flush_strategy: FlushStrategy::Periodically(PeriodicStrategy { interval: 60000, diff --git a/bottlecap/src/secrets/decrypt.rs b/bottlecap/src/secrets/decrypt.rs index 66dc9ebc0..d6cab57f0 100644 --- a/bottlecap/src/secrets/decrypt.rs +++ b/bottlecap/src/secrets/decrypt.rs @@ -18,87 +18,97 @@ use tracing::debug; use tracing::error; pub async fn resolve_secrets(config: Arc, aws_config: Arc) -> Option { - let api_key_candidate = - if !config.api_key_secret_arn.is_empty() || !config.kms_api_key.is_empty() { - let before_decrypt = Instant::now(); + let api_key_candidate = if !config.api_key_secret_arn.is_empty() + || !config.kms_api_key.is_empty() + || !config.api_key_ssm_arn.is_empty() + { + let before_decrypt = Instant::now(); + + let builder = match create_reqwest_client_builder() { + Ok(builder) => builder, + Err(err) => { + error!("Error creating reqwest client builder: {}", err); + return None; + } + }; - let builder = match create_reqwest_client_builder() { - Ok(builder) => builder, - Err(err) => { - error!("Error creating reqwest client builder: {}", err); - return None; - } - }; + let client = match builder.build() { + Ok(client) => client, + Err(err) => { + error!("Error creating reqwest client: {}", err); + return None; + } + }; - let client = match builder.build() { - Ok(client) => client, + let mut aws_credentials = AwsCredentials::from_env(); + + if aws_credentials.aws_secret_access_key.is_empty() + && aws_credentials.aws_access_key_id.is_empty() + && !aws_credentials + .aws_container_credentials_full_uri + .is_empty() + && !aws_credentials.aws_container_authorization_token.is_empty() + { + // We're in Snap Start + let credentials = match get_snapstart_credentials(&aws_credentials, &client).await { + Ok(credentials) => credentials, Err(err) => { - error!("Error creating reqwest client: {}", err); + error!("Error getting Snap Start credentials: {}", err); return None; } }; + aws_credentials.aws_access_key_id = credentials["AccessKeyId"] + .as_str() + .unwrap_or_default() + .to_string(); + aws_credentials.aws_secret_access_key = credentials["SecretAccessKey"] + .as_str() + .unwrap_or_default() + .to_string(); + aws_credentials.aws_session_token = credentials["Token"] + .as_str() + .unwrap_or_default() + .to_string(); + } - let mut aws_credentials = AwsCredentials::from_env(); - - if aws_credentials.aws_secret_access_key.is_empty() - && aws_credentials.aws_access_key_id.is_empty() - && !aws_credentials - .aws_container_credentials_full_uri - .is_empty() - && !aws_credentials.aws_container_authorization_token.is_empty() - { - // We're in Snap Start - let credentials = match get_snapstart_credentials(&aws_credentials, &client).await { - Ok(credentials) => credentials, - Err(err) => { - error!("Error getting Snap Start credentials: {}", err); - return None; - } - }; - aws_credentials.aws_access_key_id = credentials["AccessKeyId"] - .as_str() - .unwrap_or_default() - .to_string(); - aws_credentials.aws_secret_access_key = credentials["SecretAccessKey"] - .as_str() - .unwrap_or_default() - .to_string(); - aws_credentials.aws_session_token = credentials["Token"] - .as_str() - .unwrap_or_default() - .to_string(); - } - - let decrypted_key = if config.kms_api_key.is_empty() { - decrypt_aws_sm( - &client, - config.api_key_secret_arn.clone(), - aws_config, - &aws_credentials, - ) - .await - } else { - decrypt_aws_kms( - &client, - config.kms_api_key.clone(), - aws_config, - &aws_credentials, - ) - .await - }; + let decrypted_key = if !config.kms_api_key.is_empty() { + decrypt_aws_kms( + &client, + config.kms_api_key.clone(), + aws_config, + &aws_credentials, + ) + .await + } else if !config.api_key_secret_arn.is_empty() { + decrypt_aws_sm( + &client, + config.api_key_secret_arn.clone(), + aws_config, + &aws_credentials, + ) + .await + } else { + decrypt_aws_ssm( + &client, + config.api_key_ssm_arn.clone(), + aws_config, + &aws_credentials, + ) + .await + }; - debug!("Decrypt took {} ms", before_decrypt.elapsed().as_millis()); + debug!("Decrypt took {} ms", before_decrypt.elapsed().as_millis()); - match decrypted_key { - Ok(key) => Some(key), - Err(err) => { - error!("Error decrypting key: {}", err); - None - } + match decrypted_key { + Ok(key) => Some(key), + Err(err) => { + error!("Error decrypting key: {}", err); + None } - } else { - Some(config.api_key.clone()) - }; + } + } else { + Some(config.api_key.clone()) + }; clean_api_key(api_key_candidate) } @@ -215,6 +225,39 @@ async fn decrypt_aws_sm( } } +async fn decrypt_aws_ssm( + client: &Client, + parameter_arn: String, + aws_config: Arc, + aws_credentials: &AwsCredentials, +) -> Result> { + let json_body = &serde_json::json!({ "Name": parameter_arn, "WithDecryption": true }); + let parameter_region = parameter_arn + .split(':') + .nth(3) + .unwrap_or(&aws_config.region) + .to_string(); + let headers = build_get_secret_signed_headers( + aws_config, + aws_credentials, + parameter_region, + RequestArgs { + service: "ssm".to_string(), + body: json_body, + time: Utc::now(), + x_amz_target: "AmazonSSM.GetParameter".to_string(), + }, + ); + + let v = request(json_body, headers?, client).await?; + if let Some(parameter) = v["Parameter"].as_object() { + if let Some(value) = parameter["Value"].as_str() { + return Ok(value.to_string()); + } + } + Err(Error::new(std::io::ErrorKind::InvalidData, v.to_string()).into()) +} + async fn get_snapstart_credentials( aws_credentials: &AwsCredentials, client: &Client, @@ -483,4 +526,122 @@ mod tests { assert_eq!(headers.get(k).expect("cannot get header"), v); } } + + #[test] + #[allow(clippy::unwrap_used)] + fn test_ssm_parameter_headers() { + let time = Utc.from_utc_datetime( + &NaiveDateTime::parse_from_str("2024-05-30 09:10:11", "%Y-%m-%d %H:%M:%S").unwrap(), + ); + let headers = build_get_secret_signed_headers( + Arc::new(AwsConfig { + region: "us-east-1".to_string(), + aws_lwa_proxy_lambda_runtime_api: Some("***".into()), + function_name: "arn:some-function".to_string(), + sandbox_init_time: Instant::now(), + runtime_api: String::new(), + exec_wrapper: None, + }), + &AwsCredentials{ + aws_access_key_id: "AKIDEXAMPLE".to_string(), + aws_secret_access_key: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY".to_string(), + aws_session_token: "AQoDYXdzEJr...".to_string(), + aws_container_authorization_token: String::new(), + aws_container_credentials_full_uri: String::new(), + }, + "us-east-1".to_string(), + RequestArgs { + service: "ssm".to_string(), + body: &serde_json::json!({ "Name": "arn:aws:ssm:us-east-1:account-id:parameter/my-parameter", "WithDecryption": true }), + time, + x_amz_target: "AmazonSSM.GetParameter".to_string(), + }, + ).unwrap(); + + let mut expected_headers = HeaderMap::new(); + expected_headers.insert("authorization", HeaderValue::from_str("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240530/us-east-1/ssm/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token;x-amz-target, Signature=f45eebacb9f2f575cf1a235c65c97adc1c0d6ca21174044baa0c6457d4518d64").unwrap()); + expected_headers.insert( + "host", + HeaderValue::from_str("ssm.us-east-1.amazonaws.com").unwrap(), + ); + expected_headers.insert( + "content-type", + HeaderValue::from_str("application/x-amz-json-1.1").unwrap(), + ); + expected_headers.insert( + "x-amz-date", + HeaderValue::from_str("20240530T091011Z").unwrap(), + ); + expected_headers.insert( + "x-amz-target", + HeaderValue::from_str("AmazonSSM.GetParameter").unwrap(), + ); + expected_headers.insert( + "x-amz-security-token", + HeaderValue::from_str("AQoDYXdzEJr...").unwrap(), + ); + + for (k, v) in &expected_headers { + assert_eq!(headers.get(k).expect("cannot get header"), v); + } + } + + #[test] + #[allow(clippy::unwrap_used)] + fn test_cross_region_ssm_parameter() { + let time = Utc.from_utc_datetime( + &NaiveDateTime::parse_from_str("2024-05-30 09:10:11", "%Y-%m-%d %H:%M:%S").unwrap(), + ); + let headers = build_get_secret_signed_headers( + Arc::new(AwsConfig { + region: "us-east-1".to_string(), + aws_lwa_proxy_lambda_runtime_api: Some("***".into()), + function_name: "arn:some-function".to_string(), + sandbox_init_time: Instant::now(), + runtime_api: String::new(), + exec_wrapper: None, + }), + &AwsCredentials{ + aws_access_key_id: "AKIDEXAMPLE".to_string(), + aws_secret_access_key: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY".to_string(), + aws_session_token: "AQoDYXdzEJr...".to_string(), + aws_container_authorization_token: String::new(), + aws_container_credentials_full_uri: String::new(), + }, + "eu-west-1".to_string(), + RequestArgs { + service: "ssm".to_string(), + body: &serde_json::json!({ "Name": "arn:aws:ssm:eu-west-1:account-id:parameter/my-parameter", "WithDecryption": true }), + time, + x_amz_target: "AmazonSSM.GetParameter".to_string(), + }, + ).unwrap(); + + let mut expected_headers = HeaderMap::new(); + expected_headers.insert("authorization", HeaderValue::from_str("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240530/eu-west-1/ssm/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token;x-amz-target, Signature=4bdd2a27b632981c73a70b52ea84e92d3e26fe2c768e9cf18b96630188748b7a").unwrap()); + expected_headers.insert( + "host", + HeaderValue::from_str("ssm.eu-west-1.amazonaws.com").unwrap(), + ); + expected_headers.insert( + "content-type", + HeaderValue::from_str("application/x-amz-json-1.1").unwrap(), + ); + expected_headers.insert( + "x-amz-date", + HeaderValue::from_str("20240530T091011Z").unwrap(), + ); + expected_headers.insert( + "x-amz-target", + HeaderValue::from_str("AmazonSSM.GetParameter").unwrap(), + ); + expected_headers.insert( + "x-amz-security-token", + HeaderValue::from_str("AQoDYXdzEJr...").unwrap(), + ); + + for (k, v) in &expected_headers { + assert_eq!(headers.get(k).expect("cannot get header"), v); + } + } }