mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
app-server: move transport into dedicated crate (#20545)
## Why `codex-app-server` currently owns both request-processing code and transport implementation details. Splitting the transport layer into its own crate makes that boundary explicit, reduces the amount of transport-specific dependency surface carried by `codex-app-server`, and gives future transport work a narrower place to evolve. ## What changed - Added `codex-app-server-transport` and moved the existing transport tree into it, including stdio, unix socket, websocket, remote-control transport, and websocket auth. - Moved shared transport-facing message types into the new crate so both the transport implementation and `codex-app-server` use the same definitions. - Kept processor-facing connection state and outbound routing in `codex-app-server`, with the routing tests moved next to that local wrapper. - Updated workspace metadata, Bazel crate metadata, and `codex-app-server` dependencies for the new crate boundary. ## Validation - `cargo metadata --locked --no-deps` - `git diff --check` - Attempted `cargo test -p codex-app-server-transport`, `cargo test -p codex-app-server`, `just fix -p codex-app-server-transport`, and `just fix -p codex-app-server`; all were blocked before compilation by the existing `packageproxy` resolution failure for locked `rustls-webpki = 0.103.13`. - Attempted Bazel build / lockfile validation; those were blocked by external fetch failures against BuildBuddy / GitHub while resolving `v8`.
This commit is contained in:
Ruslan Nigmatullin
committed by
GitHub
Unverified
GitHub
parent
5744b85b9a
commit
41e171fcf2
@@ -1,5 +1,4 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicI64;
|
||||
use std::sync::atomic::Ordering;
|
||||
@@ -15,7 +14,6 @@ use codex_app_server_protocol::ServerRequestPayload;
|
||||
use codex_otel::span_w3c_trace_context;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use serde::Serialize;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
@@ -26,22 +24,17 @@ use tracing::warn;
|
||||
use crate::error_code::INTERNAL_ERROR_CODE;
|
||||
use crate::error_code::internal_error;
|
||||
use crate::server_request_error::TURN_TRANSITION_PENDING_REQUEST_ERROR_REASON;
|
||||
pub(crate) use codex_app_server_transport::ConnectionId;
|
||||
pub(crate) use codex_app_server_transport::OutgoingError;
|
||||
pub(crate) use codex_app_server_transport::OutgoingMessage;
|
||||
pub(crate) use codex_app_server_transport::OutgoingResponse;
|
||||
pub(crate) use codex_app_server_transport::QueuedOutgoingMessage;
|
||||
|
||||
#[cfg(test)]
|
||||
use codex_protocol::account::PlanType;
|
||||
|
||||
pub(crate) type ClientRequestResult = std::result::Result<Result, JSONRPCErrorError>;
|
||||
|
||||
/// Stable identifier for a transport connection.
|
||||
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
|
||||
pub(crate) struct ConnectionId(pub(crate) u64);
|
||||
|
||||
impl fmt::Display for ConnectionId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}", self.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Stable identifier for a client request scoped to a transport connection.
|
||||
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
|
||||
pub(crate) struct ConnectionRequestId {
|
||||
@@ -96,21 +89,6 @@ pub(crate) enum OutgoingEnvelope {
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct QueuedOutgoingMessage {
|
||||
pub(crate) message: OutgoingMessage,
|
||||
pub(crate) write_complete_tx: Option<oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
impl QueuedOutgoingMessage {
|
||||
pub(crate) fn new(message: OutgoingMessage) -> Self {
|
||||
Self {
|
||||
message,
|
||||
write_complete_tx: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends messages to the client and manages request callbacks.
|
||||
pub(crate) struct OutgoingMessageSender {
|
||||
next_server_request_id: AtomicI64,
|
||||
@@ -665,30 +643,6 @@ impl OutgoingMessageSender {
|
||||
}
|
||||
}
|
||||
|
||||
/// Outgoing message from the server to the client.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(untagged)]
|
||||
pub(crate) enum OutgoingMessage {
|
||||
Request(ServerRequest),
|
||||
/// AppServerNotification is specific to the case where this is run as an
|
||||
/// "app server" as opposed to an MCP server.
|
||||
AppServerNotification(ServerNotification),
|
||||
Response(OutgoingResponse),
|
||||
Error(OutgoingError),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub(crate) struct OutgoingResponse {
|
||||
pub id: RequestId,
|
||||
pub result: Result,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
pub(crate) struct OutgoingError {
|
||||
pub error: JSONRPCErrorError,
|
||||
pub id: RequestId,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
use crate::message_processor::ConnectionSessionState;
|
||||
use crate::outgoing_message::OutgoingEnvelope;
|
||||
use codex_app_server_protocol::ExperimentalApi;
|
||||
use codex_app_server_protocol::ServerRequest;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::sync::RwLock;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::Ordering;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
pub use codex_app_server_transport::AppServerTransport;
|
||||
pub(crate) use codex_app_server_transport::CHANNEL_CAPACITY;
|
||||
pub(crate) use codex_app_server_transport::ConnectionId;
|
||||
pub(crate) use codex_app_server_transport::ConnectionOrigin;
|
||||
pub(crate) use codex_app_server_transport::OutgoingMessage;
|
||||
pub(crate) use codex_app_server_transport::QueuedOutgoingMessage;
|
||||
pub(crate) use codex_app_server_transport::RemoteControlHandle;
|
||||
pub(crate) use codex_app_server_transport::TransportEvent;
|
||||
pub use codex_app_server_transport::app_server_control_socket_path;
|
||||
pub use codex_app_server_transport::auth;
|
||||
pub(crate) use codex_app_server_transport::start_control_socket_acceptor;
|
||||
pub(crate) use codex_app_server_transport::start_remote_control;
|
||||
pub(crate) use codex_app_server_transport::start_stdio_connection;
|
||||
pub(crate) use codex_app_server_transport::start_websocket_acceptor;
|
||||
|
||||
pub(crate) struct ConnectionState {
|
||||
pub(crate) outbound_initialized: Arc<AtomicBool>,
|
||||
pub(crate) outbound_experimental_api_enabled: Arc<AtomicBool>,
|
||||
pub(crate) outbound_opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
|
||||
pub(crate) session: Arc<ConnectionSessionState>,
|
||||
}
|
||||
|
||||
impl ConnectionState {
|
||||
pub(crate) fn new(
|
||||
origin: ConnectionOrigin,
|
||||
outbound_initialized: Arc<AtomicBool>,
|
||||
outbound_experimental_api_enabled: Arc<AtomicBool>,
|
||||
outbound_opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
outbound_initialized,
|
||||
outbound_experimental_api_enabled,
|
||||
outbound_opted_out_notification_methods,
|
||||
session: Arc::new(ConnectionSessionState::new(origin)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct OutboundConnectionState {
|
||||
pub(crate) initialized: Arc<AtomicBool>,
|
||||
pub(crate) experimental_api_enabled: Arc<AtomicBool>,
|
||||
pub(crate) opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
|
||||
pub(crate) writer: mpsc::Sender<QueuedOutgoingMessage>,
|
||||
disconnect_sender: Option<CancellationToken>,
|
||||
}
|
||||
|
||||
impl OutboundConnectionState {
|
||||
pub(crate) fn new(
|
||||
writer: mpsc::Sender<QueuedOutgoingMessage>,
|
||||
initialized: Arc<AtomicBool>,
|
||||
experimental_api_enabled: Arc<AtomicBool>,
|
||||
opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
|
||||
disconnect_sender: Option<CancellationToken>,
|
||||
) -> Self {
|
||||
Self {
|
||||
initialized,
|
||||
experimental_api_enabled,
|
||||
opted_out_notification_methods,
|
||||
writer,
|
||||
disconnect_sender,
|
||||
}
|
||||
}
|
||||
|
||||
fn can_disconnect(&self) -> bool {
|
||||
self.disconnect_sender.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn request_disconnect(&self) {
|
||||
if let Some(disconnect_sender) = &self.disconnect_sender {
|
||||
disconnect_sender.cancel();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn should_skip_notification_for_connection(
|
||||
connection_state: &OutboundConnectionState,
|
||||
message: &OutgoingMessage,
|
||||
) -> bool {
|
||||
let Ok(opted_out_notification_methods) = connection_state.opted_out_notification_methods.read()
|
||||
else {
|
||||
warn!("failed to read outbound opted-out notifications");
|
||||
return false;
|
||||
};
|
||||
match message {
|
||||
OutgoingMessage::AppServerNotification(notification) => {
|
||||
if notification.experimental_reason().is_some()
|
||||
&& !connection_state
|
||||
.experimental_api_enabled
|
||||
.load(Ordering::Acquire)
|
||||
{
|
||||
return true;
|
||||
}
|
||||
let method = notification.to_string();
|
||||
opted_out_notification_methods.contains(method.as_str())
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn disconnect_connection(
|
||||
connections: &mut HashMap<ConnectionId, OutboundConnectionState>,
|
||||
connection_id: ConnectionId,
|
||||
) -> bool {
|
||||
if let Some(connection_state) = connections.remove(&connection_id) {
|
||||
connection_state.request_disconnect();
|
||||
return true;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
async fn send_message_to_connection(
|
||||
connections: &mut HashMap<ConnectionId, OutboundConnectionState>,
|
||||
connection_id: ConnectionId,
|
||||
message: OutgoingMessage,
|
||||
write_complete_tx: Option<tokio::sync::oneshot::Sender<()>>,
|
||||
) -> bool {
|
||||
let Some(connection_state) = connections.get(&connection_id) else {
|
||||
warn!("dropping message for disconnected connection: {connection_id:?}");
|
||||
return false;
|
||||
};
|
||||
let message = filter_outgoing_message_for_connection(connection_state, message);
|
||||
if should_skip_notification_for_connection(connection_state, &message) {
|
||||
return false;
|
||||
}
|
||||
|
||||
let writer = connection_state.writer.clone();
|
||||
let queued_message = QueuedOutgoingMessage {
|
||||
message,
|
||||
write_complete_tx,
|
||||
};
|
||||
if connection_state.can_disconnect() {
|
||||
match writer.try_send(queued_message) {
|
||||
Ok(()) => false,
|
||||
Err(mpsc::error::TrySendError::Full(_)) => {
|
||||
warn!(
|
||||
"disconnecting slow connection after outbound queue filled: {connection_id:?}"
|
||||
);
|
||||
disconnect_connection(connections, connection_id)
|
||||
}
|
||||
Err(mpsc::error::TrySendError::Closed(_)) => {
|
||||
disconnect_connection(connections, connection_id)
|
||||
}
|
||||
}
|
||||
} else if writer.send(queued_message).await.is_err() {
|
||||
disconnect_connection(connections, connection_id)
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
fn filter_outgoing_message_for_connection(
|
||||
connection_state: &OutboundConnectionState,
|
||||
message: OutgoingMessage,
|
||||
) -> OutgoingMessage {
|
||||
let experimental_api_enabled = connection_state
|
||||
.experimental_api_enabled
|
||||
.load(Ordering::Acquire);
|
||||
match message {
|
||||
OutgoingMessage::Request(ServerRequest::CommandExecutionRequestApproval {
|
||||
request_id,
|
||||
mut params,
|
||||
}) => {
|
||||
if !experimental_api_enabled {
|
||||
params.strip_experimental_fields();
|
||||
}
|
||||
OutgoingMessage::Request(ServerRequest::CommandExecutionRequestApproval {
|
||||
request_id,
|
||||
params,
|
||||
})
|
||||
}
|
||||
_ => message,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn route_outgoing_envelope(
|
||||
connections: &mut HashMap<ConnectionId, OutboundConnectionState>,
|
||||
envelope: OutgoingEnvelope,
|
||||
) {
|
||||
match envelope {
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message,
|
||||
write_complete_tx,
|
||||
} => {
|
||||
let _ =
|
||||
send_message_to_connection(connections, connection_id, message, write_complete_tx)
|
||||
.await;
|
||||
}
|
||||
OutgoingEnvelope::Broadcast { message } => {
|
||||
let target_connections: Vec<ConnectionId> = connections
|
||||
.iter()
|
||||
.filter_map(|(connection_id, connection_state)| {
|
||||
if connection_state.initialized.load(Ordering::Acquire)
|
||||
&& !should_skip_notification_for_connection(connection_state, &message)
|
||||
{
|
||||
Some(*connection_id)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
for connection_id in target_connections {
|
||||
let _ = send_message_to_connection(
|
||||
connections,
|
||||
connection_id,
|
||||
message.clone(),
|
||||
/*write_complete_tx*/ None,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "transport_tests.rs"]
|
||||
mod tests;
|
||||
@@ -1,751 +0,0 @@
|
||||
use anyhow::Context;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::http::StatusCode;
|
||||
use axum::http::header::AUTHORIZATION;
|
||||
use clap::Args;
|
||||
use clap::ValueEnum;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use constant_time_eq::constant_time_eq_32;
|
||||
use jsonwebtoken::Algorithm;
|
||||
use jsonwebtoken::DecodingKey;
|
||||
use jsonwebtoken::Validation;
|
||||
use jsonwebtoken::decode;
|
||||
use serde::Deserialize;
|
||||
use sha2::Digest;
|
||||
use sha2::Sha256;
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use std::net::SocketAddr;
|
||||
use std::path::Path;
|
||||
use std::path::PathBuf;
|
||||
use time::OffsetDateTime;
|
||||
|
||||
const DEFAULT_MAX_CLOCK_SKEW_SECONDS: u64 = 30;
|
||||
const MIN_SIGNED_BEARER_SECRET_BYTES: usize = 32;
|
||||
const INVALID_AUTHORIZATION_HEADER_MESSAGE: &str = "invalid authorization header";
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Args)]
|
||||
pub struct AppServerWebsocketAuthArgs {
|
||||
/// Websocket auth mode for non-loopback listeners.
|
||||
#[arg(long = "ws-auth", value_name = "MODE", value_enum)]
|
||||
pub ws_auth: Option<WebsocketAuthCliMode>,
|
||||
|
||||
/// Absolute path to the capability-token file.
|
||||
#[arg(long = "ws-token-file", value_name = "PATH")]
|
||||
pub ws_token_file: Option<PathBuf>,
|
||||
|
||||
/// Hex-encoded SHA-256 digest of the capability token.
|
||||
#[arg(long = "ws-token-sha256", value_name = "HEX")]
|
||||
pub ws_token_sha256: Option<String>,
|
||||
|
||||
/// Absolute path to the shared secret file for signed JWT bearer tokens.
|
||||
#[arg(long = "ws-shared-secret-file", value_name = "PATH")]
|
||||
pub ws_shared_secret_file: Option<PathBuf>,
|
||||
|
||||
/// Expected issuer for signed JWT bearer tokens.
|
||||
#[arg(long = "ws-issuer", value_name = "ISSUER")]
|
||||
pub ws_issuer: Option<String>,
|
||||
|
||||
/// Expected audience for signed JWT bearer tokens.
|
||||
#[arg(long = "ws-audience", value_name = "AUDIENCE")]
|
||||
pub ws_audience: Option<String>,
|
||||
|
||||
/// Maximum clock skew when validating signed JWT bearer tokens.
|
||||
#[arg(long = "ws-max-clock-skew-seconds", value_name = "SECONDS")]
|
||||
pub ws_max_clock_skew_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
|
||||
pub enum WebsocketAuthCliMode {
|
||||
CapabilityToken,
|
||||
SignedBearerToken,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AppServerWebsocketAuthSettings {
|
||||
pub config: Option<AppServerWebsocketAuthConfig>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AppServerWebsocketAuthConfig {
|
||||
CapabilityToken {
|
||||
source: AppServerWebsocketCapabilityTokenSource,
|
||||
},
|
||||
SignedBearerToken {
|
||||
shared_secret_file: AbsolutePathBuf,
|
||||
issuer: Option<String>,
|
||||
audience: Option<String>,
|
||||
max_clock_skew_seconds: u64,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AppServerWebsocketCapabilityTokenSource {
|
||||
TokenFile { token_file: AbsolutePathBuf },
|
||||
TokenSha256 { token_sha256: [u8; 32] },
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub(crate) struct WebsocketAuthPolicy {
|
||||
pub(crate) mode: Option<WebsocketAuthMode>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(crate) enum WebsocketAuthMode {
|
||||
CapabilityToken {
|
||||
token_sha256: [u8; 32],
|
||||
},
|
||||
SignedBearerToken {
|
||||
shared_secret: Vec<u8>,
|
||||
issuer: Option<String>,
|
||||
audience: Option<String>,
|
||||
max_clock_skew_seconds: i64,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct WebsocketAuthError {
|
||||
status_code: StatusCode,
|
||||
message: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct JwtClaims {
|
||||
exp: i64,
|
||||
nbf: Option<i64>,
|
||||
iss: Option<String>,
|
||||
aud: Option<JwtAudienceClaim>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum JwtAudienceClaim {
|
||||
Single(String),
|
||||
Multiple(Vec<String>),
|
||||
}
|
||||
|
||||
impl WebsocketAuthError {
|
||||
pub(crate) fn status_code(&self) -> StatusCode {
|
||||
self.status_code
|
||||
}
|
||||
|
||||
pub(crate) fn message(&self) -> &'static str {
|
||||
self.message
|
||||
}
|
||||
}
|
||||
|
||||
impl AppServerWebsocketAuthArgs {
|
||||
pub fn try_into_settings(self) -> anyhow::Result<AppServerWebsocketAuthSettings> {
|
||||
let normalize = |value: Option<String>| {
|
||||
value.and_then(|value| {
|
||||
let trimmed = value.trim();
|
||||
(!trimmed.is_empty()).then(|| trimmed.to_string())
|
||||
})
|
||||
};
|
||||
|
||||
let config = match self.ws_auth {
|
||||
Some(WebsocketAuthCliMode::CapabilityToken) => {
|
||||
if self.ws_shared_secret_file.is_some()
|
||||
|| self.ws_issuer.is_some()
|
||||
|| self.ws_audience.is_some()
|
||||
|| self.ws_max_clock_skew_seconds.is_some()
|
||||
{
|
||||
anyhow::bail!(
|
||||
"`--ws-shared-secret-file`, `--ws-issuer`, `--ws-audience`, and `--ws-max-clock-skew-seconds` require `--ws-auth signed-bearer-token`"
|
||||
);
|
||||
}
|
||||
let source = match (self.ws_token_file, self.ws_token_sha256) {
|
||||
(Some(_), Some(_)) => {
|
||||
anyhow::bail!(
|
||||
"`--ws-token-file` and `--ws-token-sha256` are mutually exclusive"
|
||||
);
|
||||
}
|
||||
(Some(token_file), None) => {
|
||||
AppServerWebsocketCapabilityTokenSource::TokenFile {
|
||||
token_file: absolute_path_arg("--ws-token-file", token_file)?,
|
||||
}
|
||||
}
|
||||
(None, Some(token_sha256)) => {
|
||||
AppServerWebsocketCapabilityTokenSource::TokenSha256 {
|
||||
token_sha256: sha256_digest_arg("--ws-token-sha256", &token_sha256)?,
|
||||
}
|
||||
}
|
||||
(None, None) => {
|
||||
anyhow::bail!(
|
||||
"`--ws-token-file` or `--ws-token-sha256` is required when `--ws-auth capability-token` is set"
|
||||
);
|
||||
}
|
||||
};
|
||||
Some(AppServerWebsocketAuthConfig::CapabilityToken { source })
|
||||
}
|
||||
Some(WebsocketAuthCliMode::SignedBearerToken) => {
|
||||
if self.ws_token_file.is_some() || self.ws_token_sha256.is_some() {
|
||||
anyhow::bail!(
|
||||
"`--ws-token-file` and `--ws-token-sha256` require `--ws-auth capability-token`, not `signed-bearer-token`"
|
||||
);
|
||||
}
|
||||
let shared_secret_file = self.ws_shared_secret_file.context(
|
||||
"`--ws-shared-secret-file` is required when `--ws-auth signed-bearer-token` is set",
|
||||
)?;
|
||||
Some(AppServerWebsocketAuthConfig::SignedBearerToken {
|
||||
shared_secret_file: absolute_path_arg(
|
||||
"--ws-shared-secret-file",
|
||||
shared_secret_file,
|
||||
)?,
|
||||
issuer: normalize(self.ws_issuer),
|
||||
audience: normalize(self.ws_audience),
|
||||
max_clock_skew_seconds: self
|
||||
.ws_max_clock_skew_seconds
|
||||
.unwrap_or(DEFAULT_MAX_CLOCK_SKEW_SECONDS),
|
||||
})
|
||||
}
|
||||
None => {
|
||||
if self.ws_token_file.is_some()
|
||||
|| self.ws_token_sha256.is_some()
|
||||
|| self.ws_shared_secret_file.is_some()
|
||||
|| self.ws_issuer.is_some()
|
||||
|| self.ws_audience.is_some()
|
||||
|| self.ws_max_clock_skew_seconds.is_some()
|
||||
{
|
||||
anyhow::bail!(
|
||||
"websocket auth flags require `--ws-auth capability-token` or `--ws-auth signed-bearer-token`"
|
||||
);
|
||||
}
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
Ok(AppServerWebsocketAuthSettings { config })
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn policy_from_settings(
|
||||
settings: &AppServerWebsocketAuthSettings,
|
||||
) -> io::Result<WebsocketAuthPolicy> {
|
||||
let mode = match settings.config.as_ref() {
|
||||
Some(AppServerWebsocketAuthConfig::CapabilityToken { source }) => match source {
|
||||
AppServerWebsocketCapabilityTokenSource::TokenFile { token_file } => {
|
||||
let token = read_trimmed_secret(token_file.as_ref())?;
|
||||
Some(WebsocketAuthMode::CapabilityToken {
|
||||
token_sha256: sha256_digest(token.as_bytes()),
|
||||
})
|
||||
}
|
||||
AppServerWebsocketCapabilityTokenSource::TokenSha256 { token_sha256 } => {
|
||||
Some(WebsocketAuthMode::CapabilityToken {
|
||||
token_sha256: *token_sha256,
|
||||
})
|
||||
}
|
||||
},
|
||||
Some(AppServerWebsocketAuthConfig::SignedBearerToken {
|
||||
shared_secret_file,
|
||||
issuer,
|
||||
audience,
|
||||
max_clock_skew_seconds,
|
||||
}) => {
|
||||
let shared_secret = read_trimmed_secret(shared_secret_file.as_ref())?.into_bytes();
|
||||
validate_signed_bearer_secret(shared_secret_file.as_ref(), &shared_secret)?;
|
||||
let max_clock_skew_seconds = i64::try_from(*max_clock_skew_seconds).map_err(|_| {
|
||||
io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
"websocket auth clock skew must fit in a signed 64-bit integer",
|
||||
)
|
||||
})?;
|
||||
Some(WebsocketAuthMode::SignedBearerToken {
|
||||
shared_secret,
|
||||
issuer: issuer.clone(),
|
||||
audience: audience.clone(),
|
||||
max_clock_skew_seconds,
|
||||
})
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
|
||||
Ok(WebsocketAuthPolicy { mode })
|
||||
}
|
||||
|
||||
pub(crate) fn should_warn_about_unauthenticated_non_loopback_listener(
|
||||
bind_address: SocketAddr,
|
||||
policy: &WebsocketAuthPolicy,
|
||||
) -> bool {
|
||||
!bind_address.ip().is_loopback() && policy.mode.is_none()
|
||||
}
|
||||
|
||||
pub(crate) fn authorize_upgrade(
|
||||
headers: &HeaderMap,
|
||||
policy: &WebsocketAuthPolicy,
|
||||
) -> Result<(), WebsocketAuthError> {
|
||||
let Some(mode) = policy.mode.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let token = bearer_token_from_headers(headers)?;
|
||||
match mode {
|
||||
WebsocketAuthMode::CapabilityToken { token_sha256 } => {
|
||||
let actual_sha256 = sha256_digest(token.as_bytes());
|
||||
if constant_time_eq_32(token_sha256, &actual_sha256) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(unauthorized("invalid websocket bearer token"))
|
||||
}
|
||||
}
|
||||
WebsocketAuthMode::SignedBearerToken {
|
||||
shared_secret,
|
||||
issuer,
|
||||
audience,
|
||||
max_clock_skew_seconds,
|
||||
} => verify_signed_bearer_token(
|
||||
token,
|
||||
shared_secret,
|
||||
issuer.as_deref(),
|
||||
audience.as_deref(),
|
||||
*max_clock_skew_seconds,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn verify_signed_bearer_token(
|
||||
token: &str,
|
||||
shared_secret: &[u8],
|
||||
issuer: Option<&str>,
|
||||
audience: Option<&str>,
|
||||
max_clock_skew_seconds: i64,
|
||||
) -> Result<(), WebsocketAuthError> {
|
||||
let claims = decode_jwt_claims(token, shared_secret)?;
|
||||
validate_jwt_claims(&claims, issuer, audience, max_clock_skew_seconds)
|
||||
}
|
||||
|
||||
fn decode_jwt_claims(token: &str, shared_secret: &[u8]) -> Result<JwtClaims, WebsocketAuthError> {
|
||||
let mut validation = Validation::new(Algorithm::HS256);
|
||||
validation.required_spec_claims.clear();
|
||||
validation.validate_exp = false;
|
||||
validation.validate_nbf = false;
|
||||
validation.validate_aud = false;
|
||||
|
||||
decode::<JwtClaims>(token, &DecodingKey::from_secret(shared_secret), &validation)
|
||||
.map(|token_data| token_data.claims)
|
||||
.map_err(|_| unauthorized("invalid websocket jwt"))
|
||||
}
|
||||
|
||||
fn validate_jwt_claims(
|
||||
claims: &JwtClaims,
|
||||
issuer: Option<&str>,
|
||||
audience: Option<&str>,
|
||||
max_clock_skew_seconds: i64,
|
||||
) -> Result<(), WebsocketAuthError> {
|
||||
let now = OffsetDateTime::now_utc().unix_timestamp();
|
||||
if now > claims.exp.saturating_add(max_clock_skew_seconds) {
|
||||
return Err(unauthorized("expired websocket jwt"));
|
||||
}
|
||||
if let Some(nbf) = claims.nbf
|
||||
&& now < nbf.saturating_sub(max_clock_skew_seconds)
|
||||
{
|
||||
return Err(unauthorized("websocket jwt is not valid yet"));
|
||||
}
|
||||
if let Some(expected_issuer) = issuer
|
||||
&& claims.iss.as_deref() != Some(expected_issuer)
|
||||
{
|
||||
return Err(unauthorized("websocket jwt issuer mismatch"));
|
||||
}
|
||||
if let Some(expected_audience) = audience
|
||||
&& !audience_matches(claims.aud.as_ref(), expected_audience)
|
||||
{
|
||||
return Err(unauthorized("websocket jwt audience mismatch"));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn audience_matches(audience: Option<&JwtAudienceClaim>, expected_audience: &str) -> bool {
|
||||
match audience {
|
||||
Some(JwtAudienceClaim::Single(actual)) => actual == expected_audience,
|
||||
Some(JwtAudienceClaim::Multiple(actual)) => {
|
||||
actual.iter().any(|audience| audience == expected_audience)
|
||||
}
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn bearer_token_from_headers(headers: &HeaderMap) -> Result<&str, WebsocketAuthError> {
|
||||
let raw_header = headers
|
||||
.get(AUTHORIZATION)
|
||||
.ok_or_else(|| unauthorized("missing websocket bearer token"))?;
|
||||
let header = raw_header
|
||||
.to_str()
|
||||
.map_err(|_| unauthorized(INVALID_AUTHORIZATION_HEADER_MESSAGE))?;
|
||||
let Some((scheme, token)) = header.split_once(' ') else {
|
||||
return Err(unauthorized(INVALID_AUTHORIZATION_HEADER_MESSAGE));
|
||||
};
|
||||
if !scheme.eq_ignore_ascii_case("Bearer") {
|
||||
return Err(unauthorized(INVALID_AUTHORIZATION_HEADER_MESSAGE));
|
||||
}
|
||||
let token = token.trim();
|
||||
if token.is_empty() {
|
||||
return Err(unauthorized(INVALID_AUTHORIZATION_HEADER_MESSAGE));
|
||||
}
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
fn validate_signed_bearer_secret(path: &Path, shared_secret: &[u8]) -> io::Result<()> {
|
||||
if shared_secret.len() < MIN_SIGNED_BEARER_SECRET_BYTES {
|
||||
return Err(io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"signed websocket bearer secret {} must be at least {MIN_SIGNED_BEARER_SECRET_BYTES} bytes",
|
||||
path.display()
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn read_trimmed_secret(path: &std::path::Path) -> io::Result<String> {
|
||||
let raw = std::fs::read_to_string(path).map_err(|err| {
|
||||
io::Error::new(
|
||||
err.kind(),
|
||||
format!(
|
||||
"failed to read websocket auth secret {}: {err}",
|
||||
path.display()
|
||||
),
|
||||
)
|
||||
})?;
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err(io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
format!("websocket auth secret {} must not be empty", path.display()),
|
||||
));
|
||||
}
|
||||
Ok(trimmed.to_string())
|
||||
}
|
||||
|
||||
fn absolute_path_arg(flag_name: &str, path: PathBuf) -> anyhow::Result<AbsolutePathBuf> {
|
||||
AbsolutePathBuf::try_from(path).with_context(|| format!("{flag_name} must be an absolute path"))
|
||||
}
|
||||
|
||||
fn sha256_digest_arg(flag_name: &str, value: &str) -> anyhow::Result<[u8; 32]> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.len() != 64 {
|
||||
anyhow::bail!("{flag_name} must be a 64-character hex SHA-256 digest");
|
||||
}
|
||||
|
||||
let mut digest = [0u8; 32];
|
||||
for (index, pair) in trimmed.as_bytes().chunks_exact(2).enumerate() {
|
||||
let high = hex_nibble(flag_name, pair[0])?;
|
||||
let low = hex_nibble(flag_name, pair[1])?;
|
||||
digest[index] = (high << 4) | low;
|
||||
}
|
||||
Ok(digest)
|
||||
}
|
||||
|
||||
fn hex_nibble(flag_name: &str, byte: u8) -> anyhow::Result<u8> {
|
||||
match byte {
|
||||
b'0'..=b'9' => Ok(byte - b'0'),
|
||||
b'a'..=b'f' => Ok(byte - b'a' + 10),
|
||||
b'A'..=b'F' => Ok(byte - b'A' + 10),
|
||||
_ => anyhow::bail!("{flag_name} must be a 64-character hex SHA-256 digest"),
|
||||
}
|
||||
}
|
||||
|
||||
fn sha256_digest(input: &[u8]) -> [u8; 32] {
|
||||
let mut digest = [0u8; 32];
|
||||
digest.copy_from_slice(&Sha256::digest(input));
|
||||
digest
|
||||
}
|
||||
|
||||
fn unauthorized(message: &'static str) -> WebsocketAuthError {
|
||||
WebsocketAuthError {
|
||||
status_code: StatusCode::UNAUTHORIZED,
|
||||
message,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use axum::http::HeaderValue;
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||
use hmac::Hmac;
|
||||
use hmac::Mac;
|
||||
use serde_json::json;
|
||||
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
fn signed_token(shared_secret: &[u8], claims: serde_json::Value) -> String {
|
||||
let header = URL_SAFE_NO_PAD.encode(br#"{"alg":"HS256","typ":"JWT"}"#);
|
||||
let claims_segment = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap());
|
||||
let payload = format!("{header}.{claims_segment}");
|
||||
let mut mac = HmacSha256::new_from_slice(shared_secret).unwrap();
|
||||
mac.update(payload.as_bytes());
|
||||
let signature = URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
|
||||
format!("{payload}.{signature}")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warns_about_unauthenticated_non_loopback_listener() {
|
||||
let policy = WebsocketAuthPolicy::default();
|
||||
assert!(should_warn_about_unauthenticated_non_loopback_listener(
|
||||
"0.0.0.0:8765".parse().unwrap(),
|
||||
&policy,
|
||||
));
|
||||
assert!(!should_warn_about_unauthenticated_non_loopback_listener(
|
||||
"127.0.0.1:8765".parse().unwrap(),
|
||||
&policy,
|
||||
));
|
||||
assert!(!should_warn_about_unauthenticated_non_loopback_listener(
|
||||
"0.0.0.0:8765".parse().unwrap(),
|
||||
&WebsocketAuthPolicy {
|
||||
mode: Some(WebsocketAuthMode::CapabilityToken {
|
||||
token_sha256: [0u8; 32],
|
||||
}),
|
||||
},
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_token_args_require_token_file_or_hash() {
|
||||
let err = AppServerWebsocketAuthArgs {
|
||||
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
|
||||
..Default::default()
|
||||
}
|
||||
.try_into_settings()
|
||||
.expect_err("capability-token mode should require a token source");
|
||||
assert!(
|
||||
err.to_string().contains("--ws-token-file")
|
||||
&& err.to_string().contains("--ws-token-sha256"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_token_args_accept_token_hash() {
|
||||
let settings = AppServerWebsocketAuthArgs {
|
||||
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
|
||||
ws_token_sha256: Some("ab".repeat(32)),
|
||||
..Default::default()
|
||||
}
|
||||
.try_into_settings()
|
||||
.expect("capability-token hash args should parse");
|
||||
|
||||
assert_eq!(
|
||||
settings,
|
||||
AppServerWebsocketAuthSettings {
|
||||
config: Some(AppServerWebsocketAuthConfig::CapabilityToken {
|
||||
source: AppServerWebsocketCapabilityTokenSource::TokenSha256 {
|
||||
token_sha256: [0xab; 32],
|
||||
},
|
||||
}),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_token_args_reject_multiple_token_sources() {
|
||||
let err = AppServerWebsocketAuthArgs {
|
||||
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
|
||||
ws_token_file: Some(PathBuf::from("/tmp/token")),
|
||||
ws_token_sha256: Some("ab".repeat(32)),
|
||||
..Default::default()
|
||||
}
|
||||
.try_into_settings()
|
||||
.expect_err("capability-token mode should reject multiple token sources");
|
||||
assert!(
|
||||
err.to_string().contains("mutually exclusive"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_token_args_reject_malformed_token_hash() {
|
||||
let err = AppServerWebsocketAuthArgs {
|
||||
ws_auth: Some(WebsocketAuthCliMode::CapabilityToken),
|
||||
ws_token_sha256: Some("not-a-sha256".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
.try_into_settings()
|
||||
.expect_err("capability-token mode should reject malformed token hashes");
|
||||
assert!(
|
||||
err.to_string().contains("64-character hex"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capability_token_hash_policy_authorizes_matching_bearer_token() {
|
||||
let settings = AppServerWebsocketAuthSettings {
|
||||
config: Some(AppServerWebsocketAuthConfig::CapabilityToken {
|
||||
source: AppServerWebsocketCapabilityTokenSource::TokenSha256 {
|
||||
token_sha256: sha256_digest(b"super-secret-token"),
|
||||
},
|
||||
}),
|
||||
};
|
||||
let policy = policy_from_settings(&settings).expect("hash policy should build");
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer super-secret-token"),
|
||||
);
|
||||
authorize_upgrade(&headers, &policy).expect("matching token should authorize");
|
||||
|
||||
headers.insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer wrong-token"),
|
||||
);
|
||||
let err = authorize_upgrade(&headers, &policy).expect_err("wrong token should fail");
|
||||
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signed_bearer_args_require_mode_when_mode_specific_flags_are_set() {
|
||||
let err = AppServerWebsocketAuthArgs {
|
||||
ws_shared_secret_file: Some(PathBuf::from("/tmp/secret")),
|
||||
..Default::default()
|
||||
}
|
||||
.try_into_settings()
|
||||
.expect_err("mode-specific flags should require --ws-auth");
|
||||
assert!(
|
||||
err.to_string().contains("websocket auth flags require"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signed_bearer_args_default_clock_skew_and_trim_optional_claims() {
|
||||
let settings = AppServerWebsocketAuthArgs {
|
||||
ws_auth: Some(WebsocketAuthCliMode::SignedBearerToken),
|
||||
ws_shared_secret_file: Some(PathBuf::from("/tmp/secret")),
|
||||
ws_issuer: Some(" issuer ".to_string()),
|
||||
ws_audience: Some(" ".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
.try_into_settings()
|
||||
.expect("signed bearer args should parse");
|
||||
|
||||
assert_eq!(
|
||||
settings,
|
||||
AppServerWebsocketAuthSettings {
|
||||
config: Some(AppServerWebsocketAuthConfig::SignedBearerToken {
|
||||
shared_secret_file: AbsolutePathBuf::from_absolute_path("/tmp/secret")
|
||||
.expect("absolute path"),
|
||||
issuer: Some("issuer".to_string()),
|
||||
audience: None,
|
||||
max_clock_skew_seconds: DEFAULT_MAX_CLOCK_SKEW_SECONDS,
|
||||
}),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signed_bearer_token_verification_rejects_tampering() {
|
||||
let shared_secret = b"0123456789abcdef0123456789abcdef";
|
||||
let token = signed_token(
|
||||
shared_secret,
|
||||
json!({
|
||||
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
|
||||
}),
|
||||
);
|
||||
let tampered = token.replace(".eyJleHAi", ".eyJleHBi");
|
||||
let err = verify_signed_bearer_token(
|
||||
&tampered,
|
||||
shared_secret,
|
||||
/*issuer*/ None,
|
||||
/*audience*/ None,
|
||||
/*max_clock_skew_seconds*/ 30,
|
||||
)
|
||||
.expect_err("tampered jwt should fail");
|
||||
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signed_bearer_token_verification_accepts_valid_token() {
|
||||
let shared_secret = b"0123456789abcdef0123456789abcdef";
|
||||
let token = signed_token(
|
||||
shared_secret,
|
||||
json!({
|
||||
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
|
||||
"iss": "issuer",
|
||||
"aud": "audience",
|
||||
}),
|
||||
);
|
||||
verify_signed_bearer_token(
|
||||
&token,
|
||||
shared_secret,
|
||||
Some("issuer"),
|
||||
Some("audience"),
|
||||
/*max_clock_skew_seconds*/ 30,
|
||||
)
|
||||
.expect("valid signed token should verify");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signed_bearer_token_verification_accepts_multiple_audiences() {
|
||||
let shared_secret = b"0123456789abcdef0123456789abcdef";
|
||||
let token = signed_token(
|
||||
shared_secret,
|
||||
json!({
|
||||
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
|
||||
"aud": ["other-audience", "audience"],
|
||||
}),
|
||||
);
|
||||
verify_signed_bearer_token(
|
||||
&token,
|
||||
shared_secret,
|
||||
/*issuer*/ None,
|
||||
Some("audience"),
|
||||
/*max_clock_skew_seconds*/ 30,
|
||||
)
|
||||
.expect("jwt audience arrays should verify");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signed_bearer_token_verification_rejects_alg_none_tokens() {
|
||||
let claims_segment = URL_SAFE_NO_PAD.encode(
|
||||
serde_json::to_vec(&json!({
|
||||
"exp": OffsetDateTime::now_utc().unix_timestamp() + 60,
|
||||
}))
|
||||
.unwrap(),
|
||||
);
|
||||
let header_segment = URL_SAFE_NO_PAD.encode(br#"{"alg":"none","typ":"JWT"}"#);
|
||||
let token = format!("{header_segment}.{claims_segment}.");
|
||||
let err = verify_signed_bearer_token(
|
||||
&token,
|
||||
b"0123456789abcdef0123456789abcdef",
|
||||
/*issuer*/ None,
|
||||
/*audience*/ None,
|
||||
/*max_clock_skew_seconds*/ 30,
|
||||
)
|
||||
.expect_err("alg=none jwt should be rejected");
|
||||
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signed_bearer_token_verification_rejects_missing_exp() {
|
||||
let shared_secret = b"0123456789abcdef0123456789abcdef";
|
||||
let token = signed_token(
|
||||
shared_secret,
|
||||
json!({
|
||||
"iss": "issuer",
|
||||
}),
|
||||
);
|
||||
let err = verify_signed_bearer_token(
|
||||
&token,
|
||||
shared_secret,
|
||||
/*issuer*/ None,
|
||||
/*audience*/ None,
|
||||
/*max_clock_skew_seconds*/ 30,
|
||||
)
|
||||
.expect_err("jwt without exp should be rejected");
|
||||
assert_eq!(err.status_code(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_signed_bearer_secret_rejects_short_secret() {
|
||||
let err = validate_signed_bearer_secret(Path::new("/tmp/secret"), b"too-short")
|
||||
.expect_err("short shared secret should be rejected");
|
||||
assert_eq!(err.kind(), ErrorKind::InvalidInput);
|
||||
assert!(
|
||||
err.to_string().contains("must be at least 32 bytes"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,570 +0,0 @@
|
||||
use super::CHANNEL_CAPACITY;
|
||||
use super::TransportEvent;
|
||||
use super::next_connection_id;
|
||||
use super::protocol::ClientEnvelope;
|
||||
pub use super::protocol::ClientEvent;
|
||||
pub use super::protocol::ClientId;
|
||||
use super::protocol::PongStatus;
|
||||
use super::protocol::ServerEvent;
|
||||
use super::protocol::StreamId;
|
||||
use crate::outgoing_message::ConnectionId;
|
||||
use crate::outgoing_message::QueuedOutgoingMessage;
|
||||
use crate::transport::ConnectionOrigin;
|
||||
use crate::transport::remote_control::QueuedServerEnvelope;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use std::collections::HashMap;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::watch;
|
||||
use tokio::task::JoinSet;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::Instant;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
const REMOTE_CONTROL_CLIENT_IDLE_TIMEOUT: Duration = Duration::from_secs(10 * 60);
|
||||
pub(crate) const REMOTE_CONTROL_IDLE_SWEEP_INTERVAL: Duration = Duration::from_secs(30);
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct Stopped;
|
||||
|
||||
struct ClientState {
|
||||
connection_id: ConnectionId,
|
||||
disconnect_token: CancellationToken,
|
||||
last_activity_at: Instant,
|
||||
last_inbound_seq_id: Option<u64>,
|
||||
status_tx: watch::Sender<PongStatus>,
|
||||
}
|
||||
|
||||
pub(crate) struct ClientTracker {
|
||||
clients: HashMap<(ClientId, StreamId), ClientState>,
|
||||
legacy_stream_ids: HashMap<ClientId, StreamId>,
|
||||
join_set: JoinSet<(ClientId, StreamId)>,
|
||||
server_event_tx: mpsc::Sender<QueuedServerEnvelope>,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: CancellationToken,
|
||||
}
|
||||
|
||||
impl ClientTracker {
|
||||
pub(crate) fn new(
|
||||
server_event_tx: mpsc::Sender<QueuedServerEnvelope>,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: &CancellationToken,
|
||||
) -> Self {
|
||||
Self {
|
||||
clients: HashMap::new(),
|
||||
legacy_stream_ids: HashMap::new(),
|
||||
join_set: JoinSet::new(),
|
||||
server_event_tx,
|
||||
transport_event_tx,
|
||||
shutdown_token: shutdown_token.child_token(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn bookkeep_join_set(&mut self) -> Option<(ClientId, StreamId)> {
|
||||
while let Some(join_result) = self.join_set.join_next().await {
|
||||
let Ok(client_key) = join_result else {
|
||||
continue;
|
||||
};
|
||||
return Some(client_key);
|
||||
}
|
||||
futures::future::pending().await
|
||||
}
|
||||
|
||||
pub(crate) async fn shutdown(&mut self) {
|
||||
self.shutdown_token.cancel();
|
||||
|
||||
while let Some(client_key) = self.clients.keys().next().cloned() {
|
||||
let _ = self.close_client(&client_key).await;
|
||||
}
|
||||
|
||||
self.drain_join_set().await;
|
||||
}
|
||||
|
||||
async fn drain_join_set(&mut self) {
|
||||
while self.join_set.join_next().await.is_some() {}
|
||||
}
|
||||
|
||||
pub(crate) async fn handle_message(
|
||||
&mut self,
|
||||
client_envelope: ClientEnvelope,
|
||||
) -> Result<(), Stopped> {
|
||||
let ClientEnvelope {
|
||||
client_id,
|
||||
event,
|
||||
stream_id,
|
||||
seq_id,
|
||||
cursor: _,
|
||||
} = client_envelope;
|
||||
let is_legacy_stream_id = stream_id.is_none();
|
||||
let is_initialize = matches!(&event, ClientEvent::ClientMessage { message } if remote_control_message_starts_connection(message));
|
||||
let stream_id = match stream_id {
|
||||
Some(stream_id) => stream_id,
|
||||
None if is_initialize => {
|
||||
// TODO(ruslan): delete this fallback once all clients are updated to send stream_id.
|
||||
self.legacy_stream_ids
|
||||
.remove(&client_id)
|
||||
.unwrap_or_else(StreamId::new_random)
|
||||
}
|
||||
None => self
|
||||
.legacy_stream_ids
|
||||
.get(&client_id)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| {
|
||||
if matches!(&event, ClientEvent::Ping) {
|
||||
StreamId::new_random()
|
||||
} else {
|
||||
StreamId(String::new())
|
||||
}
|
||||
}),
|
||||
};
|
||||
if stream_id.0.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let client_key = (client_id.clone(), stream_id.clone());
|
||||
match event {
|
||||
ClientEvent::ClientMessage { message } => {
|
||||
if let Some(seq_id) = seq_id
|
||||
&& let Some(client) = self.clients.get(&client_key)
|
||||
&& client
|
||||
.last_inbound_seq_id
|
||||
.is_some_and(|last_seq_id| last_seq_id >= seq_id)
|
||||
&& !is_initialize
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if is_initialize && self.clients.contains_key(&client_key) {
|
||||
self.close_client(&client_key).await?;
|
||||
}
|
||||
|
||||
if let Some(connection_id) = self.clients.get_mut(&client_key).map(|client| {
|
||||
client.last_activity_at = Instant::now();
|
||||
if let Some(seq_id) = seq_id {
|
||||
client.last_inbound_seq_id = Some(seq_id);
|
||||
}
|
||||
client.connection_id
|
||||
}) {
|
||||
self.send_transport_event(TransportEvent::IncomingMessage {
|
||||
connection_id,
|
||||
message,
|
||||
})
|
||||
.await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if !is_initialize {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let connection_id = next_connection_id();
|
||||
let (writer_tx, writer_rx) =
|
||||
mpsc::channel::<QueuedOutgoingMessage>(CHANNEL_CAPACITY);
|
||||
let disconnect_token = self.shutdown_token.child_token();
|
||||
self.send_transport_event(TransportEvent::ConnectionOpened {
|
||||
connection_id,
|
||||
origin: ConnectionOrigin::RemoteControl,
|
||||
writer: writer_tx,
|
||||
disconnect_sender: Some(disconnect_token.clone()),
|
||||
})
|
||||
.await?;
|
||||
|
||||
let (status_tx, status_rx) = watch::channel(PongStatus::Active);
|
||||
self.join_set.spawn(Self::run_client_outbound(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
self.server_event_tx.clone(),
|
||||
writer_rx,
|
||||
status_rx,
|
||||
disconnect_token.clone(),
|
||||
));
|
||||
self.clients.insert(
|
||||
client_key,
|
||||
ClientState {
|
||||
connection_id,
|
||||
disconnect_token,
|
||||
last_activity_at: Instant::now(),
|
||||
last_inbound_seq_id: if is_legacy_stream_id { None } else { seq_id },
|
||||
status_tx,
|
||||
},
|
||||
);
|
||||
if is_legacy_stream_id {
|
||||
self.legacy_stream_ids.insert(client_id.clone(), stream_id);
|
||||
}
|
||||
self.send_transport_event(TransportEvent::IncomingMessage {
|
||||
connection_id,
|
||||
message,
|
||||
})
|
||||
.await
|
||||
}
|
||||
ClientEvent::ClientMessageChunk { .. } | ClientEvent::Ack { .. } => Ok(()),
|
||||
ClientEvent::Ping => {
|
||||
if let Some(client) = self.clients.get_mut(&client_key) {
|
||||
client.last_activity_at = Instant::now();
|
||||
let _ = client.status_tx.send(PongStatus::Active);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let server_event_tx = self.server_event_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let server_envelope = QueuedServerEnvelope {
|
||||
event: ServerEvent::Pong {
|
||||
status: PongStatus::Unknown,
|
||||
},
|
||||
client_id,
|
||||
stream_id,
|
||||
write_complete_tx: None,
|
||||
};
|
||||
let _ = server_event_tx.send(server_envelope).await;
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
ClientEvent::ClientClosed => self.close_client(&client_key).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_client_outbound(
|
||||
client_id: ClientId,
|
||||
stream_id: StreamId,
|
||||
server_event_tx: mpsc::Sender<QueuedServerEnvelope>,
|
||||
mut writer_rx: mpsc::Receiver<QueuedOutgoingMessage>,
|
||||
mut status_rx: watch::Receiver<PongStatus>,
|
||||
disconnect_token: CancellationToken,
|
||||
) -> (ClientId, StreamId) {
|
||||
loop {
|
||||
let (event, write_complete_tx) = tokio::select! {
|
||||
_ = disconnect_token.cancelled() => {
|
||||
break;
|
||||
}
|
||||
queued_message = writer_rx.recv() => {
|
||||
let Some(queued_message) = queued_message else {
|
||||
break;
|
||||
};
|
||||
let event = ServerEvent::ServerMessage {
|
||||
message: Box::new(queued_message.message),
|
||||
};
|
||||
(event, queued_message.write_complete_tx)
|
||||
}
|
||||
changed = status_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
break;
|
||||
}
|
||||
let event = ServerEvent::Pong { status: status_rx.borrow().clone() };
|
||||
(event, None)
|
||||
}
|
||||
};
|
||||
let send_result = tokio::select! {
|
||||
_ = disconnect_token.cancelled() => {
|
||||
break;
|
||||
}
|
||||
send_result = server_event_tx.send(QueuedServerEnvelope {
|
||||
event,
|
||||
client_id: client_id.clone(),
|
||||
stream_id: stream_id.clone(),
|
||||
write_complete_tx,
|
||||
}) => send_result,
|
||||
};
|
||||
if send_result.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
(client_id, stream_id)
|
||||
}
|
||||
|
||||
pub(crate) async fn close_expired_clients(
|
||||
&mut self,
|
||||
) -> Result<Vec<(ClientId, StreamId)>, Stopped> {
|
||||
let now = Instant::now();
|
||||
let expired_client_ids: Vec<(ClientId, StreamId)> = self
|
||||
.clients
|
||||
.iter()
|
||||
.filter_map(|(client_key, client)| {
|
||||
(!remote_control_client_is_alive(client, now)).then_some(client_key.clone())
|
||||
})
|
||||
.collect();
|
||||
for client_key in &expired_client_ids {
|
||||
self.close_client(client_key).await?;
|
||||
}
|
||||
Ok(expired_client_ids)
|
||||
}
|
||||
|
||||
pub(super) async fn close_client(
|
||||
&mut self,
|
||||
client_key: &(ClientId, StreamId),
|
||||
) -> Result<(), Stopped> {
|
||||
let Some(client) = self.clients.remove(client_key) else {
|
||||
return Ok(());
|
||||
};
|
||||
if self
|
||||
.legacy_stream_ids
|
||||
.get(&client_key.0)
|
||||
.is_some_and(|stream_id| stream_id == &client_key.1)
|
||||
{
|
||||
self.legacy_stream_ids.remove(&client_key.0);
|
||||
}
|
||||
client.disconnect_token.cancel();
|
||||
self.send_transport_event(TransportEvent::ConnectionClosed {
|
||||
connection_id: client.connection_id,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_transport_event(&self, event: TransportEvent) -> Result<(), Stopped> {
|
||||
self.transport_event_tx
|
||||
.send(event)
|
||||
.await
|
||||
.map_err(|_| Stopped)
|
||||
}
|
||||
}
|
||||
|
||||
fn remote_control_message_starts_connection(message: &JSONRPCMessage) -> bool {
|
||||
matches!(
|
||||
message,
|
||||
JSONRPCMessage::Request(codex_app_server_protocol::JSONRPCRequest { method, .. })
|
||||
if method == "initialize"
|
||||
)
|
||||
}
|
||||
|
||||
fn remote_control_client_is_alive(client: &ClientState, now: Instant) -> bool {
|
||||
now.duration_since(client.last_activity_at) < REMOTE_CONTROL_CLIENT_IDLE_TIMEOUT
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::outgoing_message::OutgoingMessage;
|
||||
use crate::transport::remote_control::protocol::ClientEnvelope;
|
||||
use crate::transport::remote_control::protocol::ClientEvent;
|
||||
use codex_app_server_protocol::ConfigWarningNotification;
|
||||
use codex_app_server_protocol::JSONRPCRequest;
|
||||
use codex_app_server_protocol::RequestId;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use tokio::time::timeout;
|
||||
|
||||
fn initialize_envelope(client_id: &str) -> ClientEnvelope {
|
||||
initialize_envelope_with_stream_id(client_id, /*stream_id*/ None)
|
||||
}
|
||||
|
||||
fn initialize_envelope_with_stream_id(
|
||||
client_id: &str,
|
||||
stream_id: Option<&str>,
|
||||
) -> ClientEnvelope {
|
||||
ClientEnvelope {
|
||||
event: ClientEvent::ClientMessage {
|
||||
message: JSONRPCMessage::Request(JSONRPCRequest {
|
||||
id: RequestId::Integer(1),
|
||||
method: "initialize".to_string(),
|
||||
params: Some(json!({
|
||||
"clientInfo": {
|
||||
"name": "remote-test-client",
|
||||
"version": "0.1.0"
|
||||
}
|
||||
})),
|
||||
trace: None,
|
||||
}),
|
||||
},
|
||||
client_id: ClientId(client_id.to_string()),
|
||||
stream_id: stream_id.map(|stream_id| StreamId(stream_id.to_string())),
|
||||
seq_id: Some(0),
|
||||
cursor: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_outbound_task_emits_connection_closed() {
|
||||
let (server_event_tx, _server_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let mut client_tracker =
|
||||
ClientTracker::new(server_event_tx, transport_event_tx, &shutdown_token);
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope("client-1"))
|
||||
.await
|
||||
.expect("initialize should open client");
|
||||
|
||||
let (connection_id, disconnect_sender) = match transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("connection opened should be sent")
|
||||
{
|
||||
TransportEvent::ConnectionOpened {
|
||||
connection_id,
|
||||
disconnect_sender: Some(disconnect_sender),
|
||||
..
|
||||
} => (connection_id, disconnect_sender),
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
match transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("initialize should be forwarded")
|
||||
{
|
||||
TransportEvent::IncomingMessage {
|
||||
connection_id: incoming_connection_id,
|
||||
..
|
||||
} => assert_eq!(incoming_connection_id, connection_id),
|
||||
other => panic!("expected incoming initialize, got {other:?}"),
|
||||
}
|
||||
|
||||
disconnect_sender.cancel();
|
||||
let closed_client_id = timeout(Duration::from_secs(1), client_tracker.bookkeep_join_set())
|
||||
.await
|
||||
.expect("bookkeeping should process the closed task")
|
||||
.expect("closed task should return client id");
|
||||
assert_eq!(closed_client_id.0, ClientId("client-1".to_string()));
|
||||
client_tracker
|
||||
.close_client(&closed_client_id)
|
||||
.await
|
||||
.expect("closed client should emit connection closed");
|
||||
|
||||
match transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("connection closed should be sent")
|
||||
{
|
||||
TransportEvent::ConnectionClosed {
|
||||
connection_id: closed_connection_id,
|
||||
} => assert_eq!(closed_connection_id, connection_id),
|
||||
other => panic!("expected connection closed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_cancels_blocked_outbound_forwarding() {
|
||||
let (server_event_tx, _server_event_rx) = mpsc::channel(1);
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let mut client_tracker =
|
||||
ClientTracker::new(server_event_tx.clone(), transport_event_tx, &shutdown_token);
|
||||
|
||||
server_event_tx
|
||||
.send(QueuedServerEnvelope {
|
||||
event: ServerEvent::Pong {
|
||||
status: PongStatus::Unknown,
|
||||
},
|
||||
client_id: ClientId("queued-client".to_string()),
|
||||
stream_id: StreamId("queued-stream".to_string()),
|
||||
write_complete_tx: None,
|
||||
})
|
||||
.await
|
||||
.expect("server event queue should accept prefill");
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope("client-1"))
|
||||
.await
|
||||
.expect("initialize should open client");
|
||||
|
||||
let writer = match transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("connection opened should be sent")
|
||||
{
|
||||
TransportEvent::ConnectionOpened { writer, .. } => writer,
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
let _ = transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("initialize should be forwarded");
|
||||
|
||||
writer
|
||||
.send(QueuedOutgoingMessage::new(
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification {
|
||||
summary: "test".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
},
|
||||
)),
|
||||
))
|
||||
.await
|
||||
.expect("writer should accept queued message");
|
||||
|
||||
timeout(Duration::from_secs(1), client_tracker.shutdown())
|
||||
.await
|
||||
.expect("shutdown should not hang on blocked server forwarding");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initialize_with_new_stream_id_opens_new_connection_for_same_client() {
|
||||
let (server_event_tx, _server_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let mut client_tracker =
|
||||
ClientTracker::new(server_event_tx, transport_event_tx, &shutdown_token);
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope_with_stream_id(
|
||||
"client-1",
|
||||
Some("stream-1"),
|
||||
))
|
||||
.await
|
||||
.expect("first initialize should open client");
|
||||
let first_connection_id = match transport_event_rx.recv().await.expect("open event") {
|
||||
TransportEvent::ConnectionOpened { connection_id, .. } => connection_id,
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
let _ = transport_event_rx.recv().await.expect("initialize event");
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope_with_stream_id(
|
||||
"client-1",
|
||||
Some("stream-2"),
|
||||
))
|
||||
.await
|
||||
.expect("second initialize should open client");
|
||||
let second_connection_id = match transport_event_rx.recv().await.expect("open event") {
|
||||
TransportEvent::ConnectionOpened { connection_id, .. } => connection_id,
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
|
||||
assert_ne!(first_connection_id, second_connection_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn legacy_initialize_without_stream_id_resets_inbound_seq_id() {
|
||||
let (server_event_tx, _server_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let mut client_tracker =
|
||||
ClientTracker::new(server_event_tx, transport_event_tx, &shutdown_token);
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope("client-1"))
|
||||
.await
|
||||
.expect("initialize should open client");
|
||||
let connection_id = match transport_event_rx.recv().await.expect("open event") {
|
||||
TransportEvent::ConnectionOpened { connection_id, .. } => connection_id,
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
let _ = transport_event_rx.recv().await.expect("initialize event");
|
||||
|
||||
client_tracker
|
||||
.handle_message(ClientEnvelope {
|
||||
event: ClientEvent::ClientMessage {
|
||||
message: JSONRPCMessage::Notification(
|
||||
codex_app_server_protocol::JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
},
|
||||
),
|
||||
},
|
||||
client_id: ClientId("client-1".to_string()),
|
||||
stream_id: None,
|
||||
seq_id: Some(0),
|
||||
cursor: None,
|
||||
})
|
||||
.await
|
||||
.expect("legacy followup should be forwarded");
|
||||
|
||||
match transport_event_rx.recv().await.expect("followup event") {
|
||||
TransportEvent::IncomingMessage {
|
||||
connection_id: incoming_connection_id,
|
||||
..
|
||||
} => assert_eq!(incoming_connection_id, connection_id),
|
||||
other => panic!("expected incoming message, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,514 +0,0 @@
|
||||
use super::protocol::EnrollRemoteServerRequest;
|
||||
use super::protocol::EnrollRemoteServerResponse;
|
||||
use super::protocol::RemoteControlTarget;
|
||||
use axum::http::HeaderMap;
|
||||
use codex_api::SharedAuthProvider;
|
||||
use codex_login::default_client::build_reqwest_client;
|
||||
use codex_state::RemoteControlEnrollmentRecord;
|
||||
use codex_state::StateRuntime;
|
||||
use gethostname::gethostname;
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
|
||||
const REMOTE_CONTROL_ENROLL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
|
||||
const REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES: usize = 4096;
|
||||
|
||||
const REQUEST_ID_HEADER: &str = "x-request-id";
|
||||
const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id";
|
||||
const CF_RAY_HEADER: &str = "cf-ray";
|
||||
pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) struct RemoteControlEnrollment {
|
||||
pub(super) account_id: String,
|
||||
pub(super) environment_id: String,
|
||||
pub(super) server_id: String,
|
||||
pub(super) server_name: String,
|
||||
}
|
||||
|
||||
pub(super) struct RemoteControlConnectionAuth {
|
||||
pub(super) auth_provider: SharedAuthProvider,
|
||||
pub(super) account_id: String,
|
||||
}
|
||||
|
||||
pub(super) async fn load_persisted_remote_control_enrollment(
|
||||
state_db: Option<&StateRuntime>,
|
||||
remote_control_target: &RemoteControlTarget,
|
||||
account_id: &str,
|
||||
app_server_client_name: Option<&str>,
|
||||
) -> io::Result<Option<RemoteControlEnrollment>> {
|
||||
let Some(state_db) = state_db else {
|
||||
return Err(io::Error::new(
|
||||
ErrorKind::NotFound,
|
||||
format!(
|
||||
"remote control enrollment cache unavailable because sqlite state db is disabled: websocket_url={}, account_id={}, app_server_client_name={:?}",
|
||||
remote_control_target.websocket_url, account_id, app_server_client_name
|
||||
),
|
||||
));
|
||||
};
|
||||
let enrollment = match state_db
|
||||
.get_remote_control_enrollment(
|
||||
&remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(enrollment) => enrollment,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"failed to load persisted remote control enrollment: websocket_url={}, account_id={}, app_server_client_name={:?}, err={err}",
|
||||
remote_control_target.websocket_url, account_id, app_server_client_name
|
||||
);
|
||||
return Err(io::Error::other(err));
|
||||
}
|
||||
};
|
||||
|
||||
match enrollment {
|
||||
Some(enrollment) => {
|
||||
info!(
|
||||
"reusing persisted remote control enrollment: websocket_url={}, account_id={}, app_server_client_name={:?}, server_id={}, environment_id={}",
|
||||
remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
enrollment.server_id,
|
||||
enrollment.environment_id
|
||||
);
|
||||
Ok(Some(RemoteControlEnrollment {
|
||||
account_id: enrollment.account_id,
|
||||
environment_id: enrollment.environment_id,
|
||||
server_id: enrollment.server_id,
|
||||
server_name: enrollment.server_name,
|
||||
}))
|
||||
}
|
||||
None => {
|
||||
info!(
|
||||
"no persisted remote control enrollment found: websocket_url={}, account_id={}, app_server_client_name={:?}",
|
||||
remote_control_target.websocket_url, account_id, app_server_client_name
|
||||
);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn update_persisted_remote_control_enrollment(
|
||||
state_db: Option<&StateRuntime>,
|
||||
remote_control_target: &RemoteControlTarget,
|
||||
account_id: &str,
|
||||
app_server_client_name: Option<&str>,
|
||||
enrollment: Option<&RemoteControlEnrollment>,
|
||||
) -> io::Result<()> {
|
||||
let Some(state_db) = state_db else {
|
||||
return Err(io::Error::new(
|
||||
ErrorKind::NotFound,
|
||||
format!(
|
||||
"remote control enrollment persistence unavailable because sqlite state db is disabled: websocket_url={}, account_id={}, app_server_client_name={:?}, has_enrollment={}",
|
||||
remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
enrollment.is_some()
|
||||
),
|
||||
));
|
||||
};
|
||||
if let &Some(enrollment) = &enrollment
|
||||
&& enrollment.account_id != account_id
|
||||
{
|
||||
return Err(io::Error::other(format!(
|
||||
"enrollment account_id does not match expected account_id `{account_id}`"
|
||||
)));
|
||||
}
|
||||
|
||||
if let Some(enrollment) = enrollment {
|
||||
state_db
|
||||
.upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord {
|
||||
websocket_url: remote_control_target.websocket_url.clone(),
|
||||
account_id: account_id.to_string(),
|
||||
app_server_client_name: app_server_client_name.map(str::to_string),
|
||||
server_id: enrollment.server_id.clone(),
|
||||
environment_id: enrollment.environment_id.clone(),
|
||||
server_name: enrollment.server_name.clone(),
|
||||
})
|
||||
.await
|
||||
.map_err(io::Error::other)?;
|
||||
info!(
|
||||
"persisted remote control enrollment: websocket_url={}, account_id={}, app_server_client_name={:?}, server_id={}, environment_id={}",
|
||||
remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
enrollment.server_id,
|
||||
enrollment.environment_id
|
||||
);
|
||||
Ok(())
|
||||
} else {
|
||||
let rows_affected = state_db
|
||||
.delete_remote_control_enrollment(
|
||||
&remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
)
|
||||
.await
|
||||
.map_err(io::Error::other)?;
|
||||
info!(
|
||||
"cleared persisted remote control enrollment: websocket_url={}, account_id={}, app_server_client_name={:?}, rows_affected={rows_affected}",
|
||||
remote_control_target.websocket_url, account_id, app_server_client_name
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn preview_remote_control_response_body(body: &[u8]) -> String {
|
||||
let body = String::from_utf8_lossy(body);
|
||||
let trimmed = body.trim();
|
||||
if trimmed.is_empty() {
|
||||
return "<empty>".to_string();
|
||||
}
|
||||
if trimmed.len() <= REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES {
|
||||
return trimmed.to_string();
|
||||
}
|
||||
|
||||
let mut cut = REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES;
|
||||
while !trimmed.is_char_boundary(cut) {
|
||||
cut = cut.saturating_sub(1);
|
||||
}
|
||||
let mut truncated = trimmed[..cut].to_string();
|
||||
truncated.push_str("...");
|
||||
truncated
|
||||
}
|
||||
|
||||
pub(crate) fn format_headers(headers: &HeaderMap) -> String {
|
||||
let request_id_str = headers
|
||||
.get(REQUEST_ID_HEADER)
|
||||
.or_else(|| headers.get(OAI_REQUEST_ID_HEADER))
|
||||
.map(|value| value.to_str().unwrap_or("<invalid utf-8>").to_owned())
|
||||
.unwrap_or_else(|| "<none>".to_owned());
|
||||
let cf_ray_str = headers
|
||||
.get(CF_RAY_HEADER)
|
||||
.map(|value| value.to_str().unwrap_or("<invalid utf-8>").to_owned())
|
||||
.unwrap_or_else(|| "<none>".to_owned());
|
||||
format!("request-id: {request_id_str}, cf-ray: {cf_ray_str}")
|
||||
}
|
||||
|
||||
pub(super) async fn enroll_remote_control_server(
|
||||
remote_control_target: &RemoteControlTarget,
|
||||
auth: &RemoteControlConnectionAuth,
|
||||
) -> io::Result<RemoteControlEnrollment> {
|
||||
let enroll_url = &remote_control_target.enroll_url;
|
||||
let server_name = gethostname().to_string_lossy().trim().to_string();
|
||||
let request = EnrollRemoteServerRequest {
|
||||
name: server_name.clone(),
|
||||
os: std::env::consts::OS,
|
||||
arch: std::env::consts::ARCH,
|
||||
app_server_version: env!("CARGO_PKG_VERSION"),
|
||||
};
|
||||
let client = build_reqwest_client();
|
||||
let mut auth_headers = HeaderMap::new();
|
||||
auth.auth_provider.add_auth_headers(&mut auth_headers);
|
||||
let http_request = client
|
||||
.post(enroll_url)
|
||||
.timeout(REMOTE_CONTROL_ENROLL_TIMEOUT)
|
||||
.headers(auth_headers)
|
||||
.header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)
|
||||
.json(&request);
|
||||
|
||||
let response = http_request.send().await.map_err(|err| {
|
||||
io::Error::other(format!(
|
||||
"failed to enroll remote control server at `{enroll_url}`: {err}"
|
||||
))
|
||||
})?;
|
||||
let headers = response.headers().clone();
|
||||
let status = response.status();
|
||||
let body = response.bytes().await.map_err(|err| {
|
||||
io::Error::other(format!(
|
||||
"failed to read remote control enrollment response from `{enroll_url}`: {err}"
|
||||
))
|
||||
})?;
|
||||
let body_preview = preview_remote_control_response_body(&body);
|
||||
if !status.is_success() {
|
||||
let headers_str = format_headers(&headers);
|
||||
let error_kind = if matches!(status.as_u16(), 401 | 403) {
|
||||
ErrorKind::PermissionDenied
|
||||
} else {
|
||||
ErrorKind::Other
|
||||
};
|
||||
return Err(io::Error::new(
|
||||
error_kind,
|
||||
format!(
|
||||
"remote control server enrollment failed at `{enroll_url}`: HTTP {status}, {headers_str}, body: {body_preview}"
|
||||
),
|
||||
));
|
||||
}
|
||||
|
||||
let enrollment = serde_json::from_slice::<EnrollRemoteServerResponse>(&body).map_err(|err| {
|
||||
let headers_str = format_headers(&headers);
|
||||
io::Error::other(format!(
|
||||
"failed to parse remote control enrollment response from `{enroll_url}`: HTTP {status}, {headers_str}, body: {body_preview}, decode error: {err}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(RemoteControlEnrollment {
|
||||
account_id: auth.account_id.clone(),
|
||||
environment_id: enrollment.environment_id,
|
||||
server_id: enrollment.server_id,
|
||||
server_name,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::transport::remote_control::protocol::normalize_remote_control_url;
|
||||
use codex_state::StateRuntime;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tempfile::TempDir;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
|
||||
async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc<StateRuntime> {
|
||||
StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string())
|
||||
.await
|
||||
.expect("state runtime should initialize")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() {
|
||||
let codex_home = TempDir::new().expect("temp dir should create");
|
||||
let state_db = remote_control_state_runtime(&codex_home).await;
|
||||
let first_target = normalize_remote_control_url("https://chatgpt.com/remote/control")
|
||||
.expect("first target should parse");
|
||||
let second_target =
|
||||
normalize_remote_control_url("https://api.chatgpt-staging.com/other/control")
|
||||
.expect("second target should parse");
|
||||
let first_enrollment = RemoteControlEnrollment {
|
||||
account_id: "account-a".to_string(),
|
||||
environment_id: "env_first".to_string(),
|
||||
server_id: "srv_e_first".to_string(),
|
||||
server_name: "first-server".to_string(),
|
||||
};
|
||||
let second_enrollment = RemoteControlEnrollment {
|
||||
account_id: "account-a".to_string(),
|
||||
environment_id: "env_second".to_string(),
|
||||
server_id: "srv_e_second".to_string(),
|
||||
server_name: "second-server".to_string(),
|
||||
};
|
||||
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
Some("desktop-client"),
|
||||
Some(&first_enrollment),
|
||||
)
|
||||
.await
|
||||
.expect("first enrollment should persist");
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&second_target,
|
||||
"account-a",
|
||||
Some("desktop-client"),
|
||||
Some(&second_enrollment),
|
||||
)
|
||||
.await
|
||||
.expect("second enrollment should persist");
|
||||
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
Some("desktop-client"),
|
||||
)
|
||||
.await
|
||||
.expect("first enrollment should load"),
|
||||
Some(first_enrollment.clone())
|
||||
);
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-b",
|
||||
Some("desktop-client"),
|
||||
)
|
||||
.await
|
||||
.expect("missing account should load"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&second_target,
|
||||
"account-a",
|
||||
Some("desktop-client"),
|
||||
)
|
||||
.await
|
||||
.expect("second enrollment should load"),
|
||||
Some(second_enrollment)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn clearing_persisted_remote_control_enrollment_removes_only_matching_entry() {
|
||||
let codex_home = TempDir::new().expect("temp dir should create");
|
||||
let state_db = remote_control_state_runtime(&codex_home).await;
|
||||
let first_target = normalize_remote_control_url("https://chatgpt.com/remote/control")
|
||||
.expect("first target should parse");
|
||||
let second_target =
|
||||
normalize_remote_control_url("https://api.chatgpt-staging.com/other/control")
|
||||
.expect("second target should parse");
|
||||
let first_enrollment = RemoteControlEnrollment {
|
||||
account_id: "account-a".to_string(),
|
||||
environment_id: "env_first".to_string(),
|
||||
server_id: "srv_e_first".to_string(),
|
||||
server_name: "first-server".to_string(),
|
||||
};
|
||||
let second_enrollment = RemoteControlEnrollment {
|
||||
account_id: "account-a".to_string(),
|
||||
environment_id: "env_second".to_string(),
|
||||
server_id: "srv_e_second".to_string(),
|
||||
server_name: "second-server".to_string(),
|
||||
};
|
||||
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
Some(&first_enrollment),
|
||||
)
|
||||
.await
|
||||
.expect("first enrollment should persist");
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&second_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
Some(&second_enrollment),
|
||||
)
|
||||
.await
|
||||
.expect("second enrollment should persist");
|
||||
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
/*enrollment*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("matching enrollment should clear");
|
||||
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("cleared enrollment should load"),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&second_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("remaining enrollment should load"),
|
||||
Some(second_enrollment)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enroll_remote_control_server_parse_failure_includes_response_body() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let remote_control_url = format!(
|
||||
"http://127.0.0.1:{}/backend-api/",
|
||||
listener
|
||||
.local_addr()
|
||||
.expect("listener should have a local addr")
|
||||
.port()
|
||||
);
|
||||
let remote_control_target =
|
||||
normalize_remote_control_url(&remote_control_url).expect("target should parse");
|
||||
let enroll_url = remote_control_target.enroll_url.clone();
|
||||
let response_body = json!({
|
||||
"error": "not enrolled",
|
||||
});
|
||||
let expected_body = response_body.to_string();
|
||||
let server_task = tokio::spawn(async move {
|
||||
let stream = accept_http_request(&listener).await;
|
||||
respond_with_json(stream, response_body).await;
|
||||
});
|
||||
|
||||
let err = enroll_remote_control_server(
|
||||
&remote_control_target,
|
||||
&RemoteControlConnectionAuth {
|
||||
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
|
||||
account_id: "account_id".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("invalid response should fail to parse");
|
||||
|
||||
server_task.await.expect("server task should succeed");
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
format!(
|
||||
"failed to parse remote control enrollment response from `{enroll_url}`: HTTP 200 OK, request-id: <none>, cf-ray: <none>, body: {expected_body}, decode error: missing field `server_id` at line 1 column {}",
|
||||
expected_body.len()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
async fn accept_http_request(listener: &TcpListener) -> TcpStream {
|
||||
let (stream, _) = timeout(Duration::from_secs(5), listener.accept())
|
||||
.await
|
||||
.expect("HTTP request should arrive in time")
|
||||
.expect("listener accept should succeed");
|
||||
let mut reader = BufReader::new(stream);
|
||||
|
||||
let mut request_line = String::new();
|
||||
reader
|
||||
.read_line(&mut request_line)
|
||||
.await
|
||||
.expect("request line should read");
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
reader
|
||||
.read_line(&mut line)
|
||||
.await
|
||||
.expect("header line should read");
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
reader.into_inner()
|
||||
}
|
||||
|
||||
async fn respond_with_json(mut stream: TcpStream, body: serde_json::Value) {
|
||||
let body = body.to_string();
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
);
|
||||
stream
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("response should write");
|
||||
stream.flush().await.expect("response should flush");
|
||||
}
|
||||
}
|
||||
@@ -1,128 +0,0 @@
|
||||
mod client_tracker;
|
||||
mod enroll;
|
||||
mod protocol;
|
||||
mod segment;
|
||||
mod websocket;
|
||||
|
||||
use crate::transport::remote_control::websocket::RemoteControlChannels;
|
||||
use crate::transport::remote_control::websocket::RemoteControlStatusPublisher;
|
||||
use crate::transport::remote_control::websocket::RemoteControlWebsocket;
|
||||
|
||||
pub use self::protocol::ClientId;
|
||||
use self::protocol::ServerEvent;
|
||||
use self::protocol::StreamId;
|
||||
use self::protocol::normalize_remote_control_url;
|
||||
use super::CHANNEL_CAPACITY;
|
||||
use super::TransportEvent;
|
||||
use super::next_connection_id;
|
||||
use codex_app_server_protocol::RemoteControlConnectionStatus;
|
||||
use codex_app_server_protocol::RemoteControlStatusChangedNotification;
|
||||
use codex_login::AuthManager;
|
||||
use codex_state::StateRuntime;
|
||||
use std::io;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::sync::watch;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
pub(super) struct QueuedServerEnvelope {
|
||||
pub(super) event: ServerEvent,
|
||||
pub(super) client_id: ClientId,
|
||||
pub(super) stream_id: StreamId,
|
||||
pub(super) write_complete_tx: Option<oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct RemoteControlHandle {
|
||||
enabled_tx: Arc<watch::Sender<bool>>,
|
||||
status_tx: Arc<watch::Sender<RemoteControlStatusChangedNotification>>,
|
||||
state_db_available: bool,
|
||||
}
|
||||
|
||||
impl RemoteControlHandle {
|
||||
pub(crate) fn set_enabled(&self, enabled: bool) {
|
||||
let requested_enabled = enabled;
|
||||
let enabled = enabled && self.state_db_available;
|
||||
if requested_enabled && !self.state_db_available {
|
||||
warn!("remote control cannot be enabled because sqlite state db is unavailable");
|
||||
}
|
||||
self.enabled_tx.send_if_modified(|state| {
|
||||
let changed = *state != enabled;
|
||||
*state = enabled;
|
||||
changed
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn status_receiver(
|
||||
&self,
|
||||
) -> watch::Receiver<RemoteControlStatusChangedNotification> {
|
||||
self.status_tx.subscribe()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn start_remote_control(
|
||||
remote_control_url: String,
|
||||
state_db: Option<Arc<StateRuntime>>,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: CancellationToken,
|
||||
app_server_client_name_rx: Option<oneshot::Receiver<String>>,
|
||||
initial_enabled: bool,
|
||||
) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> {
|
||||
let state_db_available = state_db.is_some();
|
||||
let requested_initial_enabled = initial_enabled;
|
||||
let initial_enabled = initial_enabled && state_db_available;
|
||||
if requested_initial_enabled && !state_db_available {
|
||||
warn!("remote control disabled because sqlite state db is unavailable");
|
||||
}
|
||||
let remote_control_target = if initial_enabled {
|
||||
Some(normalize_remote_control_url(&remote_control_url)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let (enabled_tx, enabled_rx) = watch::channel(initial_enabled);
|
||||
let initial_status = RemoteControlStatusChangedNotification {
|
||||
status: if initial_enabled {
|
||||
RemoteControlConnectionStatus::Connecting
|
||||
} else {
|
||||
RemoteControlConnectionStatus::Disabled
|
||||
},
|
||||
environment_id: None,
|
||||
};
|
||||
let (status_tx, _status_rx) = watch::channel(initial_status);
|
||||
let status_publisher = RemoteControlStatusPublisher::new(status_tx.clone());
|
||||
let join_handle = tokio::spawn(async move {
|
||||
RemoteControlWebsocket::new(
|
||||
remote_control_url,
|
||||
remote_control_target,
|
||||
state_db,
|
||||
auth_manager,
|
||||
RemoteControlChannels {
|
||||
transport_event_tx,
|
||||
status_publisher,
|
||||
},
|
||||
shutdown_token,
|
||||
enabled_rx,
|
||||
)
|
||||
.run(app_server_client_name_rx)
|
||||
.await;
|
||||
});
|
||||
|
||||
Ok((
|
||||
join_handle,
|
||||
RemoteControlHandle {
|
||||
enabled_tx: Arc::new(enabled_tx),
|
||||
status_tx: Arc::new(status_tx),
|
||||
state_db_available,
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod segment_tests;
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -1,277 +0,0 @@
|
||||
use crate::outgoing_message::OutgoingMessage;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use url::Host;
|
||||
use url::Url;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) struct RemoteControlTarget {
|
||||
pub(super) websocket_url: String,
|
||||
pub(super) enroll_url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(super) struct EnrollRemoteServerRequest {
|
||||
pub(super) name: String,
|
||||
pub(super) os: &'static str,
|
||||
pub(super) arch: &'static str,
|
||||
pub(super) app_server_version: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(super) struct EnrollRemoteServerResponse {
|
||||
pub(super) server_id: String,
|
||||
pub(super) environment_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct ClientId(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct StreamId(pub String);
|
||||
|
||||
impl StreamId {
|
||||
pub fn new_random() -> Self {
|
||||
Self(uuid::Uuid::now_v7().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ClientEvent {
|
||||
ClientMessage {
|
||||
message: JSONRPCMessage,
|
||||
},
|
||||
ClientMessageChunk {
|
||||
segment_id: usize,
|
||||
segment_count: usize,
|
||||
message_size_bytes: usize,
|
||||
message_chunk_base64: String,
|
||||
},
|
||||
/// Backend-generated acknowledgement for all server envelopes addressed to
|
||||
/// `client_id` and `stream_id` whose envelope `seq_id` is less than or equal
|
||||
/// to this ack's `seq_id`. Chunk acknowledgements carry `segment_id` so the
|
||||
/// sender can retain only the still-unacked wire chunks on reconnect.
|
||||
Ack {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
segment_id: Option<usize>,
|
||||
},
|
||||
Ping,
|
||||
ClientClosed,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) struct ClientEnvelope {
|
||||
#[serde(flatten)]
|
||||
pub(crate) event: ClientEvent,
|
||||
#[serde(rename = "client_id")]
|
||||
pub(crate) client_id: ClientId,
|
||||
#[serde(rename = "stream_id", skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) stream_id: Option<StreamId>,
|
||||
/// For `Ack`, this is the backend-generated per-stream cursor over
|
||||
/// `ServerEnvelope.seq_id`.
|
||||
#[serde(rename = "seq_id", skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) seq_id: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) cursor: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum PongStatus {
|
||||
Active,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ServerEvent {
|
||||
ServerMessage {
|
||||
message: Box<OutgoingMessage>,
|
||||
},
|
||||
ServerMessageChunk {
|
||||
segment_id: usize,
|
||||
segment_count: usize,
|
||||
message_size_bytes: usize,
|
||||
message_chunk_base64: String,
|
||||
},
|
||||
#[allow(dead_code)]
|
||||
Ack,
|
||||
Pong {
|
||||
status: PongStatus,
|
||||
},
|
||||
}
|
||||
|
||||
impl ServerEvent {
|
||||
pub(crate) fn segment_id(&self) -> Option<usize> {
|
||||
match self {
|
||||
Self::ServerMessageChunk { segment_id, .. } => Some(*segment_id),
|
||||
Self::ServerMessage { .. } | Self::Ack | Self::Pong { .. } => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) struct ServerEnvelope {
|
||||
#[serde(flatten)]
|
||||
pub(crate) event: ServerEvent,
|
||||
#[serde(rename = "client_id")]
|
||||
pub(crate) client_id: ClientId,
|
||||
#[serde(rename = "stream_id")]
|
||||
pub(crate) stream_id: StreamId,
|
||||
#[serde(rename = "seq_id")]
|
||||
pub(crate) seq_id: u64,
|
||||
}
|
||||
|
||||
fn is_allowed_remote_control_chatgpt_host(host: &Option<Host<&str>>) -> bool {
|
||||
let Some(Host::Domain(host)) = *host else {
|
||||
return false;
|
||||
};
|
||||
host == "chatgpt.com"
|
||||
|| host == "chatgpt-staging.com"
|
||||
|| host.ends_with(".chatgpt.com")
|
||||
|| host.ends_with(".chatgpt-staging.com")
|
||||
}
|
||||
|
||||
fn is_localhost(host: &Option<Host<&str>>) -> bool {
|
||||
match host {
|
||||
Some(Host::Domain("localhost")) => true,
|
||||
Some(Host::Ipv4(ip)) => ip.is_loopback(),
|
||||
Some(Host::Ipv6(ip)) => ip.is_loopback(),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn normalize_remote_control_url(
|
||||
remote_control_url: &str,
|
||||
) -> io::Result<RemoteControlTarget> {
|
||||
let map_url_parse_error = |err: url::ParseError| -> io::Error {
|
||||
io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
format!("invalid remote control URL `{remote_control_url}`: {err}"),
|
||||
)
|
||||
};
|
||||
let map_scheme_error = |_: ()| -> io::Error {
|
||||
io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"invalid remote control URL `{remote_control_url}`; expected HTTPS URL for chatgpt.com or chatgpt-staging.com, or HTTP/HTTPS URL for localhost"
|
||||
),
|
||||
)
|
||||
};
|
||||
|
||||
let mut remote_control_url = Url::parse(remote_control_url).map_err(map_url_parse_error)?;
|
||||
if !remote_control_url.path().ends_with('/') {
|
||||
let normalized_path = format!("{}/", remote_control_url.path());
|
||||
remote_control_url.set_path(&normalized_path);
|
||||
}
|
||||
|
||||
let enroll_url = remote_control_url
|
||||
.join("wham/remote/control/server/enroll")
|
||||
.map_err(map_url_parse_error)?;
|
||||
let mut websocket_url = remote_control_url
|
||||
.join("wham/remote/control/server")
|
||||
.map_err(map_url_parse_error)?;
|
||||
let host = enroll_url.host();
|
||||
match enroll_url.scheme() {
|
||||
"https" if is_localhost(&host) || is_allowed_remote_control_chatgpt_host(&host) => {
|
||||
websocket_url.set_scheme("wss").map_err(map_scheme_error)?;
|
||||
}
|
||||
"http" if is_localhost(&host) => {
|
||||
websocket_url.set_scheme("ws").map_err(map_scheme_error)?;
|
||||
}
|
||||
_ => return Err(map_scheme_error(())),
|
||||
}
|
||||
|
||||
Ok(RemoteControlTarget {
|
||||
websocket_url: websocket_url.to_string(),
|
||||
enroll_url: enroll_url.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn normalize_remote_control_url_accepts_chatgpt_https_urls() {
|
||||
assert_eq!(
|
||||
normalize_remote_control_url("https://chatgpt.com/backend-api")
|
||||
.expect("chatgpt.com URL should normalize"),
|
||||
RemoteControlTarget {
|
||||
websocket_url: "wss://chatgpt.com/backend-api/wham/remote/control/server"
|
||||
.to_string(),
|
||||
enroll_url: "https://chatgpt.com/backend-api/wham/remote/control/server/enroll"
|
||||
.to_string(),
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_remote_control_url("https://api.chatgpt-staging.com/backend-api")
|
||||
.expect("chatgpt-staging.com subdomain URL should normalize"),
|
||||
RemoteControlTarget {
|
||||
websocket_url:
|
||||
"wss://api.chatgpt-staging.com/backend-api/wham/remote/control/server"
|
||||
.to_string(),
|
||||
enroll_url:
|
||||
"https://api.chatgpt-staging.com/backend-api/wham/remote/control/server/enroll"
|
||||
.to_string(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_remote_control_url_accepts_localhost_urls() {
|
||||
assert_eq!(
|
||||
normalize_remote_control_url("http://localhost:8080/backend-api")
|
||||
.expect("localhost http URL should normalize"),
|
||||
RemoteControlTarget {
|
||||
websocket_url: "ws://localhost:8080/backend-api/wham/remote/control/server"
|
||||
.to_string(),
|
||||
enroll_url: "http://localhost:8080/backend-api/wham/remote/control/server/enroll"
|
||||
.to_string(),
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_remote_control_url("https://localhost:8443/backend-api")
|
||||
.expect("localhost https URL should normalize"),
|
||||
RemoteControlTarget {
|
||||
websocket_url: "wss://localhost:8443/backend-api/wham/remote/control/server"
|
||||
.to_string(),
|
||||
enroll_url: "https://localhost:8443/backend-api/wham/remote/control/server/enroll"
|
||||
.to_string(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_remote_control_url_rejects_unsupported_urls() {
|
||||
for remote_control_url in [
|
||||
"http://chatgpt.com/backend-api",
|
||||
"http://example.com/backend-api",
|
||||
"https://example.com/backend-api",
|
||||
"https://chat.openai.com/backend-api",
|
||||
"https://chatgpt.com.evil.com/backend-api",
|
||||
"https://evilchatgpt.com/backend-api",
|
||||
"https://foo.localhost/backend-api",
|
||||
] {
|
||||
let err = normalize_remote_control_url(remote_control_url)
|
||||
.expect_err("unsupported URL should be rejected");
|
||||
|
||||
assert_eq!(err.kind(), ErrorKind::InvalidInput);
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
format!(
|
||||
"invalid remote control URL `{remote_control_url}`; expected HTTPS URL for chatgpt.com or chatgpt-staging.com, or HTTP/HTTPS URL for localhost"
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,449 +0,0 @@
|
||||
use super::protocol::ClientEnvelope;
|
||||
use super::protocol::ClientEvent;
|
||||
use super::protocol::ClientId;
|
||||
use super::protocol::ServerEnvelope;
|
||||
use super::protocol::ServerEvent;
|
||||
use super::protocol::StreamId;
|
||||
use base64::DecodeSliceError;
|
||||
use base64::Engine;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use std::io::Write;
|
||||
use tokio::time::Instant;
|
||||
use tracing::warn;
|
||||
|
||||
pub(super) const REMOTE_CONTROL_SEGMENT_TARGET_BYTES: usize = 100 * 1024;
|
||||
pub(super) const REMOTE_CONTROL_SEGMENT_MAX_BYTES: usize = 150 * 1024;
|
||||
pub(super) const REMOTE_CONTROL_REASSEMBLED_MAX_BYTES: usize = 100 * 1024 * 1024;
|
||||
pub(super) const REMOTE_CONTROL_SEGMENT_COUNT_MAX: usize = 1024;
|
||||
const REMOTE_CONTROL_SEGMENT_ASSEMBLY_MAX_COUNT: usize = 128;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ClientSegmentAssembly {
|
||||
stream_id: StreamId,
|
||||
metadata: ClientSegmentMetadata,
|
||||
raw: Vec<u8>,
|
||||
next_segment_id: usize,
|
||||
last_chunk_seen_at: Instant,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct ClientSegmentMetadata {
|
||||
seq_id: u64,
|
||||
segment_count: usize,
|
||||
message_size_bytes: usize,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct ClientSegmentReassembler {
|
||||
assemblies: HashMap<ClientId, ClientSegmentAssembly>,
|
||||
}
|
||||
|
||||
pub(super) enum ClientSegmentObservation {
|
||||
Forward(Box<ClientEnvelope>),
|
||||
Pending,
|
||||
Dropped,
|
||||
}
|
||||
|
||||
impl ClientSegmentReassembler {
|
||||
pub(super) fn observe(&mut self, envelope: ClientEnvelope) -> ClientSegmentObservation {
|
||||
let ClientEvent::ClientMessageChunk {
|
||||
segment_id,
|
||||
segment_count,
|
||||
message_size_bytes,
|
||||
message_chunk_base64,
|
||||
} = &envelope.event
|
||||
else {
|
||||
return ClientSegmentObservation::Forward(Box::new(envelope));
|
||||
};
|
||||
let segment_id = *segment_id;
|
||||
let segment_count = *segment_count;
|
||||
let message_size_bytes = *message_size_bytes;
|
||||
|
||||
let Some(metadata) = ClientSegmentMetadata::from_envelope(&envelope) else {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping segmented remote-control client envelope without seq_id"
|
||||
);
|
||||
return ClientSegmentObservation::Dropped;
|
||||
};
|
||||
let Some(stream_id) = envelope.stream_id.clone() else {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping segmented remote-control client envelope without stream_id"
|
||||
);
|
||||
return ClientSegmentObservation::Dropped;
|
||||
};
|
||||
if self.should_ignore_chunk(&envelope.client_id, &stream_id, metadata.seq_id, segment_id) {
|
||||
return ClientSegmentObservation::Dropped;
|
||||
}
|
||||
if segment_count == 0
|
||||
|| segment_count > REMOTE_CONTROL_SEGMENT_COUNT_MAX
|
||||
|| segment_id >= segment_count
|
||||
|| message_size_bytes == 0
|
||||
|| message_size_bytes > REMOTE_CONTROL_REASSEMBLED_MAX_BYTES
|
||||
|| message_chunk_base64.is_empty()
|
||||
{
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping invalid segmented remote-control client envelope"
|
||||
);
|
||||
self.remove_assembly(&envelope.client_id, &stream_id);
|
||||
return ClientSegmentObservation::Dropped;
|
||||
}
|
||||
|
||||
let now = Instant::now();
|
||||
match self.assemblies.get(&envelope.client_id) {
|
||||
Some(assembly) if assembly.stream_id != stream_id => {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"resetting segmented remote-control client envelope after stream change"
|
||||
);
|
||||
self.assemblies.insert(
|
||||
envelope.client_id.clone(),
|
||||
ClientSegmentAssembly {
|
||||
stream_id: stream_id.clone(),
|
||||
metadata: metadata.clone(),
|
||||
raw: Vec::new(),
|
||||
next_segment_id: 0,
|
||||
last_chunk_seen_at: now,
|
||||
},
|
||||
);
|
||||
}
|
||||
Some(_) => {}
|
||||
None => {
|
||||
self.evict_assemblies_if_full();
|
||||
self.assemblies.insert(
|
||||
envelope.client_id.clone(),
|
||||
ClientSegmentAssembly {
|
||||
stream_id: stream_id.clone(),
|
||||
metadata: metadata.clone(),
|
||||
raw: Vec::new(),
|
||||
next_segment_id: 0,
|
||||
last_chunk_seen_at: now,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
let result = {
|
||||
let Some(assembly) = self.assemblies.get_mut(&envelope.client_id) else {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping segmented remote-control client envelope without assembly"
|
||||
);
|
||||
return ClientSegmentObservation::Dropped;
|
||||
};
|
||||
if metadata.seq_id < assembly.metadata.seq_id {
|
||||
AssemblyUpdate::Ignore
|
||||
} else if assembly.metadata != metadata {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"resetting segmented remote-control client envelope after metadata mismatch"
|
||||
);
|
||||
AssemblyUpdate::Drop
|
||||
} else if segment_id < assembly.next_segment_id {
|
||||
AssemblyUpdate::Pending
|
||||
} else if segment_id != assembly.next_segment_id {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping out-of-order segmented remote-control client envelope"
|
||||
);
|
||||
AssemblyUpdate::Drop
|
||||
} else {
|
||||
assembly.last_chunk_seen_at = now;
|
||||
let chunk_start = assembly.raw.len();
|
||||
let decoded_chunk_len = base64::decoded_len_estimate(message_chunk_base64.len());
|
||||
let chunk_end = usize::min(
|
||||
message_size_bytes,
|
||||
chunk_start.saturating_add(decoded_chunk_len),
|
||||
);
|
||||
assembly.raw.resize(chunk_end, 0);
|
||||
match base64::engine::general_purpose::STANDARD.decode_slice(
|
||||
message_chunk_base64.as_bytes(),
|
||||
&mut assembly.raw[chunk_start..],
|
||||
) {
|
||||
Ok(decoded_chunk_len) => {
|
||||
assembly.raw.truncate(chunk_start + decoded_chunk_len);
|
||||
assembly.next_segment_id += 1;
|
||||
if assembly.next_segment_id < segment_count {
|
||||
AssemblyUpdate::Pending
|
||||
} else if assembly.raw.len() != message_size_bytes {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping reassembled remote-control client envelope with mismatched size"
|
||||
);
|
||||
AssemblyUpdate::Drop
|
||||
} else {
|
||||
match serde_json::from_slice::<JSONRPCMessage>(&assembly.raw) {
|
||||
Ok(message) => AssemblyUpdate::Complete(message),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping invalid reassembled remote-control client envelope: {err}"
|
||||
);
|
||||
AssemblyUpdate::Drop
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(DecodeSliceError::OutputSliceTooSmall) => {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping segmented remote-control client envelope after size overflow"
|
||||
);
|
||||
AssemblyUpdate::Drop
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
client_id = envelope.client_id.0.as_str(),
|
||||
"dropping segmented remote-control client envelope with invalid base64: {err}"
|
||||
);
|
||||
AssemblyUpdate::Drop
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
match result {
|
||||
AssemblyUpdate::Pending => ClientSegmentObservation::Pending,
|
||||
AssemblyUpdate::Ignore => ClientSegmentObservation::Dropped,
|
||||
AssemblyUpdate::Drop => {
|
||||
self.remove_assembly(&envelope.client_id, &stream_id);
|
||||
ClientSegmentObservation::Dropped
|
||||
}
|
||||
AssemblyUpdate::Complete(message) => {
|
||||
self.remove_assembly(&envelope.client_id, &stream_id);
|
||||
ClientSegmentObservation::Forward(Box::new(ClientEnvelope {
|
||||
event: ClientEvent::ClientMessage { message },
|
||||
..envelope
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn invalidate_stream(&mut self, client_id: &ClientId, stream_id: &StreamId) {
|
||||
self.remove_assembly(client_id, stream_id);
|
||||
}
|
||||
|
||||
pub(super) fn invalidate_client(&mut self, client_id: &ClientId) {
|
||||
self.assemblies.remove(client_id);
|
||||
}
|
||||
|
||||
pub(super) fn should_ignore_chunk(
|
||||
&self,
|
||||
client_id: &ClientId,
|
||||
stream_id: &StreamId,
|
||||
seq_id: u64,
|
||||
segment_id: usize,
|
||||
) -> bool {
|
||||
self.assemblies.get(client_id).is_some_and(|assembly| {
|
||||
assembly.stream_id == *stream_id
|
||||
&& (seq_id < assembly.metadata.seq_id
|
||||
|| (seq_id == assembly.metadata.seq_id
|
||||
&& segment_id < assembly.next_segment_id))
|
||||
})
|
||||
}
|
||||
|
||||
fn remove_assembly(&mut self, client_id: &ClientId, stream_id: &StreamId) {
|
||||
if self
|
||||
.assemblies
|
||||
.get(client_id)
|
||||
.is_some_and(|assembly| &assembly.stream_id == stream_id)
|
||||
{
|
||||
self.assemblies.remove(client_id);
|
||||
}
|
||||
}
|
||||
|
||||
fn evict_assemblies_if_full(&mut self) {
|
||||
while self.assemblies.len() >= REMOTE_CONTROL_SEGMENT_ASSEMBLY_MAX_COUNT {
|
||||
let Some(client_id) = self
|
||||
.assemblies
|
||||
.iter()
|
||||
.min_by_key(|(_, assembly)| assembly.last_chunk_seen_at)
|
||||
.map(|(client_id, _)| client_id.clone())
|
||||
else {
|
||||
return;
|
||||
};
|
||||
self.assemblies.remove(&client_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum AssemblyUpdate {
|
||||
Pending,
|
||||
Ignore,
|
||||
Drop,
|
||||
Complete(JSONRPCMessage),
|
||||
}
|
||||
|
||||
impl ClientSegmentMetadata {
|
||||
fn from_envelope(envelope: &ClientEnvelope) -> Option<Self> {
|
||||
let ClientEvent::ClientMessageChunk {
|
||||
segment_count,
|
||||
message_size_bytes,
|
||||
..
|
||||
} = &envelope.event
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
Some(Self {
|
||||
seq_id: envelope.seq_id?,
|
||||
segment_count: *segment_count,
|
||||
message_size_bytes: *message_size_bytes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn split_server_envelope_for_transport(
|
||||
envelope: ServerEnvelope,
|
||||
) -> io::Result<Vec<ServerEnvelope>> {
|
||||
if !matches!(envelope.event, ServerEvent::ServerMessage { .. }) {
|
||||
return Ok(vec![envelope]);
|
||||
}
|
||||
|
||||
let envelope_size_bytes = serialized_len(&envelope)?;
|
||||
if envelope_size_bytes <= REMOTE_CONTROL_SEGMENT_MAX_BYTES {
|
||||
return Ok(vec![envelope]);
|
||||
}
|
||||
|
||||
let ServerEvent::ServerMessage { message } = envelope.event.clone() else {
|
||||
unreachable!("server message variant checked above");
|
||||
};
|
||||
let raw = serde_json::to_vec(message.as_ref()).map_err(io::Error::other)?;
|
||||
let message_size_bytes = raw.len();
|
||||
if message_size_bytes > REMOTE_CONTROL_REASSEMBLED_MAX_BYTES {
|
||||
warn!("dropping remote-control server envelope that exceeds reassembled size limit");
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let minimal_segment_count =
|
||||
usize::min(message_size_bytes.max(1), REMOTE_CONTROL_SEGMENT_COUNT_MAX);
|
||||
let minimal_chunk = &raw[..usize::min(raw.len(), 1)];
|
||||
if serialized_chunk_len(
|
||||
&envelope,
|
||||
/*segment_id*/ 0,
|
||||
minimal_segment_count,
|
||||
message_size_bytes,
|
||||
minimal_chunk,
|
||||
)? > REMOTE_CONTROL_SEGMENT_MAX_BYTES
|
||||
{
|
||||
warn!("dropping remote-control server envelope that cannot fit within segment size limit");
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut segment_count = usize::max(
|
||||
2,
|
||||
message_size_bytes.div_ceil(REMOTE_CONTROL_SEGMENT_TARGET_BYTES),
|
||||
);
|
||||
loop {
|
||||
let chunk_size = usize::max(1, message_size_bytes.div_ceil(segment_count));
|
||||
segment_count = message_size_bytes.div_ceil(chunk_size);
|
||||
let segments_fit = raw
|
||||
.chunks(chunk_size)
|
||||
.enumerate()
|
||||
.all(|(segment_id, chunk)| {
|
||||
serialized_chunk_len(
|
||||
&envelope,
|
||||
segment_id,
|
||||
segment_count,
|
||||
message_size_bytes,
|
||||
chunk,
|
||||
)
|
||||
.is_ok_and(|size| size <= REMOTE_CONTROL_SEGMENT_MAX_BYTES)
|
||||
});
|
||||
if segments_fit {
|
||||
return raw
|
||||
.chunks(chunk_size)
|
||||
.enumerate()
|
||||
.map(|(segment_id, chunk)| {
|
||||
build_chunk_envelope(
|
||||
&envelope,
|
||||
segment_id,
|
||||
segment_count,
|
||||
message_size_bytes,
|
||||
chunk,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
if chunk_size == 1 {
|
||||
warn!(
|
||||
"dropping remote-control server envelope that cannot fit within segment size limit"
|
||||
);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let next_segment_count = segment_count + 1;
|
||||
let next_chunk_size = usize::max(1, message_size_bytes.div_ceil(next_segment_count));
|
||||
segment_count = if next_chunk_size == chunk_size {
|
||||
message_size_bytes
|
||||
} else {
|
||||
next_segment_count
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
fn serialized_chunk_len(
|
||||
envelope: &ServerEnvelope,
|
||||
segment_id: usize,
|
||||
segment_count: usize,
|
||||
message_size_bytes: usize,
|
||||
chunk: &[u8],
|
||||
) -> io::Result<usize> {
|
||||
serialized_len(&build_chunk_envelope(
|
||||
envelope,
|
||||
segment_id,
|
||||
segment_count,
|
||||
message_size_bytes,
|
||||
chunk,
|
||||
)?)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct CountingWriter {
|
||||
len: usize,
|
||||
}
|
||||
|
||||
impl Write for CountingWriter {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.len += buf.len();
|
||||
Ok(buf.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn serialized_len(value: &impl serde::Serialize) -> io::Result<usize> {
|
||||
let mut writer = CountingWriter::default();
|
||||
serde_json::to_writer(&mut writer, value).map_err(io::Error::other)?;
|
||||
Ok(writer.len)
|
||||
}
|
||||
|
||||
fn build_chunk_envelope(
|
||||
envelope: &ServerEnvelope,
|
||||
segment_id: usize,
|
||||
segment_count: usize,
|
||||
message_size_bytes: usize,
|
||||
chunk: &[u8],
|
||||
) -> io::Result<ServerEnvelope> {
|
||||
if segment_count > REMOTE_CONTROL_SEGMENT_COUNT_MAX {
|
||||
return Err(io::Error::new(
|
||||
ErrorKind::InvalidData,
|
||||
"remote-control segment count exceeds maximum",
|
||||
));
|
||||
}
|
||||
Ok(ServerEnvelope {
|
||||
event: ServerEvent::ServerMessageChunk {
|
||||
segment_id,
|
||||
segment_count,
|
||||
message_size_bytes,
|
||||
message_chunk_base64: base64::engine::general_purpose::STANDARD.encode(chunk),
|
||||
},
|
||||
client_id: envelope.client_id.clone(),
|
||||
stream_id: envelope.stream_id.clone(),
|
||||
seq_id: envelope.seq_id,
|
||||
})
|
||||
}
|
||||
@@ -1,386 +0,0 @@
|
||||
use super::protocol::ClientEnvelope;
|
||||
use super::protocol::ClientEvent;
|
||||
use super::protocol::ClientId;
|
||||
use super::protocol::ServerEnvelope;
|
||||
use super::protocol::ServerEvent;
|
||||
use super::protocol::StreamId;
|
||||
use super::segment::ClientSegmentObservation;
|
||||
use super::segment::ClientSegmentReassembler;
|
||||
use super::segment::REMOTE_CONTROL_SEGMENT_MAX_BYTES;
|
||||
use super::segment::split_server_envelope_for_transport;
|
||||
use crate::outgoing_message::OutgoingMessage;
|
||||
use base64::Engine;
|
||||
use codex_app_server_protocol::ConfigWarningNotification;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use codex_app_server_protocol::JSONRPCNotification;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn reassembles_client_message_chunks() {
|
||||
let message = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
let raw = serde_json::to_vec(&message).expect("message should serialize");
|
||||
let split = raw.len() / 2;
|
||||
let client_id = ClientId("client-1".to_string());
|
||||
let stream_id = Some(StreamId("stream-1".to_string()));
|
||||
let mut reassembler = ClientSegmentReassembler::default();
|
||||
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
/*seq_id*/ 7,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[..split],
|
||||
)),
|
||||
ClientSegmentObservation::Pending
|
||||
));
|
||||
let reassembled = match reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
stream_id,
|
||||
/*seq_id*/ 7,
|
||||
/*segment_id*/ 1,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[split..],
|
||||
)) {
|
||||
ClientSegmentObservation::Forward(reassembled) => *reassembled,
|
||||
ClientSegmentObservation::Pending | ClientSegmentObservation::Dropped => {
|
||||
panic!("message should reassemble")
|
||||
}
|
||||
};
|
||||
assert_eq!(reassembled.client_id, client_id);
|
||||
assert_eq!(
|
||||
reassembled.stream_id,
|
||||
Some(StreamId("stream-1".to_string()))
|
||||
);
|
||||
assert_eq!(reassembled.seq_id, Some(7));
|
||||
assert_eq!(reassembled.cursor, None);
|
||||
match reassembled.event {
|
||||
ClientEvent::ClientMessage {
|
||||
message: reassembled_message,
|
||||
} => assert_eq!(reassembled_message, message),
|
||||
other => panic!("expected client message, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn splits_large_server_messages_into_wire_chunks() {
|
||||
let envelope = ServerEnvelope {
|
||||
event: ServerEvent::ServerMessage {
|
||||
message: Box::new(OutgoingMessage::AppServerNotification(
|
||||
ServerNotification::ConfigWarning(ConfigWarningNotification {
|
||||
summary: "x".repeat(REMOTE_CONTROL_SEGMENT_MAX_BYTES),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
}),
|
||||
)),
|
||||
},
|
||||
client_id: ClientId("client-1".to_string()),
|
||||
stream_id: StreamId("stream-1".to_string()),
|
||||
seq_id: 9,
|
||||
};
|
||||
|
||||
let segments = split_server_envelope_for_transport(envelope).expect("split should succeed");
|
||||
|
||||
assert!(segments.len() > 1);
|
||||
assert!(
|
||||
segments
|
||||
.iter()
|
||||
.all(|segment| matches!(segment.event, ServerEvent::ServerMessageChunk { .. }))
|
||||
);
|
||||
assert!(segments.iter().all(|segment| segment.seq_id == 9));
|
||||
assert!(segments.iter().all(|segment| {
|
||||
serde_json::to_vec(segment)
|
||||
.expect("segment should serialize")
|
||||
.len()
|
||||
<= REMOTE_CONTROL_SEGMENT_MAX_BYTES
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalidates_incomplete_stream_assemblies() {
|
||||
let message = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
let raw = serde_json::to_vec(&message).expect("message should serialize");
|
||||
let split = raw.len() / 2;
|
||||
let client_id = ClientId("client-1".to_string());
|
||||
let stream_id = StreamId("stream-1".to_string());
|
||||
let mut reassembler = ClientSegmentReassembler::default();
|
||||
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
Some(stream_id.clone()),
|
||||
/*seq_id*/ 7,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[..split],
|
||||
)),
|
||||
ClientSegmentObservation::Pending
|
||||
));
|
||||
reassembler.invalidate_stream(&client_id, &stream_id);
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id,
|
||||
Some(stream_id),
|
||||
/*seq_id*/ 7,
|
||||
/*segment_id*/ 1,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[split..],
|
||||
)),
|
||||
ClientSegmentObservation::Dropped
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resets_incomplete_client_assembly_when_stream_changes() {
|
||||
let message = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
let raw = serde_json::to_vec(&message).expect("message should serialize");
|
||||
let split = raw.len() / 2;
|
||||
let client_id = ClientId("client-1".to_string());
|
||||
let first_stream_id = StreamId("stream-1".to_string());
|
||||
let second_stream_id = StreamId("stream-2".to_string());
|
||||
let mut reassembler = ClientSegmentReassembler::default();
|
||||
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
Some(first_stream_id.clone()),
|
||||
/*seq_id*/ 7,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[..split],
|
||||
)),
|
||||
ClientSegmentObservation::Pending
|
||||
));
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
Some(second_stream_id.clone()),
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[..split],
|
||||
)),
|
||||
ClientSegmentObservation::Pending
|
||||
));
|
||||
let reassembled = match reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
Some(second_stream_id),
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 1,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[split..],
|
||||
)) {
|
||||
ClientSegmentObservation::Forward(reassembled) => *reassembled,
|
||||
ClientSegmentObservation::Pending | ClientSegmentObservation::Dropped => {
|
||||
panic!("replacement stream should reassemble")
|
||||
}
|
||||
};
|
||||
assert_eq!(
|
||||
reassembled.stream_id,
|
||||
Some(StreamId("stream-2".to_string()))
|
||||
);
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id,
|
||||
Some(first_stream_id),
|
||||
/*seq_id*/ 7,
|
||||
/*segment_id*/ 1,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[split..],
|
||||
)),
|
||||
ClientSegmentObservation::Dropped
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_stale_chunks_without_dropping_newer_assembly() {
|
||||
let message = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
let raw = serde_json::to_vec(&message).expect("message should serialize");
|
||||
let split = raw.len() / 2;
|
||||
let client_id = ClientId("client-1".to_string());
|
||||
let stream_id = Some(StreamId("stream-1".to_string()));
|
||||
let mut reassembler = ClientSegmentReassembler::default();
|
||||
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[..split],
|
||||
)),
|
||||
ClientSegmentObservation::Pending
|
||||
));
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
/*seq_id*/ 7,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[..split],
|
||||
)),
|
||||
ClientSegmentObservation::Dropped
|
||||
));
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id,
|
||||
stream_id,
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 1,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[split..],
|
||||
)),
|
||||
ClientSegmentObservation::Forward(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_invalid_stale_chunks_without_dropping_newer_assembly() {
|
||||
let message = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
let raw = serde_json::to_vec(&message).expect("message should serialize");
|
||||
let split = raw.len() / 2;
|
||||
let client_id = ClientId("client-1".to_string());
|
||||
let stream_id = Some(StreamId("stream-1".to_string()));
|
||||
let mut reassembler = ClientSegmentReassembler::default();
|
||||
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[..split],
|
||||
)),
|
||||
ClientSegmentObservation::Pending
|
||||
));
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
/*seq_id*/ 7,
|
||||
/*segment_id*/ 1,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
b"",
|
||||
)),
|
||||
ClientSegmentObservation::Dropped
|
||||
));
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id,
|
||||
stream_id,
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 1,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[split..],
|
||||
)),
|
||||
ClientSegmentObservation::Forward(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_invalid_duplicate_chunks_without_dropping_current_assembly() {
|
||||
let message = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
let raw = serde_json::to_vec(&message).expect("message should serialize");
|
||||
let split = raw.len() / 2;
|
||||
let client_id = ClientId("client-1".to_string());
|
||||
let stream_id = Some(StreamId("stream-1".to_string()));
|
||||
let mut reassembler = ClientSegmentReassembler::default();
|
||||
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[..split],
|
||||
)),
|
||||
ClientSegmentObservation::Pending
|
||||
));
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 0,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
b"",
|
||||
)),
|
||||
ClientSegmentObservation::Dropped
|
||||
));
|
||||
assert!(matches!(
|
||||
reassembler.observe(chunk_envelope(
|
||||
client_id,
|
||||
stream_id,
|
||||
/*seq_id*/ 8,
|
||||
/*segment_id*/ 1,
|
||||
/*segment_count*/ 2,
|
||||
raw.len(),
|
||||
&raw[split..],
|
||||
)),
|
||||
ClientSegmentObservation::Forward(_)
|
||||
));
|
||||
}
|
||||
|
||||
fn chunk_envelope(
|
||||
client_id: ClientId,
|
||||
stream_id: Option<StreamId>,
|
||||
seq_id: u64,
|
||||
segment_id: usize,
|
||||
segment_count: usize,
|
||||
message_size_bytes: usize,
|
||||
chunk: &[u8],
|
||||
) -> ClientEnvelope {
|
||||
ClientEnvelope {
|
||||
event: ClientEvent::ClientMessageChunk {
|
||||
segment_id,
|
||||
segment_count,
|
||||
message_size_bytes,
|
||||
message_chunk_base64: base64::engine::general_purpose::STANDARD.encode(chunk),
|
||||
},
|
||||
client_id,
|
||||
stream_id,
|
||||
seq_id: Some(seq_id),
|
||||
cursor: None,
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,113 +0,0 @@
|
||||
use super::CHANNEL_CAPACITY;
|
||||
use super::ConnectionOrigin;
|
||||
use super::TransportEvent;
|
||||
use super::forward_incoming_message;
|
||||
use super::next_connection_id;
|
||||
use super::serialize_outgoing_message;
|
||||
use crate::outgoing_message::QueuedOutgoingMessage;
|
||||
use codex_app_server_protocol::InitializeParams;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use codex_app_server_protocol::JSONRPCRequest;
|
||||
use std::io::ErrorKind;
|
||||
use std::io::Result as IoResult;
|
||||
use tokio::io;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::debug;
|
||||
use tracing::error;
|
||||
use tracing::info;
|
||||
|
||||
pub(crate) async fn start_stdio_connection(
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
stdio_handles: &mut Vec<JoinHandle<()>>,
|
||||
initialize_client_name_tx: oneshot::Sender<String>,
|
||||
) -> IoResult<()> {
|
||||
let connection_id = next_connection_id();
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel::<QueuedOutgoingMessage>(CHANNEL_CAPACITY);
|
||||
let writer_tx_for_reader = writer_tx.clone();
|
||||
transport_event_tx
|
||||
.send(TransportEvent::ConnectionOpened {
|
||||
connection_id,
|
||||
origin: ConnectionOrigin::Stdio,
|
||||
writer: writer_tx,
|
||||
disconnect_sender: None,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| std::io::Error::new(ErrorKind::BrokenPipe, "processor unavailable"))?;
|
||||
|
||||
let transport_event_tx_for_reader = transport_event_tx.clone();
|
||||
stdio_handles.push(tokio::spawn(async move {
|
||||
let stdin = io::stdin();
|
||||
let reader = BufReader::new(stdin);
|
||||
let mut lines = reader.lines();
|
||||
let mut initialize_client_name_tx = Some(initialize_client_name_tx);
|
||||
|
||||
loop {
|
||||
match lines.next_line().await {
|
||||
Ok(Some(line)) => {
|
||||
if let Some(client_name) = stdio_initialize_client_name(&line)
|
||||
&& let Some(initialize_client_name_tx) = initialize_client_name_tx.take()
|
||||
{
|
||||
let _ = initialize_client_name_tx.send(client_name);
|
||||
}
|
||||
if !forward_incoming_message(
|
||||
&transport_event_tx_for_reader,
|
||||
&writer_tx_for_reader,
|
||||
connection_id,
|
||||
&line,
|
||||
)
|
||||
.await
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(err) => {
|
||||
error!("Failed reading stdin: {err}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = transport_event_tx_for_reader
|
||||
.send(TransportEvent::ConnectionClosed { connection_id })
|
||||
.await;
|
||||
debug!("stdin reader finished (EOF)");
|
||||
}));
|
||||
|
||||
stdio_handles.push(tokio::spawn(async move {
|
||||
let mut stdout = io::stdout();
|
||||
while let Some(queued_message) = writer_rx.recv().await {
|
||||
let Some(mut json) = serialize_outgoing_message(queued_message.message) else {
|
||||
continue;
|
||||
};
|
||||
json.push('\n');
|
||||
if let Err(err) = stdout.write_all(json.as_bytes()).await {
|
||||
error!("Failed to write to stdout: {err}");
|
||||
break;
|
||||
}
|
||||
if let Some(write_complete_tx) = queued_message.write_complete_tx {
|
||||
let _ = write_complete_tx.send(());
|
||||
}
|
||||
}
|
||||
info!("stdout writer exited (channel closed)");
|
||||
}));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn stdio_initialize_client_name(line: &str) -> Option<String> {
|
||||
let message = serde_json::from_str::<JSONRPCMessage>(line).ok()?;
|
||||
let JSONRPCMessage::Request(JSONRPCRequest { method, params, .. }) = message else {
|
||||
return None;
|
||||
};
|
||||
if method != "initialize" {
|
||||
return None;
|
||||
}
|
||||
let params = serde_json::from_value::<InitializeParams>(params?).ok()?;
|
||||
Some(params.client_info.name)
|
||||
}
|
||||
@@ -1,165 +0,0 @@
|
||||
use std::io::ErrorKind;
|
||||
use std::io::Result as IoResult;
|
||||
use std::path::Path;
|
||||
|
||||
use super::TransportEvent;
|
||||
use crate::transport::websocket::run_websocket_connection;
|
||||
use codex_uds::UnixListener;
|
||||
use codex_uds::UnixStream;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use futures::StreamExt;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio::time::Duration;
|
||||
use tokio_tungstenite::accept_async;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::error;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
|
||||
#[cfg(unix)]
|
||||
const CONTROL_SOCKET_MODE: u32 = 0o600;
|
||||
|
||||
pub(crate) async fn start_control_socket_acceptor(
|
||||
socket_path: AbsolutePathBuf,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: CancellationToken,
|
||||
) -> IoResult<JoinHandle<()>> {
|
||||
prepare_control_socket_path(socket_path.as_path()).await?;
|
||||
let listener = UnixListener::bind(socket_path.as_path()).await?;
|
||||
let socket_guard = ControlSocketFileGuard { socket_path };
|
||||
set_control_socket_permissions(socket_guard.socket_path.as_path()).await?;
|
||||
info!(
|
||||
socket_path = %socket_guard.socket_path.display(),
|
||||
"app-server control socket listening"
|
||||
);
|
||||
|
||||
Ok(tokio::spawn(run_control_socket_acceptor(
|
||||
listener,
|
||||
transport_event_tx,
|
||||
shutdown_token,
|
||||
socket_guard,
|
||||
)))
|
||||
}
|
||||
|
||||
async fn run_control_socket_acceptor(
|
||||
mut listener: UnixListener,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: CancellationToken,
|
||||
socket_guard: ControlSocketFileGuard,
|
||||
) {
|
||||
let _socket_guard = socket_guard;
|
||||
loop {
|
||||
let stream = tokio::select! {
|
||||
_ = shutdown_token.cancelled() => {
|
||||
break;
|
||||
}
|
||||
result = listener.accept() => {
|
||||
match result {
|
||||
Ok(stream) => stream,
|
||||
Err(err) => {
|
||||
if matches!(
|
||||
err.kind(),
|
||||
ErrorKind::ConnectionAborted | ErrorKind::ConnectionReset | ErrorKind::Interrupted
|
||||
) {
|
||||
warn!("recoverable control socket accept error: {err}");
|
||||
continue;
|
||||
}
|
||||
error!("control socket accept error: {err}");
|
||||
tokio::time::sleep(Duration::from_secs(1)).await;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let transport_event_tx = transport_event_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let websocket_stream = match accept_async(stream).await {
|
||||
Ok(websocket_stream) => websocket_stream,
|
||||
Err(err) => {
|
||||
warn!("failed to upgrade control socket websocket connection: {err}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let (websocket_writer, websocket_reader) = websocket_stream.split();
|
||||
run_websocket_connection(websocket_writer, websocket_reader, transport_event_tx).await;
|
||||
});
|
||||
}
|
||||
info!("control socket acceptor shutting down");
|
||||
}
|
||||
|
||||
async fn prepare_control_socket_path(socket_path: &Path) -> IoResult<()> {
|
||||
if let Some(parent) = socket_path.parent() {
|
||||
codex_uds::prepare_private_socket_directory(parent).await?;
|
||||
}
|
||||
|
||||
match UnixStream::connect(socket_path).await {
|
||||
Ok(_stream) => {
|
||||
return Err(std::io::Error::new(
|
||||
ErrorKind::AddrInUse,
|
||||
format!(
|
||||
"app-server control socket is already in use at {}",
|
||||
socket_path.display()
|
||||
),
|
||||
));
|
||||
}
|
||||
Err(err) if err.kind() == ErrorKind::NotFound => return Ok(()),
|
||||
Err(err) if err.kind() == ErrorKind::ConnectionRefused => {}
|
||||
Err(err) => {
|
||||
if !socket_path.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
|
||||
if !socket_path.try_exists()? {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if !codex_uds::is_stale_socket_path(socket_path).await? {
|
||||
return Err(std::io::Error::new(
|
||||
ErrorKind::AlreadyExists,
|
||||
format!(
|
||||
"app-server control socket path exists and is not a socket: {}",
|
||||
socket_path.display()
|
||||
),
|
||||
));
|
||||
}
|
||||
tokio::fs::remove_file(socket_path).await
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
async fn set_control_socket_permissions(socket_path: &Path) -> IoResult<()> {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
tokio::fs::set_permissions(
|
||||
socket_path,
|
||||
std::fs::Permissions::from_mode(CONTROL_SOCKET_MODE),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
async fn set_control_socket_permissions(_socket_path: &Path) -> IoResult<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct ControlSocketFileGuard {
|
||||
socket_path: AbsolutePathBuf,
|
||||
}
|
||||
|
||||
impl Drop for ControlSocketFileGuard {
|
||||
fn drop(&mut self) {
|
||||
if let Err(err) = std::fs::remove_file(self.socket_path.as_path())
|
||||
&& err.kind() != ErrorKind::NotFound
|
||||
{
|
||||
warn!(
|
||||
socket_path = %self.socket_path.display(),
|
||||
%err,
|
||||
"failed to remove app-server control socket file"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,201 +0,0 @@
|
||||
use super::AppServerTransport;
|
||||
use super::CHANNEL_CAPACITY;
|
||||
use super::TransportEvent;
|
||||
use super::app_server_control_socket_path;
|
||||
use super::start_control_socket_acceptor;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use codex_app_server_protocol::JSONRPCNotification;
|
||||
use codex_core::config::find_codex_home;
|
||||
use codex_uds::UnixStream;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::io::Result as IoResult;
|
||||
use std::path::Path;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
use tokio_tungstenite::client_async;
|
||||
use tokio_tungstenite::tungstenite::Bytes;
|
||||
use tokio_tungstenite::tungstenite::Message as WebSocketMessage;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[test]
|
||||
fn listen_unix_socket_parses_as_unix_socket_transport() {
|
||||
assert_eq!(
|
||||
AppServerTransport::from_listen_url("unix://"),
|
||||
Ok(AppServerTransport::UnixSocket {
|
||||
socket_path: default_control_socket_path()
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn listen_unix_socket_accepts_absolute_custom_path() {
|
||||
assert_eq!(
|
||||
AppServerTransport::from_listen_url("unix:///tmp/codex.sock"),
|
||||
Ok(AppServerTransport::UnixSocket {
|
||||
socket_path: absolute_path("/tmp/codex.sock")
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn listen_unix_socket_accepts_relative_custom_path() {
|
||||
assert_eq!(
|
||||
AppServerTransport::from_listen_url("unix://codex.sock"),
|
||||
Ok(AppServerTransport::UnixSocket {
|
||||
socket_path: AbsolutePathBuf::relative_to_current_dir("codex.sock")
|
||||
.expect("relative path should resolve")
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn control_socket_acceptor_upgrades_and_forwards_websocket_text_messages_and_pings() {
|
||||
let temp_dir = tempfile::TempDir::new().expect("temp dir");
|
||||
let socket_path = test_socket_path(temp_dir.path());
|
||||
let (transport_event_tx, mut transport_event_rx) =
|
||||
mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let accept_handle = start_control_socket_acceptor(
|
||||
socket_path.clone(),
|
||||
transport_event_tx,
|
||||
shutdown_token.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("control socket acceptor should start");
|
||||
|
||||
let stream = connect_to_socket(socket_path.as_path())
|
||||
.await
|
||||
.expect("client should connect");
|
||||
let (mut websocket, response) = client_async("ws://localhost/rpc", stream)
|
||||
.await
|
||||
.expect("websocket upgrade should complete");
|
||||
assert_eq!(response.status().as_u16(), 101);
|
||||
|
||||
let opened = timeout(Duration::from_secs(1), transport_event_rx.recv())
|
||||
.await
|
||||
.expect("connection opened event should arrive")
|
||||
.expect("connection opened event");
|
||||
let connection_id = match opened {
|
||||
TransportEvent::ConnectionOpened { connection_id, .. } => connection_id,
|
||||
_ => panic!("expected connection opened event"),
|
||||
};
|
||||
|
||||
let notification = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
websocket
|
||||
.send(WebSocketMessage::Text(
|
||||
serde_json::to_string(¬ification)
|
||||
.expect("notification should serialize")
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.expect("notification should send");
|
||||
|
||||
let incoming = timeout(Duration::from_secs(1), transport_event_rx.recv())
|
||||
.await
|
||||
.expect("incoming message event should arrive")
|
||||
.expect("incoming message event");
|
||||
assert_eq!(
|
||||
match incoming {
|
||||
TransportEvent::IncomingMessage {
|
||||
connection_id: incoming_connection_id,
|
||||
message,
|
||||
} => (incoming_connection_id, message),
|
||||
_ => panic!("expected incoming message event"),
|
||||
},
|
||||
(connection_id, notification)
|
||||
);
|
||||
|
||||
websocket
|
||||
.send(WebSocketMessage::Ping(Bytes::from_static(b"check")))
|
||||
.await
|
||||
.expect("ping should send");
|
||||
let pong = timeout(Duration::from_secs(1), websocket.next())
|
||||
.await
|
||||
.expect("pong should arrive")
|
||||
.expect("pong frame")
|
||||
.expect("pong should be valid");
|
||||
assert_eq!(pong, WebSocketMessage::Pong(Bytes::from_static(b"check")));
|
||||
|
||||
websocket.close(None).await.expect("close should send");
|
||||
let closed = timeout(Duration::from_secs(1), transport_event_rx.recv())
|
||||
.await
|
||||
.expect("connection closed event should arrive")
|
||||
.expect("connection closed event");
|
||||
assert!(matches!(
|
||||
closed,
|
||||
TransportEvent::ConnectionClosed {
|
||||
connection_id: closed_connection_id,
|
||||
} if closed_connection_id == connection_id
|
||||
));
|
||||
|
||||
shutdown_token.cancel();
|
||||
accept_handle.await.expect("acceptor should join");
|
||||
assert_socket_path_removed(socket_path.as_path());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn control_socket_file_is_private_after_bind() {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
let temp_dir = tempfile::TempDir::new().expect("temp dir");
|
||||
let socket_path = test_socket_path(temp_dir.path());
|
||||
let (transport_event_tx, _transport_event_rx) =
|
||||
mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let accept_handle = start_control_socket_acceptor(
|
||||
socket_path.clone(),
|
||||
transport_event_tx,
|
||||
shutdown_token.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("control socket acceptor should start");
|
||||
|
||||
let metadata = tokio::fs::metadata(socket_path.as_path())
|
||||
.await
|
||||
.expect("socket metadata should exist");
|
||||
assert_eq!(metadata.permissions().mode() & 0o777, 0o600);
|
||||
|
||||
shutdown_token.cancel();
|
||||
accept_handle.await.expect("acceptor should join");
|
||||
}
|
||||
|
||||
fn absolute_path(path: &str) -> AbsolutePathBuf {
|
||||
AbsolutePathBuf::from_absolute_path(path).expect("absolute path")
|
||||
}
|
||||
|
||||
fn default_control_socket_path() -> AbsolutePathBuf {
|
||||
let codex_home = find_codex_home().expect("codex home");
|
||||
app_server_control_socket_path(&codex_home).expect("default control socket path")
|
||||
}
|
||||
|
||||
fn test_socket_path(temp_dir: &Path) -> AbsolutePathBuf {
|
||||
AbsolutePathBuf::from_absolute_path(
|
||||
temp_dir
|
||||
.join("app-server-control")
|
||||
.join("app-server-control.sock"),
|
||||
)
|
||||
.expect("socket path should resolve")
|
||||
}
|
||||
|
||||
async fn connect_to_socket(socket_path: &Path) -> IoResult<UnixStream> {
|
||||
UnixStream::connect(socket_path).await
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn assert_socket_path_removed(socket_path: &Path) {
|
||||
assert!(!socket_path.exists());
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn assert_socket_path_removed(_socket_path: &Path) {
|
||||
// uds_windows uses a regular filesystem path as its rendezvous point,
|
||||
// but there is no Unix socket filesystem node to assert on.
|
||||
}
|
||||
@@ -1,388 +0,0 @@
|
||||
use super::CHANNEL_CAPACITY;
|
||||
use super::ConnectionOrigin;
|
||||
use super::TransportEvent;
|
||||
use super::auth::WebsocketAuthPolicy;
|
||||
use super::auth::authorize_upgrade;
|
||||
use super::auth::should_warn_about_unauthenticated_non_loopback_listener;
|
||||
use super::forward_incoming_message;
|
||||
use super::next_connection_id;
|
||||
use super::serialize_outgoing_message;
|
||||
use crate::outgoing_message::ConnectionId;
|
||||
use crate::outgoing_message::QueuedOutgoingMessage;
|
||||
use axum::Router;
|
||||
use axum::body::Body;
|
||||
use axum::body::Bytes;
|
||||
use axum::extract::ConnectInfo;
|
||||
use axum::extract::State;
|
||||
use axum::extract::ws::Message as AxumWebSocketMessage;
|
||||
use axum::extract::ws::WebSocketUpgrade;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::http::Request;
|
||||
use axum::http::StatusCode;
|
||||
use axum::http::header::ORIGIN;
|
||||
use axum::middleware;
|
||||
use axum::middleware::Next;
|
||||
use axum::response::IntoResponse;
|
||||
use axum::response::Response;
|
||||
use axum::routing::any;
|
||||
use axum::routing::get;
|
||||
use futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use owo_colors::OwoColorize;
|
||||
use owo_colors::Stream;
|
||||
use owo_colors::Style;
|
||||
use std::io::Result as IoResult;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio_tungstenite::tungstenite::Message as TungsteniteWebSocketMessage;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::error;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
|
||||
/// WebSocket clients can briefly lag behind normal turn output bursts while the
|
||||
/// writer task is healthy, so give them more headroom than internal channels.
|
||||
const WEBSOCKET_OUTBOUND_CHANNEL_CAPACITY: usize = 32 * 1024;
|
||||
const _: () = assert!(WEBSOCKET_OUTBOUND_CHANNEL_CAPACITY > CHANNEL_CAPACITY);
|
||||
|
||||
fn colorize(text: &str, style: Style) -> String {
|
||||
text.if_supports_color(Stream::Stderr, |value| value.style(style))
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[allow(clippy::print_stderr)]
|
||||
fn print_websocket_startup_banner(addr: SocketAddr) {
|
||||
let title = colorize("codex app-server (WebSockets)", Style::new().bold().cyan());
|
||||
let listening_label = colorize("listening on:", Style::new().dimmed());
|
||||
let listen_url = colorize(&format!("ws://{addr}"), Style::new().green());
|
||||
let ready_label = colorize("readyz:", Style::new().dimmed());
|
||||
let ready_url = colorize(&format!("http://{addr}/readyz"), Style::new().green());
|
||||
let health_label = colorize("healthz:", Style::new().dimmed());
|
||||
let health_url = colorize(&format!("http://{addr}/healthz"), Style::new().green());
|
||||
let note_label = colorize("note:", Style::new().dimmed());
|
||||
eprintln!("{title}");
|
||||
eprintln!(" {listening_label} {listen_url}");
|
||||
eprintln!(" {ready_label} {ready_url}");
|
||||
eprintln!(" {health_label} {health_url}");
|
||||
if addr.ip().is_loopback() {
|
||||
eprintln!(
|
||||
" {note_label} binds localhost only (use SSH port-forwarding for remote access)"
|
||||
);
|
||||
} else {
|
||||
eprintln!(
|
||||
" {note_label} websocket auth is opt-in in this build; configure `--ws-auth ...` before real remote use"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct WebSocketListenerState {
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
auth_policy: Arc<WebsocketAuthPolicy>,
|
||||
}
|
||||
|
||||
async fn health_check_handler() -> StatusCode {
|
||||
StatusCode::OK
|
||||
}
|
||||
|
||||
async fn reject_requests_with_origin_header(
|
||||
request: Request<Body>,
|
||||
next: Next,
|
||||
) -> Result<Response, StatusCode> {
|
||||
if request.headers().contains_key(ORIGIN) {
|
||||
warn!(
|
||||
method = %request.method(),
|
||||
uri = %request.uri(),
|
||||
"rejecting websocket listener request with Origin header"
|
||||
);
|
||||
Err(StatusCode::FORBIDDEN)
|
||||
} else {
|
||||
Ok(next.run(request).await)
|
||||
}
|
||||
}
|
||||
|
||||
async fn websocket_upgrade_handler(
|
||||
websocket: WebSocketUpgrade,
|
||||
ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
|
||||
State(state): State<WebSocketListenerState>,
|
||||
headers: HeaderMap,
|
||||
) -> impl IntoResponse {
|
||||
if let Err(err) = authorize_upgrade(&headers, state.auth_policy.as_ref()) {
|
||||
warn!(
|
||||
%peer_addr,
|
||||
message = err.message(),
|
||||
"rejecting websocket client during upgrade"
|
||||
);
|
||||
return (err.status_code(), err.message()).into_response();
|
||||
}
|
||||
info!(%peer_addr, "websocket client connected");
|
||||
websocket
|
||||
.on_upgrade(move |stream| async move {
|
||||
let (websocket_writer, websocket_reader) = stream.split();
|
||||
run_websocket_connection(websocket_writer, websocket_reader, state.transport_event_tx)
|
||||
.await;
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
|
||||
pub(crate) async fn start_websocket_acceptor(
|
||||
bind_address: SocketAddr,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: CancellationToken,
|
||||
auth_policy: WebsocketAuthPolicy,
|
||||
) -> IoResult<JoinHandle<()>> {
|
||||
if should_warn_about_unauthenticated_non_loopback_listener(bind_address, &auth_policy) {
|
||||
warn!(
|
||||
%bind_address,
|
||||
"starting non-loopback websocket listener without auth; websocket auth is opt-in for now and will become the default in a future release"
|
||||
);
|
||||
}
|
||||
let listener = TcpListener::bind(bind_address).await?;
|
||||
let local_addr = listener.local_addr()?;
|
||||
print_websocket_startup_banner(local_addr);
|
||||
info!("app-server websocket listening on ws://{local_addr}");
|
||||
|
||||
let router = Router::new()
|
||||
.route("/readyz", get(health_check_handler))
|
||||
.route("/healthz", get(health_check_handler))
|
||||
.fallback(any(websocket_upgrade_handler))
|
||||
.layer(middleware::from_fn(reject_requests_with_origin_header))
|
||||
.with_state(WebSocketListenerState {
|
||||
transport_event_tx,
|
||||
auth_policy: Arc::new(auth_policy),
|
||||
});
|
||||
let server = axum::serve(
|
||||
listener,
|
||||
router.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(async move {
|
||||
shutdown_token.cancelled().await;
|
||||
});
|
||||
Ok(tokio::spawn(async move {
|
||||
if let Err(err) = server.await {
|
||||
error!("websocket acceptor failed: {err}");
|
||||
}
|
||||
info!("websocket acceptor shutting down");
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) async fn run_websocket_connection<M, SinkError, StreamError>(
|
||||
websocket_writer: impl futures::sink::Sink<M, Error = SinkError> + Send + 'static,
|
||||
websocket_reader: impl futures::stream::Stream<Item = Result<M, StreamError>> + Send + 'static,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
) where
|
||||
M: AppServerWebSocketMessage + Send + 'static,
|
||||
SinkError: Send + 'static,
|
||||
StreamError: std::fmt::Display + Send + 'static,
|
||||
{
|
||||
let connection_id = next_connection_id();
|
||||
let (writer_tx, writer_rx) =
|
||||
mpsc::channel::<QueuedOutgoingMessage>(WEBSOCKET_OUTBOUND_CHANNEL_CAPACITY);
|
||||
let writer_tx_for_reader = writer_tx.clone();
|
||||
let disconnect_token = CancellationToken::new();
|
||||
if transport_event_tx
|
||||
.send(TransportEvent::ConnectionOpened {
|
||||
connection_id,
|
||||
origin: ConnectionOrigin::WebSocket,
|
||||
writer: writer_tx,
|
||||
disconnect_sender: Some(disconnect_token.clone()),
|
||||
})
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
let (writer_control_tx, writer_control_rx) = mpsc::channel::<M>(CHANNEL_CAPACITY);
|
||||
let mut outbound_task = tokio::spawn(run_websocket_outbound_loop(
|
||||
websocket_writer,
|
||||
writer_rx,
|
||||
writer_control_rx,
|
||||
disconnect_token.clone(),
|
||||
));
|
||||
let mut inbound_task = tokio::spawn(run_websocket_inbound_loop(
|
||||
websocket_reader,
|
||||
transport_event_tx.clone(),
|
||||
writer_tx_for_reader,
|
||||
writer_control_tx,
|
||||
connection_id,
|
||||
disconnect_token.clone(),
|
||||
));
|
||||
|
||||
tokio::select! {
|
||||
_ = &mut outbound_task => {
|
||||
disconnect_token.cancel();
|
||||
inbound_task.abort();
|
||||
}
|
||||
_ = &mut inbound_task => {
|
||||
disconnect_token.cancel();
|
||||
outbound_task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
let _ = transport_event_tx
|
||||
.send(TransportEvent::ConnectionClosed { connection_id })
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(crate) enum IncomingWebSocketMessage {
|
||||
Text(String),
|
||||
Binary,
|
||||
Ping(Bytes),
|
||||
Pong,
|
||||
Close,
|
||||
}
|
||||
|
||||
/// Converts concrete WebSocket message types into the small message surface the
|
||||
/// app-server transport needs, and constructs the only outbound frames it
|
||||
/// sends directly.
|
||||
pub(crate) trait AppServerWebSocketMessage: Sized {
|
||||
fn text(text: String) -> Self;
|
||||
fn pong(payload: Bytes) -> Self;
|
||||
fn into_incoming(self) -> Option<IncomingWebSocketMessage>;
|
||||
}
|
||||
|
||||
impl AppServerWebSocketMessage for AxumWebSocketMessage {
|
||||
fn text(text: String) -> Self {
|
||||
Self::Text(text.into())
|
||||
}
|
||||
|
||||
fn pong(payload: Bytes) -> Self {
|
||||
Self::Pong(payload)
|
||||
}
|
||||
|
||||
fn into_incoming(self) -> Option<IncomingWebSocketMessage> {
|
||||
Some(match self {
|
||||
Self::Text(text) => IncomingWebSocketMessage::Text(text.to_string()),
|
||||
Self::Binary(_) => IncomingWebSocketMessage::Binary,
|
||||
Self::Ping(payload) => IncomingWebSocketMessage::Ping(payload),
|
||||
Self::Pong(_) => IncomingWebSocketMessage::Pong,
|
||||
Self::Close(_) => IncomingWebSocketMessage::Close,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl AppServerWebSocketMessage for TungsteniteWebSocketMessage {
|
||||
fn text(text: String) -> Self {
|
||||
Self::Text(text.into())
|
||||
}
|
||||
|
||||
fn pong(payload: Bytes) -> Self {
|
||||
Self::Pong(payload)
|
||||
}
|
||||
|
||||
fn into_incoming(self) -> Option<IncomingWebSocketMessage> {
|
||||
Some(match self {
|
||||
Self::Text(text) => IncomingWebSocketMessage::Text(text.to_string()),
|
||||
Self::Binary(_) => IncomingWebSocketMessage::Binary,
|
||||
Self::Ping(payload) => IncomingWebSocketMessage::Ping(payload),
|
||||
Self::Pong(_) => IncomingWebSocketMessage::Pong,
|
||||
Self::Close(_) => IncomingWebSocketMessage::Close,
|
||||
Self::Frame(_) => return None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_websocket_outbound_loop<M, SinkError>(
|
||||
websocket_writer: impl futures::sink::Sink<M, Error = SinkError> + Send + 'static,
|
||||
mut writer_rx: mpsc::Receiver<QueuedOutgoingMessage>,
|
||||
mut writer_control_rx: mpsc::Receiver<M>,
|
||||
disconnect_token: CancellationToken,
|
||||
) where
|
||||
M: AppServerWebSocketMessage + Send + 'static,
|
||||
SinkError: Send + 'static,
|
||||
{
|
||||
tokio::pin!(websocket_writer);
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = disconnect_token.cancelled() => {
|
||||
break;
|
||||
}
|
||||
message = writer_control_rx.recv() => {
|
||||
let Some(message) = message else {
|
||||
break;
|
||||
};
|
||||
if websocket_writer.send(message).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
queued_message = writer_rx.recv() => {
|
||||
let Some(queued_message) = queued_message else {
|
||||
break;
|
||||
};
|
||||
let Some(json) = serialize_outgoing_message(queued_message.message) else {
|
||||
continue;
|
||||
};
|
||||
if websocket_writer.send(M::text(json)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
if let Some(write_complete_tx) = queued_message.write_complete_tx {
|
||||
let _ = write_complete_tx.send(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_websocket_inbound_loop<M, StreamError>(
|
||||
websocket_reader: impl futures::stream::Stream<Item = Result<M, StreamError>> + Send + 'static,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
writer_tx_for_reader: mpsc::Sender<QueuedOutgoingMessage>,
|
||||
writer_control_tx: mpsc::Sender<M>,
|
||||
connection_id: ConnectionId,
|
||||
disconnect_token: CancellationToken,
|
||||
) where
|
||||
M: AppServerWebSocketMessage + Send + 'static,
|
||||
StreamError: std::fmt::Display + Send + 'static,
|
||||
{
|
||||
tokio::pin!(websocket_reader);
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = disconnect_token.cancelled() => {
|
||||
break;
|
||||
}
|
||||
incoming_message = websocket_reader.next() => {
|
||||
match incoming_message {
|
||||
Some(Ok(message)) => match message.into_incoming() {
|
||||
Some(IncomingWebSocketMessage::Text(text)) => {
|
||||
if !forward_incoming_message(
|
||||
&transport_event_tx,
|
||||
&writer_tx_for_reader,
|
||||
connection_id,
|
||||
&text,
|
||||
)
|
||||
.await
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Some(IncomingWebSocketMessage::Ping(payload)) => {
|
||||
match writer_control_tx.try_send(M::pong(payload)) {
|
||||
Ok(()) => {}
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => break,
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
|
||||
warn!("websocket control queue full while replying to ping; closing connection");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(IncomingWebSocketMessage::Pong) => {}
|
||||
Some(IncomingWebSocketMessage::Close) => break,
|
||||
Some(IncomingWebSocketMessage::Binary) => {
|
||||
warn!("dropping unsupported binary websocket message");
|
||||
}
|
||||
None => {}
|
||||
},
|
||||
None => break,
|
||||
Some(Err(err)) => {
|
||||
warn!("websocket receive error: {err}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,532 @@
|
||||
use super::*;
|
||||
use codex_app_server_protocol::ConfigWarningNotification;
|
||||
use codex_app_server_protocol::RequestId;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use codex_app_server_protocol::ThreadGoal;
|
||||
use codex_app_server_protocol::ThreadGoalStatus;
|
||||
use codex_app_server_protocol::ThreadGoalUpdatedNotification;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
|
||||
fn absolute_path(path: &str) -> AbsolutePathBuf {
|
||||
AbsolutePathBuf::from_absolute_path(path).expect("absolute path")
|
||||
}
|
||||
|
||||
fn thread_goal_updated_notification() -> ServerNotification {
|
||||
ServerNotification::ThreadGoalUpdated(ThreadGoalUpdatedNotification {
|
||||
thread_id: "thread-1".to_string(),
|
||||
turn_id: None,
|
||||
goal: ThreadGoal {
|
||||
thread_id: "thread-1".to_string(),
|
||||
objective: "ship goal mode".to_string(),
|
||||
status: ThreadGoalStatus::Active,
|
||||
token_budget: None,
|
||||
tokens_used: 0,
|
||||
time_used_seconds: 0,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn to_connection_notification_respects_opt_out_filters() {
|
||||
let connection_id = ConnectionId(7);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
let initialized = Arc::new(AtomicBool::new(true));
|
||||
let opted_out_notification_methods =
|
||||
Arc::new(RwLock::new(HashSet::from(["configWarning".to_string()])));
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
connection_id,
|
||||
OutboundConnectionState::new(
|
||||
writer_tx,
|
||||
initialized,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
opted_out_notification_methods,
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification {
|
||||
summary: "task_started".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
},
|
||||
)),
|
||||
write_complete_tx: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
writer_rx.try_recv().is_err(),
|
||||
"opted-out notification should be dropped"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn to_connection_notifications_are_dropped_for_opted_out_clients() {
|
||||
let connection_id = ConnectionId(10);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
connection_id,
|
||||
OutboundConnectionState::new(
|
||||
writer_tx,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(RwLock::new(HashSet::from(["configWarning".to_string()]))),
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification {
|
||||
summary: "task_started".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
},
|
||||
)),
|
||||
write_complete_tx: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
writer_rx.try_recv().is_err(),
|
||||
"opted-out notifications should not reach clients"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn to_connection_notifications_are_preserved_for_non_opted_out_clients() {
|
||||
let connection_id = ConnectionId(11);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
connection_id,
|
||||
OutboundConnectionState::new(
|
||||
writer_tx,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(RwLock::new(HashSet::new())),
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification {
|
||||
summary: "task_started".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
},
|
||||
)),
|
||||
write_complete_tx: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let message = writer_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("notification should reach non-opted-out clients");
|
||||
assert!(matches!(
|
||||
message.message,
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification { summary, .. }
|
||||
)) if summary == "task_started"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn experimental_notifications_are_dropped_without_capability() {
|
||||
let connection_id = ConnectionId(12);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
connection_id,
|
||||
OutboundConnectionState::new(
|
||||
writer_tx,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(false)),
|
||||
Arc::new(RwLock::new(HashSet::new())),
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::AppServerNotification(thread_goal_updated_notification()),
|
||||
write_complete_tx: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(
|
||||
writer_rx.try_recv().is_err(),
|
||||
"experimental notifications should not reach clients without capability"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn experimental_notifications_are_preserved_with_capability() {
|
||||
let connection_id = ConnectionId(13);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
connection_id,
|
||||
OutboundConnectionState::new(
|
||||
writer_tx,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(RwLock::new(HashSet::new())),
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::AppServerNotification(thread_goal_updated_notification()),
|
||||
write_complete_tx: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let message = writer_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("experimental notification should reach opted-in client");
|
||||
assert!(matches!(
|
||||
message.message,
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ThreadGoalUpdated(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn command_execution_request_approval_strips_additional_permissions_without_capability() {
|
||||
let connection_id = ConnectionId(8);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
connection_id,
|
||||
OutboundConnectionState::new(
|
||||
writer_tx,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(false)),
|
||||
Arc::new(RwLock::new(HashSet::new())),
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::Request(ServerRequest::CommandExecutionRequestApproval {
|
||||
request_id: RequestId::Integer(1),
|
||||
params: codex_app_server_protocol::CommandExecutionRequestApprovalParams {
|
||||
thread_id: "thr_123".to_string(),
|
||||
turn_id: "turn_123".to_string(),
|
||||
item_id: "call_123".to_string(),
|
||||
approval_id: None,
|
||||
reason: Some("Need extra read access".to_string()),
|
||||
network_approval_context: None,
|
||||
command: Some("cat file".to_string()),
|
||||
cwd: Some(absolute_path("/tmp")),
|
||||
command_actions: None,
|
||||
additional_permissions: Some(
|
||||
codex_app_server_protocol::AdditionalPermissionProfile {
|
||||
network: None,
|
||||
file_system: Some(
|
||||
codex_app_server_protocol::AdditionalFileSystemPermissions {
|
||||
read: Some(vec![absolute_path("/tmp/allowed")]),
|
||||
write: None,
|
||||
glob_scan_max_depth: None,
|
||||
entries: None,
|
||||
},
|
||||
),
|
||||
},
|
||||
),
|
||||
proposed_execpolicy_amendment: None,
|
||||
proposed_network_policy_amendments: None,
|
||||
available_decisions: None,
|
||||
},
|
||||
}),
|
||||
write_complete_tx: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let message = writer_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("request should be delivered to the connection");
|
||||
let json = serde_json::to_value(message.message).expect("request should serialize");
|
||||
assert_eq!(json["params"].get("additionalPermissions"), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn command_execution_request_approval_keeps_additional_permissions_with_capability() {
|
||||
let connection_id = ConnectionId(9);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
connection_id,
|
||||
OutboundConnectionState::new(
|
||||
writer_tx,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(RwLock::new(HashSet::new())),
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::Request(ServerRequest::CommandExecutionRequestApproval {
|
||||
request_id: RequestId::Integer(1),
|
||||
params: codex_app_server_protocol::CommandExecutionRequestApprovalParams {
|
||||
thread_id: "thr_123".to_string(),
|
||||
turn_id: "turn_123".to_string(),
|
||||
item_id: "call_123".to_string(),
|
||||
approval_id: None,
|
||||
reason: Some("Need extra read access".to_string()),
|
||||
network_approval_context: None,
|
||||
command: Some("cat file".to_string()),
|
||||
cwd: Some(absolute_path("/tmp")),
|
||||
command_actions: None,
|
||||
additional_permissions: Some(
|
||||
codex_app_server_protocol::AdditionalPermissionProfile {
|
||||
network: None,
|
||||
file_system: Some(
|
||||
codex_app_server_protocol::AdditionalFileSystemPermissions {
|
||||
read: Some(vec![absolute_path("/tmp/allowed")]),
|
||||
write: None,
|
||||
glob_scan_max_depth: None,
|
||||
entries: None,
|
||||
},
|
||||
),
|
||||
},
|
||||
),
|
||||
proposed_execpolicy_amendment: None,
|
||||
proposed_network_policy_amendments: None,
|
||||
available_decisions: None,
|
||||
},
|
||||
}),
|
||||
write_complete_tx: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let message = writer_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("request should be delivered to the connection");
|
||||
let json = serde_json::to_value(message.message).expect("request should serialize");
|
||||
let allowed_path = absolute_path("/tmp/allowed").to_string_lossy().into_owned();
|
||||
assert_eq!(
|
||||
json["params"]["additionalPermissions"],
|
||||
json!({
|
||||
"network": null,
|
||||
"fileSystem": {
|
||||
"read": [allowed_path],
|
||||
"write": null,
|
||||
},
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn broadcast_does_not_block_on_slow_connection() {
|
||||
let fast_connection_id = ConnectionId(1);
|
||||
let slow_connection_id = ConnectionId(2);
|
||||
|
||||
let (fast_writer_tx, mut fast_writer_rx) = mpsc::channel(1);
|
||||
let (slow_writer_tx, mut slow_writer_rx) = mpsc::channel(1);
|
||||
let fast_disconnect_token = CancellationToken::new();
|
||||
let slow_disconnect_token = CancellationToken::new();
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
fast_connection_id,
|
||||
OutboundConnectionState::new(
|
||||
fast_writer_tx,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(RwLock::new(HashSet::new())),
|
||||
Some(fast_disconnect_token.clone()),
|
||||
),
|
||||
);
|
||||
connections.insert(
|
||||
slow_connection_id,
|
||||
OutboundConnectionState::new(
|
||||
slow_writer_tx.clone(),
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(RwLock::new(HashSet::new())),
|
||||
Some(slow_disconnect_token.clone()),
|
||||
),
|
||||
);
|
||||
|
||||
let queued_message = OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification {
|
||||
summary: "already-buffered".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
},
|
||||
));
|
||||
slow_writer_tx
|
||||
.try_send(QueuedOutgoingMessage::new(queued_message))
|
||||
.expect("channel should have room");
|
||||
|
||||
let broadcast_message = OutgoingMessage::AppServerNotification(
|
||||
ServerNotification::ConfigWarning(ConfigWarningNotification {
|
||||
summary: "test".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
}),
|
||||
);
|
||||
timeout(
|
||||
Duration::from_millis(100),
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::Broadcast {
|
||||
message: broadcast_message,
|
||||
},
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("broadcast should return even when one connection is slow");
|
||||
assert!(!connections.contains_key(&slow_connection_id));
|
||||
assert!(slow_disconnect_token.is_cancelled());
|
||||
assert!(!fast_disconnect_token.is_cancelled());
|
||||
let fast_message = fast_writer_rx
|
||||
.try_recv()
|
||||
.expect("fast connection should receive the broadcast notification");
|
||||
assert!(matches!(
|
||||
fast_message.message,
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification { summary, .. }
|
||||
)) if summary == "test"
|
||||
));
|
||||
|
||||
let slow_message = slow_writer_rx
|
||||
.try_recv()
|
||||
.expect("slow connection should retain its original buffered message");
|
||||
assert!(matches!(
|
||||
slow_message.message,
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification { summary, .. }
|
||||
)) if summary == "already-buffered"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn to_connection_stdio_waits_instead_of_disconnecting_when_writer_queue_is_full() {
|
||||
let connection_id = ConnectionId(3);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
writer_tx
|
||||
.send(QueuedOutgoingMessage::new(
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification {
|
||||
summary: "queued".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
},
|
||||
)),
|
||||
))
|
||||
.await
|
||||
.expect("channel should accept the first queued message");
|
||||
|
||||
let mut connections = HashMap::new();
|
||||
connections.insert(
|
||||
connection_id,
|
||||
OutboundConnectionState::new(
|
||||
writer_tx,
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(AtomicBool::new(true)),
|
||||
Arc::new(RwLock::new(HashSet::new())),
|
||||
/*disconnect_sender*/ None,
|
||||
),
|
||||
);
|
||||
|
||||
let route_task = tokio::spawn(async move {
|
||||
route_outgoing_envelope(
|
||||
&mut connections,
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification {
|
||||
summary: "second".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
},
|
||||
)),
|
||||
write_complete_tx: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
let first = timeout(Duration::from_millis(100), writer_rx.recv())
|
||||
.await
|
||||
.expect("first queued message should be readable")
|
||||
.expect("first queued message should exist");
|
||||
timeout(Duration::from_millis(100), route_task)
|
||||
.await
|
||||
.expect("routing should finish after the first queued message is drained")
|
||||
.expect("routing task should succeed");
|
||||
|
||||
assert!(matches!(
|
||||
first.message,
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification { summary, .. }
|
||||
)) if summary == "queued"
|
||||
));
|
||||
let second = writer_rx
|
||||
.try_recv()
|
||||
.expect("second notification should be delivered once the queue has room");
|
||||
assert!(matches!(
|
||||
second.message,
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification { summary, .. }
|
||||
)) if summary == "second"
|
||||
));
|
||||
}
|
||||
Reference in New Issue
Block a user