From e45b5cadc0f8f0dfbdf28494505fb34838c801cd Mon Sep 17 00:00:00 2001
From: Peter Tripp <peter@zed.dev>
Date: Mon, 17 Feb 2025 16:42:22 +0000
Subject: [PATCH] Redact Google Gemini API keys from error messaging and log
 (#24884)

Now:
```
ERROR assistant_context_editor] error sending request for url (https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:countTokens?key=REDACTED)
```

Release Notes:

- Improved redaction of Google Gemini keys from API errors in logs
---
 Cargo.lock                                  |  1 +
 crates/reqwest_client/Cargo.toml            |  1 +
 crates/reqwest_client/src/reqwest_client.rs | 21 +++++++++++++++++++--
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 12dd44c71ac25bb77a92918225c69de72c76f92c..6b809d43c578ff44cbf8e8cadd0152510f1a26a3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -11136,6 +11136,7 @@ dependencies = [
  "gpui",
  "http_client",
  "log",
+ "regex",
  "reqwest 0.12.8",
  "serde",
  "smol",
diff --git a/crates/reqwest_client/Cargo.toml b/crates/reqwest_client/Cargo.toml
index b1b911632374e100db7dc633ae317e584fe7ace8..e02e7b7e52f54d8bd750c81377b11d2abe0b08e6 100644
--- a/crates/reqwest_client/Cargo.toml
+++ b/crates/reqwest_client/Cargo.toml
@@ -28,6 +28,7 @@ serde.workspace = true
 smol.workspace = true
 log.workspace = true
 tokio = { workspace = true, features = ["rt", "rt-multi-thread"] }
+regex.workspace = true
 reqwest.workspace = true
 
 [dev-dependencies]
diff --git a/crates/reqwest_client/src/reqwest_client.rs b/crates/reqwest_client/src/reqwest_client.rs
index 7c84ec8bab91db394d91f46dbef0e7f5745f55f1..010d0f653912196405af3af0533a4651d71b7117 100644
--- a/crates/reqwest_client/src/reqwest_client.rs
+++ b/crates/reqwest_client/src/reqwest_client.rs
@@ -1,9 +1,11 @@
-use std::{any::type_name, mem, pin::Pin, sync::OnceLock, task::Poll, time::Duration};
+use std::sync::{LazyLock, OnceLock};
+use std::{any::type_name, borrow::Cow, mem, pin::Pin, task::Poll, time::Duration};
 
 use anyhow::anyhow;
 use bytes::{BufMut, Bytes, BytesMut};
 use futures::{AsyncRead, TryStreamExt as _};
 use http_client::{http, RedirectPolicy};
+use regex::Regex;
 use reqwest::{
     header::{HeaderMap, HeaderValue},
     redirect,
@@ -12,6 +14,7 @@ use smol::future::FutureExt;
 
 const DEFAULT_CAPACITY: usize = 4096;
 static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
+static REDACT_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"key=[^&]+").unwrap());
 
 pub struct ReqwestClient {
     client: reqwest::Client,
@@ -180,6 +183,17 @@ pub fn poll_read_buf(
     Poll::Ready(Ok(n))
 }
 
+fn redact_error(mut error: reqwest::Error) -> reqwest::Error {
+    if let Some(url) = error.url_mut() {
+        if let Some(query) = url.query() {
+            if let Cow::Owned(redacted) = REDACT_REGEX.replace_all(query, "key=REDACTED") {
+                url.set_query(Some(redacted.as_str()));
+            }
+        }
+    }
+    error
+}
+
 impl http_client::HttpClient for ReqwestClient {
     fn proxy(&self) -> Option<&http::Uri> {
         self.proxy.as_ref()
@@ -217,7 +231,10 @@ impl http_client::HttpClient for ReqwestClient {
 
         let handle = self.handle.clone();
         async move {
-            let mut response = handle.spawn(async { request.send().await }).await??;
+            let mut response = handle
+                .spawn(async { request.send().await })
+                .await?
+                .map_err(redact_error)?;
 
             let headers = mem::take(response.headers_mut());
             let mut builder = http::Response::builder()