Compare commits

...

2 Commits

Author SHA1 Message Date
DuProcess 02607b49b1 Harden reconnect and stream open reliability
ci / test (push) Has been cancelled
ci / fuzz-smoke (push) Has been cancelled
ci / windows-client (push) Has been cancelled
ci / package-release (linux-x86_64, ubuntu-latest) (push) Has been cancelled
ci / package-release (macos-aarch64, macos-14) (push) Has been cancelled
ci / package-release (macos-x86_64, macos-13) (push) Has been cancelled
ci / package-release (windows-x86_64, windows-latest) (push) Has been cancelled
ci / remote-bench (push) Has been cancelled
2026-06-27 21:16:12 -04:00
DuProcess 48e3de2922 Refresh reconnect send address and guard prebuilts
ci / test (push) Has been cancelled
ci / fuzz-smoke (push) Has been cancelled
ci / windows-client (push) Has been cancelled
ci / package-release (linux-x86_64, ubuntu-latest) (push) Has been cancelled
ci / package-release (macos-aarch64, macos-14) (push) Has been cancelled
ci / package-release (macos-x86_64, macos-13) (push) Has been cancelled
ci / package-release (windows-x86_64, windows-latest) (push) Has been cancelled
ci / remote-bench (push) Has been cancelled
2026-06-27 20:11:41 -04:00
7 changed files with 573 additions and 30 deletions
Generated
+1 -1
View File
@@ -436,7 +436,7 @@ dependencies = [
[[package]]
name = "dosh"
version = "0.1.2"
version = "0.1.4"
dependencies = [
"anyhow",
"base64",
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "dosh"
version = "0.1.2"
version = "0.1.4"
edition = "2024"
license = "MIT"
+26
View File
@@ -116,6 +116,31 @@ function Verify-ArchiveChecksum($Url, $Archive) {
}
}
function Expected-ArchiveVersion($Url) {
if ($Url -match "/releases/download/v([^/]+)/") {
return $Matches[1]
}
if (-not $BinaryUrl -and -not $BinaryBase -and $BinaryVersion -ne "latest") {
return $BinaryVersion.TrimStart("v")
}
return $null
}
function Verify-ArchiveVersion($ExtractDir, $Url) {
$expected = Expected-ArchiveVersion $Url
if (-not $expected) {
return
}
$versionFile = Get-ChildItem -Path $ExtractDir -Recurse -File -Filter VERSION | Select-Object -First 1
if (-not $versionFile) {
throw "prebuilt archive missing VERSION for $Url"
}
$actual = (Get-Content $versionFile.FullName -Raw).Trim()
if ($actual -ne $expected) {
throw "prebuilt archive version mismatch for ${Url}: expected $expected, got $actual"
}
}
$bindir = Join-Path $Prefix "bin"
$configDir = Join-Path $HOME ".config\dosh"
New-Item -ItemType Directory -Force -Path $bindir, $configDir | Out-Null
@@ -137,6 +162,7 @@ function Install-Prebuilt {
Invoke-WebRequest -UseBasicParsing -Uri $url -OutFile $zip
Verify-ArchiveChecksum $url $zip
Expand-Archive -Force -Path $zip -DestinationPath $extract
Verify-ArchiveVersion $extract $url
foreach ($bin in @("dosh-client.exe", "dosh-bench.exe")) {
$found = Get-ChildItem -Path $extract -Recurse -File -Filter $bin | Select-Object -First 1
if (-not $found) {
+35
View File
@@ -268,6 +268,40 @@ verify_archive_checksum() {
fi
}
expected_archive_version() {
url="$1"
case "$url" in
*/releases/download/v*/*)
tag="${url#*/releases/download/v}"
tag="${tag%%/*}"
printf '%s\n' "$tag"
return 0
;;
esac
if [ -z "$binary_url" ] && [ -z "$binary_base" ] && [ "$binary_version" != "latest" ]; then
printf '%s\n' "${binary_version#v}"
return 0
fi
return 1
}
verify_archive_version() {
extract_dir="$1"
url="$2"
expected="$(expected_archive_version "$url" || true)"
[ -n "$expected" ] || return 0
version_file="$(find "$extract_dir" -type f -name VERSION 2>/dev/null | sed -n '1p')"
if [ -z "$version_file" ]; then
echo "prebuilt archive missing VERSION for $url" >&2
return 1
fi
actual="$(tr -d '\r\n' <"$version_file")"
if [ "$actual" != "$expected" ]; then
echo "prebuilt archive version mismatch for $url: expected $expected, got $actual" >&2
return 1
fi
}
try_install_prebuilt() {
download_url="$(release_latest_tag_download_url || release_download_url)" || return 1
tmpdir="$(mktemp -d)"
@@ -291,6 +325,7 @@ try_install_prebuilt() {
echo "prebuilt archive could not be extracted: $download_url" >&2
return 1
fi
verify_archive_version "$tmpdir/extract" "$download_url" || return 1
install_extracted_binary "$tmpdir/extract" dosh-client "$bindir/dosh-client" || return 1
ln -sf dosh-client "$bindir/dosh"
if [ "$role" = "server" ] || [ "$role" = "both" ]; then
+38
View File
@@ -30,6 +30,7 @@ base="${GITEA_URL:-https://git.palav.dev}"
repo="${GITEA_REPO:-Palav/dosh}"
token="${GITEA_TOKEN:-}"
title="${GITEA_TITLE:-$tag}"
expected_version="${tag#v}"
if [ -z "$token" ] && [ -f "$HOME/.config/dosh/gitea-token" ]; then
token="$(tr -d '\r\n' <"$HOME/.config/dosh/gitea-token")"
@@ -102,12 +103,49 @@ delete_existing_asset() {
done
}
artifact_version() {
artifact="$1"
case "$artifact" in
*.tar.gz)
tar -xOf "$artifact" dosh/VERSION 2>/dev/null | tr -d '\r\n'
;;
*.zip)
unzip -p "$artifact" dosh/VERSION 2>/dev/null | tr -d '\r\n'
;;
*)
return 1
;;
esac
}
verify_release_artifact_version() {
artifact="$1"
name="$(basename "$artifact")"
case "$name" in
*.sha256)
return 0
;;
dosh-*.tar.gz|dosh-*.zip)
actual="$(artifact_version "$artifact" || true)"
if [ -z "$actual" ]; then
echo "release artifact missing dosh/VERSION: $artifact" >&2
exit 1
fi
if [ "$actual" != "$expected_version" ]; then
echo "release artifact version mismatch: $artifact has $actual, expected $expected_version" >&2
exit 1
fi
;;
esac
}
uploaded=0
for artifact in "$@"; do
if [ ! -f "$artifact" ]; then
continue
fi
name="$(basename "$artifact")"
verify_release_artifact_version "$artifact"
delete_existing_asset "$name"
echo "uploading $name"
curl -fsS \
+206 -14
View File
@@ -203,6 +203,14 @@ struct PendingStreamChunk {
attempts: u32,
}
#[derive(Clone)]
struct PendingStreamOpen {
target_host: String,
target_port: u16,
last_sent: Instant,
attempts: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StartupGateMode {
HoldAll,
@@ -2667,7 +2675,7 @@ async fn try_resume_fresh_process(
)?;
socket.send_to(&packet, addr).await?;
let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
if packet.header.conn_id != cached.client_id {
if !is_resume_response_for_client(&packet, cached.client_id) {
return Ok(None);
}
if packet.header.kind == PacketKind::AttachReject {
@@ -2722,7 +2730,7 @@ async fn try_live_resume(
*send_seq += 1;
socket.send_to(&packet, addr).await?;
let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
if packet.header.conn_id != cached.client_id {
if !is_resume_response_for_client(&packet, cached.client_id) {
return Ok(None);
}
if packet.header.kind == PacketKind::AttachReject {
@@ -2750,6 +2758,11 @@ async fn try_live_resume(
Ok((frame, next))
}
fn is_resume_response_for_client(packet: &protocol::Packet, client_id: [u8; 16]) -> bool {
packet.header.conn_id == client_id
|| (packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16])
}
async fn recv_response_until<T, F>(socket: &UdpSocket, wait: Duration, mut accept: F) -> Result<T>
where
F: FnMut(protocol::Packet) -> Result<Option<T>>,
@@ -2797,7 +2810,7 @@ async fn run_terminal(
} else {
Some(RawMode::enter()?)
};
let addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
let mut addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
let mut send_seq = 2u64;
let mut last_packet_at = Instant::now();
let mut status_tick = tokio::time::interval(Duration::from_secs(1));
@@ -2844,6 +2857,7 @@ async fn run_terminal(
let mut opened_streams: HashSet<u64> = HashSet::new();
let mut stream_send_credit: HashMap<u64, usize> = HashMap::new();
let mut stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>> = HashMap::new();
let mut stream_pending_opens: HashMap<u64, PendingStreamOpen> = HashMap::new();
let mut stream_next_send_offset: HashMap<u64, u64> = HashMap::new();
let mut stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>> = HashMap::new();
let mut stream_next_recv_offset: HashMap<u64, u64> = HashMap::new();
@@ -2989,6 +3003,7 @@ async fn run_terminal(
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
@@ -3060,6 +3075,7 @@ async fn run_terminal(
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
@@ -3164,6 +3180,7 @@ async fn run_terminal(
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
@@ -3192,6 +3209,17 @@ async fn run_terminal(
continue;
};
last_packet_at = Instant::now();
if opened_streams.contains(&open.stream_id) {
send_stream_open_ok(
&socket,
addr,
&cred,
&mut send_seq,
open.stream_id,
)
.await?;
continue;
}
// SSH-agent stream: the server is asking us to splice this
// stream into our LOCAL ssh-agent. Only honor it when we
// actually opted in (agent_sock is Some); otherwise reject,
@@ -3340,6 +3368,7 @@ async fn run_terminal(
continue;
};
last_packet_at = Instant::now();
stream_pending_opens.remove(&ok.stream_id);
opened_streams.insert(ok.stream_id);
stream_send_credit.entry(ok.stream_id).or_insert(STREAM_INITIAL_WINDOW);
stream_next_send_offset.entry(ok.stream_id).or_insert(0);
@@ -3369,6 +3398,7 @@ async fn run_terminal(
continue;
};
last_packet_at = Instant::now();
stream_pending_opens.remove(&reject.stream_id);
if pending_socks_replies.remove(&reject.stream_id)
&& let Some(writer) = stream_writers.get_mut(&reject.stream_id)
{
@@ -3453,6 +3483,7 @@ 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);
@@ -3477,6 +3508,12 @@ async fn run_terminal(
if socks5_reply {
pending_socks_replies.insert(stream_id);
}
stream_pending_opens.insert(stream_id, PendingStreamOpen {
target_host: target_host.clone(),
target_port,
last_sent: Instant::now(),
attempts: 1,
});
send_stream_open(
&socket,
addr,
@@ -3505,6 +3542,7 @@ 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);
@@ -3534,6 +3572,7 @@ async fn run_terminal(
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
@@ -3599,6 +3638,13 @@ async fn run_terminal(
&mut send_seq,
&mut stream_sent_data,
).await?;
retransmit_stream_opens(
&socket,
addr,
&cred,
&mut send_seq,
&mut stream_pending_opens,
).await?;
}
}
}
@@ -3826,6 +3872,11 @@ async fn reconnect(
Ok(Some(frame))
}
fn refresh_live_addr(addr: &mut SocketAddr, cred: &CachedCredential) -> Result<()> {
*addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
Ok(())
}
fn queue_pending_user_input(
pending: &mut VecDeque<Vec<u8>>,
pending_bytes: &mut usize,
@@ -4126,6 +4177,38 @@ async fn retransmit_stream_data(
Ok(())
}
async fn retransmit_stream_opens(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_pending_opens: &mut HashMap<u64, PendingStreamOpen>,
) -> Result<()> {
let now = Instant::now();
let mut retransmit = Vec::new();
for (stream_id, pending) in stream_pending_opens.iter_mut() {
if now.duration_since(pending.last_sent) < Duration::from_millis(200) {
continue;
}
pending.last_sent = now;
pending.attempts = pending.attempts.saturating_add(1);
retransmit.push((*stream_id, pending.target_host.clone(), pending.target_port));
}
for (stream_id, target_host, target_port) in retransmit {
send_stream_open(
socket,
addr,
cred,
send_seq,
stream_id,
target_host,
target_port,
)
.await?;
}
Ok(())
}
fn accept_stream_data(
stream_next_recv_offset: &mut HashMap<u64, u64>,
stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
@@ -5337,25 +5420,27 @@ const TERMINAL_CLEANUP: &[u8] = concat!(
#[cfg(test)]
mod tests {
use super::{
DisconnectStatus, DynamicForward, FrameBuffer, LocalForward, MAX_PENDING_USER_INPUT_BYTES,
POST_SUBMIT_ALL_INPUT_HOLD, 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_matches_escape, is_local_status_target, latest_release_download_url,
CachedCredential, DisconnectStatus, DynamicForward, 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_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,
queue_pending_user_input, raw_contains_host_table, recv_response_until, refresh_live_addr,
render_status_clear, render_status_overlay, requested_env, resolved_startup_command,
rewrite_forward_command, selected_predict_mode, selected_udp_host, server_version_mismatch,
should_hold_during_startup_gate, should_hold_post_submit_input, split_after_command_submit,
ssh_destination_host, ssh_username, ssh_with_user, startup_command, status_ssh_target,
toml_bare_key_or_quoted, update_check_requested, valid_forward_host,
retransmit_stream_opens, rewrite_forward_command, selected_predict_mode, selected_udp_host,
server_version_mismatch, should_hold_during_startup_gate, should_hold_post_submit_input,
split_after_command_submit, ssh_destination_host, ssh_username, ssh_with_user,
startup_command, status_ssh_target, toml_bare_key_or_quoted, update_check_requested,
valid_forward_host,
};
use dosh::config::{ClientConfig, CommandExtension, HostConfig};
use dosh::native::EnvVar;
use dosh::protocol::{self, Frame, PacketKind};
use ed25519_dalek::{SigningKey, VerifyingKey};
use std::collections::VecDeque;
use std::collections::{HashMap, VecDeque};
use std::fs;
use std::time::Duration;
@@ -6052,6 +6137,113 @@ mod tests {
assert_eq!(post_submit_hold_duration(b"x"), POST_SUBMIT_ALL_INPUT_HOLD);
}
#[test]
fn reconnect_refreshes_live_send_address_from_credentials() {
let mut addr = "127.0.0.1:50000".parse().unwrap();
let cred = CachedCredential {
server: "host".to_string(),
session: "term".to_string(),
mode: "normal".to_string(),
udp_host: "127.0.0.1".to_string(),
udp_port: 50001,
client_id: [1; 16],
session_key: [2; 32],
session_key_id: [3; 16],
attach_ticket: Vec::new(),
attach_ticket_psk: [4; 32],
last_rendered_seq: 0,
};
refresh_live_addr(&mut addr, &cred).unwrap();
assert_eq!(addr.port(), 50001);
}
#[test]
fn resume_accepts_unknown_client_reject_from_legacy_zero_id_server() {
let client_id = [7u8; 16];
let matching_reject = protocol::Packet {
header: protocol::Header {
kind: PacketKind::AttachReject,
flags: 0,
conn_id: client_id,
seq: 0,
ack: 0,
session_key_id: [0u8; 16],
body_len: 0,
},
body: Vec::new(),
};
let zero_id_reject = protocol::Packet {
header: protocol::Header {
conn_id: [0u8; 16],
..matching_reject.header.clone()
},
body: Vec::new(),
};
let unrelated_frame = protocol::Packet {
header: protocol::Header {
kind: PacketKind::Frame,
conn_id: [9u8; 16],
..matching_reject.header.clone()
},
body: Vec::new(),
};
assert!(is_resume_response_for_client(&matching_reject, client_id));
assert!(is_resume_response_for_client(&zero_id_reject, client_id));
assert!(!is_resume_response_for_client(&unrelated_frame, client_id));
}
#[tokio::test]
async fn pending_stream_open_is_retransmitted_until_answered() {
let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let addr = receiver.local_addr().unwrap();
let cred = CachedCredential {
server: "host".to_string(),
session: "term".to_string(),
mode: "forward-only".to_string(),
udp_host: "127.0.0.1".to_string(),
udp_port: addr.port(),
client_id: [1; 16],
session_key: [2; 32],
session_key_id: protocol::session_key_id(&[2; 32]),
attach_ticket: Vec::new(),
attach_ticket_psk: [3; 32],
last_rendered_seq: 0,
};
let mut send_seq = 10;
let mut pending = HashMap::from([(
44,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 8080,
last_sent: std::time::Instant::now() - Duration::from_millis(250),
attempts: 1,
},
)]);
retransmit_stream_opens(&sender, addr, &cred, &mut send_seq, &mut pending)
.await
.unwrap();
let mut buf = [0u8; 2048];
let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
.await
.unwrap()
.unwrap();
let packet = protocol::decode(&buf[..n]).unwrap();
assert_eq!(packet.header.kind, PacketKind::StreamOpen);
assert_eq!(packet.header.seq, 10);
let plain =
protocol::decrypt_body(&packet, &cred.session_key, protocol::CLIENT_TO_SERVER).unwrap();
let open: protocol::StreamOpen = protocol::from_body(&plain).unwrap();
assert_eq!(open.stream_id, 44);
assert_eq!(open.target_port, 8080);
assert_eq!(send_seq, 11);
assert_eq!(pending[&44].attempts, 2);
}
#[test]
fn update_check_accepts_only_check_flag() {
assert!(!update_check_requested(&[]).unwrap());
+266 -14
View File
@@ -394,6 +394,7 @@ struct ClientState {
opened_streams: HashSet<u64>,
stream_send_credit: HashMap<u64, usize>,
stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
stream_pending_opens: HashMap<u64, PendingStreamOpen>,
stream_next_send_offset: HashMap<u64, u64>,
stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_next_recv_offset: HashMap<u64, u64>,
@@ -425,6 +426,14 @@ struct PendingStreamChunk {
attempts: u32,
}
#[derive(Clone)]
struct PendingStreamOpen {
target_host: String,
target_port: u16,
last_sent: Instant,
attempts: u32,
}
struct PendingNativeAuth {
client: dosh::native::NativeClientHello,
server: NativeServerHello,
@@ -1150,6 +1159,7 @@ async fn handle_native_user_auth(
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_pending_opens: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
@@ -1293,6 +1303,7 @@ async fn handle_bootstrap_attach(
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_pending_opens: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
@@ -1428,6 +1439,7 @@ async fn handle_ticket_attach(
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_pending_opens: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
@@ -1503,7 +1515,10 @@ async fn handle_resume(
) -> Result<()> {
let (key, session_name) = match find_client_decrypt_key(state, &packet.header) {
Ok(found) => found,
Err(_) => return send_reject(socket, peer, "unknown client").await,
Err(_) => {
return send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client")
.await;
}
};
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let req: ResumeRequest = protocol::from_body(&body)?;
@@ -1743,7 +1758,11 @@ async fn handle_stream_open(
let (key, session_name) = find_client_decrypt_key(state, &packet.header)?;
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
let open: StreamOpen = protocol::from_body(&body)?;
let allowed = {
enum StreamOpenCheck {
Duplicate,
New(Result<()>),
}
let check = {
let mut locked = state.lock().expect("server state poisoned");
let config = locked.config.clone();
let session = locked
@@ -1759,17 +1778,27 @@ async fn handle_stream_open(
}
client.endpoint = peer;
client.last_seen = Instant::now();
stream_open_allowed(&config, client, &open)
if client.opened_streams.contains(&open.stream_id) {
StreamOpenCheck::Duplicate
} else {
StreamOpenCheck::New(stream_open_allowed(&config, client, &open))
}
};
if let Err(err) = allowed {
return send_stream_open_reject(
state,
socket,
packet.header.conn_id,
open.stream_id,
err.to_string(),
)
.await;
match check {
StreamOpenCheck::Duplicate => {
return send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await;
}
StreamOpenCheck::New(Ok(())) => {}
StreamOpenCheck::New(Err(err)) => {
return send_stream_open_reject(
state,
socket,
packet.header.conn_id,
open.stream_id,
err.to_string(),
)
.await;
}
}
let target = format!("{}:{}", open.target_host, open.target_port);
@@ -1928,6 +1957,7 @@ async fn handle_stream_open_ok(
}
client.endpoint = peer;
client.last_seen = Instant::now();
client.stream_pending_opens.remove(&ok.stream_id);
client.opened_streams.insert(ok.stream_id);
client
.stream_send_credit
@@ -1960,6 +1990,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);
@@ -2018,6 +2049,7 @@ async fn handle_stream_close(
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);
@@ -2331,10 +2363,40 @@ async fn send_stream_open_to_client(
) -> Result<()> {
let body = protocol::to_body(&StreamOpen {
stream_id,
target_host,
target_host: target_host.clone(),
target_port,
})?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpen, body).await
let send = {
let mut locked = state.lock().expect("server state poisoned");
let mut found = None;
if let Some(client) = locked.client_mut(&client_id) {
client.send_seq += 1;
let packet = protocol::encode_encrypted(
PacketKind::StreamOpen,
client_id,
client.send_seq,
client.last_acked,
&client.session_key,
SERVER_TO_CLIENT,
&body,
)?;
client.stream_pending_opens.insert(
stream_id,
PendingStreamOpen {
target_host,
target_port,
last_sent: Instant::now(),
attempts: 1,
},
);
found = Some((client.endpoint, packet));
}
found
};
if let Some((endpoint, packet)) = send {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
async fn send_stream_open_reject(
@@ -2857,6 +2919,38 @@ async fn retransmit_pending(
sends.push((client.endpoint, pending.packet.clone()));
}
}
let pending_open_ids: Vec<u64> =
client.stream_pending_opens.keys().copied().collect();
let mut retransmit_opens = Vec::new();
for stream_id in pending_open_ids {
let Some(open) = client.stream_pending_opens.get_mut(&stream_id) else {
continue;
};
if now.duration_since(open.last_sent) < Duration::from_millis(200) {
continue;
}
open.last_sent = now;
open.attempts = open.attempts.saturating_add(1);
retransmit_opens.push((stream_id, open.target_host.clone(), open.target_port));
}
for (stream_id, target_host, target_port) in retransmit_opens {
client.send_seq += 1;
let body = protocol::to_body(&StreamOpen {
stream_id,
target_host,
target_port,
})?;
let packet = protocol::encode_encrypted(
PacketKind::StreamOpen,
*client_id,
client.send_seq,
client.last_acked,
&client.session_key,
SERVER_TO_CLIENT,
&body,
)?;
sends.push((client.endpoint, packet));
}
let stream_ids: Vec<u64> = client.stream_sent_data.keys().copied().collect();
let mut retransmit_chunks = Vec::new();
for stream_id in stream_ids {
@@ -3299,6 +3393,46 @@ mod tests {
assert_eq!(ServerConfig::default().client_timeout_secs, 2_592_000);
}
#[tokio::test]
async fn unknown_resume_reject_keeps_client_id() {
let (pty_tx, _rx) = mpsc::unbounded_channel();
let state = Arc::new(Mutex::new(ServerState::new(
ServerConfig::default(),
[0u8; 32],
pty_tx,
)));
let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let client_id = [9u8; 16];
let packet = protocol::Packet {
header: protocol::Header {
kind: PacketKind::ResumeRequest,
flags: 1,
conn_id: client_id,
seq: 77,
ack: 0,
session_key_id: [8u8; 16],
body_len: 0,
},
body: Vec::new(),
};
handle_resume(&state, &sender, receiver.local_addr().unwrap(), &packet)
.await
.unwrap();
let mut buf = [0u8; 1024];
let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
.await
.unwrap()
.unwrap();
let reject = protocol::decode(&buf[..n]).unwrap();
assert_eq!(reject.header.kind, PacketKind::AttachReject);
assert_eq!(reject.header.conn_id, client_id);
let body: AttachReject = protocol::from_body(&reject.body).unwrap();
assert_eq!(body.reason, "unknown client");
}
fn test_client_state(session_key: [u8; 32]) -> ClientState {
ClientState {
endpoint: "127.0.0.1:9".parse().unwrap(),
@@ -3318,6 +3452,7 @@ mod tests {
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_pending_opens: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
@@ -3389,6 +3524,121 @@ mod tests {
assert!(locked.lookup_client(&client_id).is_none());
}
#[tokio::test]
async fn duplicate_stream_open_for_open_stream_resends_ok() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
let client_id = [11u8; 16];
let session_key = [12u8; 32];
let stream_id = 99;
let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
let mut client = test_client_state(session_key);
client.endpoint = receiver.local_addr().unwrap();
client.opened_streams.insert(stream_id);
state.sessions.insert(
"test".to_string(),
Session {
pty: None,
parser: vt100::Parser::new(24, 80, 100),
clients: HashMap::from([(client_id, client)]),
output_seq: 0,
recent: VecDeque::new(),
empty_since: None,
holder_control: None,
persistent: false,
bytes_since_persist: 0,
last_persisted_seq: 0,
},
);
state.client_index.insert(client_id, "test".to_string());
let state = Arc::new(Mutex::new(state));
let raw = protocol::encode_encrypted(
PacketKind::StreamOpen,
client_id,
2,
0,
&session_key,
CLIENT_TO_SERVER,
&protocol::to_body(&StreamOpen {
stream_id,
target_host: "127.0.0.1".to_string(),
target_port: 9,
})
.unwrap(),
)
.unwrap();
let packet = protocol::decode(&raw).unwrap();
handle_stream_open(&state, &sender, receiver.local_addr().unwrap(), &packet)
.await
.unwrap();
let mut buf = [0u8; 2048];
let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
.await
.unwrap()
.unwrap();
let packet = protocol::decode(&buf[..n]).unwrap();
assert_eq!(packet.header.kind, PacketKind::StreamOpenOk);
let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
let ok: StreamOpenOk = protocol::from_body(&plain).unwrap();
assert_eq!(ok.stream_id, stream_id);
}
#[tokio::test]
async fn pending_server_stream_open_is_retransmitted() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
let mut state = ServerState::new(ServerConfig::default(), [0u8; 32], pty_tx);
let client_id = [13u8; 16];
let session_key = [14u8; 32];
let stream_id = 123;
let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let sender = Arc::new(UdpSocket::bind("127.0.0.1:0").await.unwrap());
let mut client = test_client_state(session_key);
client.endpoint = receiver.local_addr().unwrap();
client.stream_pending_opens.insert(
stream_id,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 8080,
last_sent: Instant::now() - Duration::from_millis(250),
attempts: 1,
},
);
state.sessions.insert(
"test".to_string(),
Session {
pty: None,
parser: vt100::Parser::new(24, 80, 100),
clients: HashMap::from([(client_id, client)]),
output_seq: 0,
recent: VecDeque::new(),
empty_since: None,
holder_control: None,
persistent: false,
bytes_since_persist: 0,
last_persisted_seq: 0,
},
);
state.client_index.insert(client_id, "test".to_string());
let state = Arc::new(Mutex::new(state));
retransmit_pending(&state, &sender).await.unwrap();
let mut buf = [0u8; 2048];
let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
.await
.unwrap()
.unwrap();
let packet = protocol::decode(&buf[..n]).unwrap();
assert_eq!(packet.header.kind, PacketKind::StreamOpen);
let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT).unwrap();
let open: StreamOpen = protocol::from_body(&plain).unwrap();
assert_eq!(open.stream_id, stream_id);
assert_eq!(open.target_port, 8080);
}
#[test]
fn forward_only_session_does_not_allocate_pty() {
let (pty_tx, _pty_rx) = mpsc::unbounded_channel();
@@ -3520,6 +3770,7 @@ mod tests {
opened_streams: HashSet::new(),
stream_send_credit: HashMap::new(),
stream_pending_data: HashMap::new(),
stream_pending_opens: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),
@@ -3615,6 +3866,7 @@ mod tests {
opened_streams: HashSet::from([42]),
stream_send_credit: HashMap::from([(42, 0)]),
stream_pending_data: HashMap::new(),
stream_pending_opens: HashMap::new(),
stream_next_send_offset: HashMap::new(),
stream_sent_data: HashMap::new(),
stream_next_recv_offset: HashMap::new(),