| 1 | //! One RFC 8628 device-authorization polling loop, shared by every Codewhale |
| 2 | //! device-code flow (xAI/Grok device login, Codewhale account login). |
| 3 | //! |
| 4 | //! Ported from pi (<https://github.com/badlogic/pi-mono>), MIT licensed, |
| 5 | //! Copyright (c) 2025 Mario Zechner — see |
| 6 | //! `packages/ai/src/auth/oauth/device-code.ts` for the original |
| 7 | //! `pollOAuthDeviceCodeFlow`. The accumulated behaviours carried over from it: |
| 8 | //! |
| 9 | //! * the RFC 8628 §3.2 default of 5 seconds when the server omits `interval`; |
| 10 | //! * `slow_down` handling that **prefers a server-supplied interval** over the |
| 11 | //! client-tracked one. Trusting only the client-tracked value lets WSL/VM |
| 12 | //! clock drift poll early forever; RFC 8628 §3.5's +5s step is the fallback; |
| 13 | //! * a hard deadline derived from `expires_in`, never slept past even after |
| 14 | //! `slow_down` backoff; |
| 15 | //! * a distinct timeout message when at least one `slow_down` was seen, so the |
| 16 | //! clock-drift case is diagnosable instead of looking like a plain timeout. |
| 17 | //! |
| 18 | //! The loop is generic over the poll result and does no I/O of its own: the |
| 19 | //! caller supplies the poll and the sleep. Nothing here ever holds, formats, or |
| 20 | //! logs a token — `T` is opaque to this module and is never `Debug`-printed. |
| 21 | |
| 22 | use std::time::{Duration, Instant}; |
| 23 | |
| 24 | use anyhow::{Result, bail}; |
| 25 | |
| 26 | /// RFC 8628 §3.2: when the authorization server omits `interval`, clients must |
| 27 | /// poll no faster than every 5 seconds. |
| 28 | pub const DEFAULT_POLL_INTERVAL_SECS: u64 = 5; |
| 29 | /// RFC 8628 §3.5: `slow_down` increases the polling interval by 5 seconds. |
| 30 | pub const SLOW_DOWN_STEP_SECS: u64 = 5; |
| 31 | /// Never poll faster than once a second, whatever the server asks for. |
| 32 | const MINIMUM_INTERVAL: Duration = Duration::from_secs(1); |
| 33 | |
| 34 | /// What one poll of the token endpoint told us. |
| 35 | /// |
| 36 | /// A terminal failure is reported by returning `Err` from the poll closure, so |
| 37 | /// each provider keeps its own error text. |
| 38 | pub enum DevicePollOutcome<T> { |
| 39 | /// The user approved; `T` is the provider's parsed token material. |
| 40 | Complete(T), |
| 41 | /// `authorization_pending` — keep the current interval. |
| 42 | Pending, |
| 43 | /// `slow_down` — back off. `interval_seconds` is the server's new minimum |
| 44 | /// when it supplied one (preferred over the client-tracked interval). |
| 45 | SlowDown { interval_seconds: Option<u64> }, |
| 46 | } |
| 47 | |
| 48 | /// A configured device-code polling run. Build one, then [`DeviceCodePoll::run`]. |
| 49 | pub struct DeviceCodePoll { |
| 50 | interval: Duration, |
| 51 | max_interval: Option<Duration>, |
| 52 | lifetime: Duration, |
| 53 | wait_before_first_poll: bool, |
| 54 | timeout_message: String, |
| 55 | slow_down_timeout_message: Option<String>, |
| 56 | } |
| 57 | |
| 58 | impl DeviceCodePoll { |
| 59 | /// Start a run that gives up after `lifetime` with `timeout_message`. |
| 60 | /// |
| 61 | /// The interval starts at the RFC 8628 default of 5 seconds; callers pass |
| 62 | /// the server's `interval` through [`DeviceCodePoll::interval_seconds`]. |
| 63 | #[must_use] |
| 64 | pub fn new(lifetime: Duration, timeout_message: impl Into<String>) -> Self { |
| 65 | Self { |
| 66 | interval: Duration::from_secs(DEFAULT_POLL_INTERVAL_SECS), |
| 67 | max_interval: None, |
| 68 | lifetime, |
| 69 | wait_before_first_poll: false, |
| 70 | timeout_message: timeout_message.into(), |
| 71 | slow_down_timeout_message: None, |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | /// Apply the server-advertised `interval`. `None` (or a zero/absent value, |
| 76 | /// which RFC 8628 permits) keeps the 5-second default. |
| 77 | #[must_use] |
| 78 | pub fn interval_seconds(mut self, seconds: Option<u64>) -> Self { |
| 79 | if let Some(seconds) = seconds.filter(|seconds| *seconds > 0) { |
| 80 | self.interval = self.clamp_interval(Duration::from_secs(seconds)); |
| 81 | } |
| 82 | self |
| 83 | } |
| 84 | |
| 85 | /// Cap the interval, including after `slow_down` backoff. |
| 86 | #[must_use] |
| 87 | pub fn max_interval_seconds(mut self, seconds: u64) -> Self { |
| 88 | self.max_interval = Some(Duration::from_secs(seconds.max(1))); |
| 89 | self.interval = self.clamp_interval(self.interval); |
| 90 | self |
| 91 | } |
| 92 | |
| 93 | /// Sleep one interval before the first poll. |
| 94 | /// |
| 95 | /// Device-code endpoints that answer `authorization_pending` (xAI) want |
| 96 | /// this; endpoints whose first response is already meaningful (the |
| 97 | /// Codewhale account service, which returns HTTP 202 while pending) poll |
| 98 | /// immediately and sleep afterwards. |
| 99 | #[must_use] |
| 100 | pub fn wait_before_first_poll(mut self, wait: bool) -> Self { |
| 101 | self.wait_before_first_poll = wait; |
| 102 | self |
| 103 | } |
| 104 | |
| 105 | /// Message used instead of the plain timeout message when the run saw at |
| 106 | /// least one `slow_down`. This is the WSL/VM clock-drift tell. |
| 107 | #[must_use] |
| 108 | pub fn slow_down_timeout_message(mut self, message: impl Into<String>) -> Self { |
| 109 | self.slow_down_timeout_message = Some(message.into()); |
| 110 | self |
| 111 | } |
| 112 | |
| 113 | fn clamp_interval(&self, interval: Duration) -> Duration { |
| 114 | let interval = interval.max(MINIMUM_INTERVAL); |
| 115 | match self.max_interval { |
| 116 | Some(max) => interval.min(max), |
| 117 | None => interval, |
| 118 | } |
| 119 | } |
| 120 | |
| 121 | /// Poll until the flow completes, fails, or the deadline passes. |
| 122 | /// |
| 123 | /// `sleep` is injected so tests never wait in real time. `poll` returns |
| 124 | /// `Err` for any terminal failure (denied, expired, transport error). |
| 125 | pub fn run<T, S, P>(self, mut sleep: S, mut poll: P) -> Result<T> |
| 126 | where |
| 127 | S: FnMut(Duration), |
| 128 | P: FnMut() -> Result<DevicePollOutcome<T>>, |
| 129 | { |
| 130 | let deadline = Instant::now() + self.lifetime; |
| 131 | let mut interval = self.interval; |
| 132 | let mut saw_slow_down = false; |
| 133 | |
| 134 | if self.wait_before_first_poll { |
| 135 | let remaining = deadline.saturating_duration_since(Instant::now()); |
| 136 | if remaining.is_zero() { |
| 137 | return Err(self.timed_out(saw_slow_down)); |
| 138 | } |
| 139 | sleep(interval.min(remaining)); |
| 140 | } |
| 141 | |
| 142 | while Instant::now() < deadline { |
| 143 | match poll()? { |
| 144 | DevicePollOutcome::Complete(value) => return Ok(value), |
| 145 | DevicePollOutcome::Pending => {} |
| 146 | DevicePollOutcome::SlowDown { interval_seconds } => { |
| 147 | saw_slow_down = true; |
| 148 | // Prefer the server's new minimum when it gave one: a |
| 149 | // purely client-tracked interval polls early forever when |
| 150 | // the clock drifts (WSL, suspended VMs). |
| 151 | interval = match interval_seconds.filter(|seconds| *seconds > 0) { |
| 152 | Some(seconds) => self.clamp_interval(Duration::from_secs(seconds)), |
| 153 | None => { |
| 154 | self.clamp_interval(interval + Duration::from_secs(SLOW_DOWN_STEP_SECS)) |
| 155 | } |
| 156 | }; |
| 157 | } |
| 158 | } |
| 159 | |
| 160 | // Never sleep past the code's expiry, even after slow_down backoff. |
| 161 | let remaining = deadline.saturating_duration_since(Instant::now()); |
| 162 | if remaining.is_zero() { |
| 163 | break; |
| 164 | } |
| 165 | sleep(interval.min(remaining)); |
| 166 | } |
| 167 | |
| 168 | Err(self.timed_out(saw_slow_down)) |
| 169 | } |
| 170 | |
| 171 | fn timed_out(&self, saw_slow_down: bool) -> anyhow::Error { |
| 172 | match (saw_slow_down, self.slow_down_timeout_message.as_deref()) { |
| 173 | (true, Some(message)) => anyhow::anyhow!("{message}"), |
| 174 | _ => anyhow::anyhow!("{}", self.timeout_message), |
| 175 | } |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | /// Reject a device-code verification URI that must not be handed to a browser |
| 180 | /// opener. |
| 181 | /// |
| 182 | /// Ported from pi's `validateVerificationUri` |
| 183 | /// (`packages/ai/src/auth/oauth/xai.ts`, MIT, Copyright (c) 2025 Mario |
| 184 | /// Zechner): the URI comes straight off the wire and is passed to the platform |
| 185 | /// "open this" call, so a malicious or compromised response could otherwise |
| 186 | /// launch `file:`, a custom app scheme, or a helper with attacker-chosen |
| 187 | /// arguments. pi requires `https:`; Codewhale additionally allows `http:` on a |
| 188 | /// loopback host, which is what self-hosted issuers and the device-code tests |
| 189 | /// use — matching the loopback allowance the account login already makes. |
| 190 | /// |
| 191 | /// Embedded credentials are rejected in every case. |
| 192 | pub fn validate_browser_verification_uri(raw: &str, context: &str) -> Result<String> { |
| 193 | let trimmed = raw.trim(); |
| 194 | let Ok(url) = url_scheme_and_host(trimmed) else { |
| 195 | bail!("{context} returned an unusable verification URI"); |
| 196 | }; |
| 197 | let (scheme, host, has_credentials) = url; |
| 198 | if has_credentials { |
| 199 | bail!("{context} returned a verification URI with embedded credentials"); |
| 200 | } |
| 201 | let allowed = scheme == "https" || (scheme == "http" && is_loopback_host(&host)); |
| 202 | if !allowed { |
| 203 | bail!("{context} returned an untrusted verification URI"); |
| 204 | } |
| 205 | Ok(trimmed.to_string()) |
| 206 | } |
| 207 | |
| 208 | /// Minimal scheme/host/credential split, so this module stays free of a URL |
| 209 | /// dependency (`codewhale-config` deliberately has no `reqwest`/`url`). |
| 210 | pub(crate) fn url_scheme_and_host(raw: &str) -> Result<(String, String, bool), ()> { |
| 211 | let (scheme, rest) = raw.split_once("://").ok_or(())?; |
| 212 | if scheme.is_empty() |
| 213 | || !scheme |
| 214 | .bytes() |
| 215 | .all(|b| b.is_ascii_alphanumeric() || b == b'+' || b == b'-' || b == b'.') |
| 216 | { |
| 217 | return Err(()); |
| 218 | } |
| 219 | let authority = rest |
| 220 | .split(['/', '?', '#']) |
| 221 | .next() |
| 222 | .filter(|authority| !authority.is_empty()) |
| 223 | .ok_or(())?; |
| 224 | let (credentials, hostport) = match authority.rsplit_once('@') { |
| 225 | Some((credentials, hostport)) => (!credentials.is_empty(), hostport), |
| 226 | None => (false, authority), |
| 227 | }; |
| 228 | let host = match hostport.strip_prefix('[') { |
| 229 | // IPv6 literal: [::1]:8080 |
| 230 | Some(rest) => rest.split_once(']').ok_or(())?.0.to_string(), |
| 231 | None => hostport.split(':').next().ok_or(())?.to_string(), |
| 232 | }; |
| 233 | if host.is_empty() { |
| 234 | return Err(()); |
| 235 | } |
| 236 | Ok(( |
| 237 | scheme.to_ascii_lowercase(), |
| 238 | host.to_ascii_lowercase(), |
| 239 | credentials, |
| 240 | )) |
| 241 | } |
| 242 | |
| 243 | pub(crate) fn is_loopback_host(host: &str) -> bool { |
| 244 | if host == "localhost" || host == "::1" { |
| 245 | return true; |
| 246 | } |
| 247 | host.parse::<std::net::IpAddr>() |
| 248 | .is_ok_and(|address| address.is_loopback()) |
| 249 | } |
| 250 | |
| 251 | #[cfg(test)] |
| 252 | mod tests { |
| 253 | use super::*; |
| 254 | use std::cell::RefCell; |
| 255 | |
| 256 | fn recording_sleep(log: &RefCell<Vec<Duration>>) -> impl FnMut(Duration) + '_ { |
| 257 | move |duration| log.borrow_mut().push(duration) |
| 258 | } |
| 259 | |
| 260 | #[test] |
| 261 | fn completes_on_first_poll_without_waiting() { |
| 262 | let slept = RefCell::new(Vec::new()); |
| 263 | let value = DeviceCodePoll::new(Duration::from_secs(60), "timed out") |
| 264 | .run(recording_sleep(&slept), || { |
| 265 | Ok(DevicePollOutcome::Complete("token")) |
| 266 | }) |
| 267 | .expect("first poll completes"); |
| 268 | assert_eq!(value, "token"); |
| 269 | assert!(slept.borrow().is_empty(), "no sleep before the first poll"); |
| 270 | } |
| 271 | |
| 272 | #[test] |
| 273 | fn waits_one_interval_before_the_first_poll_when_asked() { |
| 274 | let slept = RefCell::new(Vec::new()); |
| 275 | DeviceCodePoll::new(Duration::from_secs(60), "timed out") |
| 276 | .interval_seconds(Some(3)) |
| 277 | .wait_before_first_poll(true) |
| 278 | .run(recording_sleep(&slept), || { |
| 279 | Ok(DevicePollOutcome::Complete(())) |
| 280 | }) |
| 281 | .expect("completes after the initial wait"); |
| 282 | assert_eq!(slept.borrow().as_slice(), [Duration::from_secs(3)]); |
| 283 | } |
| 284 | |
| 285 | #[test] |
| 286 | fn omitted_interval_uses_the_rfc_default_of_five_seconds() { |
| 287 | let slept = RefCell::new(Vec::new()); |
| 288 | let mut polls = 0; |
| 289 | DeviceCodePoll::new(Duration::from_secs(600), "timed out") |
| 290 | .interval_seconds(None) |
| 291 | .run(recording_sleep(&slept), || { |
| 292 | polls += 1; |
| 293 | if polls == 1 { |
| 294 | Ok(DevicePollOutcome::Pending) |
| 295 | } else { |
| 296 | Ok(DevicePollOutcome::Complete(())) |
| 297 | } |
| 298 | }) |
| 299 | .expect("completes"); |
| 300 | assert_eq!(slept.borrow().as_slice(), [Duration::from_secs(5)]); |
| 301 | } |
| 302 | |
| 303 | #[test] |
| 304 | fn slow_down_without_an_interval_adds_five_seconds() { |
| 305 | let slept = RefCell::new(Vec::new()); |
| 306 | let mut polls = 0; |
| 307 | DeviceCodePoll::new(Duration::from_secs(600), "timed out") |
| 308 | .interval_seconds(Some(2)) |
| 309 | .run(recording_sleep(&slept), || { |
| 310 | polls += 1; |
| 311 | match polls { |
| 312 | 1 => Ok(DevicePollOutcome::Pending), |
| 313 | 2 => Ok(DevicePollOutcome::SlowDown { |
| 314 | interval_seconds: None, |
| 315 | }), |
| 316 | _ => Ok(DevicePollOutcome::Complete(())), |
| 317 | } |
| 318 | }) |
| 319 | .expect("completes"); |
| 320 | assert_eq!( |
| 321 | slept.borrow().as_slice(), |
| 322 | [Duration::from_secs(2), Duration::from_secs(7)] |
| 323 | ); |
| 324 | } |
| 325 | |
| 326 | #[test] |
| 327 | fn slow_down_prefers_a_server_supplied_interval() { |
| 328 | // The clock-drift fix: the server's new minimum wins over the |
| 329 | // client-tracked interval, in both directions. |
| 330 | let slept = RefCell::new(Vec::new()); |
| 331 | let mut polls = 0; |
| 332 | DeviceCodePoll::new(Duration::from_secs(600), "timed out") |
| 333 | .interval_seconds(Some(2)) |
| 334 | .run(recording_sleep(&slept), || { |
| 335 | polls += 1; |
| 336 | match polls { |
| 337 | 1 => Ok(DevicePollOutcome::SlowDown { |
| 338 | interval_seconds: Some(30), |
| 339 | }), |
| 340 | _ => Ok(DevicePollOutcome::Complete(())), |
| 341 | } |
| 342 | }) |
| 343 | .expect("completes"); |
| 344 | assert_eq!(slept.borrow().as_slice(), [Duration::from_secs(30)]); |
| 345 | } |
| 346 | |
| 347 | #[test] |
| 348 | fn interval_never_drops_below_one_second_or_exceeds_the_cap() { |
| 349 | let slept = RefCell::new(Vec::new()); |
| 350 | let mut polls = 0; |
| 351 | DeviceCodePoll::new(Duration::from_secs(600), "timed out") |
| 352 | .interval_seconds(Some(0)) |
| 353 | .max_interval_seconds(10) |
| 354 | .run(recording_sleep(&slept), || { |
| 355 | polls += 1; |
| 356 | match polls { |
| 357 | 1 => Ok(DevicePollOutcome::SlowDown { |
| 358 | interval_seconds: Some(99), |
| 359 | }), |
| 360 | _ => Ok(DevicePollOutcome::Complete(())), |
| 361 | } |
| 362 | }) |
| 363 | .expect("completes"); |
| 364 | // interval 0 falls back to the RFC default (5s), capped at 10s. |
| 365 | assert_eq!(slept.borrow().as_slice(), [Duration::from_secs(10)]); |
| 366 | } |
| 367 | |
| 368 | #[test] |
| 369 | fn never_sleeps_past_the_deadline() { |
| 370 | let slept = RefCell::new(Vec::new()); |
| 371 | let error = DeviceCodePoll::new(Duration::from_millis(30), "timed out") |
| 372 | .interval_seconds(Some(600)) |
| 373 | .run( |
| 374 | |duration| { |
| 375 | slept.borrow_mut().push(duration); |
| 376 | std::thread::sleep(duration); |
| 377 | }, |
| 378 | || Ok(DevicePollOutcome::<()>::Pending), |
| 379 | ) |
| 380 | .expect_err("deadline stops the loop"); |
| 381 | assert_eq!(error.to_string(), "timed out"); |
| 382 | for duration in slept.borrow().iter() { |
| 383 | assert!( |
| 384 | *duration <= Duration::from_millis(30), |
| 385 | "slept {duration:?} past a 30ms deadline" |
| 386 | ); |
| 387 | } |
| 388 | } |
| 389 | |
| 390 | #[test] |
| 391 | fn a_terminal_poll_error_stops_immediately() { |
| 392 | let slept = RefCell::new(Vec::new()); |
| 393 | let error = DeviceCodePoll::new(Duration::from_secs(600), "timed out") |
| 394 | .run(recording_sleep(&slept), || { |
| 395 | Err::<DevicePollOutcome<()>, _>(anyhow::anyhow!("access_denied")) |
| 396 | }) |
| 397 | .expect_err("terminal errors propagate"); |
| 398 | assert_eq!(error.to_string(), "access_denied"); |
| 399 | assert!(slept.borrow().is_empty()); |
| 400 | } |
| 401 | |
| 402 | #[test] |
| 403 | fn timing_out_after_slow_down_reports_the_clock_drift_message() { |
| 404 | let error = DeviceCodePoll::new(Duration::from_millis(5), "plain timeout") |
| 405 | .interval_seconds(Some(1)) |
| 406 | .slow_down_timeout_message("clock drift timeout") |
| 407 | .run(std::thread::sleep, || { |
| 408 | Ok(DevicePollOutcome::<()>::SlowDown { |
| 409 | interval_seconds: None, |
| 410 | }) |
| 411 | }) |
| 412 | .expect_err("deadline stops the loop"); |
| 413 | assert_eq!(error.to_string(), "clock drift timeout"); |
| 414 | } |
| 415 | |
| 416 | #[test] |
| 417 | fn timing_out_without_slow_down_reports_the_plain_message() { |
| 418 | let error = DeviceCodePoll::new(Duration::from_millis(5), "plain timeout") |
| 419 | .interval_seconds(Some(1)) |
| 420 | .slow_down_timeout_message("clock drift timeout") |
| 421 | .run(std::thread::sleep, || Ok(DevicePollOutcome::<()>::Pending)) |
| 422 | .expect_err("deadline stops the loop"); |
| 423 | assert_eq!(error.to_string(), "plain timeout"); |
| 424 | } |
| 425 | |
| 426 | #[test] |
| 427 | fn verification_uri_must_be_https_or_loopback_http() { |
| 428 | assert_eq!( |
| 429 | validate_browser_verification_uri("https://accounts.x.ai/device", "xAI").unwrap(), |
| 430 | "https://accounts.x.ai/device" |
| 431 | ); |
| 432 | assert!(validate_browser_verification_uri("http://127.0.0.1:8080/verify", "xAI").is_ok()); |
| 433 | assert!(validate_browser_verification_uri("http://localhost/verify", "xAI").is_ok()); |
| 434 | assert!(validate_browser_verification_uri("http://[::1]:9/verify", "xAI").is_ok()); |
| 435 | |
| 436 | for hostile in [ |
| 437 | "http://accounts.x.ai/device", |
| 438 | "file:///etc/passwd", |
| 439 | "javascript:alert(1)", |
| 440 | "vscode://attacker/run", |
| 441 | "data:text/html,<script>", |
| 442 | "https://", |
| 443 | "not a url", |
| 444 | "", |
| 445 | ] { |
| 446 | assert!( |
| 447 | validate_browser_verification_uri(hostile, "xAI").is_err(), |
| 448 | "accepted {hostile}" |
| 449 | ); |
| 450 | } |
| 451 | } |
| 452 | |
| 453 | #[test] |
| 454 | fn verification_uri_rejects_embedded_credentials() { |
| 455 | let error = |
| 456 | validate_browser_verification_uri("https://user:pass@accounts.x.ai/device", "xAI") |
| 457 | .expect_err("credentials must be rejected"); |
| 458 | assert!(error.to_string().contains("embedded credentials")); |
| 459 | } |
| 460 | } |
| 461 |