fix: bound websocket thread window enqueue waits

Wait for the writer to drain within the finite request deadline, terminate when a CLOSED frame cannot be queued, and select parser and wire regressions in CI.

Co-authored-by: Kalvin Chau <kalvin@block.xyz>

Signed-off-by: Kalvin Chau <kalvin@block.xyz>
This commit is contained in:
am
2026-09-23 15:30:29 -07:00
parent f225a0abed
commit 32416d88e6
4 changed files with 95 additions and 25 deletions

View File

@@ -591,6 +591,7 @@ jobs:
cargo test -p buzz-test-client --test e2e_relay invite -- --ignored --nocapture
cargo test -p buzz-test-client --test e2e_relay nip43_membership_snapshots_are_rejected -- --ignored --nocapture
cargo test -p buzz-test-client --test e2e_relay nip29_departure_wire -- --ignored --nocapture
cargo test -p buzz-test-client --test e2e_relay test_thread_window_req_returns_rows_bounds_and_eose -- --ignored --nocapture
env:
RELAY_URL: ws://localhost:3000
GIT_CREDENTIAL_NOSTR_BIN: ${{ github.workspace }}/target/ci/git-credential-nostr

View File

@@ -465,8 +465,9 @@ test-unit:
# non-postgres_tests cases only "pass" without a database by waiting out
# the ~30s sqlx acquire timeout, so they do not belong in the infra-free
# unit job either.
# The protocol parser tests are pure and protect REQ extension routing.
cargo nextest run -p buzz-relay --lib \
-E '(test(/^api::admin::/) - test(=api::admin::tests::disabled_mode_allows_unauthenticated_requests_on_the_admin_host) - test(=api::admin::tests::nip98_mode_unrostered_signer_does_not_consume_a_replay_slot)) + test(/^handlers::channel_authz::/) + test(/^handlers::moderation_authz::/) + test(/^handlers::side_effects::tests::/) + test(/^storage_sweep::tests::/)'
-E '(test(/^api::admin::/) - test(=api::admin::tests::disabled_mode_allows_unauthenticated_requests_on_the_admin_host) - test(=api::admin::tests::nip98_mode_unrostered_signer_does_not_consume_a_replay_slot)) + test(/^handlers::channel_authz::/) + test(/^handlers::moderation_authz::/) + test(/^handlers::side_effects::tests::/) + test(/^storage_sweep::tests::/) + test(/^protocol::tests::/)'
# ACP author-gate and queue tests protect the trust boundary between
# relay events and agent prompts. They are infra-free; ignored lifecycle
# tests remain excluded and run in their dedicated integration lanes.

View File

@@ -467,7 +467,36 @@ async fn websocket_batch_enforces_token_scope_and_rejects_queue_truncation() {
"complete page closed socket"
);
// Leave room for the two EVENT frames, but not for EOSE.
// A slow but draining writer must receive the entire page even when its
// queue holds fewer frames than the window produces.
let (draining, mut draining_rx) = f.ws_conn(None);
let drain = tokio::spawn(async move {
let mut frames = Vec::new();
for _ in 0..7 {
let message = draining_rx.recv().await.expect("complete window frame");
frames.push(serde_json::from_str::<Value>(message.to_text().unwrap()).unwrap());
}
frames
});
crate::handlers::req::handle_req(
"draining".into(),
vec![parsed.clone(); 3],
vec![filter.clone(); 3],
vec![None; 3],
draining.clone(),
f.state.clone(),
)
.await;
let drained = tokio::time::timeout(std::time::Duration::from_secs(2), drain)
.await
.expect("draining writer completes")
.unwrap();
assert_eq!(drained.len(), 7);
assert!(drained[..6].iter().all(|frame| frame[0] == "EVENT"));
assert_eq!(drained[6][0], "EOSE");
assert!(!draining.cancel.is_cancelled());
// A queue with room for the page but not EOSE must close, not complete.
let (terminal, mut terminal_rx) = f.ws_conn(None);
for _ in 0..2 {
assert!(terminal.send("padding".into()));
@@ -491,9 +520,7 @@ async fn websocket_batch_enforces_token_scope_and_rejects_queue_truncation() {
assert_eq!(terminal_frames.len(), 2);
assert!(terminal_frames.iter().all(|frame| frame[0] == "EVENT"));
// The full queue also drops an incorrectly attempted EOSE, so checking
// drained frames alone cannot distinguish a complete from a truncated page.
// Assert the production sender reports truncation at the source.
// The source reports truncation if delivery cannot proceed by its bound.
let (probe, _probe_rx) = f.ws_conn(None);
let requests = super::parse(&vec![filter.clone(); 3]).unwrap();
let sent = crate::handlers::req::serve_thread_windows(
@@ -503,14 +530,14 @@ async fn websocket_batch_enforces_token_scope_and_rejects_queue_truncation() {
&f.keys.public_key(),
&probe,
&f.state,
tokio::time::Instant::now() + std::time::Duration::from_secs(1),
)
.await
.unwrap();
assert!(!sent, "a failed EVENT enqueue must not authorize EOSE");
// A stalled writer cannot receive EOSE; the finite request must close.
let (slow, mut slow_rx) = f.ws_conn(None);
// Capacity four: one row + bounds per filter. Two windows fill it,
// then the third fails to enqueue. Draining later must not reveal EOSE.
crate::handlers::req::handle_req(
"full".into(),
vec![parsed.clone(); 3],
@@ -561,6 +588,23 @@ async fn websocket_batch_enforces_token_scope_and_rejects_queue_truncation() {
.await;
assert_eq!(frame(&mut mixed_rx)[0], "CLOSED");
assert!(mixed_rx.try_recv().is_err());
// A rejected window with a saturated terminal queue must not leave its
// subscription waiting on a socket that will never deliver CLOSED.
let (rejected, _rejected_rx) = f.ws_conn(None);
for _ in 0..4 {
assert!(rejected.send("padding".into()));
}
crate::handlers::req::handle_req(
"invalid-full".into(),
vec![nostr::Filter::new()],
vec![json!({"thread_window":true,"until":1})],
vec![None],
rejected.clone(),
f.state.clone(),
)
.await;
assert!(rejected.cancel.is_cancelled());
}
fn frame(rx: &mut tokio::sync::mpsc::Receiver<axum::extract::ws::Message>) -> Value {

View File

@@ -105,20 +105,18 @@ pub async fn handle_req(
};
if thread_windows.iter().any(Option::is_some) {
if thread_windows.iter().any(Option::is_none) {
conn.send(RelayMessage::closed(
send_thread_window_closed(
&conn,
&sub_id,
"invalid: thread_window cannot mix with other query modes",
));
);
return;
}
let reader = match PublicKey::from_slice(&pubkey_bytes) {
Ok(reader) => reader,
Err(error) => {
warn!(conn_id = %conn_id, %error, "Authenticated reader key became invalid");
conn.send(RelayMessage::closed(
&sub_id,
"error: invalid authenticated reader",
));
send_thread_window_closed(&conn, &sub_id, "error: invalid authenticated reader");
return;
}
};
@@ -128,8 +126,9 @@ pub async fn handle_req(
if let Some(replaced) = state.sub_registry.remove_subscription(conn_id, &sub_id) {
release_subscription_topics(&state, &conn.tenant, &replaced.scope).await;
}
let result = tokio::time::timeout(
crate::api::bridge::thread_window::DEADLINE,
let deadline = tokio::time::Instant::now() + crate::api::bridge::thread_window::DEADLINE;
let result = tokio::time::timeout_at(
deadline,
serve_thread_windows(
&sub_id,
thread_windows.iter().flatten(),
@@ -137,15 +136,13 @@ pub async fn handle_req(
&reader,
&conn,
&state,
deadline,
),
)
.await;
match result {
Ok(Ok(true)) => {
// A window is complete only if its terminal frame was queued.
// A saturated writer must close so the client can retry the
// incomplete page instead of waiting forever for EOSE.
if !conn.send(RelayMessage::eose(&sub_id)) {
if !send_thread_window_frame(&conn, RelayMessage::eose(&sub_id), deadline).await {
conn.cancel.cancel();
}
}
@@ -156,10 +153,7 @@ pub async fn handle_req(
}
Ok(Err((status, body))) => close_thread_window_error(&conn, &sub_id, status, &body),
Err(_) => {
conn.send(RelayMessage::closed(
&sub_id,
"error: thread window deadline exceeded",
));
send_thread_window_closed(&conn, &sub_id, "error: thread window deadline exceeded");
}
}
return;
@@ -571,6 +565,7 @@ pub(crate) async fn serve_thread_windows<'a>(
reader: &PublicKey,
conn: &ConnectionState,
state: &AppState,
deadline: tokio::time::Instant,
) -> Result<bool, crate::api::bridge::thread_window::Error> {
let events = crate::api::bridge::thread_window::query_batch(
state,
@@ -581,13 +576,42 @@ pub(crate) async fn serve_thread_windows<'a>(
)
.await?;
for event in events {
if !conn.send(serde_json::json!(["EVENT", sub_id, event]).to_string()) {
if !send_thread_window_frame(
conn,
serde_json::json!(["EVENT", sub_id, event]).to_string(),
deadline,
)
.await
{
return Ok(false);
}
}
Ok(true)
}
// Finite pages must not fail merely because the producer briefly outruns the
// socket writer. The same deadline bounds both query and queue delivery.
async fn send_thread_window_frame(
conn: &ConnectionState,
frame: String,
deadline: tokio::time::Instant,
) -> bool {
tokio::select! {
biased;
_ = conn.cancel.cancelled() => false,
result = tokio::time::timeout_at(
deadline,
conn.send_tx.send(axum::extract::ws::Message::Text(frame.into())),
) => matches!(result, Ok(Ok(()))),
}
}
fn send_thread_window_closed(conn: &ConnectionState, sub_id: &str, reason: &str) {
if !conn.send(RelayMessage::closed(sub_id, reason)) {
conn.cancel.cancel();
}
}
fn close_thread_window_error(
conn: &ConnectionState,
sub_id: &str,
@@ -604,7 +628,7 @@ fn close_thread_window_error(
} else {
"error"
};
conn.send(RelayMessage::closed(sub_id, &format!("{prefix}: {detail}")));
send_thread_window_closed(conn, sub_id, &format!("{prefix}: {detail}"));
}
/// FTS candidate hits fetched per page. Pages are always full regardless of