yt/auth/
token_source.rs

1use std::path::{Path, PathBuf};
2
3use chrono::Utc;
4use futures::future::BoxFuture;
5use miette::IntoDiagnostic as _;
6use secrecy::{ExposeSecret as _, SecretString};
7use sha2::{Digest as _, Sha256};
8use url::Url;
9
10use crate::api::TokenSource;
11use crate::auth::oauth::{self, Refresh};
12use crate::credentials::{self, OauthCredential, StoredCredential};
13use crate::error::{Result, YtError};
14
15/// A permanent token, from the keyring or `YOUTRACK_TOKEN`.
16pub struct StaticToken {
17    token: SecretString,
18    label: String,
19}
20
21impl StaticToken {
22    pub fn new(token: SecretString, label: impl Into<String>) -> Self {
23        crate::redact::register(token.expose_secret());
24        Self {
25            token,
26            label: label.into(),
27        }
28    }
29}
30
31impl TokenSource for StaticToken {
32    fn bearer(&self) -> BoxFuture<'_, Result<SecretString>> {
33        Box::pin(async move { Ok(SecretString::from(self.token.expose_secret().to_string())) })
34    }
35
36    fn refresh_after_unauthorized(&self) -> BoxFuture<'_, Result<bool>> {
37        Box::pin(async { Ok(false) })
38    }
39
40    fn label(&self) -> &str {
41        &self.label
42    }
43}
44
45/// A browser sign-in: refreshes the access token when it is about to expire or is rejected.
46pub struct OAuthToken {
47    http: reqwest::Client,
48    account: String,
49    label: String,
50    lock_path: PathBuf,
51    current: tokio::sync::Mutex<OauthCredential>,
52}
53
54impl OAuthToken {
55    pub fn new(
56        http: reqwest::Client,
57        account: impl Into<String>,
58        label: impl Into<String>,
59        state_dir: &Path,
60        credential: OauthCredential,
61    ) -> Self {
62        let account = account.into();
63        crate::redact::register(&credential.access_token);
64        crate::redact::register(&credential.refresh_token);
65        let digest = Sha256::digest(account.as_bytes());
66        let name: String = digest
67            .iter()
68            .take(8)
69            .map(|byte| format!("{byte:02x}"))
70            .collect();
71        Self {
72            http,
73            lock_path: state_dir.join(format!("refresh-{name}.lock")),
74            account,
75            label: label.into(),
76            current: tokio::sync::Mutex::new(credential),
77        }
78    }
79
80    /// Refreshes under a file lock shared by every yt process, then re-reads the keyring first in
81    /// case another process already refreshed.
82    async fn refresh(&self, current: &mut OauthCredential) -> Result<()> {
83        let _lock = lock(self.lock_path.clone()).await?;
84        if let Some(StoredCredential::Oauth(stored)) = credentials::load(&self.account).await?
85            && stored.access_token != current.access_token
86            && !expiring(&stored)
87        {
88            *current = stored;
89            return Ok(());
90        }
91        let endpoint = Url::parse(&current.token_endpoint).map_err(|error| YtError::OAuth {
92            message: format!("invalid token endpoint: {error}"),
93            help: None,
94        })?;
95        let refresh_token = SecretString::from(current.refresh_token.clone());
96        match oauth::refresh(
97            &self.http,
98            &endpoint,
99            &current.client_id,
100            &refresh_token,
101            &current.scope,
102        )
103        .await?
104        {
105            Refresh::Expired => Err(YtError::SessionExpired {
106                profile: self.label.clone(),
107            }
108            .into()),
109            Refresh::Fresh(tokens) => {
110                current.access_token = tokens.access_token.expose_secret().to_string();
111                if let Some(rotated) = tokens.refresh_token {
112                    current.refresh_token = rotated.expose_secret().to_string();
113                }
114                current.expires_at = Utc::now() + chrono::Duration::seconds(tokens.expires_in);
115                credentials::save(&self.account, &StoredCredential::Oauth(current.clone())).await
116            }
117        }
118    }
119}
120
121impl TokenSource for OAuthToken {
122    fn bearer(&self) -> BoxFuture<'_, Result<SecretString>> {
123        Box::pin(async move {
124            let mut current = self.current.lock().await;
125            if expiring(&current) {
126                self.refresh(&mut current).await?;
127            }
128            Ok(SecretString::from(current.access_token.clone()))
129        })
130    }
131
132    fn refresh_after_unauthorized(&self) -> BoxFuture<'_, Result<bool>> {
133        Box::pin(async move {
134            let mut current = self.current.lock().await;
135            self.refresh(&mut current).await?;
136            Ok(true)
137        })
138    }
139
140    fn label(&self) -> &str {
141        &self.label
142    }
143}
144
145fn expiring(credential: &OauthCredential) -> bool {
146    credential.expires_at - Utc::now() < chrono::Duration::seconds(60)
147}
148
149async fn lock(path: PathBuf) -> Result<std::fs::File> {
150    tokio::task::spawn_blocking(move || -> std::io::Result<std::fs::File> {
151        if let Some(parent) = path.parent() {
152            std::fs::create_dir_all(parent)?;
153        }
154        let file = std::fs::OpenOptions::new()
155            .create(true)
156            .truncate(false)
157            .write(true)
158            .open(&path)?;
159        file.lock()?;
160        Ok(file)
161    })
162    .await
163    .into_diagnostic()?
164    .into_diagnostic()
165}
166
167#[cfg(test)]
168mod tests {
169    use std::time::Duration;
170
171    use serde_json::json;
172    use wiremock::matchers::{body_string_contains, method};
173    use wiremock::{Mock, MockServer, ResponseTemplate};
174
175    use super::*;
176
177    const ACCOUNT: &str = "jane@refresh.youtrack.cloud";
178
179    fn http() -> reqwest::Client {
180        crate::http::client(&crate::http::HttpConfig {
181            timeout: Duration::from_secs(5),
182            retries: 0,
183            ca_file: None,
184        })
185        .unwrap()
186    }
187
188    fn credential(server: &MockServer, expires_in: i64) -> OauthCredential {
189        let now = Utc::now();
190        OauthCredential {
191            access_token: "old-access-0123456789".into(),
192            refresh_token: "refresh-0123456789".into(),
193            expires_at: now + chrono::Duration::seconds(expires_in),
194            client_id: "yt-cli".into(),
195            token_endpoint: format!("{}/hub/api/rest/oauth2/token", server.uri()),
196            scope: "svc".into(),
197            created_at: now,
198        }
199    }
200
201    async fn setup(server: &MockServer, expires_in: i64) -> (tempfile::TempDir, OauthCredential) {
202        keyring_core::set_default_store(keyring_core::mock::Store::new().unwrap());
203        let credential = credential(server, expires_in);
204        credentials::save(ACCOUNT, &StoredCredential::Oauth(credential.clone()))
205            .await
206            .unwrap();
207        (tempfile::tempdir().unwrap(), credential)
208    }
209
210    fn fresh_tokens() -> ResponseTemplate {
211        ResponseTemplate::new(200).set_body_json(json!({"access_token": "new-access-0123456789", "expires_in": 3600, "token_type": "Bearer"}))
212    }
213
214    #[tokio::test]
215    async fn a_valid_token_is_used_as_is() {
216        let server = MockServer::start().await;
217        Mock::given(method("POST"))
218            .respond_with(fresh_tokens())
219            .expect(0)
220            .mount(&server)
221            .await;
222        let (state, credential) = setup(&server, 3600).await;
223        let source = OAuthToken::new(http(), ACCOUNT, "acme", state.path(), credential);
224        assert_eq!(
225            source.bearer().await.unwrap().expose_secret(),
226            "old-access-0123456789"
227        );
228    }
229
230    #[tokio::test]
231    async fn an_automatically_registered_client_refreshes_with_its_document_url() {
232        let server = MockServer::start().await;
233        Mock::given(method("POST"))
234            .and(body_string_contains(
235                "client_id=https%3A%2F%2Fvaivars.github.io%2Frust-yt-cli%2Foauth-client.json",
236            ))
237            .respond_with(fresh_tokens())
238            .expect(1)
239            .mount(&server)
240            .await;
241        keyring_core::set_default_store(keyring_core::mock::Store::new().unwrap());
242        let credential = OauthCredential {
243            client_id: crate::auth::login::CIMD_CLIENT_ID.into(),
244            ..credential(&server, 10)
245        };
246        credentials::save(ACCOUNT, &StoredCredential::Oauth(credential.clone()))
247            .await
248            .unwrap();
249        let state = tempfile::tempdir().unwrap();
250        let source = OAuthToken::new(http(), ACCOUNT, "acme", state.path(), credential);
251        assert_eq!(
252            source.bearer().await.unwrap().expose_secret(),
253            "new-access-0123456789"
254        );
255    }
256
257    #[tokio::test]
258    async fn an_expiring_token_is_refreshed_and_saved() {
259        let server = MockServer::start().await;
260        Mock::given(method("POST"))
261            .and(body_string_contains("grant_type=refresh_token"))
262            .respond_with(fresh_tokens())
263            .expect(1)
264            .mount(&server)
265            .await;
266        let (state, credential) = setup(&server, 10).await;
267        let source = OAuthToken::new(http(), ACCOUNT, "acme", state.path(), credential);
268        assert_eq!(
269            source.bearer().await.unwrap().expose_secret(),
270            "new-access-0123456789"
271        );
272        let Some(StoredCredential::Oauth(saved)) = credentials::load(ACCOUNT).await.unwrap() else {
273            panic!("not saved")
274        };
275        assert_eq!(saved.access_token, "new-access-0123456789");
276        assert_eq!(saved.refresh_token, "refresh-0123456789");
277    }
278
279    #[tokio::test]
280    async fn an_invalid_grant_means_the_sign_in_expired() {
281        let server = MockServer::start().await;
282        Mock::given(method("POST"))
283            .respond_with(ResponseTemplate::new(400).set_body_json(
284                json!({"error": "invalid_grant", "error_description": "Unknown refresh token"}),
285            ))
286            .mount(&server)
287            .await;
288        let (state, credential) = setup(&server, 10).await;
289        let source = OAuthToken::new(http(), ACCOUNT, "acme", state.path(), credential);
290        let error = source.bearer().await.unwrap_err();
291        assert!(
292            matches!(error.downcast_ref::<YtError>(), Some(YtError::SessionExpired { profile }) if profile == "acme")
293        );
294        assert_eq!(crate::error::exit_code(&error), 4);
295    }
296
297    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
298    async fn two_processes_refreshing_at_once_send_one_request() {
299        let server = MockServer::start().await;
300        Mock::given(method("POST"))
301            .respond_with(fresh_tokens().set_delay(Duration::from_millis(300)))
302            .expect(1)
303            .mount(&server)
304            .await;
305        let (state, credential) = setup(&server, 10).await;
306        // Two sources with separate in-memory state stand in for two yt processes.
307        let first = OAuthToken::new(http(), ACCOUNT, "acme", state.path(), credential.clone());
308        let second = OAuthToken::new(http(), ACCOUNT, "acme", state.path(), credential);
309        let (a, b) = tokio::join!(first.bearer(), second.bearer());
310        assert_eq!(a.unwrap().expose_secret(), "new-access-0123456789");
311        assert_eq!(b.unwrap().expose_secret(), "new-access-0123456789");
312    }
313}
yt/auth/token_source.rs

yt/auth/
token_source.rs
