jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-11-01 06:51:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(relay): fence request rejection output

Browse Source 
Signed-off-by: loganj <loganj@squareup.com>
This commit is contained in:
loganj
2026-08-19 14:53:19 +00:00
parent 9e1c7abfac
commit e8e0fcc34a
2 changed files with 202 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

135
crates/buzz-relay/src/connection.rs
View File

@@ -220,12 +220,20 @@ impl ConnectionState {
}
}
pub(crate) async fn send_historical_if_permitted(
/// Queue output owned by one REQ only while its lease remains live.
///
/// The lifecycle lock covers the complete batch so a paired NOTICE/CLOSED
/// response cannot straddle a same-ID replacement or CLOSE cutover. The
/// same fence also covers historical EVENT/EOSE output.
pub(crate) async fn send_req_messages_if_permitted<I>(
&self,
sub_id: &str,
pending: &Arc<PendingSubscription>,
msg: String,
) -> bool {
messages: I,
) -> bool
where
I: IntoIterator<Item = String>,
{
let pending_subscriptions = self.pending_subscriptions.lock().await;
let permitted = pending_subscriptions
.get(sub_id)
@@ -234,10 +242,20 @@ impl ConnectionState {
return false;
}
// Keep the lifecycle lock through queue insertion. Replacement commit
// and CLOSE take the same lock, so neither can publish its cutover and
// then be followed by a stale historical frame from this lease.
self.send(msg)
// Keep the lifecycle lock through every queue insertion. Replacement
// commit and CLOSE take the same lock, so neither can publish its
// cutover and then be followed by output from this lease.
messages.into_iter().all(|msg| self.send(msg))
}
pub(crate) async fn send_historical_if_permitted(
&self,
sub_id: &str,
pending: &Arc<PendingSubscription>,
msg: String,
) -> bool {
self.send_req_messages_if_permitted(sub_id, pending, [msg])
.await
}
/// Sends a data message to this connection's outbound channel.
@@ -1110,6 +1128,109 @@ mod tests {
assert!(send_rx.try_recv().is_err());
}
#[tokio::test]
async fn stale_req_rejection_is_fenced_after_same_id_replacement_commits() {
let (send_tx, mut send_rx) = mpsc::channel(8);
let conn = Arc::new(ConnectionState {
conn_id: Uuid::new_v4(),
tenant: TenantContext::resolved(
buzz_core::CommunityId::from_uuid(Uuid::new_v4()),
"lease.test",
),
remote_addr: "127.0.0.1:1234".parse().expect("socket address"),
auth_state: RwLock::new(AuthState::Failed),
subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
pending_subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
send_tx,
ctrl_tx: mpsc::channel(1).0,
cancel: CancellationToken::new(),
backpressure_count: Arc::new(AtomicU8::new(0)),
grace_limit: 3,
});
let predecessor = conn.begin_pending_subscription("same-id").await;
let replacement = conn.begin_pending_subscription("same-id").await;
let lifecycle = conn.pending_subscriptions.lock().await;
// Model A resuming from an access-resolution error while B is at its
// registration cutover. A must wait for the lifecycle lock rather than
// queue its terminal response directly.
let stale_rejection = tokio::spawn({
let conn = Arc::clone(&conn);
let predecessor = Arc::clone(&predecessor);
async move {
conn.send_req_messages_if_permitted(
"same-id",
&predecessor,
["old-notice".into(), "old-closed".into()],
)
.await
}
});
tokio::task::yield_now().await;
replacement.commit();
assert!(conn.send("new-owner".into()));
drop(lifecycle);
assert!(!stale_rejection.await.expect("stale rejection task"));
assert_eq!(
send_rx.recv().await,
Some(WsMessage::Text("new-owner".into()))
);
assert!(send_rx.try_recv().is_err());
}
#[tokio::test]
async fn stale_req_rejection_is_fenced_after_close_acknowledgement() {
let (send_tx, mut send_rx) = mpsc::channel(8);
let conn = Arc::new(ConnectionState {
conn_id: Uuid::new_v4(),
tenant: TenantContext::resolved(
buzz_core::CommunityId::from_uuid(Uuid::new_v4()),
"lease.test",
),
remote_addr: "127.0.0.1:1234".parse().expect("socket address"),
auth_state: RwLock::new(AuthState::Failed),
subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
pending_subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
send_tx,
ctrl_tx: mpsc::channel(1).0,
cancel: CancellationToken::new(),
backpressure_count: Arc::new(AtomicU8::new(0)),
grace_limit: 3,
});
let pending = conn.begin_pending_subscription("same-id").await;
let mut lifecycle = conn.pending_subscriptions.lock().await;
let stale_rejection = tokio::spawn({
let conn = Arc::clone(&conn);
let pending = Arc::clone(&pending);
async move {
conn.send_req_messages_if_permitted(
"same-id",
&pending,
["old-notice".into(), "old-closed".into()],
)
.await
}
});
tokio::task::yield_now().await;
lifecycle
.remove("same-id")
.expect("pending owner")
.cancel_lineage();
assert!(conn.send("close-ack".into()));
drop(lifecycle);
assert!(!stale_rejection.await.expect("stale rejection task"));
assert_eq!(
send_rx.recv().await,
Some(WsMessage::Text("close-ack".into()))
);
assert!(send_rx.try_recv().is_err());
}
#[tokio::test]
async fn committed_same_id_replacement_cancels_the_previous_lease() {
let original = Arc::new(PendingSubscription::new(None));

98
crates/buzz-relay/src/handlers/req.rs
View File

@@ -81,11 +81,14 @@ pub(crate) async fn handle_req(
match &*auth {
AuthState::Authenticated(ctx) => {
if !ctx.scopes.is_empty() && !ctx.scopes.contains(&Scope::MessagesRead) {
conn.send(RelayMessage::notice("restricted: insufficient scope"));
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
Some("restricted: insufficient scope"),
"restricted: insufficient scope",
));
)
.await;
return;
}
@@ -93,23 +96,28 @@ pub(crate) async fn handle_req(
let subs = conn.subscriptions.lock().await;
if !subs.contains_key(&sub_id) && subs.len() >= MAX_SUBSCRIPTIONS {
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
None,
"error: too many subscriptions",
));
)
.await;
return;
}
(conn.conn_id, pk_bytes, ctx.channel_ids.clone())
}
_ => {
conn.send(RelayMessage::notice(
"auth-required: authenticate before subscribing",
));
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
Some("auth-required: authenticate before subscribing"),
"auth-required: not authenticated",
));
)
.await;
return;
}
}
@@ -119,10 +127,14 @@ pub(crate) async fn handle_req(
let requested_channel_ids = match extract_channel_ids_from_filters_limited(&filters) {
Ok(ids) => ids,
Err(()) => {
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
None,
"restricted: too many explicit channels",
));
)
.await;
return;
}
};
@@ -139,7 +151,7 @@ pub(crate) async fn handle_req(
Ok(ids) => ids,
Err(e) => {
warn!(conn_id = %conn_id, "Failed to get accessible channels: {e}");
conn.send(RelayMessage::closed(&sub_id, "error: database error"));
send_req_rejection(&conn, &sub_id, &pending, None, "error: database error").await;
return;
}
}
@@ -192,7 +204,8 @@ pub(crate) async fn handle_req(
}
Err(e) => {
warn!(conn_id = %conn_id, "Channel membership confirmation failed: {e}");
conn.send(RelayMessage::closed(&sub_id, "error: database error"));
send_req_rejection(&conn, &sub_id, &pending, None, "error: database error")
.await;
return;
}
}
@@ -223,10 +236,14 @@ pub(crate) async fn handle_req(
.as_ref()
.is_some_and(|authorized| authorized.is_empty())
{
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
None,
"restricted: not a channel member",
));
)
.await;
return;
}
@@ -241,24 +258,36 @@ pub(crate) async fn handle_req(
if channel_id.is_none() {
let authed_pubkey_hex = hex::encode(&pubkey_bytes);
if !p_gated_filters_authorized(&filters, &authed_pubkey_hex) {
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
None,
"restricted: p-gated events require #p matching your pubkey",
));
)
.await;
return;
}
if !engram_filters_authorized(&filters, &authed_pubkey_hex) {
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
None,
"restricted: agent-engram reads require authors=[self] or #p=[self]",
));
)
.await;
return;
}
if !author_only_filters_authorized(&filters, &authed_pubkey_hex) {
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
None,
"restricted: author-only kinds require authors=[self]",
));
)
.await;
return;
}
}
@@ -270,10 +299,14 @@ pub(crate) async fn handle_req(
let has_search = filters.iter().any(|f| f.search.is_some());
if has_search {
if filters.iter().any(|f| f.search.is_none()) {
conn.send(RelayMessage::closed(
send_req_rejection(
&conn,
&sub_id,
&pending,
None,
"error: mixed search and non-search filters not supported",
));
)
.await;
return;
}
handle_search_req(
@@ -494,6 +527,23 @@ pub(crate) async fn handle_req(
/// pages to the request.
const SEARCH_PAGE_SIZE: u32 = 100;
/// Send a terminal response owned by this REQ while its lifecycle lease is current.
/// A paired NOTICE and CLOSED are queued as one fenced batch.
async fn send_req_rejection(
conn: &ConnectionState,
sub_id: &str,
pending: &Arc<PendingSubscription>,
notice: Option<&str>,
reason: &str,
) -> bool {
let messages = notice
.into_iter()
.map(RelayMessage::notice)
.chain(std::iter::once(RelayMessage::closed(sub_id, reason)));
conn.send_req_messages_if_permitted(sub_id, pending, messages)
.await
}
/// Atomically commit a REQ into the per-connection map, fan-out registry, and
/// Redis topic refcount only while it still owns the pending lease. CLOSE uses
/// the same pending-subscription lock, so it either invalidates this lease