chore(config): create aws module (#659)

# What?

Refactors methods related to AWS config into its own module

# Motivation

Just cleaning and removing stuff from main
– [SVLS-6686](https://datadoghq.atlassian.net/browse/SVLS-6686)

[SVLS-6686]:
https://datadoghq.atlassian.net/browse/SVLS-6686?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ

Author: duncanista

bottlecap/src/bin/bottlecap/main.rs

@@ -12,7 +12,11 @@
 use bottlecap::lwa::proxy::start_lwa_proxy;
 use bottlecap::{
     base_url,
-    config::{self, get_aws_partition_by_region, AwsConfig, Config},
+    config::{
+        self,
+        aws::{build_lambda_function_arn, AwsConfig},
+        Config,
+    },
     event_bus::bus::EventBus,
     events::Event,
     lifecycle::{
@@ -179,11 +183,6 @@ async fn register(client: &Client) -> Result<RegisterResponse> {
     Ok(register_response)
 }
 
-fn build_function_arn(account_id: &str, region: &str, function_name: &str) -> String {
-    let aws_partition = get_aws_partition_by_region(region);
-    format!("arn:{aws_partition}:lambda:{region}:{account_id}:function:{function_name}")
-}
-
 #[tokio::main]
 async fn main() -> Result<()> {
     let start_time = Instant::now();
@@ -232,20 +231,10 @@ async fn main() -> Result<()> {
 }
 
 fn load_configs(start_time: Instant) -> (AwsConfig, Arc<Config>) {
-    // First load the configuration
-    let aws_config = AwsConfig {
-        region: env::var("AWS_DEFAULT_REGION").unwrap_or("us-east-1".to_string()),
-        aws_access_key_id: env::var("AWS_ACCESS_KEY_ID").unwrap_or_default(),
-        aws_secret_access_key: env::var("AWS_SECRET_ACCESS_KEY").unwrap_or_default(),
-        aws_session_token: env::var("AWS_SESSION_TOKEN").unwrap_or_default(),
-        aws_container_credentials_full_uri: env::var("AWS_CONTAINER_CREDENTIALS_FULL_URI")
-            .unwrap_or_default(),
-        aws_container_authorization_token: env::var("AWS_CONTAINER_AUTHORIZATION_TOKEN")
-            .unwrap_or_default(),
-        function_name: env::var("AWS_LAMBDA_FUNCTION_NAME").unwrap_or_default(),
-        sandbox_init_time: start_time,
-    };
-    let lambda_directory = env::var("LAMBDA_TASK_ROOT").unwrap_or_else(|_| "/var/task".to_string());
+    // First load the AWS configuration
+    let aws_config = AwsConfig::from_env(start_time);
+    let lambda_directory: String =
+        env::var("LAMBDA_TASK_ROOT").unwrap_or_else(|_| "/var/task".to_string());
     let config = match config::get_config(Path::new(&lambda_directory), &aws_config.region) {
         Ok(config) => Arc::new(config),
         Err(_e) => {
@@ -630,7 +619,7 @@ fn setup_tag_provider(
     account_id: &str,
 ) -> Arc<TagProvider> {
     let function_arn =
-        build_function_arn(account_id, &aws_config.region, &aws_config.function_name);
+        build_lambda_function_arn(account_id, &aws_config.region, &aws_config.function_name);
     let metadata_hash = hash_map::HashMap::from([(
         lambda::tags::FUNCTION_ARN_KEY.to_string(),
         function_arn.clone(),

bottlecap/src/config/aws.rs

@@ -0,0 +1,55 @@
+use std::{env, time::Instant};
+
+const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION";
+const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID";
+const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY";
+const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN";
+const AWS_CONTAINER_CREDENTIALS_FULL_URI: &str = "AWS_CONTAINER_CREDENTIALS_FULL_URI";
+const AWS_CONTAINER_AUTHORIZATION_TOKEN: &str = "AWS_CONTAINER_AUTHORIZATION_TOKEN";
+const AWS_LAMBDA_FUNCTION_NAME: &str = "AWS_LAMBDA_FUNCTION_NAME";
+
+#[allow(clippy::module_name_repetitions)]
+#[derive(Debug, Clone)]
+pub struct AwsConfig {
+    pub region: String,
+    pub aws_access_key_id: String,
+    pub aws_secret_access_key: String,
+    pub aws_session_token: String,
+    pub function_name: String,
+    pub sandbox_init_time: Instant,
+    pub aws_container_credentials_full_uri: String,
+    pub aws_container_authorization_token: String,
+}
+
+impl AwsConfig {
+    #[must_use]
+    pub fn from_env(start_time: Instant) -> Self {
+        Self {
+            region: env::var(AWS_DEFAULT_REGION).unwrap_or("us-east-1".to_string()),
+            aws_access_key_id: env::var(AWS_ACCESS_KEY_ID).unwrap_or_default(),
+            aws_secret_access_key: env::var(AWS_SECRET_ACCESS_KEY).unwrap_or_default(),
+            aws_session_token: env::var(AWS_SESSION_TOKEN).unwrap_or_default(),
+            aws_container_credentials_full_uri: env::var(AWS_CONTAINER_CREDENTIALS_FULL_URI)
+                .unwrap_or_default(),
+            aws_container_authorization_token: env::var(AWS_CONTAINER_AUTHORIZATION_TOKEN)
+                .unwrap_or_default(),
+            function_name: env::var(AWS_LAMBDA_FUNCTION_NAME).unwrap_or_default(),
+            sandbox_init_time: start_time,
+        }
+    }
+}
+
+#[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(),
+    }
+}
+
+#[must_use]
+pub fn build_lambda_function_arn(account_id: &str, region: &str, function_name: &str) -> String {
+    let aws_partition = get_aws_partition_by_region(region);
+    format!("arn:{aws_partition}:lambda:{region}:{account_id}:function:{function_name}")
+}

bottlecap/src/config/mod.rs

@@ -1,4 +1,5 @@
 pub mod apm_replace_rule;
+pub mod aws;
 pub mod env;
 pub mod flush_strategy;
 pub mod log_level;
@@ -9,7 +10,6 @@ pub mod yaml;
 
 use datadog_trace_utils::config_utils::{trace_intake_url, trace_intake_url_prefixed};
 use std::path::Path;
-use std::time::Instant;
 
 use figment::providers::{Format, Yaml};
 use figment::{providers::Env, Figment};
@@ -271,28 +271,6 @@ fn build_fqdn_logs(site: String) -> String {
     format!("https://http-intake.logs.{site}")
 }
 
-#[allow(clippy::module_name_repetitions)]
-#[derive(Debug, Clone)]
-pub struct AwsConfig {
-    pub region: String,
-    pub aws_access_key_id: String,
-    pub aws_secret_access_key: String,
-    pub aws_session_token: String,
-    pub function_name: String,
-    pub sandbox_init_time: Instant,
-    pub aws_container_credentials_full_uri: String,
-    pub aws_container_authorization_token: String,
-}
-
-#[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 datadog_trace_obfuscation::replacer::parse_rules_from_string;

bottlecap/src/lifecycle/invocation/processor.rs

@@ -13,7 +13,7 @@ use tokio::sync::{mpsc::Sender, watch};
 use tracing::{debug, warn};
 
 use crate::{
-    config::{self, AwsConfig},
+    config::{self, aws::AwsConfig},
     lifecycle::invocation::{
         base64_to_string, context::Context, context::ContextBuffer, context::ReparentingInfo,
         create_empty_span, generate_span_id, get_metadata_from_value, span_inferrer::SpanInferrer,

bottlecap/src/lifecycle/invocation/span_inferrer.rs

@@ -4,27 +4,28 @@ use datadog_trace_protobuf::pb::Span;
 use serde_json::Value;
 use tracing::debug;
 
-use crate::config::AwsConfig;
-
-use crate::lifecycle::invocation::{
-    generate_span_id,
-    triggers::{
-        api_gateway_http_event::APIGatewayHttpEvent,
-        api_gateway_rest_event::APIGatewayRestEvent,
-        api_gateway_websocket_event::APIGatewayWebSocketEvent,
-        dynamodb_event::DynamoDbRecord,
-        event_bridge_event::EventBridgeEvent,
-        kinesis_event::KinesisRecord,
-        lambda_function_url_event::LambdaFunctionUrlEvent,
-        s3_event::S3Record,
-        sns_event::{SnsEntity, SnsRecord},
-        sqs_event::{extract_trace_context_from_aws_trace_header, SqsRecord},
-        step_function_event::StepFunctionEvent,
-        Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG,
-    },
-};
 use crate::traces::span_pointers::SpanPointer;
 use crate::traces::{context::SpanContext, propagation::Propagator};
+use crate::{
+    config::aws::AwsConfig,
+    lifecycle::invocation::{
+        generate_span_id,
+        triggers::{
+            api_gateway_http_event::APIGatewayHttpEvent,
+            api_gateway_rest_event::APIGatewayRestEvent,
+            api_gateway_websocket_event::APIGatewayWebSocketEvent,
+            dynamodb_event::DynamoDbRecord,
+            event_bridge_event::EventBridgeEvent,
+            kinesis_event::KinesisRecord,
+            lambda_function_url_event::LambdaFunctionUrlEvent,
+            s3_event::S3Record,
+            sns_event::{SnsEntity, SnsRecord},
+            sqs_event::{extract_trace_context_from_aws_trace_header, SqsRecord},
+            step_function_event::StepFunctionEvent,
+            Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG,
+        },
+    },
+};
 
 #[derive(Default)]
 pub struct SpanInferrer {

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

@@ -1,4 +1,4 @@
-use crate::config::get_aws_partition_by_region;
+use crate::config::aws::get_aws_partition_by_region;
 use crate::lifecycle::invocation::{
     processor::MS_TO_NS,
     triggers::{lowercase_key, ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG},

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

@@ -1,4 +1,4 @@
-use crate::config::get_aws_partition_by_region;
+use crate::config::aws::get_aws_partition_by_region;
 use crate::lifecycle::invocation::{
     processor::MS_TO_NS,
     triggers::{lowercase_key, ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG},

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

@@ -1,5 +1,5 @@
 use crate::{
-    config::get_aws_partition_by_region,
+    config::aws::get_aws_partition_by_region,
     lifecycle::invocation::{
         processor::MS_TO_NS,
         triggers::{lowercase_key, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG},

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

@@ -1,4 +1,4 @@
-use crate::config::get_aws_partition_by_region;
+use crate::config::aws::get_aws_partition_by_region;
 use crate::lifecycle::invocation::{
     processor::MS_TO_NS,
     triggers::{

bottlecap/src/lwa/proxy.rs

@@ -397,7 +397,7 @@ fn build_forward_request(
 #[cfg(test)]
 #[allow(clippy::unwrap_used)]
 mod tests {
-    use crate::config::{AwsConfig, Config};
+    use crate::config::{aws::AwsConfig, Config};
     use crate::lifecycle::invocation::processor::Processor;
     use crate::lwa::proxy::start_lwa_proxy;
     use crate::tags::provider::Provider;