feat: add AWS delegated authentication support

.gitlab/datasources/test-suites.yaml

@@ -3,3 +3,4 @@ test_suites:
   - name: otlp
   - name: snapstart
   - name: lmi
+  - name: auth

bottlecap/src/bin/bottlecap/main.rs

@@ -150,7 +150,11 @@ async fn main() -> anyhow::Result<()> {
     let config = Arc::new(config::get_config(Path::new(&lambda_directory)));
 
     let aws_config = Arc::new(aws_config);
-    let api_key_factory = create_api_key_factory(&config, &aws_config);
+    // Build one shared reqwest::Client for metrics, logs, trace proxy flushing, and calls to
+    // Datadog APIs (e.g. delegated auth). reqwest::Client is Arc-based internally, so cloning
+    // just increments a refcount and shares the connection pool.
+    let shared_client = bottlecap::http::get_client(&config);
+    let api_key_factory = create_api_key_factory(&config, &aws_config, &shared_client);
 
     let r = response
         .await
@@ -161,6 +165,7 @@ async fn main() -> anyhow::Result<()> {
         Arc::clone(&aws_config),
         &config,
         &client,
+        shared_client,
         &r,
         Arc::clone(&api_key_factory),
         start_time,
@@ -246,17 +251,23 @@ fn get_flush_strategy_for_mode(
     }
 }
 
-fn create_api_key_factory(config: &Arc<Config>, aws_config: &Arc<AwsConfig>) -> Arc<ApiKeyFactory> {
+fn create_api_key_factory(
+    config: &Arc<Config>,
+    aws_config: &Arc<AwsConfig>,
+    client: &reqwest::Client,
+) -> Arc<ApiKeyFactory> {
     let config = Arc::clone(config);
     let aws_config = Arc::clone(aws_config);
+    let client = client.clone();
     let api_key_secret_reload_interval = config.api_key_secret_reload_interval;
 
     Arc::new(ApiKeyFactory::new_from_resolver(
         Arc::new(move || {
             let config = Arc::clone(&config);
             let aws_config = Arc::clone(&aws_config);
+            let client = client.clone();
 
-            Box::pin(async move { resolve_secrets(config, aws_config).await })
+            Box::pin(async move { resolve_secrets(config, aws_config, client).await })
         }),
         api_key_secret_reload_interval,
     ))
@@ -285,6 +296,7 @@ async fn extension_loop_active(
     aws_config: Arc<AwsConfig>,
     config: &Arc<Config>,
     client: &Client,
+    shared_client: reqwest::Client,
     r: &RegisterResponse,
     api_key_factory: Arc<ApiKeyFactory>,
     start_time: Instant,
@@ -294,11 +306,6 @@ async fn extension_loop_active(
     let account_id = r.account_id.as_ref().unwrap_or(&"none".to_string()).clone();
     let tags_provider = setup_tag_provider(&Arc::clone(&aws_config), config, &account_id);
 
-    // Build one shared reqwest::Client for metrics, logs, and trace proxy flushing.
-    // reqwest::Client is Arc-based internally, so cloning just increments a refcount
-    // and shares the connection pool.
-    let shared_client = bottlecap::http::get_client(config);
-
     let (
         logs_agent_channel,
         logs_flusher,

bottlecap/src/config/env.rs

@@ -482,6 +482,12 @@ pub struct EnvConfig {
     /// The delay between two samples of the API Security schema collection, in seconds.
     #[serde(deserialize_with = "deserialize_optional_duration_from_seconds")]
     pub api_security_sample_delay: Option<Duration>,
+
+    /// @env `DD_ORG_UUID`
+    ///
+    /// The Datadog organization UUID. When set, delegated auth is auto-enabled.
+    #[serde(deserialize_with = "deserialize_string_or_int")]
+    pub org_uuid: Option<String>,
 }
 
 #[allow(clippy::too_many_lines)]
@@ -684,6 +690,8 @@ fn merge_config(config: &mut Config, env_config: &EnvConfig) {
     merge_option_to_value!(config, env_config, appsec_waf_timeout);
     merge_option_to_value!(config, env_config, api_security_enabled);
     merge_option_to_value!(config, env_config, api_security_sample_delay);
+
+    merge_string!(config, dd_org_uuid, env_config, org_uuid);
 }
 
 #[derive(Debug, PartialEq, Clone, Copy)]
@@ -1044,6 +1052,8 @@ mod tests {
                 appsec_waf_timeout: Duration::from_secs(1),
                 api_security_enabled: false,
                 api_security_sample_delay: Duration::from_secs(60),
+
+                dd_org_uuid: String::default(),
             };
 
             assert_eq!(config, expected_config);

bottlecap/src/config/mod.rs

@@ -364,6 +364,8 @@ pub struct Config {
     pub span_dedup_timeout: Option<Duration>,
     pub api_key_secret_reload_interval: Option<Duration>,
 
+    pub dd_org_uuid: String,
+
     pub serverless_appsec_enabled: bool,
     pub appsec_rules: Option<String>,
     pub appsec_waf_timeout: Duration,
@@ -479,6 +481,8 @@ impl Default for Config {
             span_dedup_timeout: None,
             api_key_secret_reload_interval: None,
 
+            dd_org_uuid: String::default(),
+
             serverless_appsec_enabled: false,
             appsec_rules: None,
             appsec_waf_timeout: Duration::from_millis(5),

bottlecap/src/config/yaml.rs

@@ -1036,6 +1036,8 @@ api_security_sample_delay: 60 # Seconds
                 dogstatsd_so_rcvbuf: Some(1_048_576),
                 dogstatsd_buffer_size: Some(65507),
                 dogstatsd_queue_size: Some(2048),
+
+                dd_org_uuid: String::default(),
             };
 
             // Assert that

bottlecap/src/secrets/decrypt.rs

@@ -14,13 +14,19 @@ use sha2::{Digest, Sha256};
 use std::io::Error;
 use std::sync::Arc;
 use tokio::time::Instant;
-use tracing::debug;
-use tracing::error;
+use tracing::{debug, error};
 
-pub async fn resolve_secrets(config: Arc<Config>, aws_config: Arc<AwsConfig>) -> Option<String> {
+use crate::secrets::delegated_auth;
+
+pub async fn resolve_secrets(
+    config: Arc<Config>,
+    aws_config: Arc<AwsConfig>,
+    shared_client: Client,
+) -> Option<String> {
     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()
+        || !config.dd_org_uuid.is_empty()
     {
         let before_decrypt = Instant::now();
 
@@ -40,38 +46,17 @@ pub async fn resolve_secrets(config: Arc<Config>, aws_config: Arc<AwsConfig>) ->
             }
         };
 
-        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 aws_credentials = get_aws_credentials(&client).await?;
 
-        let decrypted_key = if !config.kms_api_key.is_empty() {
+        let decrypted_key = if !config.dd_org_uuid.is_empty() {
+            delegated_auth::get_delegated_api_key(
+                &config,
+                &aws_config,
+                &shared_client,
+                &aws_credentials,
+            )
+            .await
+        } else if !config.kms_api_key.is_empty() {
             decrypt_aws_kms(
                 &client,
                 config.kms_api_key.clone(),
@@ -258,6 +243,39 @@ async fn decrypt_aws_ssm(
     Err(Error::new(std::io::ErrorKind::InvalidData, v.to_string()).into())
 }
 
+async fn get_aws_credentials(client: &Client) -> Option<AwsCredentials> {
+    let mut aws_credentials = AwsCredentials::from_env();
+    // We're in SnapStart — fetch short-lived credentials from the container endpoint
+    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()
+    {
+        let credentials = match get_snapstart_credentials(&aws_credentials, client).await {
+            Ok(credentials) => credentials,
+            Err(err) => {
+                error!("Error getting SnapStart 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();
+    }
+    Some(aws_credentials)
+}
+
 async fn get_snapstart_credentials(
     aws_credentials: &AwsCredentials,
     client: &Client,