From 9081eb0f7f0af7f1291b4f90cc6ecad7d0e3d0ba Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil@katanemo.com>
Date: Fri, 8 Nov 2024 15:17:39 -0600
Subject: [PATCH] obfuscate auth header (#254)

---
 crates/common/src/lib.rs                  |  1 +
 crates/common/src/pii.rs                  | 44 +++++++++++++++++++++++
 crates/llm_gateway/src/stream_context.rs  |  3 +-
 crates/prompt_gateway/src/http_context.rs |  4 +--
 4 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 crates/common/src/pii.rs

diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs
index f2c95bc5..0984bf89 100644
--- a/crates/common/src/lib.rs
+++ b/crates/common/src/lib.rs
@@ -11,3 +11,4 @@ pub mod ratelimit;
 pub mod routing;
 pub mod stats;
 pub mod tokenizer;
+pub mod pii;
diff --git a/crates/common/src/pii.rs b/crates/common/src/pii.rs
new file mode 100644
index 00000000..5be5c295
--- /dev/null
+++ b/crates/common/src/pii.rs
@@ -0,0 +1,44 @@
+pub fn obfuscate_auth_header(headers: &mut [(String, String)]) -> &[(String, String)] {
+    headers.iter_mut().for_each(|(key, value)| {
+        if key.to_lowercase() == "authorization" {
+            if value.starts_with("Bearer ") {
+                *value = "Bearer ***".to_string();
+            } else {
+                *value = "***".to_string();
+            }
+        }
+    });
+
+    headers
+}
+
+#[cfg(test)]
+mod test {
+    use crate::pii::obfuscate_auth_header;
+
+    #[test]
+    pub fn test_obfuscate_auth_header() {
+        let mut headers = vec![("Authorization".to_string(), "Bearer 1234".to_string())];
+        obfuscate_auth_header(&mut headers);
+        assert_eq!(
+            headers,
+            vec![("Authorization".to_string(), "Bearer ***".to_string())]
+        );
+    }
+
+    #[test]
+    pub fn test_obfuscate_no_auth_header_found() {
+        let mut headers = vec![
+            (":path".to_string(), "/healthz".to_string()),
+            (":method".to_string(), "POST".to_string()),
+        ];
+        obfuscate_auth_header(&mut headers);
+        assert_eq!(
+            headers,
+            vec![
+                (":path".to_string(), "/healthz".to_string()),
+                (":method".to_string(), "POST".to_string()),
+            ]
+        );
+    }
+}
diff --git a/crates/llm_gateway/src/stream_context.rs b/crates/llm_gateway/src/stream_context.rs
index 71896108..9c3db01d 100644
--- a/crates/llm_gateway/src/stream_context.rs
+++ b/crates/llm_gateway/src/stream_context.rs
@@ -10,6 +10,7 @@ use common::consts::{
 };
 use common::errors::ServerError;
 use common::llm_providers::LlmProviders;
+use common::pii::obfuscate_auth_header;
 use common::ratelimit::Header;
 use common::{ratelimit, routing, tokenizer};
 use http::StatusCode;
@@ -153,7 +154,7 @@ impl HttpContext for StreamContext {
         debug!(
             "on_http_request_headers S[{}] req_headers={:?}",
             self.context_id,
-            self.get_http_request_headers()
+            obfuscate_auth_header(&mut self.get_http_request_headers())
         );
 
         self.request_id = self.get_http_request_header(REQUEST_ID_HEADER);
diff --git a/crates/prompt_gateway/src/http_context.rs b/crates/prompt_gateway/src/http_context.rs
index 2989e048..f14114e4 100644
--- a/crates/prompt_gateway/src/http_context.rs
+++ b/crates/prompt_gateway/src/http_context.rs
@@ -13,7 +13,7 @@ use common::{
         HEALTHZ_PATH, REQUEST_ID_HEADER, TOOL_ROLE, TRACE_PARENT_HEADER, USER_ROLE,
     },
     errors::ServerError,
-    http::{CallArgs, Client},
+    http::{CallArgs, Client}, pii::obfuscate_auth_header,
 };
 use http::StatusCode;
 use log::{debug, trace, warn};
@@ -48,7 +48,7 @@ impl HttpContext for StreamContext {
         trace!(
             "on_http_request_headers S[{}] req_headers={:?}",
             self.context_id,
-            self.get_http_request_headers()
+            obfuscate_auth_header(&mut self.get_http_request_headers())
         );
 
         self.request_id = self.get_http_request_header(REQUEST_ID_HEADER);