public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Moira Huang
2026-04-28 21:31:23 -05:00
committed by GitHub
parent c325d146ab
commit f0c8b7f723
21 changed files with 685 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/remote_server/src/auth.rs
+39
View File
@@ -0,0 +1,39 @@
use std::sync::Arc;
use warpui::r#async::BoxFuture;
type GetAuthTokenFn = dyn Fn() -> BoxFuture<'static, Option<String>> + Send + Sync;
type RemoteServerIdentityKeyFn = dyn Fn() -> String + Send + Sync;
/// App-supplied authentication context for transport-agnostic remote-server code.
///
/// Bearer tokens are delivered only through protocol messages. Identity keys
/// are non-secret stable partition keys used to select the remote daemon's
/// socket/PID directory.
///
/// This keeps the `remote_server` crate decoupled from app-side auth/server API
/// types while still allowing initial connect and reconnect handshakes to fetch
/// the current app credential.
#[derive(Clone)]
pub struct RemoteServerAuthContext {
get_auth_token: Arc<GetAuthTokenFn>,
remote_server_identity_key: Arc<RemoteServerIdentityKeyFn>,
}
impl RemoteServerAuthContext {
pub fn new(
get_auth_token: impl Fn() -> BoxFuture<'static, Option<String>> + Send + Sync + 'static,
remote_server_identity_key: impl Fn() -> String + Send + Sync + 'static,
) -> Self {
Self {
get_auth_token: Arc::new(get_auth_token),
remote_server_identity_key: Arc::new(remote_server_identity_key),
}
}
pub fn get_auth_token(&self) -> BoxFuture<'static, Option<String>> {
(self.get_auth_token)()
}
pub fn remote_server_identity_key(&self) -> String {
(self.remote_server_identity_key)()
}
}
crates/remote_server/src/client/mod.rs
+23 -6
View File
@@ -10,10 +10,10 @@ use futures::io::{AsyncRead, AsyncWrite};
use warpui::r#async::{executor, FutureExt as _};
use crate::proto::{
client_message, server_message, Abort, ClientMessage, DeleteFile, ErrorCode, Initialize,
InitializeResponse, LoadRepoMetadataDirectoryResponse, NavigatedToDirectoryResponse,
ReadFileContextRequest, ReadFileContextResponse, RunCommandRequest, RunCommandResponse,
ServerMessage, SessionBootstrapped, WriteFile,
client_message, server_message, Abort, Authenticate, ClientMessage, DeleteFile, ErrorCode,
Initialize, InitializeResponse, LoadRepoMetadataDirectoryResponse,
NavigatedToDirectoryResponse, ReadFileContextRequest, ReadFileContextResponse,
RunCommandRequest, RunCommandResponse, ServerMessage, SessionBootstrapped, WriteFile,
};
use crate::protocol::{self, ProtocolError, RequestId};
@@ -182,11 +182,16 @@ impl RemoteServerClient {
}
/// Sends an `Initialize` request and awaits the `InitializeResponse`.
pub async fn initialize(&self) -> Result<InitializeResponse, ClientError> {
pub async fn initialize(
&self,
auth_token: Option<&str>,
) -> Result<InitializeResponse, ClientError> {
let request_id = RequestId::new();
let msg = ClientMessage {
request_id: request_id.to_string(),
message: Some(client_message::Message::Initialize(Initialize {})),
message: Some(client_message::Message::Initialize(Initialize {
auth_token: auth_token.unwrap_or_default().to_owned(),
})),
};
let response = self.send_request(request_id, msg).await?;
@@ -200,6 +205,18 @@ impl RemoteServerClient {
}
}
/// Sends an `Authenticate` notification to rotate the daemon-wide
/// credential after initialization.
pub fn authenticate(&self, auth_token: &str) {
let msg = ClientMessage {
request_id: String::new(),
message: Some(client_message::Message::Authenticate(Authenticate {
auth_token: auth_token.to_owned(),
})),
};
self.send_notification(msg);
}
/// Sends a `SessionBootstrapped` notification (fire-and-forget) so the
/// server can create a `LocalCommandExecutor` for the session.
pub fn notify_session_bootstrapped(
crates/remote_server/src/client_tests.rs
+63 -3
View File
@@ -75,11 +75,69 @@ async fn initialize_round_trip() {
})
});
let resp = client.initialize().await.unwrap();
let resp = client.initialize(None).await.unwrap();
assert_eq!(resp.server_version, "test-0.1.0");
assert_eq!(resp.host_id, "test-host-id");
}
#[tokio::test]
async fn initialize_sends_empty_auth_token_when_none() {
let (client, _disconnect_rx, _executor) = setup_mock_client(|msg| {
match &msg.message {
Some(client_message::Message::Initialize(init)) => {
assert!(init.auth_token.is_empty());
}
other => panic!("Expected Initialize, got {other:?}"),
}
server_message::Message::InitializeResponse(InitializeResponse {
server_version: "test-0.1.0".to_string(),
host_id: "test-host-id".to_string(),
})
});
client.initialize(None).await.unwrap();
}
#[tokio::test]
async fn initialize_sends_auth_token_when_provided() {
let (client, _disconnect_rx, _executor) = setup_mock_client(|msg| {
match &msg.message {
Some(client_message::Message::Initialize(init)) => {
assert_eq!(init.auth_token, "secret-token");
}
other => panic!("Expected Initialize, got {other:?}"),
}
server_message::Message::InitializeResponse(InitializeResponse {
server_version: "test-0.1.0".to_string(),
host_id: "test-host-id".to_string(),
})
});
client.initialize(Some("secret-token")).await.unwrap();
}
#[tokio::test]
async fn authenticate_sends_fire_and_forget_message() {
let (client_stream, server_stream) = tokio::io::duplex(4096);
let (server_read, _server_write) = tokio::io::split(server_stream);
let (client_read, client_write) = tokio::io::split(client_stream);
let executor = executor::Background::default();
let (client, _event_rx) =
RemoteServerClient::new(client_read.compat(), client_write.compat_write(), &executor);
client.authenticate("rotated-secret");
let msg = protocol::read_client_message(&mut server_read.compat())
.await
.unwrap();
match msg.message {
Some(client_message::Message::Authenticate(auth)) => {
assert_eq!(auth.auth_token, "rotated-secret");
}
other => panic!("Expected Authenticate, got {other:?}"),
}
}
#[tokio::test]
async fn disconnected_on_closed_stream() {
let (client_stream, server_stream) = tokio::io::duplex(4096);
@@ -92,7 +150,7 @@ async fn disconnected_on_closed_stream() {
RemoteServerClient::new(client_read.compat(), client_write.compat_write(), &executor);
// An initialize call on a dead stream must complete with an error rather than hang.
let result = client.initialize().await;
let result = client.initialize(None).await;
assert!(result.is_err());
// The reader task should detect EOF and emit a Disconnected event.
@@ -148,7 +206,9 @@ async fn concurrent_in_flight_requests() {
for _ in 0..10 {
let c = std::sync::Arc::clone(&client);
handles.push(tokio::spawn(async move {
c.initialize().await.expect("concurrent initialize failed")
c.initialize(None)
.await
.expect("concurrent initialize failed")
}));
}
crates/remote_server/src/lib.rs
+1
View File
@@ -1,3 +1,4 @@
pub mod auth;
pub mod client;
pub mod host_id;
pub mod manager;
crates/remote_server/src/manager.rs
+98 -4
View File
@@ -5,6 +5,7 @@ use std::sync::Arc;
#[cfg(not(target_family = "wasm"))]
use std::time::Duration;
use crate::auth::RemoteServerAuthContext;
#[cfg(not(target_family = "wasm"))]
use crate::client::ClientEvent;
use crate::client::RemoteServerClient;
@@ -33,7 +34,9 @@ struct ReconnectParams {
host_id: HostId,
exit_status: Option<RemoteServerExitStatus>,
transport: Arc<dyn RemoteTransport>,
auth_context: Arc<RemoteServerAuthContext>,
control_path: Option<PathBuf>,
identity_key: String,
}
/// Error from [`RemoteServerManager::run_connect_and_handshake`] that
@@ -151,6 +154,12 @@ pub enum RemoteSessionState {
Connected {
client: Arc<RemoteServerClient>,
host_id: HostId,
/// Identity key that was active when this session was established.
/// Used by `rotate_auth_token` to ensure token rotation notifications
/// are only delivered to sessions that belong to the current user
/// identity, preventing a stale session for a previous identity from
/// receiving a different user's bearer token.
identity_key: String,
/// The transport's owning `Child`. See `Initializing::_child`.
#[cfg(not(target_family = "wasm"))]
_child: async_process::Child,
@@ -335,6 +344,10 @@ pub struct RemoteServerManager {
/// remote server daemon on every (re)connect. Persists until
/// `deregister_session`.
session_bootstrap_info: HashMap<SessionId, SessionBootstrapInfo>,
/// App auth context used for connection-time `Initialize` and future
/// reconnect handshakes.
#[cfg_attr(target_family = "wasm", allow(dead_code))]
auth_context: Option<Arc<RemoteServerAuthContext>>,
/// Detected remote platform per session, populated during the binary check
/// phase via `detect_platform()`. Used for telemetry.
session_platforms: HashMap<SessionId, RemotePlatform>,
@@ -354,6 +367,7 @@ impl RemoteServerManager {
spawner: ctx.spawner(),
last_navigated_path: HashMap::new(),
session_bootstrap_info: HashMap::new(),
auth_context: None,
session_platforms: HashMap::new(),
}
}
@@ -483,6 +497,7 @@ impl RemoteServerManager {
&mut self,
session_id: SessionId,
transport: T,
auth_context: Arc<RemoteServerAuthContext>,
ctx: &mut ModelContext<Self>,
) where
T: RemoteTransport + 'static,
@@ -504,6 +519,7 @@ impl RemoteServerManager {
self.sessions
.insert(session_id, RemoteSessionState::Connecting);
self.auth_context = Some(Arc::clone(&auth_context));
ctx.emit(RemoteServerManagerEvent::SessionConnecting { session_id });
let spawner = self.spawner.clone();
@@ -511,12 +527,17 @@ impl RemoteServerManager {
// Wrap the transport in an Arc so it can be stored on `Connected`
// for reconnection after a spontaneous disconnect.
let transport: Arc<dyn RemoteTransport> = Arc::new(transport);
let auth_context_for_task = Arc::clone(&auth_context);
// Capture the identity key synchronously so it travels with the
// session and can be used to filter token-rotation notifications.
let identity_key = auth_context.remote_server_identity_key();
ctx.background_executor()
.spawn(async move {
match Self::run_connect_and_handshake(
session_id,
&*transport,
&auth_context_for_task,
&spawner,
&executor,
)
@@ -525,7 +546,13 @@ impl RemoteServerManager {
Ok(host_id) => {
let _ = spawner
.spawn(move |me, ctx| {
me.mark_session_connected(session_id, host_id, transport, ctx);
me.mark_session_connected(
session_id,
host_id,
identity_key,
transport,
ctx,
);
})
.await;
}
@@ -562,13 +589,14 @@ impl RemoteServerManager {
/// 1. Calls `transport.connect()` to establish streams.
/// 2. Transitions the session to `Initializing` and starts draining the
/// event channel.
/// 3. Runs the initialize handshake.
/// 3. Runs the initialize handshake with the current auth token, if any.
///
/// Returns `Ok(host_id)` on success, or a phase-tagged error.
#[cfg(not(target_family = "wasm"))]
async fn run_connect_and_handshake(
session_id: SessionId,
transport: &dyn RemoteTransport,
auth_context: &RemoteServerAuthContext,
spawner: &ModelSpawner<Self>,
executor: &Arc<warpui::r#async::executor::Background>,
) -> Result<HostId, ConnectAndHandshakeError> {
@@ -624,8 +652,9 @@ impl RemoteServerManager {
}
// Phase 2: Initialize handshake.
let auth_token = auth_context.get_auth_token().await;
let resp = client
.initialize()
.initialize(auth_token.as_deref())
.await
.map_err(|e| ConnectAndHandshakeError::Initialize(anyhow::anyhow!("{e:#}")))?;
Ok(HostId::new(resp.host_id))
@@ -732,6 +761,42 @@ impl RemoteServerManager {
}
}
/// Rotates the daemon-wide auth credential on each connected remote host.
///
/// Only sessions whose stored `identity_key` matches the current identity
/// (from `auth_context`) receive the notification. This prevents a stale
/// session established under a previous user identity from receiving a
/// newly-rotated bearer token that belongs to a different user.
///
/// Within the matching identity, a daemon may have multiple client
/// connections. The credential is stored daemon-wide, so sending one
/// notification per connected host is sufficient.
pub fn rotate_auth_token(&self, token: String) {
let Some(ref auth_context) = self.auth_context else {
log::warn!("rotate_auth_token: no auth_context available, skipping");
return;
};
let current_identity_key = auth_context.remote_server_identity_key();
let mut authenticated_hosts = HashSet::new();
for state in self.sessions.values() {
let RemoteSessionState::Connected {
client,
host_id,
identity_key,
..
} = state
else {
continue;
};
if identity_key != &current_identity_key {
continue;
}
if authenticated_hosts.insert(host_id.clone()) {
client.authenticate(&token);
}
}
}
/// Returns the connection state for this session.
pub fn session(&self, session_id: SessionId) -> Option<&RemoteSessionState> {
self.sessions.get(&session_id)
@@ -955,6 +1020,7 @@ impl RemoteServerManager {
&mut self,
session_id: SessionId,
host_id: HostId,
identity_key: String,
transport: Arc<dyn RemoteTransport>,
ctx: &mut ModelContext<Self>,
) {
@@ -976,6 +1042,7 @@ impl RemoteServerManager {
RemoteSessionState::Connected {
client: client.clone(),
host_id: host_id.clone(),
identity_key,
_child,
control_path,
transport,
@@ -1067,6 +1134,7 @@ impl RemoteServerManager {
// with a transport available, and not being explicitly deregistered.
if let RemoteSessionState::Connected {
host_id,
identity_key,
mut _child,
control_path,
transport,
@@ -1076,7 +1144,24 @@ impl RemoteServerManager {
let exit_status = Self::capture_exit_status(&mut _child, session_id);
// Drop the old child process explicitly before reconnecting.
drop(_child);
let Some(auth_context) = self.auth_context.clone() else {
log::warn!(
"Spontaneous disconnect for session {session_id:?}, \
but no auth context is available for reconnect"
);
self.sessions
.insert(session_id, RemoteSessionState::Disconnected);
self.remove_from_host_index(&host_id, session_id);
ctx.emit(RemoteServerManagerEvent::SessionDisconnected {
session_id,
host_id: host_id.clone(),
exit_status,
});
if !self.host_to_sessions.contains_key(&host_id) {
ctx.emit(RemoteServerManagerEvent::HostDisconnected { host_id });
}
return;
};
log::info!(
"Spontaneous disconnect for session {session_id:?}, \
will attempt reconnect (transport={transport:?})"
@@ -1104,7 +1189,9 @@ impl RemoteServerManager {
host_id,
exit_status,
transport,
auth_context,
control_path,
identity_key,
},
ctx,
);
@@ -1129,7 +1216,9 @@ impl RemoteServerManager {
host_id,
exit_status,
transport,
auth_context,
control_path,
identity_key,
} = params;
log::info!(
@@ -1149,6 +1238,7 @@ impl RemoteServerManager {
let spawner = self.spawner.clone();
let executor = ctx.background_executor().clone();
let transport_clone = Arc::clone(&transport);
let auth_context_for_task = Arc::clone(&auth_context);
ctx.background_executor()
.spawn(async move {
@@ -1168,6 +1258,7 @@ impl RemoteServerManager {
match Self::run_connect_and_handshake(
session_id,
&*transport_clone,
&auth_context_for_task,
&spawner,
&executor,
)
@@ -1188,6 +1279,7 @@ impl RemoteServerManager {
me.mark_session_connected(
session_id,
new_host_id.clone(),
identity_key,
transport,
ctx,
);
@@ -1225,7 +1317,9 @@ impl RemoteServerManager {
host_id,
exit_status,
transport,
auth_context,
control_path,
identity_key,
},
ctx,
);
crates/remote_server/src/protocol_tests.rs
+6 -2
View File
@@ -10,7 +10,9 @@ use super::*;
async fn round_trip_client_message() {
let msg = ClientMessage {
request_id: "test-123".to_string(),
message: Some(client_message::Message::Initialize(Initialize {})),
message: Some(client_message::Message::Initialize(Initialize {
auth_token: String::new(),
})),
};
let mut buf = Vec::new();
@@ -119,7 +121,9 @@ async fn write_message_too_large() {
fn try_extract_request_id_from_valid_message() {
let msg = ClientMessage {
request_id: "abc-123".to_string(),
message: Some(client_message::Message::Initialize(Initialize {})),
message: Some(client_message::Message::Initialize(Initialize {
auth_token: String::new(),
})),
};
let buf = msg.encode_to_vec();
assert_eq!(try_extract_request_id(&buf), Some("abc-123".to_string()));
crates/remote_server/src/setup.rs
+32
View File
@@ -137,6 +137,38 @@ pub fn remote_server_dir() -> String {
format!("~/{warp_dir}/remote-server")
}
/// Returns a filesystem-safe directory name for a remote-server identity key.
///
/// The identity key is not secret, but it can contain bytes that are unsafe or
/// ambiguous in paths. Keep ASCII alphanumeric characters plus `-` and `_`;
/// percent-encode all other UTF-8 bytes.
pub fn remote_server_identity_dir_name(identity_key: &str) -> String {
if identity_key.is_empty() {
return "empty".to_string();
}
let mut encoded = String::with_capacity(identity_key.len());
for byte in identity_key.bytes() {
match byte {
b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'-' | b'_' => {
encoded.push(byte as char);
}
_ => encoded.push_str(&format!("%{byte:02X}")),
}
}
encoded
}
/// Returns the identity-scoped remote directory used for the daemon socket
/// and PID file.
pub fn remote_server_daemon_dir(identity_key: &str) -> String {
format!(
"{}/{}",
remote_server_dir(),
remote_server_identity_dir_name(identity_key)
)
}
/// Returns the binary name, keyed by channel.
///
/// Matches the CLI command names: `oz` (stable), `oz-preview`, `oz-dev`.