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

# Problem
When a Lambda (1) uses snap start, and (2) specifies Datadog API key
using `DD_API_KEY_SECRET_ARN`, the extension will encounter a deadlock.
For a `RwLock`, the extension first gets a read lock:

https://github.com/DataDog/datadog-lambda-extension/blob/daf633dd003447d78261e7c371838b5af21073a1/bottlecap/src/secrets/decrypt.rs#L45
then tries to get a write lock:

https://github.com/DataDog/datadog-lambda-extension/blob/daf633dd003447d78261e7c371838b5af21073a1/bottlecap/src/secrets/decrypt.rs#L65

which never finishes. This causes the function to time out.

This bug was introduced in
#717.

# This PR
Fix this bug by removing the `RwLock` usage. `AwsCredential` is only
created and used once in `resolve_secrets()`, and `resolve_secrets()` is
only called once, so there's no need to protect this struct with a lock.

# Testing
Tested on a Lambda with:
- Python 3.13 runtime
- snap start
- using `DD_API_KEY_SECRET_ARN`

Before:
- The function timed out.
- Data failed to be sent to Datadog.

After:
- The function finished without timeout.
- Data was sent to Datadog successfully.

# Notes
Jira: https://datadoghq.atlassian.net/browse/SLES-2482

Author: lym953

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<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();
     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>) {
@@ -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 })
     })))
 }
 

bottlecap/src/secrets/decrypt.rs

@@ -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();
@@ -42,36 +38,32 @@ 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_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();
@@ -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
             };