libdd_trace_utils/
config_utils.rs

1// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::trace_utils;
5use std::env;
6use tracing::{debug, error};
7
8pub const PROD_INTAKE_SUBDOMAIN: &str = "trace.agent";
9
10const TRACE_INTAKE_ROUTE: &str = "/api/v0.2/traces";
11const TRACE_STATS_INTAKE_ROUTE: &str = "/api/v0.2/stats";
12
13pub fn read_cloud_env() -> Option<(String, trace_utils::EnvironmentType)> {
14    let mut detected: Vec<(String, trace_utils::EnvironmentType)> = Vec::new();
15
16    if env::var("AWS_LAMBDA_INITIALIZATION_TYPE").is_ok() {
17        match env::var("AWS_LAMBDA_FUNCTION_NAME") {
18            Ok(name) => detected.push((name, trace_utils::EnvironmentType::LambdaFunction)),
19            Err(_) => {
20                error!("AWS Lambda environment detected but AWS_LAMBDA_FUNCTION_NAME is not set");
21            }
22        }
23    }
24
25    if env::var("FUNCTIONS_EXTENSION_VERSION").is_ok()
26        && env::var("FUNCTIONS_WORKER_RUNTIME").is_ok()
27    {
28        match env::var("WEBSITE_SITE_NAME") {
29            Ok(name) => detected.push((name, trace_utils::EnvironmentType::AzureFunction)),
30            Err(_) => {
31                error!("Azure Functions environment detected but WEBSITE_SITE_NAME is not set");
32            }
33        }
34    }
35
36    if let (Ok(name), Ok(_)) = (env::var("K_SERVICE"), env::var("FUNCTION_TARGET")) {
37        // Set by Google Cloud Functions for newer runtimes
38        detected.push((name, trace_utils::EnvironmentType::CloudFunction));
39    } else if let (Ok(name), Ok(_)) = (env::var("FUNCTION_NAME"), env::var("GCP_PROJECT")) {
40        // Set by Google Cloud Functions for older runtimes
41        detected.push((name, trace_utils::EnvironmentType::CloudFunction));
42    }
43
44    if let Ok(name) = env::var("ASCSVCRT_SPRING__APPLICATION__NAME") {
45        // Set by Azure Spring Apps
46        detected.push((name, trace_utils::EnvironmentType::AzureSpringApp));
47    }
48
49    match detected.len() {
50        0 => {
51            error!("No cloud environment detected");
52            None
53        }
54        1 => {
55            let (ref name, ref env_type) = detected[0];
56            debug!("Cloud environment detected: {env_type:?} ({name})");
57            detected.into_iter().next()
58        }
59        _ => {
60            let env_names: Vec<String> = detected
61                .iter()
62                .map(|(name, env_type)| format!("{env_type:?}({name})"))
63                .collect();
64            error!(
65                "Multiple cloud environments detected: {}",
66                env_names.join(", ")
67            );
68            None
69        }
70    }
71}
72
73pub fn trace_intake_url(site: &str) -> String {
74    construct_trace_intake_url(site, TRACE_INTAKE_ROUTE)
75}
76
77pub fn trace_intake_url_prefixed(endpoint_prefix: &str) -> String {
78    format!("{endpoint_prefix}{TRACE_INTAKE_ROUTE}")
79}
80
81pub fn trace_stats_url(site: &str) -> String {
82    construct_trace_intake_url(site, TRACE_STATS_INTAKE_ROUTE)
83}
84
85pub fn trace_stats_url_prefixed(endpoint_prefix: &str) -> String {
86    format!("{endpoint_prefix}{TRACE_STATS_INTAKE_ROUTE}")
87}
88
89fn construct_trace_intake_url(prefix: &str, route: &str) -> String {
90    format!("https://{PROD_INTAKE_SUBDOMAIN}.{prefix}{route}")
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use std::sync::Mutex;
97
98    // Mutex to ensure environment variable tests run sequentially
99    static ENV_TEST_LOCK: Mutex<()> = Mutex::new(());
100
101    fn clear_all_env_vars() {
102        unsafe {
103            env::remove_var("AWS_LAMBDA_INITIALIZATION_TYPE");
104            env::remove_var("AWS_LAMBDA_FUNCTION_NAME");
105            env::remove_var("FUNCTIONS_EXTENSION_VERSION");
106            env::remove_var("FUNCTIONS_WORKER_RUNTIME");
107            env::remove_var("WEBSITE_SITE_NAME");
108            env::remove_var("FUNCTION_NAME");
109            env::remove_var("GCP_PROJECT");
110            env::remove_var("K_SERVICE");
111            env::remove_var("FUNCTION_TARGET");
112            env::remove_var("ASCSVCRT_SPRING__APPLICATION__NAME");
113        }
114    }
115
116    #[test]
117    fn test_aws_lambda_detected() {
118        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
119        clear_all_env_vars();
120        unsafe {
121            env::set_var("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand");
122            env::set_var("AWS_LAMBDA_FUNCTION_NAME", "my-function");
123        }
124        let result = read_cloud_env();
125        assert_eq!(
126            result,
127            Some((
128                "my-function".to_string(),
129                trace_utils::EnvironmentType::LambdaFunction
130            ))
131        );
132    }
133
134    #[test]
135    fn test_aws_lambda_missing_function_name() {
136        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
137        clear_all_env_vars();
138        unsafe { env::set_var("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand") };
139        let result = read_cloud_env();
140        assert_eq!(result, None);
141    }
142
143    #[test]
144    fn test_aws_lambda_not_detected_without_init_type() {
145        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
146        clear_all_env_vars();
147        unsafe { env::set_var("AWS_LAMBDA_FUNCTION_NAME", "my-function") };
148        let result = read_cloud_env();
149        assert_eq!(result, None);
150    }
151
152    #[test]
153    fn test_azure_function_detected() {
154        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
155        clear_all_env_vars();
156        unsafe {
157            env::set_var("FUNCTIONS_EXTENSION_VERSION", "~4");
158            env::set_var("FUNCTIONS_WORKER_RUNTIME", "java");
159            env::set_var("WEBSITE_SITE_NAME", "my-azure-app");
160        }
161        let result = read_cloud_env();
162        assert_eq!(
163            result,
164            Some((
165                "my-azure-app".to_string(),
166                trace_utils::EnvironmentType::AzureFunction
167            ))
168        );
169    }
170
171    #[test]
172    fn test_azure_function_missing_site_name() {
173        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
174        clear_all_env_vars();
175        unsafe {
176            env::set_var("FUNCTIONS_EXTENSION_VERSION", "~4");
177            env::set_var("FUNCTIONS_WORKER_RUNTIME", "java");
178        }
179        let result = read_cloud_env();
180        assert_eq!(result, None);
181    }
182
183    #[test]
184    fn test_azure_function_not_detected_with_only_one_var() {
185        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
186        clear_all_env_vars();
187        unsafe { env::set_var("FUNCTIONS_EXTENSION_VERSION", "~4") };
188        let result = read_cloud_env();
189        assert_eq!(result, None);
190    }
191
192    #[test]
193    fn test_gcp_1st_gen_detected() {
194        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
195        clear_all_env_vars();
196        unsafe {
197            env::set_var("FUNCTION_NAME", "my-gcp-function");
198            env::set_var("GCP_PROJECT", "my-project");
199        }
200        let result = read_cloud_env();
201        assert_eq!(
202            result,
203            Some((
204                "my-gcp-function".to_string(),
205                trace_utils::EnvironmentType::CloudFunction
206            ))
207        );
208    }
209
210    #[test]
211    fn test_gcp_1st_gen_not_detected_without_gcp_project() {
212        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
213        clear_all_env_vars();
214        unsafe { env::set_var("FUNCTION_NAME", "my-gcp-function") };
215        let result = read_cloud_env();
216        assert_eq!(result, None);
217    }
218
219    #[test]
220    fn test_gcp_2nd_gen_detected() {
221        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
222        clear_all_env_vars();
223        unsafe {
224            env::set_var("K_SERVICE", "my-cloud-run-fn");
225            env::set_var("FUNCTION_TARGET", "myHandler");
226        }
227        let result = read_cloud_env();
228        assert_eq!(
229            result,
230            Some((
231                "my-cloud-run-fn".to_string(),
232                trace_utils::EnvironmentType::CloudFunction
233            ))
234        );
235    }
236
237    #[test]
238    fn test_gcp_2nd_gen_not_detected_without_function_target() {
239        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
240        clear_all_env_vars();
241        unsafe { env::set_var("K_SERVICE", "my-cloud-run-fn") };
242        let result = read_cloud_env();
243        assert_eq!(result, None);
244    }
245
246    #[test]
247    fn test_azure_spring_app_detected() {
248        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
249        clear_all_env_vars();
250        unsafe { env::set_var("ASCSVCRT_SPRING__APPLICATION__NAME", "my-spring-app") };
251        let result = read_cloud_env();
252        assert_eq!(
253            result,
254            Some((
255                "my-spring-app".to_string(),
256                trace_utils::EnvironmentType::AzureSpringApp
257            ))
258        );
259    }
260
261    #[test]
262    fn test_no_environment_detected() {
263        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
264        clear_all_env_vars();
265        let result = read_cloud_env();
266        assert_eq!(result, None);
267    }
268
269    #[test]
270    fn test_multiple_environments_returns_none() {
271        let _lock = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
272        clear_all_env_vars();
273        unsafe {
274            env::set_var("AWS_LAMBDA_INITIALIZATION_TYPE", "on-demand");
275            env::set_var("AWS_LAMBDA_FUNCTION_NAME", "my-lambda");
276            env::set_var("FUNCTIONS_EXTENSION_VERSION", "~4");
277            env::set_var("FUNCTIONS_WORKER_RUNTIME", "java");
278            env::set_var("WEBSITE_SITE_NAME", "my-azure-app");
279        }
280        let result = read_cloud_env();
281        assert_eq!(result, None);
282    }
283}
libdd_trace_utils/config_utils.rs

libdd_trace_utils/
config_utils.rs
