Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 46 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-client.rs
+164 -37
View File
@@ -5022,13 +5022,17 @@ async fn run_terminal(
{
let _ = writer.write_all(&[5, 5, 0, 1, 0, 0, 0, 0, 0, 0]).await;
}
opened_streams.remove(&reject.stream_id);
stream_send_credit.remove(&reject.stream_id);
stream_pending_data.remove(&reject.stream_id);
stream_next_send_offset.remove(&reject.stream_id);
stream_sent_data.remove(&reject.stream_id);
stream_next_recv_offset.remove(&reject.stream_id);
stream_recv_pending.remove(&reject.stream_id);
cleanup_stream_state(
reject.stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
);
stream_writers.remove(&reject.stream_id);
}
PacketKind::StreamData => {
@@ -5101,14 +5105,17 @@ async fn run_terminal(
};
last_packet_at = Instant::now();
pending_socks_replies.remove(&close.stream_id);
stream_pending_opens.remove(&close.stream_id);
opened_streams.remove(&close.stream_id);
stream_send_credit.remove(&close.stream_id);
stream_pending_data.remove(&close.stream_id);
stream_next_send_offset.remove(&close.stream_id);
stream_sent_data.remove(&close.stream_id);
stream_next_recv_offset.remove(&close.stream_id);
stream_recv_pending.remove(&close.stream_id);
cleanup_stream_state(
close.stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
);
stream_writers.remove(&close.stream_id);
#[cfg(unix)]
agent_writers.remove(&close.stream_id);
@@ -5160,14 +5167,17 @@ async fn run_terminal(
}
Some(ForwardEvent::Close { stream_id }) => {
pending_socks_replies.remove(&stream_id);
stream_pending_opens.remove(&stream_id);
opened_streams.remove(&stream_id);
stream_send_credit.remove(&stream_id);
stream_pending_data.remove(&stream_id);
stream_next_send_offset.remove(&stream_id);
stream_sent_data.remove(&stream_id);
stream_next_recv_offset.remove(&stream_id);
stream_recv_pending.remove(&stream_id);
cleanup_stream_state(
stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
);
stream_writers.remove(&stream_id);
#[cfg(unix)]
agent_writers.remove(&stream_id);
@@ -6140,6 +6150,28 @@ fn accept_stream_data(
(writes, consumed, *expected)
}
#[allow(clippy::too_many_arguments)]
fn cleanup_stream_state(
stream_id: u64,
opened_streams: &mut HashSet<u64>,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
stream_pending_opens: &mut HashMap<u64, PendingStreamOpen>,
stream_next_send_offset: &mut HashMap<u64, u64>,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_next_recv_offset: &mut HashMap<u64, u64>,
stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
) {
opened_streams.remove(&stream_id);
stream_send_credit.remove(&stream_id);
stream_pending_data.remove(&stream_id);
stream_pending_opens.remove(&stream_id);
stream_next_send_offset.remove(&stream_id);
stream_sent_data.remove(&stream_id);
stream_next_recv_offset.remove(&stream_id);
stream_recv_pending.remove(&stream_id);
}
fn ack_stream_data(
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_send_credit: &mut HashMap<u64, usize>,
@@ -7425,19 +7457,20 @@ mod tests {
FRAME_GAP_RESYNC_AFTER_MS, FrameBuffer, LocalForward, MAX_PENDING_USER_INPUT_BYTES,
POST_SUBMIT_ALL_INPUT_HOLD, PendingStreamOpen, PredictMode, Predictor,
RESTART_STATUS_SCRIPT, RemoteForward, STARTUP_INPUT_HOLD, SshConfig, StartupGateMode,
StatusAction, auth_allows, cache_key, cache_server_prefix, clear_cached_credentials,
ensure_tui_safe_status_overlay, input_contains_focus_in, input_matches_escape,
is_local_status_target, is_resume_response_for_client, latest_release_download_url,
load_first_native_identity_with_prompt, parse_dynamic_forward, parse_escape_key,
parse_local_forward, parse_remote_forward, parse_ssh_config, post_submit_hold_duration,
queue_pending_user_input, raw_contains_host_table, recv_response_until, refresh_live_addr,
release_tag_download_url, release_tag_from_effective_url, release_version_from_tag,
render_status_clear, render_status_overlay, requested_env, resolved_startup_command,
retransmit_stream_opens, rewrite_forward_command, selected_predict_mode, selected_udp_host,
server_version_mismatch, should_flush_terminal_input_after_contact,
should_hold_during_startup_gate, should_hold_post_submit_input,
should_repaint_idle_alternate_screen, split_after_command_submit, ssh_destination_host,
ssh_username, ssh_with_user, startup_command, status_ssh_target, strip_stale_mouse_reports,
StatusAction, auth_allows, cache_key, cache_server_prefix, cleanup_stream_state,
clear_cached_credentials, ensure_tui_safe_status_overlay, input_contains_focus_in,
input_matches_escape, is_local_status_target, is_resume_response_for_client,
latest_release_download_url, load_first_native_identity_with_prompt, parse_dynamic_forward,
parse_escape_key, parse_local_forward, parse_remote_forward, parse_ssh_config,
post_submit_hold_duration, queue_pending_user_input, raw_contains_host_table,
recv_response_until, refresh_live_addr, release_tag_download_url,
release_tag_from_effective_url, release_version_from_tag, render_status_clear,
render_status_overlay, requested_env, resolved_startup_command, retransmit_stream_opens,
rewrite_forward_command, selected_predict_mode, selected_udp_host, server_version_mismatch,
should_flush_terminal_input_after_contact, should_hold_during_startup_gate,
should_hold_post_submit_input, should_repaint_idle_alternate_screen,
split_after_command_submit, ssh_destination_host, ssh_username, ssh_with_user,
startup_command, status_ssh_target, strip_stale_mouse_reports,
strip_terminal_focus_reports, terminal_private_mode_transition, toml_bare_key_or_quoted,
update_check_requested, update_version_status, upsert_managed_block, valid_forward_host,
vscode_safe_alias,
@@ -7447,7 +7480,7 @@ mod tests {
use dosh::protocol::{self, Frame, PacketKind};
use dosh::udp::is_transient_udp_send_error;
use ed25519_dalek::{SigningKey, VerifyingKey};
use std::collections::{HashMap, VecDeque};
use std::collections::{BTreeMap, HashMap, VecDeque};
use std::fs;
use std::time::{Duration, Instant};
@@ -7493,6 +7526,100 @@ mod tests {
assert_eq!(vscode_safe_alias(":///"), "host");
}
#[test]
fn cleanup_stream_state_removes_all_per_stream_state_only_for_target() {
let stream_id = 42;
let other_stream_id = 43;
let mut opened_streams = std::collections::HashSet::from([stream_id, other_stream_id]);
let mut stream_send_credit = HashMap::from([(stream_id, 12), (other_stream_id, 34)]);
let mut stream_pending_data = HashMap::from([
(stream_id, VecDeque::from([vec![1]])),
(other_stream_id, VecDeque::from([vec![2]])),
]);
let mut stream_pending_opens = HashMap::from([
(
stream_id,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 1234,
last_sent: Instant::now(),
attempts: 1,
},
),
(
other_stream_id,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 1235,
last_sent: Instant::now(),
attempts: 1,
},
),
]);
let mut stream_next_send_offset = HashMap::from([(stream_id, 55), (other_stream_id, 66)]);
let mut stream_sent_data = HashMap::from([
(
stream_id,
BTreeMap::from([(
0,
super::PendingStreamChunk {
offset: 0,
bytes: vec![3],
last_sent: Instant::now(),
attempts: 1,
},
)]),
),
(
other_stream_id,
BTreeMap::from([(
0,
super::PendingStreamChunk {
offset: 0,
bytes: vec![4],
last_sent: Instant::now(),
attempts: 1,
},
)]),
),
]);
let mut stream_next_recv_offset = HashMap::from([(stream_id, 77), (other_stream_id, 88)]);
let mut stream_recv_pending = HashMap::from([
(stream_id, BTreeMap::from([(77, vec![5])])),
(other_stream_id, BTreeMap::from([(88, vec![6])])),
]);
cleanup_stream_state(
stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
);
assert!(!opened_streams.contains(&stream_id));
assert!(!stream_send_credit.contains_key(&stream_id));
assert!(!stream_pending_data.contains_key(&stream_id));
assert!(!stream_pending_opens.contains_key(&stream_id));
assert!(!stream_next_send_offset.contains_key(&stream_id));
assert!(!stream_sent_data.contains_key(&stream_id));
assert!(!stream_next_recv_offset.contains_key(&stream_id));
assert!(!stream_recv_pending.contains_key(&stream_id));
assert!(opened_streams.contains(&other_stream_id));
assert!(stream_send_credit.contains_key(&other_stream_id));
assert!(stream_pending_data.contains_key(&other_stream_id));
assert!(stream_pending_opens.contains_key(&other_stream_id));
assert!(stream_next_send_offset.contains_key(&other_stream_id));
assert!(stream_sent_data.contains_key(&other_stream_id));
assert!(stream_next_recv_offset.contains_key(&other_stream_id));
assert!(stream_recv_pending.contains_key(&other_stream_id));
}
#[test]
fn managed_ssh_block_is_replaced_not_duplicated() {
let dir = tempfile::tempdir().unwrap();