Skip to main content

dbx_tools_auth/
token.rs

1use oauth2::{
2    basic::BasicTokenType, AccessToken, EmptyExtraTokenFields, RefreshToken, StandardTokenResponse,
3};
4use serde::{Deserialize, Serialize};
5use time::{Duration, OffsetDateTime};
6
7use crate::{Error, Result};
8
9pub type OAuthTokenResponse = StandardTokenResponse<EmptyExtraTokenFields, BasicTokenType>;
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
12pub struct Token {
13    pub access_token: String,
14    pub token_type: String,
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub refresh_token: Option<String>,
17    #[serde(
18        rename = "expiry",
19        alias = "expires_at",
20        with = "time::serde::rfc3339::option",
21        default,
22        skip_serializing_if = "Option::is_none"
23    )]
24    pub expires_at: Option<OffsetDateTime>,
25    #[serde(default, skip_serializing_if = "Vec::is_empty")]
26    pub scopes: Vec<String>,
27}
28
29impl Token {
30    pub fn needs_refresh(&self, now: OffsetDateTime, buffer: Duration) -> bool {
31        self.expires_at.is_some_and(|expiry| expiry - now <= buffer)
32    }
33
34    pub fn is_valid(&self, now: OffsetDateTime) -> bool {
35        !self.access_token.is_empty() && self.expires_at.is_none_or(|expiry| expiry > now)
36    }
37
38    pub fn access_token(&self) -> AccessToken {
39        AccessToken::new(self.access_token.clone())
40    }
41
42    pub fn refresh_token(&self) -> Option<RefreshToken> {
43        self.refresh_token.clone().map(RefreshToken::new)
44    }
45
46    pub fn from_response(
47        response: &OAuthTokenResponse,
48        now: OffsetDateTime,
49        previous: Option<&Token>,
50    ) -> Result<Self> {
51        use oauth2::TokenResponse;
52
53        let expires_at = response
54            .expires_in()
55            .and_then(|duration| Duration::try_from(duration).ok())
56            .and_then(|duration| now.checked_add(duration));
57        let token_type = match response.token_type() {
58            BasicTokenType::Bearer => "Bearer",
59            BasicTokenType::Mac => "MAC",
60            BasicTokenType::Extension(value) => value.as_str(),
61        };
62        let refresh_token = response
63            .refresh_token()
64            .map(|token| token.secret().to_owned())
65            .or_else(|| previous.and_then(|token| token.refresh_token.clone()));
66        let scopes = response
67            .scopes()
68            .map(|scopes| {
69                scopes
70                    .iter()
71                    .map(|scope| scope.as_ref().to_owned())
72                    .collect()
73            })
74            .or_else(|| previous.map(|token| token.scopes.clone()))
75            .unwrap_or_default();
76
77        if response.access_token().secret().is_empty() {
78            return Err(Error::OAuth(
79                "token response did not contain an access token".into(),
80            ));
81        }
82        Ok(Self {
83            access_token: response.access_token().secret().to_owned(),
84            token_type: token_type.to_owned(),
85            refresh_token,
86            expires_at,
87            scopes,
88        })
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::Token;
95
96    #[test]
97    fn refresh_window_handles_full_binding_integer_range() {
98        let now = time::OffsetDateTime::now_utc();
99        let token = Token {
100            access_token: "access".into(),
101            token_type: "Bearer".into(),
102            refresh_token: None,
103            expires_at: Some(now),
104            scopes: vec![],
105        };
106        assert!(token.needs_refresh(now, time::Duration::seconds(i64::MAX)));
107        assert!(!token.needs_refresh(now, time::Duration::seconds(i64::MIN)));
108    }
109
110    #[test]
111    fn reads_databricks_cli_rfc3339_expiry() {
112        let token: Token = serde_json::from_str(
113            r#"{"access_token":"access","token_type":"Bearer","refresh_token":"refresh","expiry":"2026-07-30T18:35:12.447-04:00","expires_in":3600}"#,
114        )
115        .unwrap();
116
117        assert_eq!(
118            token.expires_at.unwrap().to_string(),
119            "2026-07-30 18:35:12.447 -04:00:00"
120        );
121    }
122}