Refactor auth state handling

Cargo.lock

@@ -8447,6 +8447,7 @@ dependencies = [
  "indicatif",
  "js-sys",
  "jsonwebtoken",
+ "parking_lot",
  "rand 0.9.2",
  "re_analytics",
  "re_log",

crates/utils/re_auth/Cargo.toml

@@ -44,6 +44,7 @@ async-trait.workspace = true
 base64.workspace = true
 http.workspace = true
 jsonwebtoken.workspace = true
+parking_lot.workspace = true
 saturating_cast.workspace = true
 thiserror.workspace = true
 tokio.workspace = true

crates/utils/re_auth/src/credentials.rs

@@ -40,18 +40,35 @@ impl CredentialsProvider for StaticCredentialsProvider {
 }
 
 #[cfg(feature = "oauth")]
-pub use oauth::CliCredentialsProvider;
+pub use oauth::{CliCredentialsProvider, subscribe_auth_changes};
 
 #[cfg(feature = "oauth")]
-mod oauth {
+pub(crate) mod oauth {
     use super::{CredentialsProvider, CredentialsProviderError, Jwt};
+    use crate::oauth;
     use crate::oauth::{Credentials, load_and_refresh_credentials};
     use tokio::sync::RwLock;
 
     // We only want to keep a single instance of credentials in memory,
     // so we store them in a static.
     static CACHE: RwLock<Option<Credentials>> = RwLock::const_new(None);
 
+    type AuthCallback = Box<dyn Fn(Option<oauth::User>) + Send>;
+    static AUTH_SUBSCRIBERS: parking_lot::Mutex<Vec<AuthCallback>> =
+        parking_lot::Mutex::new(Vec::new());
+
+    pub(crate) fn auth_update(user: Option<&oauth::User>) {
+        let subscribers = AUTH_SUBSCRIBERS.lock();
+        for sub in &*subscribers {
+            sub(user.cloned());
+        }
+    }
+
+    pub fn subscribe_auth_changes(callback: impl Fn(Option<oauth::User>) + Send + 'static) {
+        let mut subscribers = AUTH_SUBSCRIBERS.lock();
+        subscribers.push(Box::new(callback));
+    }
+
     /// Provider which uses `OAuth` credentials stored on the user's machine.
     #[derive(Debug, Default)]
     pub struct CliCredentialsProvider {
@@ -95,22 +112,26 @@ mod oauth {
                 Ok(Some(credentials)) => {
                     // Success: cache credentials and return the token.
                     let token = credentials.access_token().jwt();
+                    auth_update(Some(credentials.user()));
                     *cache = Some(credentials);
                     Ok(Some(token))
                 }
 
                 Ok(None) => {
                     re_log::debug!("no credentials available");
 
-                    // TODO(jan): we should propagate this information to the UI
+                    auth_update(None);
 
                     // There are no credentials stored on disk, so the user has not logged in yet.
                     // We represent that by saying there is no token:
                     Ok(None)
                 }
 
                 // TODO(jan): this needs to handle the case where the refresh token expired
-                Err(err) => Err(CredentialsProviderError::Custom(err.into())),
+                Err(err) => {
+                    auth_update(None);
+                    Err(CredentialsProviderError::Custom(err.into()))
+                }
             }
         }
     }

crates/utils/re_auth/src/oauth.rs

@@ -236,6 +236,8 @@ impl InMemoryCredentials {
             });
         }
 
+        crate::credentials::oauth::auth_update(Some(&self.0.user));
+
         Ok(self.0)
     }
 }

crates/viewer/re_redap_browser/src/server_modal.rs

@@ -4,9 +4,7 @@ use re_redap_client::ConnectionRegistryHandle;
 use re_ui::UiExt as _;
 use re_ui::modal::{ModalHandler, ModalWrapper};
 use re_uri::Scheme;
-use re_viewer_context::{
-    CommandSender, DisplayMode, GlobalContext, SystemCommand, SystemCommandSender as _,
-};
+use re_viewer_context::{DisplayMode, GlobalContext, SystemCommand, SystemCommandSender as _};
 use std::str::FromStr as _;
 
 use crate::{context::Context, servers::Command};
@@ -28,7 +26,6 @@ pub enum ServerModalMode {
 
 /// Authentication state for the server modal.
 struct Authentication {
-    email: Option<String>,
     token: String,
     show_token_input: bool,
     login_flow: Option<LoginFlow>,
@@ -44,22 +41,13 @@ impl Authentication {
     ///
     /// Optionally, this can be given a token, which takes
     /// precedence over stored credentials.
-    fn new(token: Option<String>, use_stored_credentials: bool) -> Self {
-        let email = if !use_stored_credentials {
-            None
-        } else {
-            re_auth::oauth::load_credentials()
-                .ok()
-                .flatten()
-                .map(|credentials| credentials.user().email.clone())
-        };
+    fn new(token: Option<String>) -> Self {
         let (token, show_token_input) = match token {
             Some(token) => (token, true),
             None => (String::new(), false),
         };
 
         Self {
-            email,
             token,
             show_token_input,
             login_flow: None,
@@ -74,8 +62,8 @@ impl Authentication {
     }
 
     fn start_login_flow(&mut self, ui: &mut egui::Ui) {
-        let login_hint = self.email.as_deref();
-        match LoginFlow::open(ui, login_hint) {
+        // TODO: Is login hint required?
+        match LoginFlow::open(ui, None) {
             Ok(flow) => {
                 self.login_flow = Some(flow);
                 self.error = None;
@@ -104,18 +92,17 @@ impl Default for ServerModal {
             mode: ServerModalMode::Add,
             scheme: Scheme::Rerun,
             host: String::new(),
-            auth: Authentication::new(None, false),
+            auth: Authentication::new(None),
             port: 443,
         }
     }
 }
 
 impl ServerModal {
     pub fn open(&mut self, mode: ServerModalMode, connection_registry: &ConnectionRegistryHandle) {
-        let use_stored_credentials = connection_registry.should_use_stored_credentials();
         *self = match mode {
             ServerModalMode::Add => {
-                let auth = Authentication::new(None, use_stored_credentials);
+                let auth = Authentication::new(None);
 
                 Self {
                     mode: ServerModalMode::Add,
@@ -129,11 +116,9 @@ impl ServerModal {
                 let credentials = connection_registry.credentials(&origin);
                 let auth = match credentials {
                     Some(re_redap_client::Credentials::Token(token)) => {
-                        Authentication::new(Some(token.to_string()), use_stored_credentials)
-                    }
-                    Some(re_redap_client::Credentials::Stored) | None => {
-                        Authentication::new(None, use_stored_credentials)
+                        Authentication::new(Some(token.to_string()))
                     }
+                    Some(re_redap_client::Credentials::Stored) | None => Authentication::new(None),
                 };
 
                 Self {
@@ -236,7 +221,7 @@ impl ServerModal {
                 ui.label("Authenticate:");
                 ui.scope(|ui| {
                     ui.shrink_width_to_current();
-                    auth_ui(ui, global_ctx.command_sender, &mut self.auth);
+                    auth_ui(ui, global_ctx, &mut self.auth);
                 });
 
                 ui.add_space(24.0);
@@ -258,7 +243,7 @@ impl ServerModal {
                         .map(Some)
                         // error is reported in the UI above
                         .map_err(|_err| ())
-                } else if self.auth.email.is_some() {
+                } else if global_ctx.logged_in() {
                     Ok(Some(re_redap_client::Credentials::Stored))
                 } else {
                     Ok(None)
@@ -312,7 +297,7 @@ impl ServerModal {
     }
 }
 
-fn auth_ui(ui: &mut egui::Ui, cmd: &CommandSender, auth: &mut Authentication) {
+fn auth_ui(ui: &mut egui::Ui, ctx: &GlobalContext<'_>, auth: &mut Authentication) {
     ui.horizontal(|ui| {
         ui.scope(|ui| {
             if auth.show_token_input {
@@ -345,10 +330,9 @@ fn auth_ui(ui: &mut egui::Ui, cmd: &CommandSender, auth: &mut Authentication) {
                 }
             } else {
                 if let Some(flow) = &mut auth.login_flow {
-                    if let Some(result) = flow.ui(ui, cmd) {
+                    if let Some(result) = flow.ui(ui, ctx.command_sender) {
                         match result {
-                            LoginFlowResult::Success(credentials) => {
-                                auth.email = Some(credentials.user().email.clone());
+                            LoginFlowResult::Success => {
                                 auth.error = None;
                                 // Clear login flow to close popup window
                                 auth.reset_login_flow();
@@ -360,9 +344,9 @@ fn auth_ui(ui: &mut egui::Ui, cmd: &CommandSender, auth: &mut Authentication) {
                             }
                         }
                     }
-                } else if let Some(email) = &auth.email {
+                } else if let Some(logged_in) = &ctx.auth_context {
                     ui.label("Continue as ");
-                    ui.label(RichText::new(email).strong().underline());
+                    ui.label(RichText::new(&logged_in.email).strong().underline());
 
                     if ui
                         .small_icon_button(&re_ui::icons::CLOSE, "Clear login status")
@@ -401,7 +385,7 @@ fn auth_ui(ui: &mut egui::Ui, cmd: &CommandSender, auth: &mut Authentication) {
     ui.horizontal(|ui| {
         ui.set_min_width(300.0);
         ui.set_width(300.0);
-        if !auth.show_token_input && auth.email.is_none() {
+        if !auth.show_token_input && !ctx.logged_in() {
             if let Some(error) = &auth.error {
                 ui.error_label(error.clone());
             }

crates/viewer/re_redap_browser/src/server_modal/login_flow.rs

@@ -4,7 +4,6 @@ mod native;
 mod web;
 
 use egui::{IntoAtoms as _, vec2};
-use re_auth::oauth::Credentials;
 use re_ui::{
     UiExt as _,
     notifications::{Notification, NotificationLevel},
@@ -23,7 +22,7 @@ pub struct LoginFlow {
 }
 
 pub enum LoginFlowResult {
-    Success(Credentials),
+    Success,
     Failure(String),
 }
 
@@ -75,7 +74,7 @@ impl LoginFlow {
                     format!("Logged in as {}", credentials.user().email),
                 )));
 
-                Some(LoginFlowResult::Success(credentials))
+                Some(LoginFlowResult::Success)
             }
 
             Ok(None) => None,

crates/viewer/re_test_context/src/lib.rs

@@ -518,6 +518,8 @@ impl TestContext {
                 display_mode: &DisplayMode::LocalRecordings(
                     store_context.recording_store_id().clone(),
                 ),
+
+                auth_context: None,
             },
             component_ui_registry: &self.component_ui_registry,
             component_fallback_registry: &self.component_fallback_registry,
@@ -753,6 +755,7 @@ impl TestContext {
                 | SystemCommand::RedoBlueprint { .. }
                 | SystemCommand::CloseAllEntries
                 | SystemCommand::SetAuthCredentials { .. }
+                | SystemCommand::OnAuthChanged(_)
                 | SystemCommand::ShowNotification { .. } => handled = false,
 
                 #[cfg(debug_assertions)]

crates/viewer/re_viewer/src/app.rs

@@ -17,8 +17,8 @@ use re_types::blueprint::components::PlayState;
 use re_ui::egui_ext::context_ext::ContextExt as _;
 use re_ui::{ContextExt as _, UICommand, UICommandSender as _, UiExt as _, notifications};
 use re_viewer_context::{
-    AppOptions, AsyncRuntimeHandle, BlueprintUndoState, CommandReceiver, CommandSender,
-    ComponentUiRegistry, DisplayMode, FallbackProviderRegistry, Item, NeedsRepaint,
+    AppOptions, AsyncRuntimeHandle, AuthContext, BlueprintUndoState, CommandReceiver,
+    CommandSender, ComponentUiRegistry, DisplayMode, FallbackProviderRegistry, Item, NeedsRepaint,
     RecordingOrTable, StorageContext, StoreContext, SystemCommand, SystemCommandSender as _,
     TableStore, TimeControlCommand, ViewClass, ViewClassRegistry, ViewClassRegistryError,
     command_channel,
@@ -180,6 +180,15 @@ impl App {
     ) -> Self {
         re_tracing::profile_function!();
 
+        {
+            let command_sender = command_channel.0.clone();
+            re_auth::credentials::subscribe_auth_changes(move |user| {
+                command_sender.send_system(SystemCommand::OnAuthChanged(
+                    user.map(|user| AuthContext { email: user.email }),
+                ));
+            });
+        }
+
         let connection_registry = connection_registry
             .unwrap_or_else(re_redap_client::ConnectionRegistry::new_with_stored_credentials);
 
@@ -1220,6 +1229,10 @@ impl App {
                 }
             }
 
+            SystemCommand::OnAuthChanged(auth) => {
+                self.state.auth_state = auth;
+            }
+
             SystemCommand::SetAuthCredentials {
                 access_token,
                 email,

crates/viewer/re_viewer/src/app_state.rs

@@ -12,7 +12,7 @@ use re_smart_channel::ReceiveSet;
 use re_types::blueprint::components::{PanelState, PlayState};
 use re_ui::{ContextExt as _, UiExt as _};
 use re_viewer_context::{
-    AppOptions, ApplicationSelectionState, AsyncRuntimeHandle, BlueprintContext,
+    AppOptions, ApplicationSelectionState, AsyncRuntimeHandle, AuthContext, BlueprintContext,
     BlueprintUndoState, CommandSender, ComponentUiRegistry, DataQueryResult, DisplayMode,
     DragAndDropManager, FallbackProviderRegistry, GlobalContext, Item, PerVisualizerInViewClass,
     SelectionChange, StorageContext, StoreContext, StoreHub, SystemCommand,
@@ -112,6 +112,10 @@ pub struct AppState {
     /// that last several frames.
     #[serde(skip)]
     pub(crate) focused_item: Option<Item>,
+
+    /// Are we logged in?
+    #[serde(skip)]
+    pub(crate) auth_state: Option<AuthContext>,
 }
 
 impl Default for AppState {
@@ -136,6 +140,7 @@ impl Default for AppState {
             view_states: Default::default(),
             selection_state: Default::default(),
             focused_item: Default::default(),
+            auth_state: Default::default(),
 
             #[cfg(feature = "testing")]
             test_hook: None,
@@ -241,6 +246,7 @@ impl AppState {
                     view_states,
                     selection_state,
                     focused_item,
+                    auth_state,
                     ..
                 } = self;
 
@@ -368,6 +374,7 @@ impl AppState {
 
                         connection_registry,
                         display_mode,
+                        auth_context: auth_state.as_ref(),
                     },
                     component_ui_registry,
                     component_fallback_registry,
@@ -413,6 +420,7 @@ impl AppState {
 
                         connection_registry,
                         display_mode,
+                        auth_context: auth_state.as_ref(),
                     },
                     component_ui_registry,
                     component_fallback_registry,

crates/viewer/re_viewer_context/src/command_sender.rs

@@ -4,7 +4,7 @@ use re_data_source::LogDataSource;
 use re_log_types::StoreId;
 use re_ui::{UICommand, UICommandSender};
 
-use crate::{RecordingOrTable, time_control::TimeControlCommand};
+use crate::{AuthContext, RecordingOrTable, time_control::TimeControlCommand};
 
 // ----------------------------------------------------------------------------
 
@@ -139,6 +139,9 @@ pub enum SystemCommand {
     #[cfg(not(target_arch = "wasm32"))]
     FileSaver(Box<dyn FnOnce() -> anyhow::Result<std::path::PathBuf> + Send + 'static>),
 
+    /// Notify about authentication changes.
+    OnAuthChanged(Option<AuthContext>),
+
     /// Set authentication credentials from an external source.
     SetAuthCredentials {
         access_token: String,