From 92272e90a7919f07178f3246ef8f82295513cfed Mon Sep 17 00:00:00 2001
From: AJ Stuyvenberg <astuyve@gmail.com>
Date: Wed, 16 Oct 2024 16:28:20 -0400
Subject: [PATCH] feat: Prefer DD_PROXY_HTTPS over HTTPS_PROXY (#673)

* feat: Prefer DD_PROXY_HTTPS over HTTPS_PROXY

* fix: no proxy on ints

* fix: clippy thx
---
 dogstatsd/src/datadog.rs            | 21 ++++++++++++++++++---
 dogstatsd/src/flusher.rs            |  9 +++++++--
 dogstatsd/tests/integration_test.rs |  1 +
 serverless/src/main.rs              |  8 +++++++-
 trace-mini-agent/src/config.rs      |  4 +++-
 5 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/dogstatsd/src/datadog.rs b/dogstatsd/src/datadog.rs
index b95882975..e73024458 100644
--- a/dogstatsd/src/datadog.rs
+++ b/dogstatsd/src/datadog.rs
@@ -8,7 +8,7 @@ use protobuf::Message;
 use reqwest;
 use serde::{Serialize, Serializer};
 use serde_json;
-use tracing::debug;
+use tracing::{debug, error};
 
 /// Interface for the `DogStatsD` metrics intake API.
 #[derive(Debug)]
@@ -20,11 +20,18 @@ pub struct DdApi {
 
 impl DdApi {
     #[must_use]
-    pub fn new(api_key: String, site: String) -> Self {
+    pub fn new(api_key: String, site: String, https_proxy: Option<String>) -> Self {
+        let client = match Self::build_client(https_proxy) {
+            Ok(client) => client,
+            Err(e) => {
+                error!("Unable to parse proxy URL, no proxy will be used. {:?}", e);
+                reqwest::Client::new()
+            }
+        };
         DdApi {
             api_key,
             fqdn_site: site,
-            client: reqwest::Client::new(),
+            client,
         }
     }
 
@@ -96,6 +103,14 @@ impl DdApi {
             }
         };
     }
+
+    fn build_client(https_proxy: Option<String>) -> Result<reqwest::Client, reqwest::Error> {
+        let mut builder = reqwest::Client::builder();
+        if let Some(proxy) = https_proxy {
+            builder = builder.proxy(reqwest::Proxy::https(proxy)?);
+        }
+        builder.build()
+    }
 }
 
 #[derive(Debug, Serialize, Clone, Copy)]
diff --git a/dogstatsd/src/flusher.rs b/dogstatsd/src/flusher.rs
index 7969c7f86..601d223b7 100644
--- a/dogstatsd/src/flusher.rs
+++ b/dogstatsd/src/flusher.rs
@@ -18,8 +18,13 @@ pub fn build_fqdn_metrics(site: String) -> String {
 
 #[allow(clippy::await_holding_lock)]
 impl Flusher {
-    pub fn new(api_key: String, aggregator: Arc<Mutex<Aggregator>>, site: String) -> Self {
-        let dd_api = datadog::DdApi::new(api_key, site);
+    pub fn new(
+        api_key: String,
+        aggregator: Arc<Mutex<Aggregator>>,
+        site: String,
+        https_proxy: Option<String>,
+    ) -> Self {
+        let dd_api = datadog::DdApi::new(api_key, site, https_proxy);
         Flusher { dd_api, aggregator }
     }
 
diff --git a/dogstatsd/tests/integration_test.rs b/dogstatsd/tests/integration_test.rs
index 1acb6543e..7ffce5411 100644
--- a/dogstatsd/tests/integration_test.rs
+++ b/dogstatsd/tests/integration_test.rs
@@ -41,6 +41,7 @@ async fn dogstatsd_server_ships_series() {
         "mock-api-key".to_string(),
         Arc::clone(&metrics_aggr),
         mock_server.url(),
+        None,
     );
 
     let server_address = "127.0.0.1:18125";
diff --git a/serverless/src/main.rs b/serverless/src/main.rs
index ff09303dd..4352d1cb8 100644
--- a/serverless/src/main.rs
+++ b/serverless/src/main.rs
@@ -40,6 +40,9 @@ pub async fn main() {
         .map(|val| val.to_lowercase() != "false")
         .unwrap_or(true);
 
+    let https_proxy = env::var("DD_PROXY_HTTPS")
+        .or_else(|_| env::var("HTTPS_PROXY"))
+        .ok();
     debug!("Starting serverless trace mini agent");
 
     let mini_agent_version = env!("CARGO_PKG_VERSION").to_string();
@@ -79,7 +82,8 @@ pub async fn main() {
 
     let mut metrics_flusher = if dd_use_dogstatsd {
         debug!("Starting dogstatsd");
-        let (_, metrics_flusher) = start_dogstatsd(dd_dogstatsd_port, dd_api_key, dd_site).await;
+        let (_, metrics_flusher) =
+            start_dogstatsd(dd_dogstatsd_port, dd_api_key, dd_site, https_proxy).await;
         info!("dogstatsd-udp: starting to listen on port {dd_dogstatsd_port}");
         metrics_flusher
     } else {
@@ -104,6 +108,7 @@ async fn start_dogstatsd(
     port: u16,
     dd_api_key: Option<String>,
     dd_site: String,
+    https_proxy: Option<String>,
 ) -> (CancellationToken, Option<Flusher>) {
     let metrics_aggr = Arc::new(Mutex::new(
         MetricsAggregator::new(EMPTY_TAGS, CONTEXTS).expect("Failed to create metrics aggregator"),
@@ -131,6 +136,7 @@ async fn start_dogstatsd(
                 dd_api_key,
                 Arc::clone(&metrics_aggr),
                 build_fqdn_metrics(dd_site),
+                https_proxy,
             );
             Some(metrics_flusher)
         }
diff --git a/trace-mini-agent/src/config.rs b/trace-mini-agent/src/config.rs
index bcd4a666f..d93a7457f 100644
--- a/trace-mini-agent/src/config.rs
+++ b/trace-mini-agent/src/config.rs
@@ -90,7 +90,9 @@ impl Config {
                 ..Default::default()
             },
             obfuscation_config,
-            proxy_url: env::var("HTTPS_PROXY").ok(),
+            proxy_url: env::var("DD_PROXY_HTTPS")
+                .or_else(|_| env::var("HTTPS_PROXY"))
+                .ok(),
         })
     }
 }