Skip to content

[SVLS-7704] add support for SSM Parameter API key #924

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 19, 2025
+242 −71
Merged

[SVLS-7704] add support for SSM Parameter API key #924

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions bottlecap/src/config/env.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>,
/// @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<String>,
/// @env `DD_SERVERLESS_LOGS_ENABLED`
///
/// Enable logs for AWS Lambda. Default is `true`.
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions bottlecap/src/config/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions bottlecap/src/config/yaml.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
303 changes: 232 additions & 71 deletions bottlecap/src/secrets/decrypt.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,87 +18,97 @@ use tracing::debug;
use tracing::error;

pub async fn resolve_secrets(config: Arc<Config>, aws_config: Arc<AwsConfig>) -> Option<String> {
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()
Copy link
Contributor

@duncanista duncanista Nov 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe change it to a method in aws_credentials?

aws_credentials.is_snap_start()

&& 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)
}
Expand Down Expand Up @@ -215,6 +225,39 @@ async fn decrypt_aws_sm(
}
}

async fn decrypt_aws_ssm(
client: &Client,
parameter_arn: String,
aws_config: Arc<AwsConfig>,
aws_credentials: &AwsCredentials,
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
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,
Expand Down Expand Up @@ -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...<remainder of session token>".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...<remainder of session token>").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...<remainder of session token>".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...<remainder of session token>").unwrap(),
);

for (k, v) in &expected_headers {
assert_eq!(headers.get(k).expect("cannot get header"), v);
}
}
}
Loading