Auth for Jupyter Notebooks (#11855)

### Related

Related to
https://linear.app/rerun/issue/RR-2861/authentication-in-jupyter-notebooks

### What

Implements authentication in Jupyter notebooks.

- By running ``rr.login()`` the user can log in.
- Uses `re_auth` module, starts a web server in the background to listen
to events.
- Automatically passes credentials to the wasm viewer.

Author: aedm

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::callback_server::OauthCallbackServer;
-use crate::oauth::api::{AuthenticateWithCode, Pkce, send_async};
-use crate::oauth::{self, Credentials};
+use crate::oauth;
+use crate::oauth::login_flow::OauthLoginFlowState;
 
 pub struct LoginOptions {
     pub open_browser: bool,
@@ -42,92 +42,47 @@ 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 pkce = Pkce::new();
-    let server = OauthCallbackServer::new(&pkce, login_hint.as_deref())?;
-    p.inc(1);
+    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,
+    };
 
-    // 2. Open authorization URL in browser
-    let login_url = server.get_login_url();
+    let progress_bar = ProgressBar::new_spinner();
 
-    // 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.
+    // 2. Open authorization URL in browser
+    let login_url = login_flow.get_login_url();
     if options.open_browser {
-        p.println("Opening login page in your browser.");
-        p.println("Once you've logged in, the process will continue here.");
-        p.println(format!(
+        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!(
             "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 {
-        p.println("Open the following page in your browser:");
-        p.println(login_url);
+        progress_bar.println("Open the following page in your browser:");
+        progress_bar.println(login_url);
     }
-    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);
+
+    // 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;
         }
+        progress_bar.inc(1);
+        std::thread::sleep(Duration::from_millis(10));
     };
 
-    // 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();
+    progress_bar.finish_and_clear();
 
     println!(
         "Success! You are now logged in as {}",

crates/utils/re_auth/src/lib.rs

@@ -9,6 +9,7 @@
 
 #[cfg(not(target_arch = "wasm32"))]
 mod error;
+
 #[cfg(not(target_arch = "wasm32"))]
 mod provider;
 
@@ -31,6 +32,8 @@ 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,6 +5,9 @@ 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.
@@ -53,6 +56,9 @@ pub enum CredentialsRefreshError {
 
     #[error("failed to deserialize credentials: {0}")]
     MalformedToken(#[from] MalformedTokenError),
+
+    #[error("no refresh token available")]
+    NoRefreshToken,
 }
 
 /// Refresh credentials if they are expired.
@@ -73,7 +79,11 @@ pub async fn refresh_credentials(
         -credentials.access_token().remaining_duration_secs()
     );
 
-    let response = api::refresh(&credentials.refresh_token).await?;
+    let Some(refresh_token) = &credentials.refresh_token else {
+        return Err(CredentialsRefreshError::NoRefreshToken);
+    };
+
+    let response = api::refresh(refresh_token).await?;
     let credentials = Credentials::from_auth_response(response)?
         .ensure_stored()
         .map_err(|err| CredentialsRefreshError::Store(err.0))?;
@@ -181,7 +191,12 @@ pub enum VerifyError {
 #[derive(Debug, serde::Serialize, serde::Deserialize)]
 pub struct Credentials {
     user: User,
-    refresh_token: RefreshToken,
+
+    // 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>,
+
     access_token: AccessToken,
 }
 
@@ -209,14 +224,42 @@ impl Credentials {
     pub fn from_auth_response(
         res: api::RefreshResponse,
     ) -> Result<InMemoryCredentials, MalformedTokenError> {
-        let access_token = AccessToken::unverified(Jwt(res.access_token))?;
+        let access_token = AccessToken::try_from_unverified_jwt(Jwt(res.access_token))?;
         Ok(InMemoryCredentials(Self {
             user: res.user,
-            refresh_token: RefreshToken(res.refresh_token),
+            refresh_token: Some(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
     }
@@ -266,25 +309,11 @@ impl AccessToken {
     /// Construct an [`AccessToken`] without verifying it.
     ///
     /// The token should come from a trusted source, like the Rerun auth API.
-    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)?;
-
+    pub(crate) fn try_from_unverified_jwt(jwt: Jwt) -> Result<Self, MalformedTokenError> {
+        let claims = jwt.unverified_claims()?;
         Ok(Self {
             token: jwt.0,
-            expires_at: payload.exp,
+            expires_at: claims.exp,
         })
     }
 }

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

@@ -0,0 +1,92 @@
+use crate::callback_server::Error;
+use crate::callback_server::OauthCallbackServer;
+use crate::oauth;
+use crate::oauth::Credentials;
+use crate::oauth::api::AuthenticateWithCode;
+use crate::oauth::api::Pkce;
+use crate::oauth::api::send_async;
+
+pub enum OauthLoginFlowState {
+    AlreadyLoggedIn(Credentials),
+    LoginFlowStarted(OauthLoginFlow),
+}
+
+pub struct OauthLoginFlow {
+    pub server: OauthCallbackServer,
+    pub login_hint: Option<String>,
+    pkce: Pkce,
+}
+
+impl OauthLoginFlow {
+    /// Login to Rerun using Authorization Code flow.
+    ///
+    /// This first checks if valid credentials already exist locally,
+    /// and doesn't perform the login flow if so, unless `force_login` is set to `true`.
+    pub async fn init(force_login: bool) -> Result<OauthLoginFlowState, Error> {
+        let mut login_hint = None;
+        if !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) => {
+                            return Ok(OauthLoginFlowState::AlreadyLoggedIn(credentials));
+                        }
+                        Err(err) => {
+                            // Credentials are bad, login again.
+                            re_log::debug!("refreshing credentials failed: {err}");
+                        }
+                    }
+                }
+
+                Ok(None) => {
+                    // No credentials yet, login as usual.
+                }
+
+                Err(err) => {
+                    re_log::debug!(
+                        "validating credentials failed, logging user in again anyway. reason: {err}"
+                    );
+                }
+            }
+        }
+
+        // Start web server that listens for the authorization code received from the auth server.
+        let pkce = Pkce::new();
+        let server = OauthCallbackServer::new(&pkce, login_hint.as_deref())?;
+
+        Ok(OauthLoginFlowState::LoginFlowStarted(Self {
+            server,
+            pkce,
+            login_hint,
+        }))
+    }
+
+    pub fn get_login_url(&self) -> &str {
+        self.server.get_login_url()
+    }
+
+    /// Polls the web server for the authorization code received from the auth server.
+    ///
+    /// This will not block, and will return `None` if no authorization code has been received yet.
+    pub async fn poll(&self) -> Result<Option<Credentials>, Error> {
+        // 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.
+        let Some(code) = self.server.check_for_browser_response()? else {
+            return Ok(None);
+        };
+
+        // Exchange code for credentials.
+        let auth = send_async(AuthenticateWithCode::new(&code, &self.pkce))
+            .await
+            .map_err(|err| Error::Generic(err.into()))?;
+
+        // Store and return credentials
+        let credentials = Credentials::from_auth_response(auth.into())?.ensure_stored()?;
+        Ok(Some(credentials))
+    }
+}

crates/utils/re_auth/src/token.rs

@@ -1,5 +1,9 @@
+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}")]
@@ -15,6 +19,22 @@ 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 {