yapi_ai/auth/
device.rs

1//! RFC 8628 device authorization polling. Port of `oauth/device-code.ts`.
2
3use std::future::Future;
4use std::time::{Duration, Instant};
5
6use serde_json::Value;
7use tokio_util::sync::CancellationToken;
8
9use super::AuthError;
10
11/// RFC 8628's grant type for exchanging a device code.
12pub(crate) const DEVICE_CODE_GRANT: &str = "urn:ietf:params:oauth:grant-type:device_code";
13
14const TIMEOUT_MESSAGE: &str = "Device flow timed out";
15const SLOW_DOWN_TIMEOUT_MESSAGE: &str = "Device flow timed out after one or more slow_down responses. This is often caused by clock drift in WSL or VM environments. Please sync or restart the VM clock and try again.";
16const MINIMUM_INTERVAL: Duration = Duration::from_secs(1);
17/// RFC 8628 section 3.2: the default when the server omits `interval`.
18const DEFAULT_INTERVAL_SECONDS: f64 = 5.0;
19/// RFC 8628 section 3.5: `slow_down` adds five seconds.
20const SLOW_DOWN_INCREMENT: Duration = Duration::from_secs(5);
21
22/// An http(s) URL, the only kind opened in the browser.
23pub(crate) fn trusted_url(value: &Value) -> Option<String> {
24    let url = url::Url::parse(value.as_str()?).ok()?;
25    matches!(url.scheme(), "http" | "https").then(|| url.to_string())
26}
27
28/// A finite number above zero.
29pub(crate) fn positive(value: &Value) -> Option<f64> {
30    value
31        .as_f64()
32        .filter(|value| value.is_finite() && *value > 0.0)
33}
34
35/// The result of one poll.
36#[derive(Debug, PartialEq)]
37pub(crate) enum Poll<T> {
38    /// Not yet authorized.
39    Pending,
40    /// Poll less often. The server may name the new interval.
41    SlowDown(Option<f64>),
42    /// The flow failed with this message.
43    Failed(String),
44    /// Authorized.
45    Complete(T),
46}
47
48fn seconds(value: f64) -> Duration {
49    Duration::try_from_secs_f64(value.max(0.0)).unwrap_or(Duration::MAX)
50}
51
52async fn sleep(duration: Duration, cancel: &CancellationToken) -> Result<(), AuthError> {
53    cancel
54        .run_until_cancelled(tokio::time::sleep(duration))
55        .await
56        .ok_or(AuthError::Cancelled)
57}
58
59/// Polls until `poll` completes, fails, the code expires or `cancel` fires.
60pub(crate) async fn poll_device_code<T, F, Fut>(
61    interval_seconds: Option<f64>,
62    expires_in_seconds: Option<f64>,
63    wait_before_first_poll: bool,
64    cancel: &CancellationToken,
65    mut poll: F,
66) -> Result<T, AuthError>
67where
68    F: FnMut() -> Fut,
69    Fut: Future<Output = Result<Poll<T>, AuthError>>,
70{
71    let deadline = expires_in_seconds.map(|expires| Instant::now() + seconds(expires));
72    let remaining = || {
73        deadline.map_or(Duration::MAX, |deadline| {
74            deadline.saturating_duration_since(Instant::now())
75        })
76    };
77    let mut interval =
78        seconds(interval_seconds.unwrap_or(DEFAULT_INTERVAL_SECONDS)).max(MINIMUM_INTERVAL);
79    let mut slow_downs = 0;
80    if wait_before_first_poll && !remaining().is_zero() {
81        sleep(interval.min(remaining()), cancel).await?;
82    }
83    while !remaining().is_zero() {
84        if cancel.is_cancelled() {
85            return Err(AuthError::Cancelled);
86        }
87        match poll().await? {
88            Poll::Complete(value) => return Ok(value),
89            Poll::Failed(message) => return Err(AuthError::Failed(message)),
90            Poll::SlowDown(server_interval) => {
91                slow_downs += 1;
92                interval = match server_interval {
93                    Some(value) if value.is_finite() && value > 0.0 => seconds(value),
94                    _ => interval + SLOW_DOWN_INCREMENT,
95                }
96                .max(MINIMUM_INTERVAL);
97            }
98            Poll::Pending => {}
99        }
100        if remaining().is_zero() {
101            break;
102        }
103        sleep(interval.min(remaining()), cancel).await?;
104    }
105    Err(AuthError::failed(if slow_downs > 0 {
106        SLOW_DOWN_TIMEOUT_MESSAGE
107    } else {
108        TIMEOUT_MESSAGE
109    }))
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    #[tokio::test(start_paused = true)]
117    async fn honors_slow_down() {
118        let cancel = CancellationToken::new();
119        let mut calls = 0;
120        let started = tokio::time::Instant::now();
121        let value = poll_device_code(Some(1.0), Some(60.0), true, &cancel, || {
122            calls += 1;
123            let result = match calls {
124                1 => Poll::Pending,
125                2 => Poll::SlowDown(None),
126                _ => Poll::Complete("token"),
127            };
128            async move { Ok(result) }
129        })
130        .await
131        .unwrap();
132        assert_eq!(value, "token");
133        // 1s before the first poll, 1s after pending, 6s after slow_down.
134        assert_eq!(started.elapsed().as_secs(), 8);
135    }
136
137    #[tokio::test(start_paused = true)]
138    async fn times_out() {
139        let cancel = CancellationToken::new();
140        let result: Result<(), _> =
141            poll_device_code(Some(1.0), Some(3.0), false, &cancel, || async {
142                Ok(Poll::SlowDown(Some(2.0)))
143            })
144            .await;
145        assert_eq!(result, Err(AuthError::failed(SLOW_DOWN_TIMEOUT_MESSAGE)));
146        let failed: Result<(), _> = poll_device_code(None, None, false, &cancel, || async {
147            Ok(Poll::Failed("Device flow failed: access_denied".into()))
148        })
149        .await;
150        assert_eq!(
151            failed,
152            Err(AuthError::failed("Device flow failed: access_denied"))
153        );
154    }
155}
yapi_ai/auth/device.rs

yapi_ai/auth/
device.rs
