Skip to main content

authkestra_engine/flow/
device_flow.rs

1use crate::auth::{
2    error::{AuthError, OAuthErrorResponse},
3    state::OAuthToken,
4};
5use serde::{Deserialize, Serialize};
6use std::thread::sleep;
7use std::time::Duration;
8
9/// Represents the response from the device authorization endpoint.
10/// Defined in RFC 8628 Section 3.2.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct DeviceAuthorizationResponse {
13    /// The device verification code.
14    pub device_code: String,
15    /// The end-user verification code.
16    pub user_code: String,
17    /// The end-user verification URI on the authorization server.
18    pub verification_uri: String,
19    /// A verification URI that includes the "user_code" (or other information)
20    /// to optimize the end-user interaction.
21    pub verification_uri_complete: Option<String>,
22    /// The lifetime in seconds of the "device_code" and "user_code".
23    pub expires_in: u64,
24    /// The minimum amount of time in seconds that the client SHOULD wait
25    /// between polling requests to the token endpoint.
26    pub interval: Option<u64>,
27}
28
29/// Orchestrates the Device Authorization Flow (RFC 8628).
30pub struct DeviceFlow {
31    client_id: String,
32    device_authorization_url: String,
33    token_url: String,
34    http_client: reqwest::Client,
35}
36
37impl DeviceFlow {
38    /// Creates a new `DeviceFlow` instance.
39    pub fn new(client_id: String, device_authorization_url: String, token_url: String) -> Self {
40        Self {
41            client_id,
42            device_authorization_url,
43            token_url,
44            http_client: reqwest::Client::new(),
45        }
46    }
47
48    /// Initiates the device authorization request.
49    /// Returns a `DeviceAuthorizationResponse` which contains the codes and URIs
50    /// to be displayed to the user.
51    pub async fn initiate_device_authorization(
52        &self,
53        scopes: &[&str],
54    ) -> Result<DeviceAuthorizationResponse, AuthError> {
55        let scope_param = scopes.join(" ");
56
57        let response = self
58            .http_client
59            .post(&self.device_authorization_url)
60            .header("Accept", "application/json")
61            .form(&[("client_id", &self.client_id), ("scope", &scope_param)])
62            .send()
63            .await
64            .map_err(|_| AuthError::Network)?;
65
66        if !response.status().is_success() {
67            let error_text = response.text().await.unwrap_or_default();
68            return Err(AuthError::Provider(format!(
69                "Device authorization request failed: {error_text}"
70            )));
71        }
72
73        let response_text = response.text().await.map_err(|e| {
74            AuthError::Provider(format!(
75                "Failed to read device authorization response body: {e}"
76            ))
77        })?;
78
79        println!("Raw device authorization response body: {response_text}");
80
81        serde_json::from_str::<DeviceAuthorizationResponse>(&response_text).map_err(|e| {
82            AuthError::Provider(format!(
83                "Failed to parse device authorization response: {e}"
84            ))
85        })
86    }
87
88    /// Polls the token endpoint until an access token is granted or an error occurs.
89    /// This function respects the `interval` specified by the provider and handles
90    /// common device flow errors like `authorization_pending` and `slow_down`.
91    pub async fn poll_for_token(
92        &self,
93        device_code: &str,
94        interval: Option<u64>,
95    ) -> Result<OAuthToken, AuthError> {
96        let mut current_interval = interval.unwrap_or(5);
97
98        loop {
99            let response = self
100                .http_client
101                .post(&self.token_url)
102                .header("Accept", "application/json")
103                .form(&[
104                    ("client_id", &self.client_id),
105                    ("device_code", &device_code.to_string()),
106                    (
107                        "grant_type",
108                        &"urn:ietf:params:oauth:grant-type:device_code".to_string(),
109                    ),
110                ])
111                .send()
112                .await
113                .map_err(|_| AuthError::Network)?;
114
115            let response_text = response.text().await.map_err(|e| {
116                AuthError::Provider(format!("Failed to read token response body: {e}"))
117            })?;
118
119            // Attempt to deserialize as OAuthErrorResponse first
120            if let Ok(oauth_error) = serde_json::from_str::<OAuthErrorResponse>(&response_text) {
121                match oauth_error.error.as_str() {
122                    "authorization_pending" => {
123                        // Continue polling
124                    }
125                    "slow_down" => {
126                        current_interval += 5;
127                    }
128                    "access_denied" => {
129                        return Err(AuthError::Provider("Access denied by user".into()));
130                    }
131                    "expired_token" => {
132                        return Err(AuthError::Provider("Device code expired".into()));
133                    }
134                    _ => {
135                        let error_description = oauth_error
136                            .error_description
137                            .unwrap_or_else(|| "No description provided".to_string());
138                        return Err(AuthError::Provider(format!(
139                            "OAuth error: {} - {}",
140                            oauth_error.error, error_description
141                        )));
142                    }
143                }
144            } else if let Ok(oauth_token) = serde_json::from_str::<OAuthToken>(&response_text) {
145                // If OAuthErrorResponse deserialization fails, attempt OAuthToken
146                return Ok(oauth_token);
147            } else {
148                // If both deserialization attempts fail
149                return Err(AuthError::Provider(
150                    "Failed to parse token response as either OAuthToken or OAuthErrorResponse"
151                        .into(),
152                ));
153            }
154
155            sleep(Duration::from_secs(current_interval));
156        }
157    }
158}