openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-02 19:43:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Keep Noise relay streams alive after repeated handshake failures (#46031)

Browse Source 
## Why

Exhausting the handshake failure budget closed the physical relay, disconnecting authenticated streams along with failed attempts.

## What changed

After eight failed handshakes, pause new handshake admission for 10 seconds. Reset incoming handshake attempts before parsing them while existing streams and pending validations continue. Failures during the cooldown do not extend it, and the failure budget resets when handshake admission resumes.

## Testing

Update relay tests to cover duplicate handshakes and early data without disconnecting the relay. Verify that encrypted traffic on an established stream continues during cooldown, rejected attempts skip authorization checks, and new handshakes succeed after cooldown expires.

GitOrigin-RevId: 0a5098ed1b3c31466f9c091f46cbb20415c7aa1d
This commit is contained in:
viyatb-oai
2026-09-16 21:21:45 +00:00
committed by copyberry
parent 2b2b0fa870
commit 105fe8761c
2 changed files with 263 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/exec-server/src/relay.rs
+46 -29
View File
@@ -12,6 +12,7 @@ use prost::Message as ProstMessage;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tokio::time::Instant;
use tokio::time::timeout;
use tokio_tungstenite::tungstenite::Message;
use tracing::debug;
@@ -50,6 +51,7 @@ use crate::websocket_pong_watchdog::WebSocketPongWatchdog;
const RELAY_MESSAGE_FRAME_VERSION: u32 = 1;
const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128;
const MAX_FAILED_NOISE_HANDSHAKES: usize = 8;
const NOISE_HANDSHAKE_FAILURE_COOLDOWN: Duration = Duration::from_secs(10);
const MAX_HARNESS_KEY_AUTHORIZATION_BYTES: usize = 4096;
const MAX_PENDING_HANDSHAKE_VALIDATIONS: usize = 32;
const HARNESS_KEY_VALIDATION_TIMEOUT: Duration = Duration::from_secs(10);
@@ -575,6 +577,7 @@ where
let mut pending_handshakes: HashMap<String, PendingHandshake> = HashMap::new();
let mut validation_tasks: JoinSet<HarnessKeyValidationResult> = JoinSet::new();
let mut failed_handshakes = 0usize;
let mut handshake_cooldown_until = None;
let mut next_validation_id = 0u64;
let mut disconnect_reason = RendezvousDisconnectReason::LocalShutdown;
@@ -633,10 +636,10 @@ where
"Noise harness authorization failure details"
);
send_reset(&physical_outgoing_tx, validation_result.stream_id);
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
warn!("closing Noise relay after repeated handshake failures");
break;
}
record_failed_handshake(
&mut failed_handshakes,
&mut handshake_cooldown_until,
);
continue;
}
if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS {
@@ -653,10 +656,10 @@ where
Err(error) => {
warn!("failed to complete Noise relay handshake: {error}");
send_reset(&physical_outgoing_tx, validation_result.stream_id);
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
warn!("closing Noise relay after repeated handshake failures");
break;
}
record_failed_handshake(
&mut failed_handshakes,
&mut handshake_cooldown_until,
);
continue;
}
};
@@ -746,6 +749,16 @@ where
let stream_id = frame.stream_id.clone();
match kind {
RelayFrameBodyKind::Handshake => {
// Stop admitting work before parsing another hybrid handshake.
// Existing streams and already-admitted validations keep running.
if let Some(deadline) = handshake_cooldown_until {
if Instant::now() < deadline {
send_reset(&physical_outgoing_tx, stream_id);
continue;
}
failed_handshakes = 0;
handshake_cooldown_until = None;
}
// Reject duplicate or busy streams before paying for a hybrid
// handshake. Malformed attempts that reach cryptography are
// covered by the connection-wide failure budget below.
@@ -756,10 +769,7 @@ where
// Removing pending state makes the in-flight validation result stale.
if pending_handshakes.remove(&stream_id).is_some() {
send_reset(&physical_outgoing_tx, stream_id);
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
warn!("closing Noise relay after repeated handshake failures");
break;
}
record_failed_handshake(&mut failed_handshakes, &mut handshake_cooldown_until);
continue;
}
if streams.len() >= MAX_ACTIVE_NOISE_RELAY_STREAMS {
@@ -788,10 +798,10 @@ where
Err(error) => {
warn!("failed to read Noise relay handshake request: {error}");
send_reset(&physical_outgoing_tx, stream_id);
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
warn!("closing Noise relay after repeated handshake failures");
break;
}
record_failed_handshake(
&mut failed_handshakes,
&mut handshake_cooldown_until,
);
continue;
}
};
@@ -815,10 +825,7 @@ where
};
let Some(authorization) = authorization else {
send_reset(&physical_outgoing_tx, stream_id);
if failed_handshake_budget_exhausted(&mut failed_handshakes) {
warn!("closing Noise relay after repeated handshake failures");
break;
}
record_failed_handshake(&mut failed_handshakes, &mut handshake_cooldown_until);
continue;
};
let harness_public_key = pending.initiator_public_key.clone();
@@ -860,11 +867,11 @@ where
let canceled_pending_handshake =
pending_handshakes.remove(&stream_id).is_some();
send_reset(&physical_outgoing_tx, stream_id);
if canceled_pending_handshake
&& failed_handshake_budget_exhausted(&mut failed_handshakes)
{
warn!("closing Noise relay after repeated handshake failures");
break;
if canceled_pending_handshake {
record_failed_handshake(
&mut failed_handshakes,
&mut handshake_cooldown_until,
);
}
continue;
};
@@ -909,11 +916,21 @@ where
/// Charge one failed authenticated-channel attempt to this physical relay.
///
/// Closing after a small fixed budget prevents a peer that has not been
/// authorized from triggering unbounded hybrid handshakes or registry checks.
fn failed_handshake_budget_exhausted(failed_handshakes: &mut usize) -> bool {
/// Pause new handshakes after a small fixed budget without disconnecting
/// authenticated streams. In-flight validation failures do not extend the
/// cooldown, so they cannot indefinitely prevent new streams from connecting.
fn record_failed_handshake(
failed_handshakes: &mut usize,
handshake_cooldown_until: &mut Option<Instant>,
) {
if handshake_cooldown_until.is_some() {
return;
}
*failed_handshakes += 1;
*failed_handshakes >= MAX_FAILED_NOISE_HANDSHAKES
if *failed_handshakes >= MAX_FAILED_NOISE_HANDSHAKES {
*handshake_cooldown_until = Some(Instant::now() + NOISE_HANDSHAKE_FAILURE_COOLDOWN);
warn!("pausing Noise relay handshakes after repeated failures");
}
}
/// Responder state held while registry authorization is pending.
codex-rs/exec-server/src/relay_noise_tests.rs
+217 -26
View File
@@ -10,9 +10,12 @@ use codex_exec_server_protocol::RequestId;
use futures::SinkExt;
use futures::StreamExt;
use pretty_assertions::assert_eq;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::net::TcpListener;
use tokio::sync::Notify;
use tokio::time::timeout;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message;
@@ -21,6 +24,7 @@ use tokio_util::task::AbortOnDropHandle;
use super::HarnessKeyValidator;
use super::MAX_FAILED_NOISE_HANDSHAKES;
use super::MAX_HARNESS_KEY_AUTHORIZATION_BYTES;
use super::NOISE_HANDSHAKE_FAILURE_COOLDOWN;
use super::RendezvousDisconnectReason;
use super::run_multiplexed_environment;
use crate::ExecServerError;
@@ -31,6 +35,8 @@ use crate::noise_channel::NoiseChannelIdentity;
use crate::noise_channel::NoiseChannelPublicKey;
use crate::noise_channel::noise_channel_prologue;
use crate::noise_relay::NoiseHarnessConnectionArgs;
use crate::noise_relay::message_framing::JsonRpcMessageDecoder;
use crate::noise_relay::message_framing::frame_jsonrpc_message;
use crate::noise_relay::noise_harness_connection_from_websocket_with_readiness;
use crate::noise_relay::stream_handler::NoiseOutboundMessage;
use crate::noise_relay::stream_handler::NoiseStreamConnection;
@@ -44,6 +50,51 @@ use crate::server::ConnectionProcessor;
const ENVIRONMENT_ID: &str = "environment-1";
const EXECUTOR_REGISTRATION_ID: &str = "registration-1";
async fn next_relay_frame<S>(websocket: &mut WebSocketStream<S>) -> Result<RelayMessageFrame>
where
S: AsyncRead + AsyncWrite + Unpin,
{
loop {
match websocket.next().await {
Some(Ok(Message::Binary(payload))) => {
return Ok(decode_relay_message_frame(payload.as_ref())?);
}
Some(Ok(Message::Ping(payload))) => websocket.send(Message::Pong(payload)).await?,
Some(Ok(Message::Pong(_) | Message::Frame(_))) => {}
Some(Ok(message)) => anyhow::bail!("expected relay frame, got {message:?}"),
Some(Err(error)) => return Err(error.into()),
None => anyhow::bail!("environment closed before sending relay frame"),
}
}
}
#[derive(Clone)]
struct EchoMessages;
impl NoiseStreamHandler for EchoMessages {
type Incoming = JSONRPCMessage;
type Outgoing = JSONRPCMessage;
fn decode(payload: bytes::Bytes) -> Result<Self::Incoming, ExecServerError> {
Ok(serde_json::from_slice(&payload)?)
}
fn encode(message: Self::Outgoing) -> Result<NoiseOutboundMessage, ExecServerError> {
ConnectionProcessor::encode(message)
}
async fn run_connection(
self,
mut connection: NoiseStreamConnection<Self::Incoming, Self::Outgoing>,
) {
while let Some(message) = connection.incoming_rx.recv().await {
if connection.outgoing_tx.send(message).await.is_err() {
break;
}
}
}
}
#[derive(Clone)]
struct ObservedRegistration(
ConnectionProcessor,
@@ -304,7 +355,7 @@ async fn pending_harness_key_validation_does_not_block_new_handshakes() -> Resul
}
#[tokio::test]
async fn duplicate_handshakes_exhaust_failure_budget() -> Result<()> {
async fn duplicate_handshakes_pause_admission_without_disconnecting() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let websocket_url = format!("ws://{}", listener.local_addr()?);
let harness_connection = tokio::spawn(connect_async(websocket_url));
@@ -366,19 +417,11 @@ async fn duplicate_handshakes_exhaust_failure_budget() -> Result<()> {
harness_websocket
.send(Message::Binary(encoded.clone().into()))
.await?;
let payload = timeout(Duration::from_secs(1), async {
loop {
match harness_websocket.next().await {
Some(Ok(Message::Binary(payload))) => break Ok(payload),
Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {}
Some(Ok(message)) => anyhow::bail!("expected reset frame, got {message:?}"),
Some(Err(error)) => break Err(error.into()),
None => anyhow::bail!("environment closed before sending reset"),
}
}
})
let reset = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
let reset = decode_relay_message_frame(payload.as_ref())?;
assert_eq!(reset.stream_id, stream_id);
assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset);
}
@@ -392,10 +435,31 @@ async fn duplicate_handshakes_exhaust_failure_budget() -> Result<()> {
}
})
.await?;
harness_websocket
.send(Message::Binary(encoded.clone().into()))
.await?;
let reset = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset);
// The next valid request is rejected before starting another registry check.
harness_websocket
.send(Message::Binary(encoded.into()))
.await?;
timeout(Duration::from_secs(1), environment_task).await??;
let reset = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset);
assert_eq!(calls.load(Ordering::SeqCst), MAX_FAILED_NOISE_HANDSHAKES);
harness_websocket.close(None).await?;
assert_eq!(
timeout(Duration::from_secs(1), environment_task).await??,
RendezvousDisconnectReason::PeerClose,
);
release.notify_waiters();
Ok(())
}
@@ -457,7 +521,13 @@ async fn oversized_harness_authorization_is_rejected_before_validation() -> Resu
}
#[tokio::test]
async fn repeated_malformed_handshakes_close_the_physical_relay() -> Result<()> {
async fn malformed_handshakes_preserve_active_streams_and_admission_recovers() -> Result<()> {
tokio::time::pause();
// Socket I/O must not auto-advance the paused cooldown or keepalive clocks.
let (_clock_guard, clock_released) = std::sync::mpsc::channel::<()>();
tokio::task::spawn_blocking(move || {
let _ = clock_released.recv();
});
let listener = TcpListener::bind("127.0.0.1:0").await?;
let websocket_url = format!("ws://{}", listener.local_addr()?);
let harness_connection = tokio::spawn(connect_async(websocket_url));
@@ -467,20 +537,46 @@ async fn repeated_malformed_handshakes_close_the_physical_relay() -> Result<()>
let environment_identity = NoiseChannelIdentity::generate()?;
let harness_identity = NoiseChannelIdentity::generate()?;
let environment_task = tokio::spawn(run_multiplexed_environment(
let calls = Arc::new(AtomicUsize::new(0));
let release = Arc::new(Notify::new());
let environment_task = AbortOnDropHandle::new(tokio::spawn(run_multiplexed_environment(
environment_websocket,
ConnectionProcessor::new(ExecServerRuntimePaths::new(
std::env::current_exe()?,
/*codex_linux_sandbox_exe*/ None,
)?),
EchoMessages,
ENVIRONMENT_ID.to_string(),
EXECUTOR_REGISTRATION_ID.to_string(),
environment_identity.clone(),
BlockingValidator {
calls: Arc::new(AtomicUsize::new(0)),
release: Arc::new(Notify::new()),
calls: Arc::clone(&calls),
release: Arc::clone(&release),
},
));
)));
let healthy_stream_id = "healthy";
let prologue =
noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, healthy_stream_id);
let (handshake, request) = InitiatorHandshake::start(
&harness_identity,
&environment_identity.public_key(),
&prologue,
b"authorization",
)?;
release.notify_one();
harness_websocket
.send(Message::Binary(
encode_relay_message_frame(&RelayMessageFrame::handshake(
healthy_stream_id.to_string(),
request,
))
.into(),
))
.await?;
let response = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
assert_eq!(response.stream_id, healthy_stream_id);
let mut healthy_transport = handshake.finish(&response.into_handshake_payload()?)?;
for attempt in 0..MAX_FAILED_NOISE_HANDSHAKES {
let stream_id = format!("malformed-{attempt}");
@@ -497,14 +593,99 @@ async fn repeated_malformed_handshakes_close_the_physical_relay() -> Result<()>
harness_websocket
.send(Message::Binary(encode_relay_message_frame(&frame).into()))
.await?;
let reset = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
assert_eq!(
reset,
RelayMessageFrame::reset(
frame.stream_id,
crate::noise_relay::NOISE_RELAY_RESET_REASON.to_string(),
)
);
}
timeout(Duration::from_secs(1), environment_task).await??;
let new_stream_id = "after-cooldown";
let prologue = noise_channel_prologue(ENVIRONMENT_ID, EXECUTOR_REGISTRATION_ID, new_stream_id);
let (new_handshake, request) = InitiatorHandshake::start(
&harness_identity,
&environment_identity.public_key(),
&prologue,
b"authorization",
)?;
let new_frame = RelayMessageFrame::handshake(new_stream_id.to_string(), request);
// Check both immediate rejection and a later attempt that must not extend
// the cooldown. Encrypted traffic on the healthy stream still round-trips.
for seq in 0..2 {
harness_websocket
.send(Message::Binary(
encode_relay_message_frame(&new_frame).into(),
))
.await?;
let reset = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
assert_eq!(reset.stream_id, new_stream_id);
assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset);
assert_eq!(calls.load(Ordering::SeqCst), 1);
let message = JSONRPCMessage::Response(JSONRPCResponse {
id: RequestId::Integer(i64::from(seq)),
result: serde_json::json!({"healthy": true}),
});
let frame = RelayMessageFrame::data(
healthy_stream_id.to_string(),
seq,
healthy_transport.encrypt(&frame_jsonrpc_message(&message)?)?,
/*trace*/ None,
);
harness_websocket
.send(Message::Binary(encode_relay_message_frame(&frame).into()))
.await?;
let response = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
assert_eq!(response.stream_id, healthy_stream_id);
let data = response.into_data()?;
assert_eq!(data.seq, seq);
let messages =
JsonRpcMessageDecoder::default().push(&healthy_transport.decrypt(&data.payload)?)?;
assert_eq!(messages, vec![message]);
tokio::time::advance(NOISE_HANDSHAKE_FAILURE_COOLDOWN / 2).await;
}
release.notify_one();
harness_websocket
.send(Message::Binary(
encode_relay_message_frame(&new_frame).into(),
))
.await?;
let response = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
assert_eq!(response.stream_id, new_stream_id);
new_handshake.finish(&response.into_handshake_payload()?)?;
assert_eq!(calls.load(Ordering::SeqCst), 2);
harness_websocket.close(None).await?;
assert_eq!(
timeout(Duration::from_secs(1), environment_task).await??,
RendezvousDisconnectReason::PeerClose,
);
Ok(())
}
#[tokio::test]
async fn repeated_early_data_during_validation_closes_the_physical_relay() -> Result<()> {
async fn early_data_during_validation_pauses_admission_without_disconnecting() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let websocket_url = format!("ws://{}", listener.local_addr()?);
let harness_connection = tokio::spawn(connect_async(websocket_url));
@@ -546,8 +727,18 @@ async fn repeated_early_data_during_validation_closes_the_physical_relay() -> Re
.send(Message::Binary(encode_relay_message_frame(&frame).into()))
.await?;
}
let reset = timeout(
Duration::from_secs(1),
next_relay_frame(&mut harness_websocket),
)
.await??;
assert_eq!(reset.validate()?, RelayFrameBodyKind::Reset);
}
timeout(Duration::from_secs(1), environment_task).await??;
harness_websocket.close(None).await?;
assert_eq!(
timeout(Duration::from_secs(1), environment_task).await??,
RendezvousDisconnectReason::PeerClose,
);
Ok(())
}