mirror of
https://github.com/openai/codex.git
synced 2026-09-09 15:58:47 +00:00
fix: require remote control token bundle
Co-authored-by: Codex noreply@openai.com
This commit is contained in:
@@ -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!(
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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()),
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user