From 264613501077350e68876bb0b98657ebef721412 Mon Sep 17 00:00:00 2001 From: Yiming Luo Date: Thu, 18 Sep 2025 15:34:01 -0400 Subject: [PATCH 1/2] fix: Fix dead lock when snap start and using secret for api key --- bottlecap/src/bin/bottlecap/main.rs | 18 +++++++----------- bottlecap/src/secrets/decrypt.rs | 26 ++++++++++++-------------- 2 files changed, 19 insertions(+), 25 deletions(-) diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 5bc4aeb72..80f86b290 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -23,7 +23,7 @@ use bottlecap::{ }, config::{ self, Config, - aws::{AwsConfig, AwsCredentials, build_lambda_function_arn}, + aws::{AwsConfig, build_lambda_function_arn}, flush_strategy::FlushStrategy, }, event_bus::{Event, EventBus}, @@ -85,7 +85,7 @@ use dogstatsd::{ use reqwest::Client; use std::{collections::hash_map, env, path::Path, sync::Arc}; use tokio::time::{Duration, Instant}; -use tokio::{sync::Mutex as TokioMutex, sync::RwLock, sync::mpsc::Sender, task::JoinHandle}; +use tokio::{sync::Mutex as TokioMutex, sync::mpsc::Sender, task::JoinHandle}; use tokio_util::sync::CancellationToken; use tracing::{debug, error}; use tracing_subscriber::EnvFilter; @@ -236,7 +236,7 @@ impl PendingFlushHandles { async fn main() -> anyhow::Result<()> { let start_time = Instant::now(); init_ustr(); - let (aws_config, aws_credentials, config) = load_configs(start_time); + let (aws_config, config) = load_configs(start_time); enable_logging_subsystem(&config); log_fips_status(&aws_config.region); @@ -254,7 +254,7 @@ async fn main() -> anyhow::Result<()> { .map_err(|e| anyhow::anyhow!("Failed to register extension: {e:?}"))?; let aws_config = Arc::new(aws_config); - let api_key_factory = create_api_key_factory(&config, &aws_config, aws_credentials); + let api_key_factory = create_api_key_factory(&config, &aws_config); match extension_loop_active( Arc::clone(&aws_config), @@ -285,14 +285,13 @@ fn init_ustr() { }); } -fn load_configs(start_time: Instant) -> (AwsConfig, AwsCredentials, Arc) { +fn load_configs(start_time: Instant) -> (AwsConfig, Arc) { // First load the AWS configuration let aws_config = AwsConfig::from_env(start_time); - let aws_credentials = AwsCredentials::from_env(); let lambda_directory: String = env::var("LAMBDA_TASK_ROOT").unwrap_or_else(|_| "/var/task".to_string()); let config = Arc::new(config::get_config(Path::new(&lambda_directory))); - (aws_config, aws_credentials, config) + (aws_config, config) } fn enable_logging_subsystem(config: &Arc) { @@ -321,18 +320,15 @@ fn enable_logging_subsystem(config: &Arc) { fn create_api_key_factory( config: &Arc, aws_config: &Arc, - aws_credentials: AwsCredentials, ) -> Arc { let config = Arc::clone(config); let aws_config = Arc::clone(aws_config); - let aws_credentials = Arc::new(RwLock::new(aws_credentials)); Arc::new(ApiKeyFactory::new_from_resolver(Arc::new(move || { let config = Arc::clone(&config); let aws_config = Arc::clone(&aws_config); - let aws_credentials = Arc::clone(&aws_credentials); - Box::pin(async move { resolve_secrets(config, aws_config, aws_credentials).await }) + Box::pin(async move { resolve_secrets(config, aws_config).await }) }))) } diff --git a/bottlecap/src/secrets/decrypt.rs b/bottlecap/src/secrets/decrypt.rs index ac1df3e5e..b22b063c4 100644 --- a/bottlecap/src/secrets/decrypt.rs +++ b/bottlecap/src/secrets/decrypt.rs @@ -13,14 +13,13 @@ use serde_json::Value; use sha2::{Digest, Sha256}; use std::io::Error; use std::sync::Arc; -use tokio::{sync::RwLock, time::Instant}; +use tokio::time::Instant; use tracing::debug; use tracing::error; pub async fn resolve_secrets( config: Arc, aws_config: Arc, - aws_credentials: Arc>, ) -> Option { let api_key_candidate = if !config.api_key_secret_arn.is_empty() || !config.kms_api_key.is_empty() { @@ -42,36 +41,35 @@ pub async fn resolve_secrets( } }; - let aws_credentials_read = aws_credentials.read().await; + let mut aws_credentials = AwsCredentials::from_env(); - if aws_credentials_read.aws_secret_access_key.is_empty() - && aws_credentials_read.aws_access_key_id.is_empty() - && !aws_credentials_read + 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_read + && !aws_credentials .aws_container_authorization_token .is_empty() { // We're in Snap Start let credentials = - match get_snapstart_credentials(&aws_credentials_read, &client).await { + match get_snapstart_credentials(&aws_credentials, &client).await { Ok(credentials) => credentials, Err(err) => { error!("Error getting Snap Start credentials: {}", err); return None; } }; - let mut aws_credentials_write = aws_credentials.write().await; - aws_credentials_write.aws_access_key_id = credentials["AccessKeyId"] + aws_credentials.aws_access_key_id = credentials["AccessKeyId"] .as_str() .unwrap_or_default() .to_string(); - aws_credentials_write.aws_secret_access_key = credentials["SecretAccessKey"] + aws_credentials.aws_secret_access_key = credentials["SecretAccessKey"] .as_str() .unwrap_or_default() .to_string(); - aws_credentials_write.aws_session_token = credentials["Token"] + aws_credentials.aws_session_token = credentials["Token"] .as_str() .unwrap_or_default() .to_string(); @@ -82,7 +80,7 @@ pub async fn resolve_secrets( &client, config.api_key_secret_arn.clone(), aws_config, - &aws_credentials_read, + &aws_credentials, ) .await } else { @@ -90,7 +88,7 @@ pub async fn resolve_secrets( &client, config.kms_api_key.clone(), aws_config, - &aws_credentials_read, + &aws_credentials, ) .await }; From bc99e51c49a2ce1dac43c0860782a84f3f570cae Mon Sep 17 00:00:00 2001 From: Yiming Luo Date: Thu, 18 Sep 2025 15:34:26 -0400 Subject: [PATCH 2/2] fmt --- bottlecap/src/bin/bottlecap/main.rs | 5 +---- bottlecap/src/secrets/decrypt.rs | 24 +++++++++--------------- 2 files changed, 10 insertions(+), 19 deletions(-) diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 80f86b290..b29e55966 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -317,10 +317,7 @@ fn enable_logging_subsystem(config: &Arc) { debug!("Logging subsystem enabled"); } -fn create_api_key_factory( - config: &Arc, - aws_config: &Arc, -) -> Arc { +fn create_api_key_factory(config: &Arc, aws_config: &Arc) -> Arc { let config = Arc::clone(config); let aws_config = Arc::clone(aws_config); diff --git a/bottlecap/src/secrets/decrypt.rs b/bottlecap/src/secrets/decrypt.rs index b22b063c4..47b71d174 100644 --- a/bottlecap/src/secrets/decrypt.rs +++ b/bottlecap/src/secrets/decrypt.rs @@ -17,10 +17,7 @@ use tokio::time::Instant; use tracing::debug; use tracing::error; -pub async fn resolve_secrets( - config: Arc, - aws_config: Arc, -) -> Option { +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(); @@ -48,19 +45,16 @@ pub async fn resolve_secrets( && !aws_credentials .aws_container_credentials_full_uri .is_empty() - && !aws_credentials - .aws_container_authorization_token - .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; - } - }; + 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()