Revert "Auth for Jupyter Notebooks (#11855)"

This reverts commit c3e5829.

Author: ntjohnson1

crates/utils/re_auth/src/cli.rs

@@ -2,10 +2,10 @@ use std::time::Duration;
 
 use indicatif::ProgressBar;
 
-use crate::OauthLoginFlow;
 pub use crate::callback_server::Error;
-use crate::oauth;
-use crate::oauth::login_flow::OauthLoginFlowState;
+use crate::callback_server::OauthCallbackServer;
+use crate::oauth::api::{AuthenticateWithCode, Pkce, send_async};
+use crate::oauth::{self, Credentials};
 
 pub struct LoginOptions {
     pub open_browser: bool,
@@ -42,47 +42,92 @@ pub async fn token() -> Result<(), Error> {
 /// This first checks if valid credentials already exist locally,
 /// and doesn't perform the login flow if so, unless `options.force_login` is set to `true`.
 pub async fn login(options: LoginOptions) -> Result<(), Error> {
+    let mut login_hint = None;
+    if !options.force_login {
+        // NOTE: If the loading fails for whatever reason, we debug log the error
+        // and have the user login again as if nothing happened.
+        match oauth::load_credentials() {
+            Ok(Some(credentials)) => {
+                login_hint = Some(credentials.user().email.clone());
+                match oauth::refresh_credentials(credentials).await {
+                    Ok(credentials) => {
+                        println!("You're already logged in as: {}", credentials.user().email);
+                        println!("Note: We've refreshed your credentials.");
+                        println!("Note: Run `rerun auth login --force` to login again.");
+                        return Ok(());
+                    }
+                    Err(err) => {
+                        re_log::debug!("refreshing credentials failed: {err}");
+                        // Credentials are bad, login again.
+                        // fallthrough
+                    }
+                }
+            }
+
+            Ok(None) => {
+                // No credentials yet, login as usual.
+                // fallthrough
+            }
+
+            Err(err) => {
+                re_log::debug!(
+                    "validating credentials failed, logging user in again anyway. reason: {err}"
+                );
+                // fallthrough
+            }
+        }
+    }
+
+    let p = ProgressBar::new_spinner();
+
     // Login process:
 
     // 1. Start web server listening for token
-    let login_flow = match OauthLoginFlow::init(options.force_login).await? {
-        OauthLoginFlowState::AlreadyLoggedIn(credentials) => {
-            println!("You're already logged in as: {}", credentials.user().email);
-            println!("Note: We've refreshed your credentials.");
-            println!("Note: Run `rerun auth login --force` to login again.");
-            return Ok(());
-        }
-        OauthLoginFlowState::LoginFlowStarted(login_flow) => login_flow,
-    };
-
-    let progress_bar = ProgressBar::new_spinner();
+    let pkce = Pkce::new();
+    let server = OauthCallbackServer::new(&pkce, login_hint.as_deref())?;
+    p.inc(1);
 
     // 2. Open authorization URL in browser
-    let login_url = login_flow.get_login_url();
+    let login_url = server.get_login_url();
+
+    // Once the user opens the link, they are redirected to the login UI.
+    // If they were already logged in, it will immediately redirect them
+    // to the login callback with an authorization code.
+    // That code is then sent by our callback page back to the web server here.
     if options.open_browser {
-        progress_bar.println("Opening login page in your browser.");
-        progress_bar.println("Once you've logged in, the process will continue here.");
-        progress_bar.println(format!(
+        p.println("Opening login page in your browser.");
+        p.println("Once you've logged in, the process will continue here.");
+        p.println(format!(
             "Alternatively, manually open this url: {login_url}"
         ));
         webbrowser::open(login_url).ok(); // Ok to ignore error here. The user can just open the above url themselves.
     } else {
-        progress_bar.println("Open the following page in your browser:");
-        progress_bar.println(login_url);
+        p.println("Open the following page in your browser:");
+        p.println(login_url);
     }
-    progress_bar.inc(1);
-
-    // 3. Wait for login to finish
-    progress_bar.set_message("Waiting for browser…");
-    let credentials = loop {
-        if let Some(code) = login_flow.poll().await? {
-            break code;
+    p.inc(1);
+
+    // 3. Wait for callback
+    p.set_message("Waiting for browser…");
+    let code = loop {
+        match server.check_for_browser_response()? {
+            None => {
+                p.inc(1);
+                std::thread::sleep(Duration::from_millis(10));
+            }
+            Some(response) => break response,
         }
-        progress_bar.inc(1);
-        std::thread::sleep(Duration::from_millis(10));
     };
 
-    progress_bar.finish_and_clear();
+    // 4. Exchange code for credentials
+    let auth = send_async(AuthenticateWithCode::new(&code, &pkce))
+        .await
+        .map_err(|err| Error::Generic(err.into()))?;
+
+    // 5. Store credentials
+    let credentials = Credentials::from_auth_response(auth.into())?.ensure_stored()?;
+
+    p.finish_and_clear();
 
     println!(
         "Success! You are now logged in as {}",

crates/utils/re_auth/src/lib.rs

@@ -9,7 +9,6 @@
 
 #[cfg(not(target_arch = "wasm32"))]
 mod error;
-
 #[cfg(not(target_arch = "wasm32"))]
 mod provider;
 
@@ -32,8 +31,6 @@ pub use token::{Jwt, TokenError};
 
 #[cfg(not(target_arch = "wasm32"))]
 pub use error::Error;
-#[cfg(all(feature = "oauth", not(target_arch = "wasm32")))]
-pub use oauth::login_flow::OauthLoginFlow;
 #[cfg(not(target_arch = "wasm32"))]
 pub use provider::{Claims, RedapProvider, SecretKey, VerificationOptions};
 #[cfg(not(target_arch = "wasm32"))]

crates/utils/re_auth/src/oauth.rs

@@ -5,9 +5,6 @@ use crate::Jwt;
 pub mod api;
 mod storage;
 
-#[cfg(not(target_arch = "wasm32"))]
-pub mod login_flow;
-
 /// Tokens with fewer than this number of seconds left before expiration
 /// are considered expired. This ensures tokens don't become expired
 /// during network transit.
@@ -56,9 +53,6 @@ pub enum CredentialsRefreshError {
 
     #[error("failed to deserialize credentials: {0}")]
     MalformedToken(#[from] MalformedTokenError),
-
-    #[error("no refresh token available")]
-    NoRefreshToken,
 }
 
 /// Refresh credentials if they are expired.
@@ -79,11 +73,7 @@ pub async fn refresh_credentials(
         -credentials.access_token().remaining_duration_secs()
     );
 
-    let Some(refresh_token) = &credentials.refresh_token else {
-        return Err(CredentialsRefreshError::NoRefreshToken);
-    };
-
-    let response = api::refresh(refresh_token).await?;
+    let response = api::refresh(&credentials.refresh_token).await?;
     let credentials = Credentials::from_auth_response(response)?
         .ensure_stored()
         .map_err(|err| CredentialsRefreshError::Store(err.0))?;
@@ -191,12 +181,7 @@ pub enum VerifyError {
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
 pub struct Credentials {
     user: User,
-
-    // Refresh token is optional because it may not be available in some cases,
-    // like the Jupyter notebook Wasm viewer. In that case, the SDK handles
-    // token refreshes.
-    refresh_token: Option<RefreshToken>,
-
+    refresh_token: RefreshToken,
     access_token: AccessToken,
 }
 
@@ -224,42 +209,14 @@ impl Credentials {
     pub fn from_auth_response(
         res: api::RefreshResponse,
     ) -> Result<InMemoryCredentials, MalformedTokenError> {
-        let access_token = AccessToken::try_from_unverified_jwt(Jwt(res.access_token))?;
+        let access_token = AccessToken::unverified(Jwt(res.access_token))?;
         Ok(InMemoryCredentials(Self {
             user: res.user,
-            refresh_token: Some(RefreshToken(res.refresh_token)),
+            refresh_token: RefreshToken(res.refresh_token),
             access_token,
         }))
     }
 
-    /// Creates credentials from raw token strings.
-    ///
-    /// Warning: it does not check the signature of the access token.
-    pub fn try_new(
-        access_token: String,
-        refresh_token: Option<String>,
-        email: String,
-    ) -> Result<InMemoryCredentials, MalformedTokenError> {
-        // TODO(aedm): check signature of the JWT token
-        let claims = Jwt(access_token.clone()).unverified_claims()?;
-
-        let user = User {
-            id: claims.sub,
-            email,
-        };
-        let access_token = AccessToken {
-            token: access_token,
-            expires_at: claims.exp,
-        };
-        let refresh_token = refresh_token.map(RefreshToken);
-
-        Ok(InMemoryCredentials(Self {
-            user,
-            access_token,
-            refresh_token,
-        }))
-    }
-
     pub fn access_token(&self) -> &AccessToken {
         &self.access_token
     }
@@ -309,11 +266,25 @@ impl AccessToken {
     /// Construct an [`AccessToken`] without verifying it.
     ///
     /// The token should come from a trusted source, like the Rerun auth API.
-    pub(crate) fn try_from_unverified_jwt(jwt: Jwt) -> Result<Self, MalformedTokenError> {
-        let claims = jwt.unverified_claims()?;
+    pub(crate) fn unverified(jwt: Jwt) -> Result<Self, MalformedTokenError> {
+        use base64::prelude::*;
+
+        let (_header, rest) = jwt
+            .as_str()
+            .split_once('.')
+            .ok_or(MalformedTokenError::MissingHeaderPayloadSeparator)?;
+        let (payload, _signature) = rest
+            .split_once('.')
+            .ok_or(MalformedTokenError::MissingPayloadSignatureSeparator)?;
+        let payload = BASE64_URL_SAFE_NO_PAD
+            .decode(payload)
+            .map_err(MalformedTokenError::Base64)?;
+        let payload: RerunCloudClaims =
+            serde_json::from_slice(&payload).map_err(MalformedTokenError::Serde)?;
+
         Ok(Self {
             token: jwt.0,
-            expires_at: claims.exp,
+            expires_at: payload.exp,
         })
     }
 }

crates/utils/re_auth/src/oauth/login_flow.rs

@@ -1,9 +1,5 @@
-use base64::Engine as _;
-use base64::prelude::BASE64_URL_SAFE_NO_PAD;
 use jsonwebtoken::decode_header;
 
-use crate::oauth::{MalformedTokenError, RerunCloudClaims};
-
 #[derive(Debug, thiserror::Error)]
 pub enum TokenError {
     #[error("token does not seem to be a valid JWT token: {0}")]
@@ -19,22 +15,6 @@ impl Jwt {
     pub fn as_str(&self) -> &str {
         &self.0
     }
-
-    pub fn unverified_claims(&self) -> Result<RerunCloudClaims, MalformedTokenError> {
-        let (_header, rest) = self
-            .as_str()
-            .split_once('.')
-            .ok_or(MalformedTokenError::MissingHeaderPayloadSeparator)?;
-        let (payload, _signature) = rest
-            .split_once('.')
-            .ok_or(MalformedTokenError::MissingPayloadSignatureSeparator)?;
-        let payload = BASE64_URL_SAFE_NO_PAD
-            .decode(payload)
-            .map_err(MalformedTokenError::Base64)?;
-        let claims: RerunCloudClaims =
-            serde_json::from_slice(&payload).map_err(MalformedTokenError::Serde)?;
-        Ok(claims)
-    }
 }
 
 impl TryFrom<String> for Jwt {