Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 58 Packages Projects Releases 51 Wiki Activity

Harden stream cleanup state

Browse Source
This commit is contained in:
DuProcess
2026-07-09 19:27:45 -04:00
parent bab3f777dc
commit 1466516053
2 changed files with 279 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dosh-server.rs
+115 -22
View File
@@ -2133,11 +2133,7 @@ async fn handle_stream_open_reject(
}
client.endpoint = peer;
client.last_seen = Instant::now();
client.stream_writers.remove(&reject.stream_id);
client.stream_pending_opens.remove(&reject.stream_id);
client.opened_streams.remove(&reject.stream_id);
client.stream_send_credit.remove(&reject.stream_id);
client.stream_pending_data.remove(&reject.stream_id);
cleanup_client_stream(client, reject.stream_id);
Ok(())
}
@@ -2192,15 +2188,7 @@ async fn handle_stream_close(
// Connection migration on any authenticated, fresh packet (spec §11).
client.endpoint = peer;
client.last_seen = Instant::now();
client.stream_writers.remove(&close.stream_id);
client.stream_pending_opens.remove(&close.stream_id);
client.opened_streams.remove(&close.stream_id);
client.stream_send_credit.remove(&close.stream_id);
client.stream_pending_data.remove(&close.stream_id);
client.stream_next_send_offset.remove(&close.stream_id);
client.stream_sent_data.remove(&close.stream_id);
client.stream_next_recv_offset.remove(&close.stream_id);
client.stream_recv_pending.remove(&close.stream_id);
cleanup_client_stream(client, close.stream_id);
Ok(())
}
@@ -3280,6 +3268,18 @@ fn accept_stream_data(client: &mut ClientState, data: StreamData) -> (Vec<Vec<u8
(writes, consumed, *expected)
}
fn cleanup_client_stream(client: &mut ClientState, stream_id: u64) {
client.stream_writers.remove(&stream_id);
client.stream_pending_opens.remove(&stream_id);
client.opened_streams.remove(&stream_id);
client.stream_send_credit.remove(&stream_id);
client.stream_pending_data.remove(&stream_id);
client.stream_next_send_offset.remove(&stream_id);
client.stream_sent_data.remove(&stream_id);
client.stream_next_recv_offset.remove(&stream_id);
client.stream_recv_pending.remove(&stream_id);
}
fn ack_stream_data(client: &mut ClientState, stream_id: u64, received_offset: u64) {
let Some(sent) = client.stream_sent_data.get_mut(&stream_id) else {
return;
@@ -3340,14 +3340,7 @@ async fn send_stream_close_to_client(
{
let mut locked = state.lock().expect("server state poisoned");
if let Some(client) = locked.client_mut(&client_id) {
client.stream_writers.remove(&stream_id);
client.opened_streams.remove(&stream_id);
client.stream_send_credit.remove(&stream_id);
client.stream_pending_data.remove(&stream_id);
client.stream_next_send_offset.remove(&stream_id);
client.stream_sent_data.remove(&stream_id);
client.stream_next_recv_offset.remove(&stream_id);
client.stream_recv_pending.remove(&stream_id);
cleanup_client_stream(client, stream_id);
}
}
let body = protocol::to_body(&StreamClose { stream_id })?;
@@ -4053,6 +4046,106 @@ mod tests {
}
}
#[test]
fn cleanup_client_stream_removes_all_per_stream_state_only_for_target() {
let stream_id = 42;
let other_stream_id = 43;
let mut client = test_client_state([7u8; 32]);
let (writer_tx, _writer_rx) = mpsc::channel(1);
client.stream_writers.insert(stream_id, writer_tx);
client.opened_streams.insert(stream_id);
client.opened_streams.insert(other_stream_id);
client.stream_send_credit.insert(stream_id, 12);
client.stream_send_credit.insert(other_stream_id, 34);
client
.stream_pending_data
.insert(stream_id, VecDeque::from([vec![1]]));
client
.stream_pending_data
.insert(other_stream_id, VecDeque::from([vec![2]]));
client.stream_pending_opens.insert(
stream_id,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 1234,
last_sent: Instant::now(),
attempts: 1,
},
);
client.stream_pending_opens.insert(
other_stream_id,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 1235,
last_sent: Instant::now(),
attempts: 1,
},
);
client.stream_next_send_offset.insert(stream_id, 55);
client.stream_next_send_offset.insert(other_stream_id, 66);
client.stream_sent_data.insert(
stream_id,
BTreeMap::from([(
0,
PendingStreamChunk {
offset: 0,
bytes: vec![3],
last_sent: Instant::now(),
attempts: 1,
},
)]),
);
client.stream_sent_data.insert(
other_stream_id,
BTreeMap::from([(
0,
PendingStreamChunk {
offset: 0,
bytes: vec![4],
last_sent: Instant::now(),
attempts: 1,
},
)]),
);
client.stream_next_recv_offset.insert(stream_id, 77);
client.stream_next_recv_offset.insert(other_stream_id, 88);
client
.stream_recv_pending
.insert(stream_id, BTreeMap::from([(77, vec![5])]));
client
.stream_recv_pending
.insert(other_stream_id, BTreeMap::from([(88, vec![6])]));
cleanup_client_stream(&mut client, stream_id);
assert!(!client.stream_writers.contains_key(&stream_id));
assert!(!client.opened_streams.contains(&stream_id));
assert!(!client.stream_send_credit.contains_key(&stream_id));
assert!(!client.stream_pending_data.contains_key(&stream_id));
assert!(!client.stream_pending_opens.contains_key(&stream_id));
assert!(!client.stream_next_send_offset.contains_key(&stream_id));
assert!(!client.stream_sent_data.contains_key(&stream_id));
assert!(!client.stream_next_recv_offset.contains_key(&stream_id));
assert!(!client.stream_recv_pending.contains_key(&stream_id));
assert!(client.opened_streams.contains(&other_stream_id));
assert!(client.stream_send_credit.contains_key(&other_stream_id));
assert!(client.stream_pending_data.contains_key(&other_stream_id));
assert!(client.stream_pending_opens.contains_key(&other_stream_id));
assert!(
client
.stream_next_send_offset
.contains_key(&other_stream_id)
);
assert!(client.stream_sent_data.contains_key(&other_stream_id));
assert!(
client
.stream_next_recv_offset
.contains_key(&other_stream_id)
);
assert!(client.stream_recv_pending.contains_key(&other_stream_id));
}
#[test]
fn client_index_stays_in_sync_with_session_clients() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();