fix(bottlecap): try kms decrypt without context as in goagent

bottlecap/src/bin/bottlecap/main.rs

@@ -50,6 +50,7 @@ use dogstatsd::{
     dogstatsd::{DogStatsD, DogStatsDConfig},
     flusher::{build_fqdn_metrics, Flusher as MetricsFlusher},
 };
+use lazy_static::lazy_static;
 use reqwest::Client;
 use serde::Deserialize;
 use std::{
@@ -69,6 +70,11 @@ use tokio_util::sync::CancellationToken;
 use tracing::{debug, error};
 use tracing_subscriber::EnvFilter;
 
+lazy_static! {
+    static ref API_KEY_REGEX: regex::Regex =
+        regex::Regex::new(r"^[a-f0-9]{32}$").expect("Invalid regex for DD API KEY");
+}
+
 #[derive(Clone, Deserialize)]
 #[serde(rename_all = "camelCase")]
 struct RegisterResponse {
@@ -175,7 +181,9 @@ async fn main() -> Result<()> {
         .await
         .map_err(|e| Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
 
-    if let Some(resolved_api_key) = resolve_secrets(Arc::clone(&config), &aws_config).await {
+    if let Some(resolved_api_key) =
+        clean_api_key(resolve_secrets(Arc::clone(&config), &aws_config).await)
+    {
         match extension_loop_active(&aws_config, &config, &client, &r, resolved_api_key).await {
             Ok(()) => {
                 debug!("Extension loop completed successfully");
@@ -194,6 +202,17 @@ async fn main() -> Result<()> {
     }
 }
 
+fn clean_api_key(maybe_key: Option<String>) -> Option<String> {
+    if let Some(key) = maybe_key {
+        let clean_key = key.trim_end_matches('\n').replace(' ', "").to_string();
+        if API_KEY_REGEX.is_match(&clean_key) {
+            return Some(clean_key);
+        }
+        error!("API key has invalid format");
+    }
+    None
+}
+
 fn load_configs() -> (AwsConfig, Arc<Config>) {
     // First load the configuration
     let aws_config = AwsConfig {

bottlecap/src/secrets/decrypt.rs

@@ -59,10 +59,13 @@ async fn decrypt_aws_kms(
     kms_key: String,
     aws_config: &AwsConfig,
 ) -> Result<String, Box<dyn std::error::Error>> {
+    // When the API key is encrypted using the AWS console, the function name is added as an
+    // encryption context. When the API key is encrypted using the AWS CLI, no encryption context
+    // is added. We need to try decrypting the API key both with and without the encryption context.
+
     let json_body = &serde_json::json!({
-        "CiphertextBlob": kms_key,
-        "encryptionContext": { "LambdaFunctionName": aws_config.function_name }}
-    );
+        "CiphertextBlob": kms_key
+    });
 
     let headers = build_get_secret_signed_headers(
         aws_config,
@@ -80,7 +83,29 @@ async fn decrypt_aws_kms(
         let secret_string = String::from_utf8(BASE64_STANDARD.decode(secret_string_b64)?)?;
         Ok(secret_string)
     } else {
-        Err(Error::new(std::io::ErrorKind::InvalidData, v.to_string()).into())
+        let json_body = &serde_json::json!({
+            "CiphertextBlob": kms_key,
+            "encryptionContext": { "LambdaFunctionName": aws_config.function_name }}
+        );
+
+        let headers = build_get_secret_signed_headers(
+            aws_config,
+            RequestArgs {
+                service: "kms".to_string(),
+                body: json_body,
+                time: Utc::now(),
+                x_amz_target: "TrentService.Decrypt".to_string(),
+            },
+        );
+
+        let v = request(json_body, headers?, client).await?;
+
+        if let Some(secret_string_b64) = v["Plaintext"].as_str() {
+            let secret_string = String::from_utf8(BASE64_STANDARD.decode(secret_string_b64)?)?;
+            Ok(secret_string)
+        } else {
+            Err(Error::new(std::io::ErrorKind::InvalidData, v.to_string()).into())
+        }
     }
 }