fix: set right domain and arn by region on secrets manager (#511)

* check whether the region is in China and use the appropriated domain

* correct arn for lambda in chinese regions

* fix: typo in china arn

* fix: reuse function to detect right aws partition and support gov too

* nest and rearrange imports

* fix imports again

Author: alexgallotta

bottlecap/src/bin/bottlecap/main.rs

@@ -11,7 +11,7 @@
 
 use bottlecap::{
     base_url,
-    config::{self, flush_strategy::FlushStrategy, AwsConfig, Config},
+    config::{self, flush_strategy::FlushStrategy, get_aws_partition_by_region, AwsConfig, Config},
     event_bus::bus::EventBus,
     events::Event,
     lifecycle::{
@@ -43,16 +43,15 @@ use bottlecap::{
 };
 use datadog_trace_obfuscation::obfuscation_config;
 use decrypt::resolve_secrets;
-use dogstatsd::metric::{SortedTags, EMPTY_TAGS};
 use dogstatsd::{
     aggregator::Aggregator as MetricsAggregator,
     constants::CONTEXTS,
     dogstatsd::{DogStatsD, DogStatsDConfig},
     flusher::{build_fqdn_metrics, Flusher as MetricsFlusher},
+    metric::{SortedTags, EMPTY_TAGS},
 };
 use reqwest::Client;
 use serde::Deserialize;
-use std::time::Duration;
 use std::{
     collections::{hash_map, HashMap},
     env,
@@ -61,11 +60,11 @@ use std::{
     path::Path,
     process::Command,
     sync::{Arc, Mutex},
+    time::Duration,
     time::Instant,
 };
 use telemetry::listener::TelemetryListenerConfig;
-use tokio::sync::mpsc::Sender;
-use tokio::sync::Mutex as TokioMutex;
+use tokio::{sync::mpsc::Sender, sync::Mutex as TokioMutex};
 use tokio_util::sync::CancellationToken;
 use tracing::{debug, error};
 use tracing_subscriber::EnvFilter;
@@ -156,7 +155,8 @@ async fn register(client: &reqwest::Client) -> Result<RegisterResponse> {
 }
 
 fn build_function_arn(account_id: &str, region: &str, function_name: &str) -> String {
-    format!("arn:aws:lambda:{region}:{account_id}:function:{function_name}")
+    let aws_partition = get_aws_partition_by_region(region);
+    format!("arn:{aws_partition}:lambda:{region}:{account_id}:function:{function_name}")
 }
 
 #[tokio::main]

bottlecap/src/config/mod.rs

@@ -311,6 +311,15 @@ pub struct AwsConfig {
     pub sandbox_init_time: Instant,
 }
 
+#[must_use]
+pub fn get_aws_partition_by_region(region: &str) -> String {
+    match region {
+        r if r.starts_with("us-gov-") => "aws-us-gov".to_string(),
+        r if r.starts_with("cn-") => "aws-cn".to_string(),
+        _ => "aws".to_string(),
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
     use super::*;

bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs

@@ -1,17 +1,14 @@
+use crate::config::get_aws_partition_by_region;
+use crate::lifecycle::invocation::{
+    processor::MS_TO_NS,
+    triggers::{lowercase_key, ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG},
+};
 use datadog_trace_protobuf::pb::Span;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::collections::HashMap;
 use tracing::debug;
 
-use crate::lifecycle::invocation::{
-    processor::MS_TO_NS,
-    triggers::{
-        get_aws_partition_by_region, lowercase_key, ServiceNameResolver, Trigger,
-        FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
-    },
-};
-
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct APIGatewayHttpEvent {
     #[serde(rename = "routeKey")]

bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs

@@ -1,17 +1,14 @@
+use crate::config::get_aws_partition_by_region;
+use crate::lifecycle::invocation::{
+    processor::MS_TO_NS,
+    triggers::{lowercase_key, ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG},
+};
 use datadog_trace_protobuf::pb::Span;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::collections::HashMap;
 use tracing::debug;
 
-use crate::lifecycle::invocation::{
-    processor::MS_TO_NS,
-    triggers::{
-        get_aws_partition_by_region, lowercase_key, ServiceNameResolver, Trigger,
-        FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
-    },
-};
-
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct APIGatewayRestEvent {
     #[serde(deserialize_with = "lowercase_key")]

bottlecap/src/lifecycle/invocation/triggers/mod.rs

@@ -57,15 +57,6 @@ pub trait Trigger: ServiceNameResolver {
     }
 }
 
-#[must_use]
-pub fn get_aws_partition_by_region(region: &str) -> String {
-    match region {
-        r if r.starts_with("us-gov-") => "aws-us-gov".to_string(),
-        r if r.starts_with("cn-") => "aws-cn".to_string(),
-        _ => "aws".to_string(),
-    }
-}
-
 /// Serialize a `HashMap` with lowercase keys
 ///
 pub fn lowercase_key<'de, D, V>(deserializer: D) -> Result<HashMap<String, V>, D::Error>

bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs

@@ -1,19 +1,18 @@
-use datadog_trace_protobuf::pb::Span;
-use serde::{Deserialize, Serialize};
-use serde_json::Value;
-use std::collections::HashMap;
-use tracing::debug;
-
+use crate::config::get_aws_partition_by_region;
 use crate::lifecycle::invocation::{
     processor::MS_TO_NS,
     triggers::{
         event_bridge_event::EventBridgeEvent,
-        get_aws_partition_by_region,
         sns_event::{SnsEntity, SnsRecord},
         ServiceNameResolver, Trigger, DATADOG_CARRIER_KEY, FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
     },
 };
 use crate::traces::context::{Sampling, SpanContext};
+use datadog_trace_protobuf::pb::Span;
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::collections::HashMap;
+use tracing::debug;
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct SqsEvent {

bottlecap/src/secrets/decrypt.rs

@@ -170,10 +170,14 @@ fn build_get_secret_signed_headers(
 ) -> Result<HeaderMap, Box<dyn std::error::Error>> {
     let amz_date = header_values.time.format("%Y%m%dT%H%M%SZ").to_string();
     let date_stamp = header_values.time.format("%Y%m%d").to_string();
-    let host = format!(
-        "{}.{}.amazonaws.com",
-        header_values.service, aws_config.region
-    );
+
+    let domain = if aws_config.region.starts_with("cn-") {
+        "amazonaws.com.cn"
+    } else {
+        "amazonaws.com"
+    };
+
+    let host = format!("{}.{}.{}", header_values.service, aws_config.region, domain);
 
     let canonical_uri = "/";
     let canonical_querystring = "";