From 9d538136b5d86a1eb0a11ef469729b7304db24fb Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Sat, 21 Mar 2026 14:39:52 -0700 Subject: [PATCH] fix(oauth): reject malformed ic2.* states in decode_hosted_oauth_state (#1441) (#1454) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(oauth): reject malformed ic2.* states instead of falling through to legacy handler (#1441) When decode_hosted_oauth_state() encountered a versioned state (ic2.*) that failed to fully parse (bad base64, invalid JSON, missing separator), it silently fell through to legacy handling which used the full malformed envelope as the flow_id. This never matched the raw nonce stored in pending_oauth_flows, breaking the OAuth callback. Restructure the versioned decode path so any ic2.* state must parse as a valid envelope or return Err — never fall through to legacy handling. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(oauth): address PR review — avoid alloc in strip_prefix, strengthen JSON parse test - Replace `strip_prefix(&format!(...))` with a `HOSTED_STATE_PREFIX_DOT` constant to avoid per-call allocation. - Fix "valid base64 but not JSON" test to compute the correct checksum so it actually exercises the JSON parse error path instead of stopping at the checksum check. Co-Authored-By: Claude Opus 4.6 (1M context) * fix: add missing fallback_deliverable field in job_monitor tests The SseEvent::JobResult struct gained a fallback_deliverable field in the structured fallback deliverables feature, but the job_monitor test constructors were not updated. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(oauth): remove HOSTED_STATE_PREFIX_DOT to avoid drift with HOSTED_STATE_PREFIX concat! requires literals and cannot reference const items, so a separate _DOT constant would duplicate the prefix string. Revert to deriving the dotted prefix via format!() — both encode and decode now use the same single HOSTED_STATE_PREFIX constant, keeping them mechanically consistent. Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- src/cli/oauth_defaults.rs | 101 +++++++++++++++++++++++++++++++------- 1 file changed, 83 insertions(+), 18 deletions(-) diff --git a/src/cli/oauth_defaults.rs b/src/cli/oauth_defaults.rs index 874cff98..b4e93704 100644 --- a/src/cli/oauth_defaults.rs +++ b/src/cli/oauth_defaults.rs @@ -579,23 +579,27 @@ pub fn encode_hosted_oauth_state(flow_id: &str, instance_name: Option<&str>) -> /// Decode hosted OAuth state in either the new versioned format or the /// legacy `instance:nonce`/`nonce` forms. pub fn decode_hosted_oauth_state(state: &str) -> Result { - if let Some(rest) = state.strip_prefix(&format!("{HOSTED_STATE_PREFIX}.")) - && let Some((payload_b64, checksum)) = rest.rsplit_once('.') - && let Ok(payload_json) = URL_SAFE_NO_PAD.decode(payload_b64) - { + if let Some(rest) = state.strip_prefix(&format!("{HOSTED_STATE_PREFIX}.")) { + let (payload_b64, checksum) = rest + .rsplit_once('.') + .ok_or("Hosted OAuth versioned state missing checksum separator")?; + let payload_json = URL_SAFE_NO_PAD + .decode(payload_b64) + .map_err(|e| format!("Hosted OAuth versioned state base64 decode failed: {e}"))?; let expected_checksum = hosted_state_checksum(&payload_json); if checksum != expected_checksum { return Err("Hosted OAuth state checksum mismatch".to_string()); } - if let Ok(payload) = serde_json::from_slice::(&payload_json) - && !payload.flow_id.trim().is_empty() - { - return Ok(DecodedHostedOAuthState { - flow_id: payload.flow_id, - instance_name: payload.instance_name.filter(|v| !v.is_empty()), - is_legacy: false, - }); + let payload: HostedOAuthStatePayload = serde_json::from_slice(&payload_json) + .map_err(|e| format!("Hosted OAuth versioned state JSON parse failed: {e}"))?; + if payload.flow_id.trim().is_empty() { + return Err("Hosted OAuth versioned state has empty flow_id".to_string()); } + return Ok(DecodedHostedOAuthState { + flow_id: payload.flow_id, + instance_name: payload.instance_name.filter(|v| !v.is_empty()), + is_legacy: false, + }); } if let Some((instance_name, flow_id)) = state.split_once(':') { @@ -1187,14 +1191,14 @@ mod tests { } #[test] - fn test_decode_hosted_oauth_state_falls_back_for_non_envelope_ic2_prefix() { + fn test_decode_hosted_oauth_state_rejects_non_envelope_ic2_prefix() { use crate::cli::oauth_defaults::decode_hosted_oauth_state; - let decoded = - decode_hosted_oauth_state("ic2.provider-owned-state").expect("prefixed fallback"); - assert_eq!(decoded.flow_id, "ic2.provider-owned-state"); - assert_eq!(decoded.instance_name, None); - assert!(decoded.is_legacy); + // "ic2." prefix must parse as a valid versioned envelope — never fall + // through to legacy handling, which would use the full malformed + // envelope as the flow_id and break OAuth callback lookup (#1441). + decode_hosted_oauth_state("ic2.provider-owned-state") + .expect_err("ic2-prefixed non-envelope state should fail"); } #[test] @@ -1244,4 +1248,65 @@ mod tests { assert!(result.url.contains("code_challenge=")); assert!(result.code_verifier.is_some()); } + + /// Malformed `ic2.*` states must return Err, never fall through to legacy + /// handling where the full envelope would be used as the flow_id (#1441). + #[test] + fn test_decode_versioned_state_rejects_malformed_envelopes() { + use crate::cli::oauth_defaults::decode_hosted_oauth_state; + + // Missing checksum separator (no second dot after prefix) + let err = + decode_hosted_oauth_state("ic2.nodots").expect_err("missing separator should fail"); + assert!( + err.contains("checksum separator"), + "unexpected error: {err}" + ); + + // Bad base64 payload + let err = decode_hosted_oauth_state("ic2.!!!badbase64!!!.fakechecksum") + .expect_err("bad base64 should fail"); + assert!(err.contains("base64"), "unexpected error: {err}"); + + // Valid base64 but not JSON: use correct checksum so we exercise JSON parsing + use base64::Engine; + use sha2::Digest; + let not_json_bytes = b"not json"; + let not_json_b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(not_json_bytes); + let digest = sha2::Sha256::digest(not_json_bytes); + let checksum = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(&digest[..super::HOSTED_STATE_CHECKSUM_BYTES]); + let err = decode_hosted_oauth_state(&format!("ic2.{not_json_b64}.{checksum}")) + .expect_err("non-JSON payload should fail with JSON parse error"); + assert!( + err.contains("JSON"), + "unexpected error (expected JSON parse failure): {err}" + ); + } + + /// Round-trip: encode_hosted_oauth_state(nonce) → decode → flow_id == nonce. + /// Ensures the registration key and lookup key are always identical (#1441). + #[test] + fn test_oauth_flow_key_round_trip_consistency() { + use crate::cli::oauth_defaults::{decode_hosted_oauth_state, encode_hosted_oauth_state}; + + let nonce = "test-nonce-abc123"; + let encoded = encode_hosted_oauth_state(nonce, Some("my-instance")); + let decoded = decode_hosted_oauth_state(&encoded).expect("round-trip decode"); + + assert_eq!( + decoded.flow_id, nonce, + "flow_id must match the original nonce" + ); + assert_eq!(decoded.instance_name.as_deref(), Some("my-instance")); + assert!(!decoded.is_legacy); + + // Also test without instance name + let encoded_no_instance = encode_hosted_oauth_state(nonce, None); + let decoded_no_instance = + decode_hosted_oauth_state(&encoded_no_instance).expect("round-trip without instance"); + assert_eq!(decoded_no_instance.flow_id, nonce); + assert_eq!(decoded_no_instance.instance_name, None); + assert!(!decoded_no_instance.is_legacy); + } }