From 4dcd6a2c1a3ec5a781a4ac3138ce6baf350c15d8 Mon Sep 17 00:00:00 2001
From: haru0017 <hammerhead0017@gmail.com>
Date: Sun, 31 May 2026 20:24:31 +0900
Subject: [PATCH] fix(auth): support --subdomain login for orgs on non-US1
 regions

Use the OAuth callback's domain parameter to determine the correct
token exchange endpoint instead of the configured site.

- Propagate domain through CallbackResult
- Use effective_site for token exchange, token storage, and session storage
- Fall back to configured site when domain is absent

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 src/auth/callback.rs | 15 ++++++++++++---
 src/commands/auth.rs | 11 +++++++----
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/auth/callback.rs b/src/auth/callback.rs
index ada9784..d55cef0 100644
--- a/src/auth/callback.rs
+++ b/src/auth/callback.rs
@@ -19,6 +19,8 @@ pub struct CallbackResult {
     pub dd_oid: Option<String>,
     /// Display name of the consented org (`dd_org_name`).
     pub dd_org_name: Option<String>,
+    /// Actual Datadog site/region from the OAuth callback (`domain`).
+    pub domain: Option<String>,
 }
 
 #[cfg(not(target_arch = "wasm32"))]
@@ -199,6 +201,7 @@ async fn accept_loop(
             error_description,
             dd_oid,
             dd_org_name,
+            domain,
         };
         if let Some(tx) = result_tx.lock().unwrap().take() {
             let _ = tx.send(result);
@@ -294,6 +297,7 @@ fn parse_callback_url(input: &str) -> Result<CallbackResult> {
     let mut error_description = None;
     let mut dd_oid = None;
     let mut dd_org_name = None;
+    let mut domain = None;
     for (k, v) in url.query_pairs() {
         match k.as_ref() {
             "code" => code = Some(v.into_owned()),
@@ -302,6 +306,7 @@ fn parse_callback_url(input: &str) -> Result<CallbackResult> {
             "error_description" => error_description = Some(v.into_owned()),
             "dd_oid" => dd_oid = Some(v.into_owned()),
             "dd_org_name" => dd_org_name = Some(v.into_owned()),
+            "domain" => domain = Some(v.into_owned()),
             _ => {}
         }
     }
@@ -315,6 +320,7 @@ fn parse_callback_url(input: &str) -> Result<CallbackResult> {
         error_description,
         dd_oid,
         dd_org_name,
+        domain,
     })
 }
 
@@ -366,17 +372,19 @@ mod tests {
         assert!(r.error_description.is_none());
         assert!(r.dd_oid.is_none());
         assert!(r.dd_org_name.is_none());
+        assert!(r.domain.is_none());
     }
 
     #[test]
-    fn parse_callback_url_extracts_dd_oid_and_org_name() {
+    fn parse_callback_url_extracts_dd_oid_org_name_and_domain() {
         // Real callback shape from a Datadog OAuth flow — the issuer appends
-        // dd_oid and (URL-encoded) dd_org_name alongside code and state.
+        // dd_oid, (URL-encoded) dd_org_name, and domain alongside code and state.
         let r = parse_callback_url(
             "http://127.0.0.1:8000/oauth/callback\
              ?code=abc&state=xyz\
              &dd_oid=00000000-1111-2222-3333-444444444444\
-             &dd_org_name=Datadog+HQ",
+             &dd_org_name=Datadog+HQ\
+             &domain=us3.datadoghq.com",
         )
         .unwrap();
         assert_eq!(
@@ -384,6 +392,7 @@ mod tests {
             Some("00000000-1111-2222-3333-444444444444")
         );
         assert_eq!(r.dd_org_name.as_deref(), Some("Datadog HQ"));
+        assert_eq!(r.domain.as_deref(), Some("us3.datadoghq.com"));
     }
 
     #[test]
diff --git a/src/commands/auth.rs b/src/commands/auth.rs
index 0a171b2..545ecc8 100644
--- a/src/commands/auth.rs
+++ b/src/commands/auth.rs
@@ -171,9 +171,12 @@ pub async fn login(
         bail!("OAuth state mismatch (possible CSRF attack)");
     }
 
-    // 7. Exchange code for tokens
+    // 7. Exchange code for tokens.
+    // Use the actual site from the callback (e.g. "us3.datadoghq.com") when
+    // available, so the token exchange targets the correct region.
+    let effective_site = result.domain.as_deref().unwrap_or(site);
     eprintln!("🔄 Exchanging authorization code for tokens...");
-    let tokens = dcr_client
+    let tokens = dcr::DcrClient::new(effective_site)
         .exchange_code(&result.code, &redirect_uri, &challenge.verifier, &creds)
         .await?;
 
@@ -221,7 +224,7 @@ pub async fn login(
     let saved_org_label = org_suffix(saved_org);
 
     let location = with_storage(|store| {
-        store.save_tokens(site, saved_org, &tokens)?;
+        store.save_tokens(effective_site, saved_org, &tokens)?;
         Ok(store.storage_location())
     })?;
 
@@ -235,7 +238,7 @@ pub async fn login(
         .map(String::from)
         .or_else(|| org_uuid.map(String::from));
     storage::save_session(&storage::SessionEntry {
-        site: site.clone(),
+        site: effective_site.to_string(),
         org: saved_org.map(String::from),
         org_uuid: saved_org_uuid,
     })?;