[APP-3801] implement remote environments auth (#9331)

Author: moirahuang

app/src/lib.rs

@@ -565,12 +565,12 @@ pub fn run() -> Result<()> {
                 }
             }
             #[cfg(not(target_family = "wasm"))]
-            warp_cli::Command::Worker(warp_cli::WorkerCommand::RemoteServerProxy) => {
-                return crate::remote_server::run_proxy();
+            warp_cli::Command::Worker(warp_cli::WorkerCommand::RemoteServerProxy(args)) => {
+                return crate::remote_server::run_proxy(args.identity_key.clone());
             }
             #[cfg(not(target_family = "wasm"))]
-            warp_cli::Command::Worker(warp_cli::WorkerCommand::RemoteServerDaemon) => {
-                return crate::remote_server::run_daemon();
+            warp_cli::Command::Worker(warp_cli::WorkerCommand::RemoteServerDaemon(args)) => {
+                return crate::remote_server::run_daemon(args.identity_key.clone());
             }
             #[cfg(not(target_family = "wasm"))]
             warp_cli::Command::Worker(warp_cli::WorkerCommand::RipgrepSearch {
@@ -1253,6 +1253,8 @@ fn initialize_app(
     ctx.add_singleton_model(|_ctx| SyncedInputState::new());
 
     ctx.add_singleton_model(remote_server::manager::RemoteServerManager::new);
+    #[cfg(not(target_family = "wasm"))]
+    remote_server::wire_auth_token_rotation(ctx);
 
     log::info!(
         "Starting warp with channel state {} and version {:?}",

app/src/remote_server/auth_context.rs

@@ -0,0 +1,49 @@
+use std::sync::Arc;
+
+use remote_server::auth::RemoteServerAuthContext;
+use warpui::r#async::BoxFuture;
+
+use crate::auth::auth_state::AuthState;
+use crate::server::server_api::auth::AuthClient;
+
+/// Builds the app-wide auth context used by remote-server connections.
+pub fn server_api_auth_context(
+    auth_state: Arc<AuthState>,
+    auth_client: Arc<dyn AuthClient>,
+) -> RemoteServerAuthContext {
+    let token_auth_state = auth_state.clone();
+    let token_auth_client = auth_client;
+    let identity_auth_state = auth_state;
+
+    RemoteServerAuthContext::new(
+        move || -> BoxFuture<'static, Option<String>> {
+            if !use_authenticated_user_identity(&token_auth_state) {
+                return Box::pin(async { None });
+            }
+
+            let auth_client = token_auth_client.clone();
+            Box::pin(async move {
+                match auth_client.get_or_refresh_access_token().await {
+                    Ok(token) => token.bearer_token(),
+                    Err(_) => None,
+                }
+            })
+        },
+        move || remote_server_identity_key(&identity_auth_state),
+    )
+}
+
+fn use_authenticated_user_identity(auth_state: &AuthState) -> bool {
+    auth_state.is_logged_in() && !auth_state.is_user_anonymous().unwrap_or(true)
+}
+
+fn remote_server_identity_key(auth_state: &AuthState) -> String {
+    if use_authenticated_user_identity(auth_state) {
+        auth_state
+            .user_id()
+            .map(|uid| uid.as_string())
+            .unwrap_or_else(|| auth_state.anonymous_id())
+    } else {
+        auth_state.anonymous_id()
+    }
+}

app/src/remote_server/auth_provider.rs

@@ -1,7 +1,15 @@
+#[cfg(not(target_family = "wasm"))]
+use crate::server::server_api::{ServerApiEvent, ServerApiProvider};
+#[cfg(not(target_family = "wasm"))]
+use remote_server::manager::RemoteServerManager;
+#[cfg(not(target_family = "wasm"))]
+use warpui::SingletonEntity;
 // Re-export everything from the `remote_server` crate so existing
 // `crate::remote_server::*` imports in `app` continue to work.
 pub use remote_server::*;
 
+#[cfg(not(target_family = "wasm"))]
+pub mod auth_context;
 #[cfg(not(target_family = "wasm"))]
 pub mod server_model;
 #[cfg(not(target_family = "wasm"))]
@@ -11,23 +19,23 @@ pub mod unix;
 
 /// Run the `remote-server-proxy` subcommand.
 #[cfg(unix)]
-pub fn run_proxy() -> anyhow::Result<()> {
-    unix::run_proxy()
+pub fn run_proxy(identity_key: String) -> anyhow::Result<()> {
+    unix::run_proxy(identity_key)
 }
 
 #[cfg(not(unix))]
-pub fn run_proxy() -> anyhow::Result<()> {
+pub fn run_proxy(_identity_key: String) -> anyhow::Result<()> {
     anyhow::bail!("remote-server-proxy is not supported on this platform")
 }
 
 /// Run the `remote-server-daemon` subcommand.
 #[cfg(unix)]
-pub fn run_daemon() -> anyhow::Result<()> {
-    unix::run_daemon()
+pub fn run_daemon(identity_key: String) -> anyhow::Result<()> {
+    unix::run_daemon(identity_key)
 }
 
 #[cfg(not(unix))]
-pub fn run_daemon() -> anyhow::Result<()> {
+pub fn run_daemon(_identity_key: String) -> anyhow::Result<()> {
     anyhow::bail!("remote-server-daemon is not supported on this platform")
 }
 
@@ -79,3 +87,17 @@ pub(super) fn run_daemon_app(
     })?;
     Ok(())
 }
+
+/// Forwards app auth-token rotation events to the remote-server manager.
+#[cfg(not(target_family = "wasm"))]
+pub fn wire_auth_token_rotation(ctx: &mut warpui::AppContext) {
+    let server_api = ServerApiProvider::handle(ctx);
+    let manager = RemoteServerManager::handle(ctx);
+    ctx.subscribe_to_model(&server_api, move |_, event, ctx| {
+        if let ServerApiEvent::AccessTokenRefreshed { token } = event {
+            manager.update(ctx, |manager, _| {
+                manager.rotate_auth_token(token.clone());
+            });
+        }
+    });
+}

app/src/remote_server/mod.rs

@@ -18,12 +18,12 @@ use warp_util::file::FileId;
 
 use super::proto::{
     client_message, delete_file_response, run_command_response, server_message,
-    write_file_response, Abort, ClientMessage, DeleteFile, DeleteFileResponse, DeleteFileSuccess,
-    ErrorCode, ErrorResponse, FailedFileRead, FileContextProto, FileOperationError,
-    InitializeResponse, NavigatedToDirectory, NavigatedToDirectoryResponse,
-    ReadFileContextResponse, RunCommandError, RunCommandErrorCode, RunCommandRequest,
-    RunCommandResponse, RunCommandSuccess, ServerMessage, SessionBootstrapped, WriteFile,
-    WriteFileResponse, WriteFileSuccess,
+    write_file_response, Abort, Authenticate, ClientMessage, DeleteFile, DeleteFileResponse,
+    DeleteFileSuccess, ErrorCode, ErrorResponse, FailedFileRead, FileContextProto,
+    FileOperationError, Initialize, InitializeResponse, NavigatedToDirectory,
+    NavigatedToDirectoryResponse, ReadFileContextResponse, RunCommandError, RunCommandErrorCode,
+    RunCommandRequest, RunCommandResponse, RunCommandSuccess, ServerMessage, SessionBootstrapped,
+    WriteFile, WriteFileResponse, WriteFileSuccess,
 };
 
 /// How long the daemon waits with no connections before exiting.
@@ -164,6 +164,13 @@ pub struct ServerModel {
     executors: HashMap<SessionId, Arc<LocalCommandExecutor>>,
     /// Tracks in-flight file write/delete operations and handles cleanup.
     pending_file_ops: PendingFileOps,
+    /// Daemon-wide bearer credential for the identity-scoped daemon.
+    ///
+    /// The token is written by Initialize when the client supplies a
+    /// non-empty credential, or by Authenticate during token rotation. It is
+    /// intentionally retained across proxy connection teardown and cleared
+    /// only by daemon process exit.
+    auth_token: Option<String>,
 }
 
 impl Entity for ServerModel {
@@ -188,6 +195,7 @@ impl ServerModel {
             host_id,
             executors: HashMap::new(),
             pending_file_ops: PendingFileOps::new(),
+            auth_token: None,
         };
         // Subscribe to FileModel and RepoMetadataModel events
         // file operation results and repo metadata pushes are forwarded to all
@@ -377,7 +385,13 @@ impl ServerModel {
         let request_id = RequestId::from(msg.request_id);
 
         let outcome = match msg.message {
-            Some(client_message::Message::Initialize(_)) => self.handle_initialize(&request_id),
+            Some(client_message::Message::Initialize(msg)) => {
+                self.handle_initialize(msg, &request_id)
+            }
+            Some(client_message::Message::Authenticate(msg)) => {
+                self.handle_authenticate(msg);
+                return;
+            }
             Some(client_message::Message::SessionBootstrapped(msg)) => {
                 self.handle_session_bootstrapped(msg);
                 return;
@@ -497,8 +511,11 @@ impl ServerModel {
     }
 
     /// Handles `Initialize` by returning the server version and host id.
-    fn handle_initialize(&self, request_id: &RequestId) -> HandlerOutcome {
+    fn handle_initialize(&mut self, msg: Initialize, request_id: &RequestId) -> HandlerOutcome {
         log::info!("Handling Initialize (request_id={request_id})");
+        if !msg.auth_token.is_empty() {
+            self.auth_token = Some(msg.auth_token);
+        }
         let server_version = ChannelState::app_version()
             .unwrap_or(env!("CARGO_PKG_VERSION"))
             .to_string();
@@ -510,6 +527,20 @@ impl ServerModel {
         ))
     }
 
+    /// Handles `Authenticate` by replacing the daemon-wide credential.
+    /// This is a notification — no response is sent.
+    fn handle_authenticate(&mut self, msg: Authenticate) {
+        if msg.auth_token.is_empty() {
+            log::warn!("Received Authenticate notification with empty auth token; ignoring");
+            return;
+        }
+        self.auth_token = Some(msg.auth_token);
+    }
+
+    pub fn auth_token(&self) -> Option<&str> {
+        self.auth_token.as_deref()
+    }
+
     /// Handles `Abort` by cancelling the in-progress request it targets.
     /// This is a notification — no response is sent.
     fn handle_abort(&mut self, abort: Abort, request_id: &RequestId) {
@@ -1074,3 +1105,7 @@ fn file_context_result_to_proto(result: ReadFileContextResult) -> ReadFileContex
         failed_files,
     }
 }
+
+#[cfg(test)]
+#[path = "server_model_tests.rs"]
+mod tests;

app/src/remote_server/server_model.rs

@@ -0,0 +1,97 @@
+use std::collections::HashMap;
+
+use super::super::proto::{Authenticate, Initialize};
+use super::super::protocol::RequestId;
+use super::{PendingFileOps, ServerModel};
+
+fn test_model() -> ServerModel {
+    ServerModel {
+        connection_senders: HashMap::new(),
+        snapshot_sent_roots_by_connection: HashMap::new(),
+        grace_timer_cancel: None,
+        in_progress: HashMap::new(),
+        host_id: "test-host-id".to_string(),
+        executors: HashMap::new(),
+        pending_file_ops: PendingFileOps::new(),
+        auth_token: None,
+    }
+}
+
+fn request_id() -> RequestId {
+    RequestId::from("test-request".to_string())
+}
+
+#[test]
+fn fresh_model_starts_without_auth_token() {
+    let model = test_model();
+
+    assert_eq!(model.auth_token(), None);
+}
+
+#[test]
+fn initialize_with_auth_token_stores_token() {
+    let mut model = test_model();
+
+    model.handle_initialize(
+        Initialize {
+            auth_token: "initial-token".to_string(),
+        },
+        &request_id(),
+    );
+
+    assert_eq!(model.auth_token(), Some("initial-token"));
+}
+
+#[test]
+fn empty_initialize_preserves_existing_auth_token() {
+    let mut model = test_model();
+    model.handle_initialize(
+        Initialize {
+            auth_token: "initial-token".to_string(),
+        },
+        &request_id(),
+    );
+
+    model.handle_initialize(
+        Initialize {
+            auth_token: String::new(),
+        },
+        &request_id(),
+    );
+
+    assert_eq!(model.auth_token(), Some("initial-token"));
+}
+
+#[test]
+fn authenticate_with_auth_token_replaces_auth_token() {
+    let mut model = test_model();
+    model.handle_initialize(
+        Initialize {
+            auth_token: "initial-token".to_string(),
+        },
+        &request_id(),
+    );
+
+    model.handle_authenticate(Authenticate {
+        auth_token: "rotated-token".to_string(),
+    });
+
+    assert_eq!(model.auth_token(), Some("rotated-token"));
+}
+
+#[test]
+fn empty_authenticate_preserves_existing_auth_token() {
+    let mut model = test_model();
+    model.handle_initialize(
+        Initialize {
+            auth_token: "initial-token".to_string(),
+        },
+        &request_id(),
+    );
+
+    model.handle_authenticate(Authenticate {
+        auth_token: String::new(),
+    });
+
+    assert_eq!(model.auth_token(), Some("initial-token"));
+}