feat: Support periodic reload for api key secret

bottlecap/Cargo.toml

@@ -61,8 +61,8 @@ datadog-trace-protobuf = { git = "https://github.com/DataDog/libdatadog", rev =
 datadog-trace-utils = { git = "https://github.com/DataDog/libdatadog", rev = "9405db9cb4ef733f3954c3ee77ce71a502e98e50" , features = ["mini_agent"] }
 datadog-trace-normalization = { git = "https://github.com/DataDog/libdatadog",  rev = "9405db9cb4ef733f3954c3ee77ce71a502e98e50" }
 datadog-trace-obfuscation = { git = "https://github.com/DataDog/libdatadog", rev = "9405db9cb4ef733f3954c3ee77ce71a502e98e50"  }
-dogstatsd = { git = "https://github.com/DataDog/serverless-components", rev = "b0b0cb9310d8d8f2038c00a46e3267e21dc3e287", default-features = false }
-datadog-fips = { git = "https://github.com/DataDog/serverless-components", rev = "fa1d2f4ea2c4c2596144a1f362935e56cf0cb3c7", default-features = false }
+dogstatsd = { git = "https://github.com/DataDog/serverless-components", rev = "936b3440a1ffc3dd68d040354b721a3042aad47d", default-features = false }
+datadog-fips = { git = "https://github.com/DataDog/serverless-components", rev = "936b3440a1ffc3dd68d040354b721a3042aad47d", default-features = false }
 libddwaf = { version = "1.28.1", git = "https://github.com/DataDog/libddwaf-rust", rev = "d1534a158d976bd4f747bf9fcc58e0712d2d17fc", default-features = false, features = ["serde"] }
 
 [dev-dependencies]

bottlecap/src/bin/bottlecap/main.rs

@@ -342,13 +342,17 @@ fn enable_logging_subsystem() {
 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 api_key_reload_interval = config.api_key_reload_interval;
 
-    Arc::new(ApiKeyFactory::new_from_resolver(Arc::new(move || {
-        let config = Arc::clone(&config);
-        let aws_config = Arc::clone(&aws_config);
+    Arc::new(ApiKeyFactory::new_from_resolver(
+        Arc::new(move || {
+            let config = Arc::clone(&config);
+            let aws_config = Arc::clone(&aws_config);
 
-        Box::pin(async move { resolve_secrets(config, aws_config).await })
-    })))
+            Box::pin(async move { resolve_secrets(config, aws_config).await })
+        }),
+        api_key_reload_interval,
+    ))
 }
 
 async fn extension_loop_idle(

bottlecap/src/config/env.rs

@@ -13,7 +13,8 @@ use crate::{
         deserialize_apm_filter_tags, deserialize_array_from_comma_separated_string,
         deserialize_key_value_pairs, deserialize_option_lossless,
         deserialize_optional_bool_from_anything, deserialize_optional_duration_from_microseconds,
-        deserialize_optional_duration_from_seconds, deserialize_optional_string,
+        deserialize_optional_duration_from_seconds,
+        deserialize_optional_duration_from_seconds_ignore_zero, deserialize_optional_string,
         deserialize_string_or_int,
         flush_strategy::FlushStrategy,
         log_level::LogLevel,
@@ -395,6 +396,13 @@ pub struct EnvConfig {
     /// Default is `false`.
     #[serde(deserialize_with = "deserialize_optional_bool_from_anything")]
     pub compute_trace_stats_on_extension: Option<bool>,
+    /// @env `DD_API_KEY_RELOAD_INTERVAL`
+    ///
+    /// The interval at which the Datadog API key is reloaded, in seconds.
+    /// If None, the API key will not be reloaded.
+    /// Default is `None`.
+    #[serde(deserialize_with = "deserialize_optional_duration_from_seconds_ignore_zero")]
+    pub api_key_reload_interval: Option<Duration>,
     /// @env `DD_SERVERLESS_APPSEC_ENABLED`
     ///
     /// Enable Application and API Protection (AAP), previously known as AppSec/ASM, for AWS Lambda.
@@ -598,6 +606,7 @@ fn merge_config(config: &mut Config, env_config: &EnvConfig) {
     merge_option_to_value!(config, env_config, capture_lambda_payload);
     merge_option_to_value!(config, env_config, capture_lambda_payload_max_depth);
     merge_option_to_value!(config, env_config, compute_trace_stats_on_extension);
+    merge_option!(config, env_config, api_key_reload_interval);
     merge_option_to_value!(config, env_config, serverless_appsec_enabled);
     merge_option!(config, env_config, appsec_rules);
     merge_option_to_value!(config, env_config, appsec_waf_timeout);
@@ -791,6 +800,7 @@ mod tests {
             jail.set_env("DD_CAPTURE_LAMBDA_PAYLOAD", "true");
             jail.set_env("DD_CAPTURE_LAMBDA_PAYLOAD_MAX_DEPTH", "5");
             jail.set_env("DD_COMPUTE_TRACE_STATS_ON_EXTENSION", "true");
+            jail.set_env("DD_API_KEY_RELOAD_INTERVAL", "10");
             jail.set_env("DD_SERVERLESS_APPSEC_ENABLED", "true");
             jail.set_env("DD_APPSEC_RULES", "/path/to/rules.json");
             jail.set_env("DD_APPSEC_WAF_TIMEOUT", "1000000"); // Microseconds
@@ -940,6 +950,7 @@ mod tests {
                 capture_lambda_payload: true,
                 capture_lambda_payload_max_depth: 5,
                 compute_trace_stats_on_extension: true,
+                api_key_reload_interval: Some(Duration::from_secs(10)),
                 serverless_appsec_enabled: true,
                 appsec_rules: Some("/path/to/rules.json".to_string()),
                 appsec_waf_timeout: Duration::from_secs(1),

bottlecap/src/config/mod.rs

@@ -343,6 +343,7 @@ pub struct Config {
     pub capture_lambda_payload: bool,
     pub capture_lambda_payload_max_depth: u32,
     pub compute_trace_stats_on_extension: bool,
+    pub api_key_reload_interval: Option<Duration>,
 
     pub serverless_appsec_enabled: bool,
     pub appsec_rules: Option<String>,
@@ -444,6 +445,7 @@ impl Default for Config {
             capture_lambda_payload: false,
             capture_lambda_payload_max_depth: 10,
             compute_trace_stats_on_extension: false,
+            api_key_reload_interval: None,
 
             serverless_appsec_enabled: false,
             appsec_rules: None,
@@ -745,6 +747,17 @@ pub fn deserialize_optional_duration_from_seconds<'de, D: Deserializer<'de>>(
     deserializer.deserialize_any(DurationVisitor)
 }
 
+// Like deserialize_optional_duration_from_seconds(), but return None if the value is 0
+pub fn deserialize_optional_duration_from_seconds_ignore_zero<'de, D: Deserializer<'de>>(
+    deserializer: D,
+) -> Result<Option<Duration>, D::Error> {
+    let duration: Option<Duration> = deserialize_optional_duration_from_seconds(deserializer)?;
+    if duration.is_some_and(|d| d.as_secs() == 0) {
+        return Ok(None);
+    }
+    Ok(duration)
+}
+
 #[cfg_attr(coverage_nightly, coverage(off))] // Test modules skew coverage metrics
 #[cfg(test)]
 pub mod tests {
@@ -1353,4 +1366,26 @@ pub mod tests {
             }
         );
     }
+
+    #[test]
+    fn test_parse_duration_from_seconds_ignore_zero() {
+        #[derive(Deserialize, Debug, PartialEq, Eq)]
+        struct Value {
+            #[serde(default)]
+            #[serde(deserialize_with = "deserialize_optional_duration_from_seconds_ignore_zero")]
+            duration: Option<Duration>,
+        }
+
+        assert_eq!(
+            serde_json::from_str::<Value>(r#"{"duration":1}"#).expect("failed to parse JSON"),
+            Value {
+                duration: Some(Duration::from_secs(1))
+            }
+        );
+
+        assert_eq!(
+            serde_json::from_str::<Value>(r#"{"duration":0}"#).expect("failed to parse JSON"),
+            Value { duration: None }
+        );
+    }
 }

bottlecap/src/config/yaml.rs

@@ -8,7 +8,8 @@ use crate::{
         deserialize_apm_replace_rules, deserialize_key_value_pair_array_to_hashmap,
         deserialize_option_lossless, deserialize_optional_bool_from_anything,
         deserialize_optional_duration_from_microseconds,
-        deserialize_optional_duration_from_seconds, deserialize_optional_string,
+        deserialize_optional_duration_from_seconds,
+        deserialize_optional_duration_from_seconds_ignore_zero, deserialize_optional_string,
         deserialize_processing_rules, deserialize_string_or_int,
         flush_strategy::FlushStrategy,
         log_level::LogLevel,
@@ -111,6 +112,8 @@ pub struct YamlConfig {
     pub capture_lambda_payload_max_depth: Option<u32>,
     #[serde(deserialize_with = "deserialize_optional_bool_from_anything")]
     pub compute_trace_stats_on_extension: Option<bool>,
+    #[serde(deserialize_with = "deserialize_optional_duration_from_seconds_ignore_zero")]
+    pub api_key_reload_interval: Option<Duration>,
     #[serde(deserialize_with = "deserialize_optional_bool_from_anything")]
     pub serverless_appsec_enabled: Option<bool>,
     #[serde(deserialize_with = "deserialize_optional_string")]
@@ -675,6 +678,7 @@ fn merge_config(config: &mut Config, yaml_config: &YamlConfig) {
     merge_option_to_value!(config, yaml_config, capture_lambda_payload);
     merge_option_to_value!(config, yaml_config, capture_lambda_payload_max_depth);
     merge_option_to_value!(config, yaml_config, compute_trace_stats_on_extension);
+    merge_option!(config, yaml_config, api_key_reload_interval);
     merge_option_to_value!(config, yaml_config, serverless_appsec_enabled);
     merge_option!(config, yaml_config, appsec_rules);
     merge_option_to_value!(config, yaml_config, appsec_waf_timeout);
@@ -844,6 +848,7 @@ lambda_proc_enhanced_metrics: false
 capture_lambda_payload: true
 capture_lambda_payload_max_depth: 5
 compute_trace_stats_on_extension: true
+api_key_reload_interval: 0
 serverless_appsec_enabled: true
 appsec_rules: "/path/to/rules.json"
 appsec_waf_timeout: 1000000 # Microseconds
@@ -976,6 +981,7 @@ api_security_sample_delay: 60 # Seconds
                 capture_lambda_payload: true,
                 capture_lambda_payload_max_depth: 5,
                 compute_trace_stats_on_extension: true,
+                api_key_reload_interval: None,
 
                 serverless_appsec_enabled: true,
                 appsec_rules: Some("/path/to/rules.json".to_string()),

bottlecap/src/logs/flusher.rs

@@ -60,7 +60,7 @@ impl Flusher {
                 if batch.is_empty() {
                     continue;
                 }
-                let req = self.create_request(batch.clone(), api_key).await;
+                let req = self.create_request(batch.clone(), api_key.as_str()).await;
                 set.spawn(async move { Self::send(req).await });
             }
         }

bottlecap/src/traces/proxy_flusher.rs

@@ -103,7 +103,7 @@ impl Flusher {
         } else {
             let mut aggregator = self.aggregator.lock().await;
             for pr in aggregator.get_batch() {
-                requests.push(self.create_request(pr, api_key).await);
+                requests.push(self.create_request(pr, api_key.as_str()).await);
             }
         }
 

bottlecap/src/traces/stats_flusher.rs

@@ -102,7 +102,8 @@ impl StatsFlusher for ServerlessStatsFlusher {
         let start = std::time::Instant::now();
 
         let resp =
-            stats_utils::send_stats_payload(serialized_stats_payload, endpoint, api_key).await;
+            stats_utils::send_stats_payload(serialized_stats_payload, endpoint, api_key.as_str())
+                .await;
         let elapsed = start.elapsed();
         debug!(
             "Stats request to {} took {} ms",

bottlecap/src/traces/trace_flusher.rs

@@ -129,7 +129,7 @@ impl TraceFlusher for ServerlessTraceFlusher {
         for trace_builders in all_batches {
             let traces: Vec<_> = trace_builders
                 .into_iter()
-                .map(|builder| builder.with_api_key(api_key))
+                .map(|builder| builder.with_api_key(api_key.as_str()))
                 .map(SendDataBuilder::build)
                 .collect();