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