diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 1123a9436..91d399a62 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/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 { } 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] diff --git a/bottlecap/src/config/mod.rs b/bottlecap/src/config/mod.rs index e7775e263..598ba1e21 100644 --- a/bottlecap/src/config/mod.rs +++ b/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::*; diff --git a/bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs b/bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs index 69673cc2f..349d8866c 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs +++ b/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")] diff --git a/bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs b/bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs index 5e725163d..191de03ca 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs +++ b/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")] diff --git a/bottlecap/src/lifecycle/invocation/triggers/mod.rs b/bottlecap/src/lifecycle/invocation/triggers/mod.rs index c089aff75..afb875280 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/mod.rs +++ b/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, D::Error> diff --git a/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs b/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs index a4bf0e44f..aa0271cc4 100644 --- a/bottlecap/src/lifecycle/invocation/triggers/sqs_event.rs +++ b/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 { diff --git a/bottlecap/src/secrets/decrypt.rs b/bottlecap/src/secrets/decrypt.rs index 25bf92920..36442f164 100644 --- a/bottlecap/src/secrets/decrypt.rs +++ b/bottlecap/src/secrets/decrypt.rs @@ -170,10 +170,14 @@ fn build_get_secret_signed_headers( ) -> Result> { 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 = "";