Skip to content

fix: Fix deadlock when snap start and using secret for api key #853

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
lym953 merged 2 commits into from
Sep 19, 2025
Merged

fix: Fix deadlock when snap start and using secret for api key #853

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2646135
fix: Fix dead lock when snap start and using secret for api key
lym953 Sep 18, 2025
bc99e51
fmt
lym953 Sep 18, 2025
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
23 changes: 8 additions & 15 deletions bottlecap/src/bin/bottlecap/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand All @@ -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),
Expand Down Expand Up @@ -285,14 +285,13 @@ fn init_ustr() {
});
}

fn load_configs(start_time: Instant) -> (AwsConfig, AwsCredentials, Arc<Config>) {
fn load_configs(start_time: Instant) -> (AwsConfig, Arc<Config>) {
// First load the AWS configuration
let aws_config = AwsConfig::from_env(start_time);
let aws_credentials = AwsCredentials::from_env();
Copy link
Copy Markdown
Contributor Author

@lym953 lym953 Sep 18, 2025

Choose a reason for hiding this comment

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

No need to load credentials from main()

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<Config>) {
Expand All @@ -318,21 +317,15 @@ fn enable_logging_subsystem(config: &Arc<Config>) {
debug!("Logging subsystem enabled");
}

fn create_api_key_factory(
config: &Arc<Config>,
aws_config: &Arc<AwsConfig>,
aws_credentials: AwsCredentials,
) -> Arc<ApiKeyFactory> {
fn create_api_key_factory(config: &Arc<Config>, aws_config: &Arc<AwsConfig>) -> Arc<ApiKeyFactory> {
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 })
})))
}

Expand Down
46 changes: 19 additions & 27 deletions bottlecap/src/secrets/decrypt.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,11 @@ 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<Config>,
aws_config: Arc<AwsConfig>,
aws_credentials: Arc<RwLock<AwsCredentials>>,
) -> Option<String> {
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();
Expand All @@ -42,36 +38,32 @@ pub async fn resolve_secrets(
}
};

let aws_credentials_read = aws_credentials.read().await;
let mut aws_credentials = AwsCredentials::from_env();
Copy link
Copy Markdown
Contributor Author

@lym953 lym953 Sep 18, 2025

Choose a reason for hiding this comment

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

Load it as needed in resolve_secrets(), which is called lazily.


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_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_read, &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"]
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_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();
Expand All @@ -82,15 +74,15 @@ pub async fn resolve_secrets(
&client,
config.api_key_secret_arn.clone(),
aws_config,
&aws_credentials_read,
&aws_credentials,
)
.await
} else {
decrypt_aws_kms(
&client,
config.kms_api_key.clone(),
aws_config,
&aws_credentials_read,
&aws_credentials,
)
.await
};
Expand Down
Loading