diff --git a/src/auth/callback.rs b/src/auth/callback.rs index 2a491b4..3eaa996 100644 --- a/src/auth/callback.rs +++ b/src/auth/callback.rs @@ -15,6 +15,10 @@ pub struct CallbackResult { pub state: String, pub error: Option, pub error_description: Option, + /// Authoritative org UUID from the OAuth server (`dd_oid`). + pub dd_oid: Option, + /// Display name of the consented org (`dd_org_name`). + pub dd_org_name: Option, } #[cfg(not(target_arch = "wasm32"))] @@ -162,11 +166,25 @@ async fn accept_loop( let state = params.get("state").cloned().unwrap_or_default(); let error = params.get("error").cloned(); let error_description = params.get("error_description").cloned(); + let dd_oid = params.get("dd_oid").cloned(); + let dd_org_name = params.get("dd_org_name").cloned(); + // `site` (full URL) and `domain` (bare host) are sent alongside the + // OAuth response so the success page can render a link back. + let site = params.get("site").cloned(); + let domain = params.get("domain").cloned(); let (status, body) = if error.is_some() { ("400 Bad Request", error_page(&error, &error_description)) } else { - ("200 OK", success_page()) + ( + "200 OK", + success_page( + site.as_deref(), + domain.as_deref(), + dd_org_name.as_deref(), + dd_oid.as_deref(), + ), + ) }; let response = format!( "HTTP/1.1 {status}\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", @@ -179,6 +197,8 @@ async fn accept_loop( state, error, error_description, + dd_oid, + dd_org_name, }; if let Some(tx) = result_tx.lock().unwrap().take() { let _ = tx.send(result); @@ -188,14 +208,62 @@ async fn accept_loop( } #[cfg(not(target_arch = "wasm32"))] -fn success_page() -> String { - r#" +fn success_page( + site: Option<&str>, + domain: Option<&str>, + org_name: Option<&str>, + org_uuid: Option<&str>, +) -> String { + // Only accept https `site` as the link href; anything else falls back to + // plain text so a malformed or hostile callback can't inject a different + // scheme into the rendered HTML. + let provider_html = match (site.filter(|s| s.starts_with("https://")), domain) { + (Some(href), Some(host)) => format!( + r#"Datadog ({})"#, + html_escape(href), + html_escape(host), + ), + (Some(href), None) => format!(r#"{0}"#, html_escape(href)), + (None, Some(host)) => format!("Datadog ({})", html_escape(host)), + (None, None) => "Datadog".to_string(), + }; + + let org_line = org_name + .map(|n| { + format!( + r#"

in org {}

"#, + html_escape(n) + ) + }) + .unwrap_or_default(); + let uuid_line = org_uuid + .map(|u| format!(r#"

({})

"#, html_escape(u))) + .unwrap_or_default(); + + format!( + r#" Pup - Authentication Successful - +

Authentication Successful

-

You can close this window and return to pup.

"#.to_string() +

Connected to {provider_html}

+{org_line}{uuid_line}

You can close this window and return to pup.

"# + ) +} + +#[cfg(not(target_arch = "wasm32"))] +fn html_escape(s: &str) -> String { + s.replace('&', "&") + .replace('<', "<") + .replace('>', ">") + .replace('"', """) + .replace('\'', "'") } #[cfg(not(target_arch = "wasm32"))] @@ -224,12 +292,16 @@ fn parse_callback_url(input: &str) -> Result { let mut state = None; let mut error = None; let mut error_description = None; + let mut dd_oid = None; + let mut dd_org_name = None; for (k, v) in url.query_pairs() { match k.as_ref() { "code" => code = Some(v.into_owned()), "state" => state = Some(v.into_owned()), "error" => error = Some(v.into_owned()), "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()), _ => {} } } @@ -241,6 +313,8 @@ fn parse_callback_url(input: &str) -> Result { state: state.unwrap_or_default(), error, error_description, + dd_oid, + dd_org_name, }) } @@ -290,6 +364,26 @@ mod tests { assert_eq!(r.state, "xyz789"); assert!(r.error.is_none()); assert!(r.error_description.is_none()); + assert!(r.dd_oid.is_none()); + assert!(r.dd_org_name.is_none()); + } + + #[test] + fn parse_callback_url_extracts_dd_oid_and_org_name() { + // Real callback shape from a Datadog OAuth flow ā€” the issuer appends + // dd_oid and (URL-encoded) dd_org_name 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", + ) + .unwrap(); + assert_eq!( + r.dd_oid.as_deref(), + Some("00000000-1111-2222-3333-444444444444") + ); + assert_eq!(r.dd_org_name.as_deref(), Some("Datadog HQ")); } #[test] @@ -460,6 +554,77 @@ mod tests { ); } + #[test] + fn success_page_renders_org_name_and_uuid() { + let html = success_page( + Some("https://app.tatooine.example"), + Some("tatooine.example"), + Some("Mos Eisley Cantina"), + Some("00000000-1111-2222-3333-444444444444"), + ); + assert!(html.contains("Authentication Successful")); + assert!(html.contains("Connected to Datadog")); + assert!(html.contains(r#"href="https://app.tatooine.example""#)); + assert!(html.contains("Mos Eisley Cantina")); + assert!(html.contains("in org")); + // UUID is rendered inline as fine-print so operators can read or copy + // it without any interaction. + assert!(html.contains(r#"

no 'in org' line"); + assert!( + !html.contains(r#"

no fine-print line" + ); + } + + #[test] + fn success_page_escapes_org_name() { + // Defense in depth: a hostile dd_org_name shouldn't be able to inject + // script tags into the success page rendered by the local listener. + let html = success_page(None, None, Some(""), None); + assert!(!html.contains("")); + assert!(html.contains("<script>alert(1)</script>")); + } + + #[test] + fn success_page_rejects_non_https_site_as_link() { + // An unexpected `site` value (anything other than https://...) must not + // become a clickable link. The bare host falls back to plain text. + let html = success_page( + Some("javascript:alert(1)"), + Some("deathstar.tatooine.example"), + None, + None, + ); + assert!(!html.contains("javascript:alert(1)")); + assert!(!html.contains("href=\"javascript")); + assert!(html.contains("deathstar.tatooine.example")); + } + + #[test] + fn success_page_uses_vanity_subdomain_in_link() { + // SAML/SSO vanity subdomains route through their own host; the link + // should reflect whatever site the issuer reported, not always app.. + let html = success_page( + Some("https://yavin4.tatooine.example"), + Some("tatooine.example"), + Some("Rebel Alliance"), + None, + ); + assert!(html.contains(r#"href="https://yavin4.tatooine.example""#)); + assert!(html.contains("Rebel Alliance")); + } + #[tokio::test] async fn read_callback_url_passes_through_oauth_error_redirect() { let r = read_callback_url_from_reader(reader( diff --git a/src/auth/dcr.rs b/src/auth/dcr.rs index 1c66a47..088fed5 100644 --- a/src/auth/dcr.rs +++ b/src/auth/dcr.rs @@ -189,7 +189,10 @@ impl DcrClient { }) } - /// Build the authorization URL for the browser. + /// Build the authorization URL for the browser. `org_uuid` is appended as + /// `dd_oid` when set; callers should coerce empty strings to `None` + /// upstream so this function doesn't have to second-guess them. + #[allow(clippy::too_many_arguments)] pub fn build_authorization_url( &self, client_id: &str, @@ -198,17 +201,22 @@ impl DcrClient { challenge: &super::pkce::PkceChallenge, scopes: &[&str], subdomain: Option<&str>, + org_uuid: Option<&str>, ) -> String { let scope = scopes.join(" "); - let params = url::form_urlencoded::Serializer::new(String::new()) + let mut serializer = url::form_urlencoded::Serializer::new(String::new()); + serializer .append_pair("response_type", "code") .append_pair("client_id", client_id) .append_pair("redirect_uri", redirect_uri) .append_pair("state", state) .append_pair("scope", &scope) .append_pair("code_challenge", &challenge.challenge) - .append_pair("code_challenge_method", &challenge.method) - .finish(); + .append_pair("code_challenge_method", &challenge.method); + if let Some(uuid) = org_uuid { + serializer.append_pair("dd_oid", uuid); + } + let params = serializer.finish(); // Use custom subdomain for SAML/SSO auth, otherwise use standard app.{site}. // The subdomain replaces the `app` prefix on whichever site is in play, so a @@ -248,6 +256,7 @@ mod tests { &challenge(), &["dashboards_read"], None, + None, ); assert!( url.starts_with("https://app.datadoghq.com/oauth2/v1/authorize?"), @@ -268,6 +277,7 @@ mod tests { &challenge(), &["dashboards_read"], Some("dd"), + None, ); assert!( url.starts_with("https://dd.datad0g.com/oauth2/v1/authorize?"), @@ -289,6 +299,7 @@ mod tests { &challenge(), &["dashboards_read"], Some("acme"), + None, ); assert!( url.starts_with("https://acme.datadoghq.eu/oauth2/v1/authorize?"), @@ -310,6 +321,7 @@ mod tests { &challenge(), &["dashboards_read"], Some(""), + None, ); assert!( url.starts_with("https://app.datadoghq.com/oauth2/v1/authorize?"), @@ -327,6 +339,7 @@ mod tests { &challenge(), &["dashboards_read", "metrics_read"], None, + None, ); assert!(url.contains("response_type=code")); assert!(url.contains("client_id=client123")); @@ -336,4 +349,37 @@ mod tests { // Scopes are joined with a space, then URL-encoded as `+` or `%20`. assert!(url.contains("scope=dashboards_read") && url.contains("metrics_read")); } + + #[test] + fn build_authorization_url_appends_dd_oid_when_org_uuid_set() { + let client = DcrClient::new("datadoghq.com"); + let url = client.build_authorization_url( + "client123", + "http://127.0.0.1:8000/oauth/callback", + "state", + &challenge(), + &["dashboards_read"], + None, + Some("00000000-1111-2222-3333-444444444444"), + ); + assert!( + url.contains("dd_oid=00000000-1111-2222-3333-444444444444"), + "expected dd_oid query param, got: {url}" + ); + } + + #[test] + fn build_authorization_url_omits_dd_oid_when_unset() { + let client = DcrClient::new("datadoghq.com"); + let url = client.build_authorization_url( + "client123", + "http://127.0.0.1:8000/oauth/callback", + "state", + &challenge(), + &["dashboards_read"], + None, + None, + ); + assert!(!url.contains("dd_oid"), "got: {url}"); + } } diff --git a/src/auth/storage.rs b/src/auth/storage.rs index 8ac9a7f..25a6bda 100644 --- a/src/auth/storage.rs +++ b/src/auth/storage.rs @@ -11,6 +11,11 @@ use super::types::{ClientCredentials, TokenSet}; pub struct SessionEntry { pub site: String, pub org: Option, + /// Authoritative org UUID (`dd_oid`) returned by the OAuth issuer. + /// `#[serde(default)]` lets sessions.json files written before this field + /// existed deserialize cleanly (treated as None). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub org_uuid: Option, } // --------------------------------------------------------------------------- @@ -878,17 +883,13 @@ pub fn list_sessions() -> Result> { } } -/// Upsert a session entry into the registry. +/// Upsert a session entry into the registry. Dedups on `(site, org)`; the +/// new entry's `org_uuid` wins so re-auth refreshes the stored UUID. #[cfg(not(target_arch = "wasm32"))] -pub fn save_session(site: &str, org: Option<&str>) -> Result<()> { +pub fn save_session(entry: &SessionEntry) -> Result<()> { let mut sessions = list_sessions()?; - let entry = SessionEntry { - site: site.to_string(), - org: org.map(String::from), - }; - // Dedup: remove any existing entry with same site+org, then append sessions.retain(|s| !(s.site == entry.site && s.org == entry.org)); - sessions.push(entry); + sessions.push(entry.clone()); write_sessions(&sessions) } @@ -900,6 +901,15 @@ pub fn remove_session(site: &str, org: Option<&str>) -> Result<()> { write_sessions(&sessions) } +/// Look up a single session entry by `(site, org)`. +#[cfg(not(target_arch = "wasm32"))] +pub fn find_session(site: &str, org: Option<&str>) -> Option { + list_sessions() + .ok()? + .into_iter() + .find(|s| s.site == site && s.org.as_deref() == org) +} + /// Look up the site for a named org session. Returns None if no session exists /// for that org, or if multiple sessions share the same org name on different /// sites (ambiguous ā€” caller must pass DD_SITE explicitly). On the ambiguous @@ -1241,6 +1251,45 @@ mod tests { .is_none()); } + // --- SessionEntry serde ------------------------------------------------- + + #[test] + fn test_session_entry_legacy_json_deserializes_with_no_org_uuid() { + // sessions.json files written before the org_uuid field existed must + // continue to deserialize. #[serde(default)] gives them None. + let legacy = r#"[{"site":"datadoghq.com","org":"prod-child"}]"#; + let parsed: Vec = serde_json::from_str(legacy).unwrap(); + assert_eq!(parsed.len(), 1); + assert_eq!(parsed[0].site, "datadoghq.com"); + assert_eq!(parsed[0].org.as_deref(), Some("prod-child")); + assert!(parsed[0].org_uuid.is_none()); + } + + #[test] + fn test_session_entry_roundtrip_with_uuid() { + let entry = SessionEntry { + site: "datadoghq.com".into(), + org: Some("prod-child".into()), + org_uuid: Some("00000000-1111-2222-3333-444444444444".into()), + }; + let json = serde_json::to_string(&entry).unwrap(); + let parsed: SessionEntry = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, entry); + } + + #[test] + fn test_session_entry_omits_uuid_when_none() { + // skip_serializing_if keeps existing on-disk shapes byte-stable for + // sessions that were never tagged with a UUID. + let entry = SessionEntry { + site: "datadoghq.com".into(), + org: None, + org_uuid: None, + }; + let json = serde_json::to_string(&entry).unwrap(); + assert!(!json.contains("org_uuid"), "got: {json}"); + } + // --- Session registry --------------------------------------------------- #[test] @@ -1259,8 +1308,18 @@ mod tests { let tmp = TempDir::new("sess_save"); std::env::set_var("PUP_CONFIG_DIR", tmp.path()); - save_session("datadoghq.com", None).unwrap(); - save_session("datadoghq.com", Some("prod-child")).unwrap(); + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: None, + org_uuid: None, + }) + .unwrap(); + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: Some("prod-child".into()), + org_uuid: None, + }) + .unwrap(); let sessions = list_sessions().unwrap(); std::env::remove_var("PUP_CONFIG_DIR"); @@ -1279,8 +1338,18 @@ mod tests { let tmp = TempDir::new("sess_dedup"); std::env::set_var("PUP_CONFIG_DIR", tmp.path()); - save_session("datadoghq.com", Some("prod")).unwrap(); - save_session("datadoghq.com", Some("prod")).unwrap(); // duplicate + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: Some("prod".into()), + org_uuid: None, + }) + .unwrap(); + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: Some("prod".into()), + org_uuid: None, + }) + .unwrap(); // duplicate let sessions = list_sessions().unwrap(); std::env::remove_var("PUP_CONFIG_DIR"); @@ -1293,8 +1362,18 @@ mod tests { let tmp = TempDir::new("sess_remove"); std::env::set_var("PUP_CONFIG_DIR", tmp.path()); - save_session("datadoghq.com", None).unwrap(); - save_session("datadoghq.com", Some("prod")).unwrap(); + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: None, + org_uuid: None, + }) + .unwrap(); + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: Some("prod".into()), + org_uuid: None, + }) + .unwrap(); remove_session("datadoghq.com", Some("prod")).unwrap(); let sessions = list_sessions().unwrap(); std::env::remove_var("PUP_CONFIG_DIR"); @@ -1319,8 +1398,18 @@ mod tests { let tmp = TempDir::new("find_sess_unique"); std::env::set_var("PUP_CONFIG_DIR", tmp.path()); - save_session("custom.datadoghq.com", Some("prod-child")).unwrap(); - save_session("datadoghq.com", None).unwrap(); + save_session(&SessionEntry { + site: "custom.datadoghq.com".into(), + org: Some("prod-child".into()), + org_uuid: None, + }) + .unwrap(); + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: None, + org_uuid: None, + }) + .unwrap(); let result = find_session_site("prod-child"); std::env::remove_var("PUP_CONFIG_DIR"); @@ -1333,7 +1422,12 @@ mod tests { let tmp = TempDir::new("find_sess_none"); std::env::set_var("PUP_CONFIG_DIR", tmp.path()); - save_session("datadoghq.com", Some("prod-child")).unwrap(); + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: Some("prod-child".into()), + org_uuid: None, + }) + .unwrap(); let result = find_session_site("nonexistent"); std::env::remove_var("PUP_CONFIG_DIR"); @@ -1348,8 +1442,18 @@ mod tests { // Same org name registered against two different sites ā†’ caller must // disambiguate via DD_SITE rather than us picking one. - save_session("datadoghq.com", Some("shared-name")).unwrap(); - save_session("datadoghq.eu", Some("shared-name")).unwrap(); + save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: Some("shared-name".into()), + org_uuid: None, + }) + .unwrap(); + save_session(&SessionEntry { + site: "datadoghq.eu".into(), + org: Some("shared-name".into()), + org_uuid: None, + }) + .unwrap(); let result = find_session_site("shared-name"); std::env::remove_var("PUP_CONFIG_DIR"); @@ -1363,7 +1467,12 @@ mod tests { std::env::set_var("PUP_CONFIG_DIR", tmp.path()); // The unnamed (org=None) session must not match any --org lookup. - save_session("datadoghq.eu", None).unwrap(); + save_session(&SessionEntry { + site: "datadoghq.eu".into(), + org: None, + org_uuid: None, + }) + .unwrap(); let result = find_session_site("anything"); std::env::remove_var("PUP_CONFIG_DIR"); diff --git a/src/commands/auth.rs b/src/commands/auth.rs index 15116ff..b597ecc 100644 --- a/src/commands/auth.rs +++ b/src/commands/auth.rs @@ -14,24 +14,40 @@ where f(&mut **store) } +/// Format the trailing " (org: )" segment used in user-facing log lines. +/// Returns an empty string when there's no org so the surrounding messages stay +/// clean for the default-org case. +fn org_suffix(org: Option<&str>) -> String { + org.map(|o| format!(" (org: {o})")).unwrap_or_default() +} + #[cfg(not(target_arch = "wasm32"))] pub async fn login( cfg: &Config, scopes: Vec, subdomain: Option<&str>, callback_port: Option, + org_uuid: Option<&str>, ) -> Result<()> { use crate::auth::{dcr, pkce}; let site = &cfg.site; let org = cfg.org.as_deref(); + // Resolve effective org_uuid: CLI flag wins; otherwise recall the UUID + // stored on the matching `(site, org)` session so re-auth keeps emitting + // `dd_oid` without re-passing the flag. + let stored_session = storage::find_session(site, org); + let effective_org_uuid: Option = org_uuid + .map(String::from) + .or_else(|| stored_session.as_ref().and_then(|s| s.org_uuid.clone())); + // 1. Start callback server. `callback_port` (from --callback-port or // PUP_OAUTH_CALLBACK_PORT) pins the port deterministically for SSH // port-forwarded workflows; otherwise we scan the DCR allowlist. let mut server = crate::auth::callback::CallbackServer::new(callback_port).await?; let redirect_uri = server.redirect_uri(); - let org_label = org.map(|o| format!(" (org: {o})")).unwrap_or_default(); + let org_label = org_suffix(org); eprintln!("\nšŸ” Starting OAuth2 login for site: {site}{org_label}\n"); if let Some(sub) = subdomain { // Compose against the actual site, not a hardcoded prod host. Mirrors @@ -92,7 +108,11 @@ pub async fn login( &challenge, &scope_strs, subdomain, + effective_org_uuid.as_deref(), ); + if let Some(uuid) = effective_org_uuid.as_deref() { + eprintln!("šŸŽÆ Hinting org UUID (dd_oid): {uuid}"); + } // 5. Open browser eprintln!("\nšŸŒ Opening browser for authentication..."); @@ -145,31 +165,154 @@ pub async fn login( .exchange_code(&result.code, &redirect_uri, &challenge.verifier, &creds) .await?; + // 8. Resolve the save target ā€” relabel under the actual returned org on + // dd_oid mismatch. The OAuth code is single-use, so refusing to save + // would throw away the user's click. + // + // Filter callback `dd_oid` through a shape check so a tampered URL + // pasted into the stdin fallback path can't seed an arbitrary string + // into the session registry. + let validated_dd_oid = result.dd_oid.as_deref().filter(|u| looks_like_uuid(u)); + let saved_org_owned: Option = match resolve_save_target( + org, + effective_org_uuid.as_deref(), + validated_dd_oid, + result.dd_org_name.as_deref(), + ) { + SaveTarget::Requested(label) => label, + SaveTarget::Reassigned { + label, + requested_uuid, + actual_uuid, + actual_name, + } => { + let actual_name_suffix = actual_name + .as_deref() + .map(|n| format!(" ({n})")) + .unwrap_or_default(); + let cleanup_hint = match org { + Some(o) => format!( + "\n Your existing '{o}' session is unchanged. \ + Run `pup auth logout --org {o}` to clear it if it's stale." + ), + None => String::new(), + }; + eprintln!( + "āš ļø Requested org UUID {requested_uuid} but OAuth returned \ + {actual_uuid}{actual_name_suffix}. Saving token under \ + \"{label}\" instead of the requested label.{cleanup_hint}" + ); + Some(label) + } + }; + let saved_org = saved_org_owned.as_deref(); + let saved_org_label = org_suffix(saved_org); + let location = with_storage(|store| { - store.save_tokens(site, org, &tokens)?; + store.save_tokens(site, saved_org, &tokens)?; Ok(store.storage_location()) })?; - // Register this session in the session registry - storage::save_session(site, org)?; + // Persist the callback's confirmed `dd_oid` if we have one. Otherwise + // fall back only to the user's explicit `--org-uuid` flag ā€” never to a + // previously-recalled hint, since carrying that forward unconfirmed + // would let a stale UUID survive a login that actually landed in a + // different org. Users who want to re-establish a hint can pass the + // flag again on the next login. + let saved_org_uuid = validated_dd_oid + .map(String::from) + .or_else(|| org_uuid.map(String::from)); + storage::save_session(&storage::SessionEntry { + site: site.clone(), + org: saved_org.map(String::from), + org_uuid: saved_org_uuid, + })?; let expires_at = chrono::DateTime::from_timestamp(tokens.issued_at + tokens.expires_in, 0) .map(|dt| dt.with_timezone(&chrono::Local).to_rfc3339()) .unwrap_or_else(|| format!("in {} hours", tokens.expires_in / 3600)); - eprintln!("\nāœ… Login successful{org_label}!"); + eprintln!("\nāœ… Login successful{saved_org_label}!"); eprintln!(" Access token expires: {expires_at}"); eprintln!(" Token stored in: {location}"); Ok(()) } +/// Loose UUID shape check (8-4-4-4-12 hex with dashes, ASCII only). The +/// callback's `dd_oid` is unsigned metadata ā€” in the stdin-paste fallback +/// path a tampered URL could carry an arbitrary string. Rejecting anything +/// that doesn't look like a UUID keeps such values from being persisted into +/// `sessions.json` or rendered as a future hint. +#[cfg(not(target_arch = "wasm32"))] +fn looks_like_uuid(s: &str) -> bool { + s.len() == 36 + && s.chars().enumerate().all(|(i, c)| match i { + 8 | 13 | 18 | 23 => c == '-', + _ => c.is_ascii_hexdigit(), + }) +} + +/// What `resolve_save_target` decided about where to save the new token. +#[cfg(not(target_arch = "wasm32"))] +#[derive(Debug)] +enum SaveTarget { + /// dd_oid hint matched, no hint was sent, or there's no `--org` to + /// relabel away from ā€” save under the requested label (or the default + /// session when `--org` was unset). + Requested(Option), + /// dd_oid mismatch on a named-org login ā€” save under `label` and warn + /// the user. Carries the requested vs actual UUIDs so the caller can + /// render a useful warning. + Reassigned { + label: String, + requested_uuid: String, + actual_uuid: String, + actual_name: Option, + }, +} + +/// Pure: pick the `(site, org)` save target for a finished login. Mismatch +/// detection compares hinted vs callback UUIDs case-insensitively (UUIDs are +/// hex). Relabel only triggers for named-org logins ā€” when `--org` is unset +/// the default session is the source of truth and gets overwritten with the +/// new UUID directly, so plain `pup auth login` stays idempotent. +/// +/// On relabel, the new label always embeds an 8-char UUID prefix so the +/// stored entry is unambiguous even if the issuer's display name collides +/// with another existing session on the same site. +#[cfg(not(target_arch = "wasm32"))] +fn resolve_save_target( + requested_org: Option<&str>, + requested_uuid: Option<&str>, + actual_uuid: Option<&str>, + actual_org_name: Option<&str>, +) -> SaveTarget { + match (requested_org, requested_uuid, actual_uuid) { + (Some(_), Some(req), Some(actual)) if !req.eq_ignore_ascii_case(actual) => { + let prefix: String = actual.chars().take(8).collect(); + let label = match actual_org_name { + Some(name) if !name.is_empty() => format!("{name} [{prefix}]"), + _ => prefix, + }; + SaveTarget::Reassigned { + label, + requested_uuid: req.to_string(), + actual_uuid: actual.to_string(), + actual_name: actual_org_name.filter(|n| !n.is_empty()).map(String::from), + } + } + _ => SaveTarget::Requested(requested_org.map(String::from)), + } +} + #[cfg(target_arch = "wasm32")] pub async fn login( _cfg: &Config, _scopes: Vec, _subdomain: Option<&str>, _callback_port: Option, + _org_uuid: Option<&str>, ) -> Result<()> { bail!( "OAuth login is not available in WASM builds.\n\ @@ -192,7 +335,7 @@ pub async fn logout(cfg: &Config) -> Result<()> { Ok(()) })?; storage::remove_session(site, org)?; - let org_label = org.map(|o| format!(" (org: {o})")).unwrap_or_default(); + let org_label = org_suffix(org); eprintln!("Logged out from {site}{org_label}. Tokens removed."); Ok(()) } @@ -225,7 +368,7 @@ pub fn status(cfg: &Config) -> Result<()> { ("valid".to_string(), format!("{mins}m{secs}s")) }; - let org_label = org.map(|o| format!(" (org: {o})")).unwrap_or_default(); + let org_label = org_suffix(org); if tokens.is_expired() { eprintln!("āš ļø Token expired for site: {site}{org_label}"); } else { @@ -256,7 +399,7 @@ pub fn status(cfg: &Config) -> Result<()> { println!("{}", serde_json::to_string_pretty(&json).unwrap()); } None => { - let org_label = org.map(|o| format!(" (org: {o})")).unwrap_or_default(); + let org_label = org_suffix(org); eprintln!("āŒ Not authenticated for site: {site}{org_label}"); let json = serde_json::json!({ "authenticated": false, @@ -311,7 +454,7 @@ pub async fn refresh(cfg: &Config) -> Result<()> { anyhow::anyhow!("no client credentials found for site {site} ā€” run 'pup auth login' first") })?; - let org_label = org.map(|o| format!(" (org: {o})")).unwrap_or_default(); + let org_label = org_suffix(org); eprintln!("šŸ”„ Refreshing access token for site: {site}{org_label}..."); let dcr_client = dcr::DcrClient::new(site); @@ -370,6 +513,7 @@ pub fn list(cfg: &Config) -> Result<()> { "expires_at": expires_at, "has_refresh": !t.refresh_token.is_empty(), "org": s.org, + "org_uuid": s.org_uuid, "scopes": scopes, "site": s.site, "status": status, @@ -379,6 +523,7 @@ pub fn list(cfg: &Config) -> Result<()> { "expires_at": null, "has_refresh": false, "org": s.org, + "org_uuid": s.org_uuid, "scopes": [], "site": s.site, "status": "no token", @@ -398,6 +543,7 @@ pub fn list(_cfg: &Config) -> Result<()> { #[cfg(all(test, not(target_arch = "wasm32")))] mod tests { use super::*; + use crate::auth::storage::SessionEntry; use crate::config::{Config, OutputFormat}; /// Temporary directory that removes itself on drop. Mirrors the helper @@ -506,8 +652,18 @@ mod tests { std::env::set_var("PUP_CONFIG_DIR", tmp.path()); std::env::set_var("DD_TOKEN_STORAGE", "file"); - storage::save_session("datadoghq.com", None).unwrap(); - storage::save_session("datadoghq.com", Some("prod-child")).unwrap(); + storage::save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: None, + org_uuid: None, + }) + .unwrap(); + storage::save_session(&SessionEntry { + site: "datadoghq.com".into(), + org: Some("prod-child".into()), + org_uuid: None, + }) + .unwrap(); let cfg = base_config(); let result = list(&cfg); @@ -591,6 +747,141 @@ mod tests { ); } + // ------------------------------------------------------------------ + // resolve_save_target ā€” pure function, no env / storage state. + // ------------------------------------------------------------------ + + #[test] + fn resolve_save_target_no_hint_keeps_requested_org() { + let t = resolve_save_target(Some("prod-child"), None, None, None); + assert!(matches!(t, SaveTarget::Requested(Some(s)) if s == "prod-child")); + } + + #[test] + fn resolve_save_target_hint_matches_keeps_requested_org() { + let uuid = "00000000-1111-2222-3333-444444444444"; + let t = resolve_save_target( + Some("prod-child"), + Some(uuid), + Some(uuid), + Some("Mos Eisley Cantina"), + ); + assert!(matches!(t, SaveTarget::Requested(Some(s)) if s == "prod-child")); + } + + #[test] + fn resolve_save_target_hint_matches_case_insensitive() { + // UUIDs are hex; normalise case so a mismatch in the issuer's + // canonicalisation doesn't trip the warn-and-relabel path. + let upper = "AABBCCDD-1111-2222-3333-EEFFAABBCCDD"; + let lower = "aabbccdd-1111-2222-3333-eeffaabbccdd"; + let t = resolve_save_target(Some("prod-child"), Some(upper), Some(lower), None); + assert!(matches!(t, SaveTarget::Requested(Some(s)) if s == "prod-child")); + } + + #[test] + fn resolve_save_target_mismatch_combines_name_and_uuid_prefix() { + // The relabel embeds the UUID prefix so a save under (site, "Other Org") + // can never silently overwrite an unrelated existing "Other Org" session + // on the same site that has a different underlying UUID. + let req = "00000000-1111-2222-3333-444444444444"; + let act = "11111111-2222-3333-4444-555555555555"; + let t = resolve_save_target(Some("prod-child"), Some(req), Some(act), Some("Other Org")); + match t { + SaveTarget::Reassigned { + label, + requested_uuid, + actual_uuid, + actual_name, + } => { + assert_eq!(label, "Other Org [11111111]"); + assert_eq!(requested_uuid, req); + assert_eq!(actual_uuid, act); + assert_eq!(actual_name.as_deref(), Some("Other Org")); + } + other => panic!("expected Reassigned, got {other:?}"), + } + } + + #[test] + fn resolve_save_target_mismatch_falls_back_to_uuid_prefix_only() { + // No dd_org_name in the callback (older issuer or unusual flow): the + // 8-char UUID prefix becomes the entire label so it's still stable + // and distinguishable. + let req = "00000000-1111-2222-3333-444444444444"; + let act = "11111111-2222-3333-4444-555555555555"; + let t = resolve_save_target(Some("prod-child"), Some(req), Some(act), None); + match t { + SaveTarget::Reassigned { + label, actual_name, .. + } => { + assert_eq!(label, "11111111"); + assert!(actual_name.is_none()); + } + other => panic!("expected Reassigned, got {other:?}"), + } + } + + #[test] + fn resolve_save_target_mismatch_treats_empty_org_name_as_missing() { + // An empty `dd_org_name=""` shouldn't produce labels like " [11111111]" + // with a leading space ā€” fall back to the UUID-only label instead. + let req = "00000000-1111-2222-3333-444444444444"; + let act = "11111111-2222-3333-4444-555555555555"; + let t = resolve_save_target(Some("prod-child"), Some(req), Some(act), Some("")); + match t { + SaveTarget::Reassigned { label, .. } => { + assert_eq!(label, "11111111"); + } + other => panic!("expected Reassigned, got {other:?}"), + } + } + + #[test] + fn resolve_save_target_actual_uuid_missing_keeps_requested() { + // If the issuer didn't echo dd_oid, we have no comparison to make; + // trust the user's --org label and the in-flight UUID we sent. + let req = "00000000-1111-2222-3333-444444444444"; + let t = resolve_save_target(Some("prod-child"), Some(req), None, None); + assert!(matches!(t, SaveTarget::Requested(Some(s)) if s == "prod-child")); + } + + #[test] + fn resolve_save_target_default_org_mismatch_stays_default() { + // Default-org login with a stale stored UUID hint that no longer + // matches the issuer's response: `pup auth login` is supposed to be + // idempotent for the default session, so we don't relabel onto a + // phantom slot. The default session gets refreshed in place; the + // stored UUID hint is updated by the caller from the callback. + let req = "00000000-1111-2222-3333-444444444444"; + let act = "11111111-2222-3333-4444-555555555555"; + let t = resolve_save_target(None, Some(req), Some(act), Some("Other Org")); + assert!(matches!(t, SaveTarget::Requested(None))); + } + + // looks_like_uuid ----------------------------------------------------------- + + #[test] + fn looks_like_uuid_accepts_canonical_form() { + assert!(looks_like_uuid("00000000-1111-2222-3333-444444444444")); + assert!(looks_like_uuid("aabbccdd-1111-2222-3333-eeffaabbccdd")); + assert!(looks_like_uuid("AABBCCDD-1111-2222-3333-EEFFAABBCCDD")); + } + + #[test] + fn looks_like_uuid_rejects_garbage() { + assert!(!looks_like_uuid("")); + assert!(!looks_like_uuid("not-a-uuid")); + // Wrong length + assert!(!looks_like_uuid("00000000-1111-2222-3333-44444444444")); + // Non-hex character + assert!(!looks_like_uuid("0000000g-1111-2222-3333-444444444444")); + // Missing dash positions + assert!(!looks_like_uuid("000000001111-2222-3333-44444444444444")); + // Hostile injection attempt + assert!(!looks_like_uuid(" ")); + } + #[tokio::test] async fn test_logout_with_org_removes_session_entry() { // save a session, then logout should remove just that org's entry @@ -603,8 +894,18 @@ mod tests { std::env::set_var("DD_TOKEN_STORAGE", "file"); let site = "logout-session.example.invalid"; - storage::save_session(site, None).unwrap(); - storage::save_session(site, Some("keep-me")).unwrap(); + storage::save_session(&SessionEntry { + site: site.into(), + org: None, + org_uuid: None, + }) + .unwrap(); + storage::save_session(&SessionEntry { + site: site.into(), + org: Some("keep-me".into()), + org_uuid: None, + }) + .unwrap(); let mut cfg = base_config(); cfg.site = site.into(); diff --git a/src/config.rs b/src/config.rs index 188c0e2..48c462a 100644 --- a/src/config.rs +++ b/src/config.rs @@ -482,6 +482,7 @@ fn env_bool(key: &str) -> bool { #[cfg(test)] mod tests { use super::*; + use crate::auth::storage::SessionEntry; use crate::test_utils::ENV_LOCK; fn make_cfg(api_key: Option<&str>, app_key: Option<&str>, token: Option<&str>) -> Config { @@ -937,7 +938,12 @@ mod tests { std::fs::create_dir_all(&tmp).unwrap(); std::env::set_var("PUP_CONFIG_DIR", &tmp); - crate::auth::storage::save_session("custom.datadoghq.com", Some("prod-child")).unwrap(); + crate::auth::storage::save_session(&SessionEntry { + site: "custom.datadoghq.com".into(), + org: Some("prod-child".into()), + org_uuid: None, + }) + .unwrap(); std::env::set_var("DD_ORG", "prod-child"); let cfg = Config::from_env().unwrap(); @@ -967,7 +973,12 @@ mod tests { std::fs::create_dir_all(&tmp).unwrap(); std::env::set_var("PUP_CONFIG_DIR", &tmp); - crate::auth::storage::save_session("custom.datadoghq.com", Some("prod-child")).unwrap(); + crate::auth::storage::save_session(&SessionEntry { + site: "custom.datadoghq.com".into(), + org: Some("prod-child".into()), + org_uuid: None, + }) + .unwrap(); std::env::set_var("DD_ORG", "prod-child"); std::env::set_var("DD_SITE", "datadoghq.eu"); @@ -998,8 +1009,18 @@ mod tests { std::fs::create_dir_all(&tmp).unwrap(); std::env::set_var("PUP_CONFIG_DIR", &tmp); - crate::auth::storage::save_session("a.datadoghq.com", Some("org-a")).unwrap(); - crate::auth::storage::save_session("b.datadoghq.com", Some("org-b")).unwrap(); + crate::auth::storage::save_session(&SessionEntry { + site: "a.datadoghq.com".into(), + org: Some("org-a".into()), + org_uuid: None, + }) + .unwrap(); + crate::auth::storage::save_session(&SessionEntry { + site: "b.datadoghq.com".into(), + org: Some("org-b".into()), + org_uuid: None, + }) + .unwrap(); // Simulate the post-from_env state where we resolved org-a's site via // the registry (site_explicit=false because the user did not set DD_SITE). @@ -1040,7 +1061,12 @@ mod tests { std::fs::create_dir_all(&tmp).unwrap(); std::env::set_var("PUP_CONFIG_DIR", &tmp); - crate::auth::storage::save_session("session.datadoghq.com", Some("org-a")).unwrap(); + crate::auth::storage::save_session(&SessionEntry { + site: "session.datadoghq.com".into(), + org: Some("org-a".into()), + org_uuid: None, + }) + .unwrap(); let mut cfg = Config { api_key: None, diff --git a/src/main.rs b/src/main.rs index 2eb191a..d3a2260 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9239,6 +9239,13 @@ enum AuthActions { /// PUP_OAUTH_CALLBACK_PORT if unset. #[arg(long, value_name = "PORT")] callback_port: Option, + /// Hint the target org by UUID (sent as `dd_oid`). Skips the org switcher when the + /// existing browser session already matches and pre-routes SAML/SSO routing for + /// first-time logins. The value is persisted with the session so subsequent + /// `pup auth login` invocations re-emit it automatically. The server validates + /// the UUID; malformed values are rejected there. + #[arg(long, value_name = "UUID")] + org_uuid: Option, }, /// Logout and clear tokens Logout, @@ -13991,6 +13998,7 @@ async fn main_inner() -> anyhow::Result<()> { site, subdomain, callback_port, + org_uuid, } => { if let Some(s) = site { cfg.set_site_explicit(s); @@ -14002,7 +14010,17 @@ async fn main_inner() -> anyhow::Result<()> { let resolved = resolve_login_scopes(scopes.as_deref(), cfg.org.as_deref(), is_read_only); let resolved_port = resolve_callback_port(callback_port)?; - commands::auth::login(&cfg, resolved, subdomain.as_deref(), resolved_port).await? + // Coerce empty `--org-uuid ""` to no-hint so callees treat + // it the same as omitting the flag, not as `dd_oid=`. + let org_uuid_hint = org_uuid.as_deref().filter(|s| !s.is_empty()); + commands::auth::login( + &cfg, + resolved, + subdomain.as_deref(), + resolved_port, + org_uuid_hint, + ) + .await? } AuthActions::Logout => commands::auth::logout(&cfg).await?, AuthActions::Status { site } => {