fix: More robust token validation

crates/turborepo-auth/src/auth/login.rs

@@ -4,13 +4,10 @@ pub use error::Error;
 use reqwest::Url;
 use tokio::sync::OnceCell;
 use tracing::warn;
-use turborepo_api_client::{Client, TokenClient};
-use turborepo_ui::start_spinner;
+use turborepo_api_client::{CacheClient, Client, TokenClient};
+use turborepo_ui::{start_spinner, BOLD, UI};
 
-use crate::{
-    auth::{check_user_token, extract_vercel_token},
-    error, ui, LoginOptions, Token,
-};
+use crate::{auth::extract_vercel_token, error, ui, LoginOptions, Token};
 
 const DEFAULT_HOST_NAME: &str = "127.0.0.1";
 const DEFAULT_PORT: u16 = 9789;
@@ -21,7 +18,9 @@ const DEFAULT_PORT: u16 = 9789;
 ///
 /// First checks if an existing option has been passed in, then if the login is
 /// to Vercel, checks if the user has a Vercel CLI token on disk.
-pub async fn login<T: Client + TokenClient>(options: &LoginOptions<'_, T>) -> Result<Token, Error> {
+pub async fn login<T: Client + TokenClient + CacheClient>(
+    options: &LoginOptions<'_, T>,
+) -> Result<Token, Error> {
     let LoginOptions {
         api_client,
         ui,
@@ -31,13 +30,31 @@ pub async fn login<T: Client + TokenClient>(options: &LoginOptions<'_, T>) -> Re
         existing_token,
     } = *options; // Deref or we get double references for each of these
 
+    // I created a closure that gives back a closure since the `is_valid` checks do
+    // a call to get the user, so instead of doing that multiple times we have
+    // `is_valid` give back the user email.
+    //
+    // In the future I want to make the Token have some non-skewable information and
+    // be able to get rid of this, but it works for now.
+    let valid_token_callback = |message: &str, ui: &UI| {
+        let message = message.to_string();
+        let ui = *ui;
+        move |user_email: &str| {
+            println!("{}", ui.apply(BOLD.apply_to(message)));
+            ui::print_cli_authorized(user_email, &ui);
+        }
+    };
     // Check if passed in token exists first.
     if let Some(token) = existing_token {
-        if Token::existing(token.to_string())
-            .is_valid(api_client)
+        let token = Token::existing(token.into());
+        if token
+            .is_valid(
+                api_client,
+                Some(valid_token_callback("Existing token found!", ui)),
+            )
             .await?
         {
-            return check_user_token(token, ui, api_client, "Existing token found!").await;
+            return Ok(token);
         }
     }
 
@@ -46,7 +63,16 @@ pub async fn login<T: Client + TokenClient>(options: &LoginOptions<'_, T>) -> Re
         // The extraction can return an error, but we don't want to fail the login if
         // the token is not found.
         if let Ok(token) = extract_vercel_token() {
-            return check_user_token(&token, ui, api_client, "Existing Vercel token found!").await;
+            let token = Token::existing(token);
+            if token
+                .is_valid(
+                    api_client,
+                    Some(valid_token_callback("Existing Vercel token found!", ui)),
+                )
+                .await?
+            {
+                return Ok(token);
+            }
         }
     }
 
@@ -104,11 +130,12 @@ mod tests {
     use std::{assert_matches::assert_matches, sync::atomic::AtomicUsize};
 
     use async_trait::async_trait;
-    use reqwest::{RequestBuilder, Response};
+    use reqwest::{Method, RequestBuilder, Response};
     use turborepo_api_client::Client;
     use turborepo_ui::UI;
     use turborepo_vercel_api::{
-        Membership, Role, SpacesResponse, Team, TeamsResponse, User, UserResponse, VerifiedSsoUser,
+        CachingStatus, CachingStatusResponse, Membership, Role, SpacesResponse, Team,
+        TeamsResponse, User, UserResponse, VerifiedSsoUser,
     };
     use turborepo_vercel_api_mock::start_test_server;
 
@@ -263,6 +290,60 @@ mod tests {
         }
     }
 
+    #[async_trait]
+    impl CacheClient for MockApiClient {
+        async fn get_artifact(
+            &self,
+            _hash: &str,
+            _token: &str,
+            _team_id: Option<&str>,
+            _team_slug: Option<&str>,
+            _method: Method,
+        ) -> Result<Option<Response>, turborepo_api_client::Error> {
+            unimplemented!("get_artifact")
+        }
+        async fn put_artifact(
+            &self,
+            _hash: &str,
+            _artifact_body: &[u8],
+            _duration: u64,
+            _tag: Option<&str>,
+            _token: &str,
+            _team_id: Option<&str>,
+            _team_slug: Option<&str>,
+        ) -> Result<(), turborepo_api_client::Error> {
+            unimplemented!("set_artifact")
+        }
+        async fn fetch_artifact(
+            &self,
+            _hash: &str,
+            _token: &str,
+            _team_id: Option<&str>,
+            _team_slug: Option<&str>,
+        ) -> Result<Option<Response>, turborepo_api_client::Error> {
+            unimplemented!("fetch_artifact")
+        }
+        async fn artifact_exists(
+            &self,
+            _hash: &str,
+            _token: &str,
+            _team_id: Option<&str>,
+            _team_slug: Option<&str>,
+        ) -> Result<Option<Response>, turborepo_api_client::Error> {
+            unimplemented!("artifact_exists")
+        }
+        async fn get_caching_status(
+            &self,
+            _token: &str,
+            _team_id: Option<&str>,
+            _team_slug: Option<&str>,
+        ) -> Result<CachingStatusResponse, turborepo_api_client::Error> {
+            Ok(CachingStatusResponse {
+                status: CachingStatus::Enabled,
+            })
+        }
+    }
+
     #[tokio::test]
     async fn test_login() {
         let port = port_scanner::request_open_port().unwrap();

crates/turborepo-auth/src/auth/mod.rs

@@ -5,17 +5,17 @@ mod sso;
 pub use login::*;
 pub use logout::*;
 pub use sso::*;
-use turborepo_api_client::{Client, TokenClient};
-use turborepo_ui::{BOLD, UI};
+use turborepo_api_client::{CacheClient, Client, TokenClient};
+use turborepo_ui::UI;
 
-use crate::{ui, LoginServer, Token};
+use crate::LoginServer;
 
 const VERCEL_TOKEN_DIR: &str = "com.vercel.cli";
 const VERCEL_TOKEN_FILE: &str = "auth.json";
 
 pub struct LoginOptions<'a, T>
 where
-    T: Client + TokenClient,
+    T: Client + TokenClient + CacheClient,
 {
     pub ui: &'a UI,
     pub login_url: &'a str,
@@ -27,7 +27,7 @@ where
 }
 impl<'a, T> LoginOptions<'a, T>
 where
-    T: Client + TokenClient,
+    T: Client + TokenClient + CacheClient,
 {
     pub fn new(
         ui: &'a UI,
@@ -46,52 +46,6 @@ where
     }
 }
 
-async fn check_user_token(
-    token: &str,
-    ui: &UI,
-    api_client: &(impl Client + TokenClient),
-    message: &str,
-) -> Result<Token, Error> {
-    let response_user = api_client.get_user(token).await?;
-    println!("{}", ui.apply(BOLD.apply_to(message)));
-    ui::print_cli_authorized(&response_user.user.email, ui);
-    Ok(Token::Existing(token.to_string()))
-}
-
-async fn check_sso_token(
-    token: &str,
-    sso_team: &str,
-    ui: &UI,
-    api_client: &(impl Client + TokenClient),
-    message: &str,
-) -> Result<Token, Error> {
-    let (result_user, result_teams) =
-        tokio::join!(api_client.get_user(token), api_client.get_teams(token),);
-
-    let token = Token::existing(token.into());
-
-    match (result_user, result_teams) {
-        (Ok(response_user), Ok(response_teams)) => {
-            if response_teams
-                .teams
-                .iter()
-                .any(|team| team.slug == sso_team)
-            {
-                if token.is_valid(api_client).await? {
-                    println!("{}", ui.apply(BOLD.apply_to(message)));
-                    ui::print_cli_authorized(&response_user.user.email, ui);
-                    Ok(token)
-                } else {
-                    Err(Error::SSOTokenExpired(sso_team.to_string()))
-                }
-            } else {
-                Err(Error::SSOTeamNotFound(sso_team.to_string()))
-            }
-        }
-        (Err(e), _) | (_, Err(e)) => Err(Error::APIError(e)),
-    }
-}
-
 fn extract_vercel_token() -> Result<String, Error> {
     let vercel_config_dir =
         turborepo_dirs::vercel_config_dir().ok_or_else(|| Error::ConfigDirNotFound)?;