Split stream writes into packet chunks

src/bin/dosh-client.rs

@@ -7966,31 +7966,34 @@ async fn queue_or_send_stream_data(
     stream_next_send_offset: &mut HashMap<u64, u64>,
     stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
 ) -> Result<()> {
-    if !opened_streams.contains(&stream_id)
-        || stream_send_credit.get(&stream_id).copied().unwrap_or(0) < bytes.len()
-        || stream_pending_data
-            .get(&stream_id)
-            .is_some_and(|pending| !pending.is_empty())
-    {
-        stream_pending_data
-            .entry(stream_id)
-            .or_default()
-            .push_back(bytes);
-        return Ok(());
-    }
-    *stream_send_credit.entry(stream_id).or_default() -= bytes.len();
-    let offset = *stream_next_send_offset.entry(stream_id).or_default();
-    stream_next_send_offset.insert(stream_id, offset.saturating_add(bytes.len() as u64));
-    stream_sent_data.entry(stream_id).or_default().insert(
-        offset,
-        PendingStreamChunk {
+    for chunk in dosh::transport::split_stream_data_bytes(bytes, STREAM_INITIAL_WINDOW) {
+        if !opened_streams.contains(&stream_id)
+            || stream_send_credit.get(&stream_id).copied().unwrap_or(0) < chunk.len()
+            || stream_pending_data
+                .get(&stream_id)
+                .is_some_and(|pending| !pending.is_empty())
+        {
+            stream_pending_data
+                .entry(stream_id)
+                .or_default()
+                .push_back(chunk);
+            continue;
+        }
+        *stream_send_credit.entry(stream_id).or_default() -= chunk.len();
+        let offset = *stream_next_send_offset.entry(stream_id).or_default();
+        stream_next_send_offset.insert(stream_id, offset.saturating_add(chunk.len() as u64));
+        stream_sent_data.entry(stream_id).or_default().insert(
             offset,
-            bytes: bytes.clone(),
-            last_sent: Instant::now(),
-            attempts: 1,
-        },
-    );
-    send_stream_data(socket, addr, cred, send_seq, stream_id, offset, bytes).await
+            PendingStreamChunk {
+                offset,
+                bytes: chunk.clone(),
+                last_sent: Instant::now(),
+                attempts: 1,
+            },
+        );
+        send_stream_data(socket, addr, cred, send_seq, stream_id, offset, chunk).await?;
+    }
+    Ok(())
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -8008,9 +8011,22 @@ async fn flush_stream_pending_data(
     let Some(pending) = stream_pending_data.get_mut(&stream_id) else {
         return Ok(());
     };
-    while let Some(bytes) = pending.front() {
+    let chunk_limit = dosh::transport::stream_data_chunk_limit(STREAM_INITIAL_WINDOW);
+    while let Some(front_len) = pending.front().map(Vec::len) {
+        if front_len > chunk_limit {
+            let Some(bytes) = pending.pop_front() else {
+                break;
+            };
+            for chunk in dosh::transport::split_stream_data_bytes(bytes, STREAM_INITIAL_WINDOW)
+                .into_iter()
+                .rev()
+            {
+                pending.push_front(chunk);
+            }
+            continue;
+        }
         let credit = stream_send_credit.get(&stream_id).copied().unwrap_or(0);
-        if credit < bytes.len() {
+        if credit < front_len {
             break;
         }
         let Some(bytes) = pending.pop_front() else {
@@ -9605,9 +9621,9 @@ mod tests {
         LOCAL_SLEEP_REPAINT_RETRY_WINDOW, LocalForward, MAX_PENDING_USER_INPUT_BYTES,
         NativeIdentityContext, POST_RECONNECT_STALE_INPUT_GRACE, POST_SUBMIT_ALL_INPUT_HOLD,
         PendingStreamOpen, PredictMode, Predictor, RESTART_STATUS_SCRIPT, RemoteForward,
-        STALE_TERMINAL_INPUT_AFTER, STARTUP_INPUT_HOLD, SshConfig, SshPathTokenContext,
-        StartupGateMode, StatusAction, UpdateOptions, UpdateRole, auth_allows, cache_key,
-        cache_server_prefix, cleanup_stream_state, clear_cached_credentials,
+        STALE_TERMINAL_INPUT_AFTER, STARTUP_INPUT_HOLD, STREAM_INITIAL_WINDOW, SshConfig,
+        SshPathTokenContext, StartupGateMode, StatusAction, UpdateOptions, UpdateRole, auth_allows,
+        cache_key, cache_server_prefix, cleanup_stream_state, clear_cached_credentials,
         effective_update_artifact_tag, ensure_tui_safe_status_overlay, expand_ssh_path_tokens,
         input_contains_focus_in, input_matches_escape, is_local_status_target,
         is_resume_response_for_client, latest_release_download_url,
@@ -9615,21 +9631,22 @@ mod tests {
         newest_client_trace_path_from, parse_dynamic_forward, parse_escape_key,
         parse_local_forward, parse_remote_forward, parse_single_remote_path, parse_ssh_config,
         parse_trace_line, parse_trace_options, parse_trace_report_options, parse_trace_summary,
-        parse_update_options, post_submit_hold_duration, queue_pending_user_input,
-        queue_stale_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, retire_stream_state, retransmit_stream_opens,
-        rewrite_forward_command, sanitize_trace_name, selected_predict_mode, selected_udp_host,
-        server_version_mismatch, should_flush_terminal_input_after_contact,
-        should_health_log_client_start, should_hold_during_startup_gate,
-        should_hold_post_submit_input, should_reconnect_before_input_for_local_sleep,
-        should_repaint_idle_terminal, should_strip_unowned_terminal_reports,
-        split_after_command_submit, split_trace_tokens, ssh_command_target, ssh_config_uses_proxy,
-        ssh_destination_host, ssh_username, ssh_with_user, startup_command, status_ssh_target,
-        strip_stale_mouse_reports, strip_terminal_focus_reports, strip_unowned_terminal_reports,
-        summarize_trace_file, summarize_trace_file_with_mode, terminal_private_mode_transition,
-        toml_bare_key_or_quoted, top_trace_events, trace_report_warnings, unix_update_script,
+        parse_update_options, post_submit_hold_duration, queue_or_send_stream_data,
+        queue_pending_user_input, queue_stale_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, retire_stream_state,
+        retransmit_stream_opens, rewrite_forward_command, sanitize_trace_name,
+        selected_predict_mode, selected_udp_host, server_version_mismatch,
+        should_flush_terminal_input_after_contact, should_health_log_client_start,
+        should_hold_during_startup_gate, should_hold_post_submit_input,
+        should_reconnect_before_input_for_local_sleep, should_repaint_idle_terminal,
+        should_strip_unowned_terminal_reports, split_after_command_submit, split_trace_tokens,
+        ssh_command_target, ssh_config_uses_proxy, ssh_destination_host, ssh_username,
+        ssh_with_user, startup_command, status_ssh_target, strip_stale_mouse_reports,
+        strip_terminal_focus_reports, strip_unowned_terminal_reports, summarize_trace_file,
+        summarize_trace_file_with_mode, terminal_private_mode_transition, toml_bare_key_or_quoted,
+        top_trace_events, trace_report_warnings, unix_update_script,
         update_binary_version_for_installer, update_installer_url, update_version_status,
         upsert_managed_block, valid_forward_host, vscode_safe_alias, wake_repaint_retry_deadline,
         windows_update_script,
@@ -10860,6 +10877,85 @@ mod tests {
         assert_eq!(pending[&44].attempts, 2);
     }
 
+    #[tokio::test]
+    async fn client_stream_send_splits_large_writes() {
+        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 stream_id = 44;
+        let opened_streams = HashSet::from([stream_id]);
+        let mut stream_send_credit = HashMap::from([(stream_id, STREAM_INITIAL_WINDOW)]);
+        let mut stream_pending_data = HashMap::new();
+        let mut stream_next_send_offset = HashMap::new();
+        let mut stream_sent_data = HashMap::new();
+        let bytes = vec![9; dosh::transport::MAX_STREAM_DATA_BYTES * 2 + 11];
+
+        queue_or_send_stream_data(
+            &sender,
+            addr,
+            &cred,
+            &mut send_seq,
+            stream_id,
+            bytes,
+            &opened_streams,
+            &mut stream_send_credit,
+            &mut stream_pending_data,
+            &mut stream_next_send_offset,
+            &mut stream_sent_data,
+        )
+        .await
+        .unwrap();
+
+        assert!(stream_pending_data.is_empty());
+        assert_eq!(stream_sent_data[&stream_id].len(), 3);
+        assert_eq!(
+            stream_next_send_offset[&stream_id],
+            (dosh::transport::MAX_STREAM_DATA_BYTES * 2 + 11) as u64
+        );
+        assert_eq!(send_seq, 13);
+
+        let mut seen = Vec::new();
+        let mut buf = vec![0u8; u16::MAX as usize];
+        for _ in 0..3 {
+            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::StreamData);
+            let plain =
+                protocol::decrypt_body(&packet, &cred.session_key, protocol::CLIENT_TO_SERVER)
+                    .unwrap();
+            let data: protocol::StreamData = protocol::from_body(&plain).unwrap();
+            seen.push((data.offset, data.bytes.len()));
+        }
+        assert_eq!(
+            seen,
+            vec![
+                (0, dosh::transport::MAX_STREAM_DATA_BYTES),
+                (
+                    dosh::transport::MAX_STREAM_DATA_BYTES as u64,
+                    dosh::transport::MAX_STREAM_DATA_BYTES
+                ),
+                ((dosh::transport::MAX_STREAM_DATA_BYTES * 2) as u64, 11)
+            ]
+        );
+    }
+
     #[tokio::test]
     async fn client_stream_retransmit_uses_observed_rtt() {
         let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();