fix: require remote control token bundle

Co-authored-by: Codex noreply@openai.com
This commit is contained in:
viyatb-oai
2026-04-20 20:03:11 -07:00
parent 6539a56c15
commit d7390cbf52
4 changed files with 184 additions and 233 deletions

View File

@@ -48,7 +48,7 @@ impl RemoteControlServerToken {
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct RemoteControlEnrollmentResult {
pub(super) enrollment: RemoteControlEnrollment,
pub(super) server_token: Option<RemoteControlServerToken>,
pub(super) server_token: RemoteControlServerToken,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -286,23 +286,9 @@ pub(super) async fn enroll_remote_control_server(
fn remote_control_server_token_from_response(
enrollment: &EnrollRemoteServerResponse,
) -> io::Result<Option<RemoteControlServerToken>> {
let Some(remote_control_token) = enrollment.remote_control_token.as_ref() else {
return Ok(None);
};
let expires_at = enrollment.expires_at.as_ref().ok_or_else(|| {
io::Error::new(
ErrorKind::InvalidData,
"remote control enrollment response included a token without expires_at",
)
})?;
let scopes = enrollment.scopes.as_ref().ok_or_else(|| {
io::Error::new(
ErrorKind::InvalidData,
"remote control enrollment response included a token without scopes",
)
})?;
if !scopes
) -> io::Result<RemoteControlServerToken> {
if !enrollment
.scopes
.iter()
.any(|scope| scope == REMOTE_CONTROL_SERVER_WEBSOCKET_SCOPE)
{
@@ -311,7 +297,7 @@ fn remote_control_server_token_from_response(
"remote control enrollment response token is missing server websocket scope",
));
}
let expires_at = DateTime::parse_from_rfc3339(expires_at)
let expires_at = DateTime::parse_from_rfc3339(&enrollment.expires_at)
.map_err(|err| {
io::Error::new(
ErrorKind::InvalidData,
@@ -320,10 +306,10 @@ fn remote_control_server_token_from_response(
})?
.with_timezone(&Utc);
Ok(Some(RemoteControlServerToken {
bearer_token: remote_control_token.clone(),
Ok(RemoteControlServerToken {
bearer_token: enrollment.remote_control_token.clone(),
expires_at,
}))
})
}
#[cfg(test)]
@@ -497,32 +483,16 @@ mod tests {
);
}
#[test]
fn remote_control_server_token_from_response_accepts_legacy_response_without_token() {
assert_eq!(
remote_control_server_token_from_response(&EnrollRemoteServerResponse {
server_id: "srv_e_test".to_string(),
environment_id: "env_test".to_string(),
remote_control_token: None,
expires_at: None,
scopes: None,
})
.expect("legacy response should parse"),
None
);
}
#[test]
fn remote_control_server_token_from_response_parses_scoped_token() {
let server_token = remote_control_server_token_from_response(&EnrollRemoteServerResponse {
server_id: "srv_e_test".to_string(),
environment_id: "env_test".to_string(),
remote_control_token: Some("remote-control-token".to_string()),
expires_at: Some("2026-04-09T12:00:00Z".to_string()),
scopes: Some(vec!["remote_control_server_websocket".to_string()]),
remote_control_token: "remote-control-token".to_string(),
expires_at: "2026-04-09T12:00:00Z".to_string(),
scopes: vec!["remote_control_server_websocket".to_string()],
})
.expect("token response should parse")
.expect("token should be present");
.expect("token response should parse");
assert_eq!(server_token.bearer_token, "remote-control-token");
assert!(

View File

@@ -29,12 +29,9 @@ pub(super) struct EnrollRemoteServerRequest {
pub(super) struct EnrollRemoteServerResponse {
pub(super) server_id: String,
pub(super) environment_id: String,
#[serde(default)]
pub(super) remote_control_token: Option<String>,
#[serde(default)]
pub(super) expires_at: Option<String>,
#[serde(default)]
pub(super) scopes: Option<Vec<String>>,
pub(super) remote_control_token: String,
pub(super) expires_at: String,
pub(super) scopes: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]

View File

@@ -114,6 +114,20 @@ fn remote_control_url_for_listener(listener: &TcpListener) -> String {
format!("http://{addr}/backend-api/")
}
fn remote_control_enroll_response(
server_id: impl Into<String>,
environment_id: impl Into<String>,
remote_control_token: impl Into<String>,
) -> serde_json::Value {
json!({
"server_id": server_id.into(),
"environment_id": environment_id.into(),
"remote_control_token": remote_control_token.into(),
"expires_at": (chrono::Utc::now() + chrono::Duration::minutes(10)).to_rfc3339(),
"scopes": ["remote_control_server_websocket"],
})
}
#[tokio::test]
async fn remote_control_transport_manages_virtual_clients_and_routes_messages() {
let listener = TcpListener::bind("127.0.0.1:0")
@@ -142,7 +156,7 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages()
);
respond_with_json(
enroll_request.stream,
json!({ "server_id": "srv_e_test", "environment_id": "env_test" }),
remote_control_enroll_response("srv_e_test", "env_test", "remote-control-token"),
)
.await;
let mut websocket = accept_remote_control_connection(&listener).await;
@@ -408,7 +422,7 @@ async fn remote_control_transport_reconnects_after_disconnect() {
);
respond_with_json(
enroll_request.stream,
json!({ "server_id": "srv_e_test", "environment_id": "env_test" }),
remote_control_enroll_response("srv_e_test", "env_test", "remote-control-token"),
)
.await;
let mut first_websocket = accept_remote_control_connection(&listener).await;
@@ -547,7 +561,7 @@ async fn remote_control_handle_set_enabled_stops_and_restarts_connections() {
);
respond_with_json(
enroll_request.stream,
json!({ "server_id": "srv_e_test", "environment_id": "env_test" }),
remote_control_enroll_response("srv_e_test", "env_test", "remote-control-token"),
)
.await;
let mut first_websocket = accept_remote_control_connection(&listener).await;
@@ -596,7 +610,7 @@ async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() {
let enroll_request = accept_http_request(&listener).await;
respond_with_json(
enroll_request.stream,
json!({ "server_id": "srv_e_test", "environment_id": "env_test" }),
remote_control_enroll_response("srv_e_test", "env_test", "remote-control-token"),
)
.await;
let mut first_websocket = accept_remote_control_connection(&listener).await;
@@ -786,7 +800,7 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
);
respond_with_json(
enroll_request.stream,
json!({ "server_id": "srv_e_test", "environment_id": "env_test" }),
remote_control_enroll_response("srv_e_test", "env_test", "remote-control-token"),
)
.await;
@@ -798,7 +812,7 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
);
assert_eq!(
handshake_request.headers.get("authorization"),
Some(&"Bearer Access Token".to_string())
Some(&"Bearer remote-control-token".to_string())
);
assert_eq!(
handshake_request
@@ -1044,7 +1058,7 @@ async fn remote_control_renews_server_token_with_existing_enrollment_ids() {
}
#[tokio::test]
async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling() {
async fn remote_control_http_mode_renews_token_for_persisted_enrollment_before_connecting() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
@@ -1053,6 +1067,7 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling
let state_db = remote_control_state_runtime(&codex_home).await;
let remote_control_target =
normalize_remote_control_url(&remote_control_url).expect("target should parse");
let expected_server_name = gethostname().to_string_lossy().trim().to_string();
let persisted_enrollment = RemoteControlEnrollment {
account_id: "account_id".to_string(),
environment_id: "env_persisted".to_string(),
@@ -1084,6 +1099,33 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling
.await
.expect("remote control should start");
let renew_request = accept_http_request(&listener).await;
assert_eq!(
renew_request.request_line,
"POST /backend-api/wham/remote/control/server/enroll HTTP/1.1"
);
assert_eq!(
serde_json::from_str::<serde_json::Value>(&renew_request.body)
.expect("renew enroll body should deserialize"),
json!({
"name": expected_server_name,
"os": std::env::consts::OS,
"arch": std::env::consts::ARCH,
"app_server_version": env!("CARGO_PKG_VERSION"),
"server_id": persisted_enrollment.server_id.clone(),
"environment_id": persisted_enrollment.environment_id.clone(),
})
);
respond_with_json(
renew_request.stream,
remote_control_enroll_response(
persisted_enrollment.server_id.clone(),
persisted_enrollment.environment_id.clone(),
"remote-control-token",
),
)
.await;
let (handshake_request, _websocket) = accept_remote_control_backend_connection(&listener).await;
assert_eq!(
handshake_request.path,
@@ -1093,6 +1135,10 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling
handshake_request.headers.get("x-codex-server-id"),
Some(&persisted_enrollment.server_id)
);
assert_eq!(
handshake_request.headers.get("authorization"),
Some(&"Bearer remote-control-token".to_string())
);
assert_eq!(
load_persisted_remote_control_enrollment(
Some(state_db.as_ref()),
@@ -1101,7 +1147,10 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling
/*app_server_client_name*/ None,
)
.await,
Some(persisted_enrollment)
Some(RemoteControlEnrollment {
server_name: expected_server_name,
..persisted_enrollment
})
);
shutdown_token.cancel();
@@ -1156,11 +1205,30 @@ async fn remote_control_stdio_mode_waits_for_client_name_before_connecting() {
.expect_err("remote control should wait for the stdio client name");
let _ = app_server_client_name_tx.send(app_server_client_name.to_string());
let renew_request = accept_http_request(&listener).await;
assert_eq!(
renew_request.request_line,
"POST /backend-api/wham/remote/control/server/enroll HTTP/1.1"
);
respond_with_json(
renew_request.stream,
remote_control_enroll_response(
persisted_enrollment.server_id.clone(),
persisted_enrollment.environment_id.clone(),
"remote-control-token",
),
)
.await;
let (handshake_request, _websocket) = accept_remote_control_backend_connection(&listener).await;
assert_eq!(
handshake_request.headers.get("x-codex-server-id"),
Some(&persisted_enrollment.server_id)
);
assert_eq!(
handshake_request.headers.get("authorization"),
Some(&"Bearer remote-control-token".to_string())
);
shutdown_token.cancel();
let _ = remote_task.await;
@@ -1226,10 +1294,11 @@ async fn remote_control_waits_for_account_id_before_enrolling() {
);
respond_with_json(
enroll_request.stream,
json!({
"server_id": expected_enrollment.server_id,
"environment_id": expected_enrollment.environment_id,
}),
remote_control_enroll_response(
expected_enrollment.server_id.clone(),
expected_enrollment.environment_id.clone(),
"remote-control-token",
),
)
.await;
@@ -1264,7 +1333,7 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404()
account_id: "account_id".to_string(),
environment_id: "env_refreshed".to_string(),
server_id: "srv_e_refreshed".to_string(),
server_name: expected_server_name,
server_name: expected_server_name.clone(),
};
update_persisted_remote_control_enrollment(
Some(state_db.as_ref()),
@@ -1291,6 +1360,33 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404()
.await
.expect("remote control should start");
let stale_token_request = accept_http_request(&listener).await;
assert_eq!(
stale_token_request.request_line,
"POST /backend-api/wham/remote/control/server/enroll HTTP/1.1"
);
assert_eq!(
serde_json::from_str::<serde_json::Value>(&stale_token_request.body)
.expect("stale token enroll body should deserialize"),
json!({
"name": expected_server_name,
"os": std::env::consts::OS,
"arch": std::env::consts::ARCH,
"app_server_version": env!("CARGO_PKG_VERSION"),
"server_id": stale_enrollment.server_id.clone(),
"environment_id": stale_enrollment.environment_id.clone(),
})
);
respond_with_json(
stale_token_request.stream,
remote_control_enroll_response(
stale_enrollment.server_id.clone(),
stale_enrollment.environment_id.clone(),
"stale-remote-control-token",
),
)
.await;
let websocket_request = accept_http_request(&listener).await;
assert_eq!(
websocket_request.request_line,
@@ -1300,6 +1396,10 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404()
websocket_request.headers.get("x-codex-server-id"),
Some(&stale_enrollment.server_id)
);
assert_eq!(
websocket_request.headers.get("authorization"),
Some(&"Bearer stale-remote-control-token".to_string())
);
respond_with_status(websocket_request.stream, "404 Not Found", "").await;
let enroll_request = accept_http_request(&listener).await;
@@ -1307,12 +1407,23 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404()
enroll_request.request_line,
"POST /backend-api/wham/remote/control/server/enroll HTTP/1.1"
);
assert_eq!(
serde_json::from_str::<serde_json::Value>(&enroll_request.body)
.expect("refreshed enroll body should deserialize"),
json!({
"name": refreshed_enrollment.server_name.clone(),
"os": std::env::consts::OS,
"arch": std::env::consts::ARCH,
"app_server_version": env!("CARGO_PKG_VERSION"),
})
);
respond_with_json(
enroll_request.stream,
json!({
"server_id": refreshed_enrollment.server_id,
"environment_id": refreshed_enrollment.environment_id,
}),
remote_control_enroll_response(
refreshed_enrollment.server_id.clone(),
refreshed_enrollment.environment_id.clone(),
"refreshed-remote-control-token",
),
)
.await;
@@ -1321,6 +1432,10 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404()
handshake_request.headers.get("x-codex-server-id"),
Some(&refreshed_enrollment.server_id)
);
assert_eq!(
handshake_request.headers.get("authorization"),
Some(&"Bearer refreshed-remote-control-token".to_string())
);
assert_eq!(
load_persisted_remote_control_enrollment(
Some(state_db.as_ref()),

View File

@@ -690,7 +690,7 @@ fn build_remote_control_websocket_request(
websocket_url: &str,
enrollment: &RemoteControlEnrollment,
auth: &RemoteControlConnectionAuth,
server_token: Option<&RemoteControlServerToken>,
server_token: &RemoteControlServerToken,
subscribe_cursor: Option<&str>,
) -> io::Result<tungstenite::http::Request<()>> {
let mut request = websocket_url.into_client_request().map_err(|err| {
@@ -711,9 +711,7 @@ fn build_remote_control_websocket_request(
"x-codex-protocol-version",
REMOTE_CONTROL_PROTOCOL_VERSION,
)?;
let authorization_header_value = server_token
.map(|server_token| format!("Bearer {}", server_token.bearer_token))
.unwrap_or_else(|| auth.authorization_header_value.clone());
let authorization_header_value = format!("Bearer {}", server_token.bearer_token);
set_remote_control_header(headers, "authorization", &authorization_header_value)?;
set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?;
if auth.is_fedramp_account {
@@ -867,11 +865,12 @@ async fn connect_remote_control_websocket_with_options(
.await;
}
let should_refresh_server_token = enrollment_state.server_token_refresh_required
|| enrollment_state
.server_token
.as_ref()
.is_some_and(|server_token| server_token.expires_soon(Utc::now()));
let should_refresh_server_token = enrollment_state.enrollment.is_some()
&& (enrollment_state.server_token_refresh_required
|| enrollment_state
.server_token
.as_ref()
.is_none_or(|server_token| server_token.expires_soon(Utc::now())));
if should_refresh_server_token {
info!(
"remote control server token needs refresh; renewing enrollment token: websocket_url={}, account_id={}",
@@ -935,18 +934,21 @@ async fn connect_remote_control_websocket_with_options(
new_enrollment.environment_id
);
enrollment_state.enrollment = Some(new_enrollment);
enrollment_state.server_token = enrollment_result.server_token;
enrollment_state.server_token = Some(enrollment_result.server_token);
enrollment_state.server_token_refresh_required = false;
}
let enrollment_ref = enrollment_state.enrollment.as_ref().ok_or_else(|| {
io::Error::other("missing remote control enrollment after enrollment step")
})?;
let server_token = enrollment_state.server_token.as_ref().ok_or_else(|| {
io::Error::other("missing remote control server token after enrollment step")
})?;
let request = build_remote_control_websocket_request(
&remote_control_target.websocket_url,
enrollment_ref,
&auth,
enrollment_state.server_token.as_ref(),
server_token,
subscribe_cursor,
)?;
@@ -982,43 +984,16 @@ async fn connect_remote_control_websocket_with_options(
tungstenite::Error::Http(response)
if matches!(response.status().as_u16(), 401 | 403) =>
{
if enrollment_state.server_token.is_some() {
info!(
"remote control websocket token auth failed with HTTP {}; renewing token before reconnecting",
response.status()
);
enrollment_state.server_token = None;
enrollment_state.server_token_refresh_required = true;
return Err(io::Error::other(format!(
"remote control websocket token auth failed with HTTP {}; re-enrolling",
response.status()
)));
} else if recover_remote_control_auth(auth_recovery).await {
return Err(io::Error::other(format!(
"remote control websocket auth failed with HTTP {}; retrying after auth recovery",
response.status()
)));
} else {
info!(
"remote control websocket rejected legacy account auth with HTTP {}; clearing enrollment before re-enrolling",
response.status()
);
if let Err(clear_err) = update_persisted_remote_control_enrollment(
state_db,
remote_control_target,
&auth.account_id,
app_server_client_name,
/*enrollment*/ None,
)
.await
{
warn!(
"failed to clear legacy remote control enrollment after auth failure: {clear_err}"
);
}
enrollment_state.enrollment = None;
enrollment_state.server_token_refresh_required = false;
}
info!(
"remote control websocket token auth failed with HTTP {}; renewing token before reconnecting",
response.status()
);
enrollment_state.server_token = None;
enrollment_state.server_token_refresh_required = true;
return Err(io::Error::other(format!(
"remote control websocket token auth failed with HTTP {}; re-enrolling",
response.status()
)));
}
_ => {}
}
@@ -1170,8 +1145,15 @@ mod tests {
}
}
fn remote_control_server_token(bearer_token: &str) -> RemoteControlServerToken {
RemoteControlServerToken {
bearer_token: bearer_token.to_string(),
expires_at: Utc::now() + chrono::Duration::minutes(10),
}
}
#[test]
fn build_remote_control_websocket_request_uses_server_token_when_available() {
fn build_remote_control_websocket_request_uses_server_token() {
let enrollment = RemoteControlEnrollment {
account_id: "account_id".to_string(),
environment_id: "env_test".to_string(),
@@ -1183,16 +1165,13 @@ mod tests {
account_id: "account_id".to_string(),
is_fedramp_account: false,
};
let server_token = RemoteControlServerToken {
bearer_token: "remote-control-token".to_string(),
expires_at: Utc::now() + chrono::Duration::minutes(10),
};
let server_token = remote_control_server_token("remote-control-token");
let request = build_remote_control_websocket_request(
"ws://localhost:8080/backend-api/wham/remote/control/server",
&enrollment,
&auth,
Some(&server_token),
&server_token,
/*subscribe_cursor*/ None,
)
.expect("websocket request should build");
@@ -1213,40 +1192,9 @@ mod tests {
);
}
#[test]
fn build_remote_control_websocket_request_falls_back_to_account_token() {
let enrollment = RemoteControlEnrollment {
account_id: "account_id".to_string(),
environment_id: "env_test".to_string(),
server_id: "srv_e_test".to_string(),
server_name: "test-server".to_string(),
};
let auth = RemoteControlConnectionAuth {
authorization_header_value: "AgentAssertion assertion".to_string(),
account_id: "account_id".to_string(),
is_fedramp_account: false,
};
let request = build_remote_control_websocket_request(
"ws://localhost:8080/backend-api/wham/remote/control/server",
&enrollment,
&auth,
/*server_token*/ None,
/*subscribe_cursor*/ None,
)
.expect("websocket request should build");
assert_eq!(
request
.headers()
.get("authorization")
.expect("authorization header should exist"),
"AgentAssertion assertion"
);
}
#[test]
fn build_remote_control_websocket_request_includes_fedramp_header() {
let server_token = remote_control_server_token("remote-control-token");
let request = build_remote_control_websocket_request(
"ws://127.0.0.1/backend-api/wham/remote/control/server",
&RemoteControlEnrollment {
@@ -1260,7 +1208,7 @@ mod tests {
account_id: "account_id".to_string(),
is_fedramp_account: true,
},
/*server_token*/ None,
&server_token,
/*subscribe_cursor*/ None,
)
.expect("request should build");
@@ -1302,10 +1250,7 @@ mod tests {
};
let mut enrollment_state = RemoteControlEnrollmentState {
enrollment: Some(enrollment.clone()),
server_token: Some(RemoteControlServerToken {
bearer_token: "remote-control-token".to_string(),
expires_at: Utc::now() + chrono::Duration::minutes(10),
}),
server_token: Some(remote_control_server_token("remote-control-token")),
..Default::default()
};
@@ -1368,7 +1313,7 @@ mod tests {
server_id: "srv_e_test".to_string(),
server_name: "test-server".to_string(),
}),
server_token: None,
server_token: Some(remote_control_server_token("remote-control-token")),
..Default::default()
};
@@ -1391,82 +1336,6 @@ mod tests {
assert_eq!(err.to_string(), expected_error);
}
#[tokio::test]
async fn connect_remote_control_websocket_recovers_after_unauthorized_reload() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let remote_control_url = remote_control_url_for_listener(&listener);
let remote_control_target =
normalize_remote_control_url(&remote_control_url).expect("target should parse");
let codex_home = TempDir::new().expect("temp dir should create");
save_auth(
codex_home.path(),
&remote_control_auth_dot_json("stale-token"),
AuthCredentialsStoreMode::File,
)
.expect("stale auth should save");
let state_db = remote_control_state_runtime(&codex_home).await;
let auth_manager = AuthManager::shared(
codex_home.path().to_path_buf(),
/*enable_codex_api_key_env*/ false,
AuthCredentialsStoreMode::File,
);
let mut auth_recovery = auth_manager.unauthorized_recovery();
let mut enrollment_state = RemoteControlEnrollmentState {
enrollment: Some(RemoteControlEnrollment {
account_id: "account_id".to_string(),
environment_id: "env_test".to_string(),
server_id: "srv_e_test".to_string(),
server_name: "test-server".to_string(),
}),
server_token: None,
..Default::default()
};
save_auth(
codex_home.path(),
&remote_control_auth_dot_json("fresh-token"),
AuthCredentialsStoreMode::File,
)
.expect("fresh auth should save");
let server_task = tokio::spawn(async move {
let (stream, request_line) = accept_http_request(&listener).await;
assert_eq!(
request_line,
"GET /backend-api/wham/remote/control/server HTTP/1.1"
);
respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await;
});
let err = connect_remote_control_websocket(
&remote_control_target,
Some(state_db.as_ref()),
&auth_manager,
&mut auth_recovery,
&mut enrollment_state,
/*subscribe_cursor*/ None,
/*app_server_client_name*/ None,
)
.await
.expect_err("unauthorized response should fail the websocket connect");
server_task.await.expect("server task should succeed");
assert_eq!(
err.to_string(),
"remote control websocket auth failed with HTTP 401 Unauthorized; retrying after auth recovery"
);
assert_eq!(
auth_manager
.auth()
.await
.expect("auth should remain available")
.get_token()
.expect("token should be readable"),
"fresh-token"
);
}
#[tokio::test]
async fn connect_remote_control_websocket_recovers_after_unauthorized_enrollment() {
let listener = TcpListener::bind("127.0.0.1:0")