Reduce russh write-path copies with direct Bytes sends

russh/src/channels/mod.rs

@@ -322,6 +322,11 @@ impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> ChannelWriteHalf<
         self.send_data(None, data).await
     }
 
+    /// Send owned bytes to a channel without copying them into the `AsyncWrite` path.
+    pub async fn data_bytes(&self, data: impl Into<Bytes>) -> Result<(), Error> {
+        self.send_bytes(None, data.into()).await
+    }
+
     /// Send data to a channel. The number of bytes added to the
     /// "sending pipeline" (to be processed by the event loop) is
     /// returned.
@@ -333,6 +338,15 @@ impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> ChannelWriteHalf<
         self.send_data(Some(ext), data).await
     }
 
+    /// Send owned extended data to a channel without copying it into the `AsyncWrite` path.
+    pub async fn extended_data_bytes(
+        &self,
+        ext: u32,
+        data: impl Into<Bytes>,
+    ) -> Result<(), Error> {
+        self.send_bytes(Some(ext), data.into()).await
+    }
+
     async fn send_data<R: tokio::io::AsyncRead + Unpin>(
         &self,
         ext: Option<u32>,
@@ -345,6 +359,49 @@ impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> ChannelWriteHalf<
         Ok(())
     }
 
+    async fn reserve_writable_chunk(&self, remaining: usize) -> Result<usize, Error> {
+        if self.max_packet_size == 0 {
+            return Err(Error::Inconsistent);
+        }
+        loop {
+            let mut window_size = self.window_size.value.lock().await;
+            let writable = (self.max_packet_size as usize)
+                .min(*window_size as usize)
+                .min(remaining);
+            if writable > 0 {
+                *window_size -= writable as u32;
+                if *window_size > 0 {
+                    self.window_size.notifier.notify_one();
+                }
+                return Ok(writable);
+            }
+            let notified = self.window_size.notifier.notified();
+            drop(window_size);
+            notified.await;
+        }
+    }
+
+    async fn send_bytes(&self, ext: Option<u32>, data: Bytes) -> Result<(), Error> {
+        if data.is_empty() {
+            return Ok(());
+        }
+
+        let mut offset = 0;
+        while offset < data.len() {
+            let writable = self.reserve_writable_chunk(data.len() - offset).await?;
+            let end = offset + writable;
+            let chunk = data.slice(offset..end);
+            let msg = match ext {
+                None => ChannelMsg::Data { data: chunk },
+                Some(ext) => ChannelMsg::ExtendedData { data: chunk, ext },
+            };
+            self.send_msg(msg).await?;
+            offset = end;
+        }
+
+        Ok(())
+    }
+
     pub async fn eof(&self) -> Result<(), Error> {
         self.send_msg(ChannelMsg::Eof).await
     }
@@ -556,6 +613,11 @@ impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> Channel<S> {
         self.write_half.data(data).await
     }
 
+    /// Send owned bytes to a channel without copying them into the `AsyncWrite` path.
+    pub async fn data_bytes(&self, data: impl Into<Bytes>) -> Result<(), Error> {
+        self.write_half.data_bytes(data).await
+    }
+
     /// Send data to a channel. The number of bytes added to the
     /// "sending pipeline" (to be processed by the event loop) is
     /// returned.
@@ -567,6 +629,15 @@ impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> Channel<S> {
         self.write_half.extended_data(ext, data).await
     }
 
+    /// Send owned extended data to a channel without copying it into the `AsyncWrite` path.
+    pub async fn extended_data_bytes(
+        &self,
+        ext: u32,
+        data: impl Into<Bytes>,
+    ) -> Result<(), Error> {
+        self.write_half.extended_data_bytes(ext, data).await
+    }
+
     pub async fn eof(&self) -> Result<(), Error> {
         self.write_half.eof().await
     }
@@ -625,3 +696,148 @@ impl<S: From<(ChannelId, ChannelMsg)> + Send + Sync + 'static> Channel<S> {
         self.write_half.make_writer_ext(ext)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use tokio::sync::mpsc;
+
+    use super::*;
+
+    fn test_write_half(
+        window_size: WindowSizeRef,
+        max_packet_size: u32,
+    ) -> (
+        ChannelWriteHalf<(ChannelId, ChannelMsg)>,
+        mpsc::Receiver<(ChannelId, ChannelMsg)>,
+    ) {
+        let (sender, receiver) = mpsc::channel(8);
+        (
+            ChannelWriteHalf {
+                id: ChannelId(7),
+                sender,
+                max_packet_size,
+                window_size,
+            },
+            receiver,
+        )
+    }
+
+    #[tokio::test]
+    async fn data_bytes_sends_one_owned_message_when_window_permits() {
+        let payload = Bytes::from_static(b"owned data");
+        let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 1024);
+
+        write_half.data_bytes(payload.clone()).await.unwrap();
+
+        match receiver.recv().await.unwrap() {
+            (ChannelId(7), ChannelMsg::Data { data }) => {
+                assert_eq!(data, payload);
+                assert_eq!(data.as_ptr(), payload.as_ptr());
+            }
+            msg => panic!("unexpected message: {msg:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn data_bytes_splits_by_max_packet_size_without_copying() {
+        let payload = Bytes::from_static(b"abcdefghij");
+        let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 4);
+
+        write_half.data_bytes(payload.clone()).await.unwrap();
+
+        for (range, expected) in [
+            (0..4, &b"abcd"[..]),
+            (4..8, &b"efgh"[..]),
+            (8..10, &b"ij"[..]),
+        ] {
+            match receiver.recv().await.unwrap() {
+                (ChannelId(7), ChannelMsg::Data { data }) => {
+                    assert_eq!(data.as_ref(), expected);
+                    assert_eq!(data.as_ptr(), payload.slice(range).as_ptr());
+                }
+                msg => panic!("unexpected message: {msg:?}"),
+            }
+        }
+        assert!(receiver.try_recv().is_err());
+    }
+
+    #[tokio::test]
+    async fn extended_data_bytes_preserves_extension_code() {
+        let payload = Bytes::from_static(b"stderr");
+        let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 1024);
+
+        write_half
+            .extended_data_bytes(1, payload.clone())
+            .await
+            .unwrap();
+
+        match receiver.recv().await.unwrap() {
+            (ChannelId(7), ChannelMsg::ExtendedData { data, ext }) => {
+                assert_eq!(ext, 1);
+                assert_eq!(data, payload);
+                assert_eq!(data.as_ptr(), payload.as_ptr());
+            }
+            msg => panic!("unexpected message: {msg:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn data_bytes_empty_payload_sends_nothing() {
+        let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 1024);
+
+        write_half.data_bytes(Bytes::new()).await.unwrap();
+
+        assert!(receiver.try_recv().is_err());
+    }
+
+    #[tokio::test]
+    async fn data_bytes_waits_for_window_update() {
+        let window_size = WindowSizeRef::new(0);
+        let (write_half, mut receiver) = test_write_half(window_size.clone(), 1024);
+        let send = tokio::spawn(async move {
+            write_half
+                .data_bytes(Bytes::from_static(b"after-window"))
+                .await
+                .unwrap();
+        });
+
+        tokio::task::yield_now().await;
+        assert!(!send.is_finished());
+
+        window_size.update(1024).await;
+        send.await.unwrap();
+
+        match receiver.recv().await.unwrap() {
+            (ChannelId(7), ChannelMsg::Data { data }) => {
+                assert_eq!(data.as_ref(), b"after-window");
+            }
+            msg => panic!("unexpected message: {msg:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn data_bytes_rejects_zero_max_packet_size() {
+        let (write_half, mut receiver) = test_write_half(WindowSizeRef::new(1024), 0);
+
+        let result = write_half.data_bytes(Bytes::from_static(b"owned")).await;
+
+        assert!(matches!(result, Err(Error::Inconsistent)));
+        assert!(receiver.try_recv().is_err());
+    }
+
+    #[tokio::test]
+    async fn channel_data_bytes_forwards_to_write_half() {
+        let (sender, mut receiver) = mpsc::channel(8);
+        let (channel, _reference) =
+            Channel::<(ChannelId, ChannelMsg)>::new(ChannelId(9), sender, 1024, 1024, 8);
+
+        channel.data_bytes(Bytes::from_static(b"channel")).await.unwrap();
+
+        match receiver.recv().await.unwrap() {
+            (ChannelId(9), ChannelMsg::Data { data }) => {
+                assert_eq!(data.as_ref(), b"channel");
+            }
+            msg => panic!("unexpected message: {msg:?}"),
+        }
+    }
+}

russh/src/cipher/mod.rs

@@ -209,6 +209,45 @@ pub(crate) trait SealingKey {
 
     fn seal(&mut self, seqn: u32, plaintext_in_ciphertext_out: &mut [u8], tag_out: &mut [u8]);
 
+    #[allow(clippy::indexing_slicing)] // PacketWriter reserves and sizes the packet buffer first
+    fn finish_packet(&mut self, offset: usize, payload_len: usize, buffer: &mut SSHBuffer) {
+        let payload_start = offset + PACKET_LENGTH_LEN + PADDING_LENGTH_LEN;
+        let payload_end = payload_start + payload_len;
+
+        trace!("writing, seqn = {:?}", buffer.seqn.0);
+        let padding_length = self.padding_length(&buffer.buffer[payload_start..payload_end]);
+        trace!("padding length {padding_length:?}");
+        let packet_length = PADDING_LENGTH_LEN + payload_len + padding_length;
+        trace!("packet_length {packet_length:?}");
+
+        // Maximum packet length:
+        // https://tools.ietf.org/html/rfc4253#section-6.1
+        assert!(packet_length <= u32::MAX as usize);
+        BigEndian::write_u32(
+            &mut buffer.buffer[offset..offset + PACKET_LENGTH_LEN],
+            packet_length as u32,
+        );
+
+        assert!(padding_length <= u8::MAX as usize);
+        buffer.buffer[offset + PACKET_LENGTH_LEN] = padding_length as u8;
+        buffer.buffer.resize(payload_end + padding_length, 0);
+        #[allow(clippy::indexing_slicing)] // length checked
+        self.fill_padding(&mut buffer.buffer[payload_end..]);
+        let tag_offset = buffer.buffer.len();
+        buffer.buffer.resize(tag_offset + self.tag_len(), 0);
+
+        #[allow(clippy::indexing_slicing)] // length checked
+        let (plaintext, tag) =
+            buffer.buffer[offset..].split_at_mut(PACKET_LENGTH_LEN + packet_length);
+
+        self.seal(buffer.seqn.0, plaintext, tag);
+
+        buffer.bytes += payload_len;
+        // Sequence numbers are on 32 bits and wrap.
+        // https://tools.ietf.org/html/rfc4253#section-6.4
+        buffer.seqn += Wrapping(1);
+    }
+
     fn write(&mut self, payload: &[u8], buffer: &mut SSHBuffer) {
         // https://tools.ietf.org/html/rfc4253#section-6
         //
@@ -225,8 +264,9 @@ pub(crate) trait SealingKey {
         // Maximum packet length:
         // https://tools.ietf.org/html/rfc4253#section-6.1
         assert!(packet_length <= u32::MAX as usize);
-        #[allow(clippy::unwrap_used)] // length checked
-        (packet_length as u32).encode(&mut buffer.buffer).unwrap();
+        buffer
+            .buffer
+            .extend_from_slice(&(packet_length as u32).to_be_bytes());
 
         assert!(padding_length <= u8::MAX as usize);
         buffer.buffer.push(padding_length as u8);

russh/src/client/encrypted.rs

@@ -581,8 +581,11 @@ impl Session {
                     }
                 }
 
-                if let Some(ref mut enc) = self.common.encrypted {
-                    new_size -= enc.flush_pending(channel_num)? as u32;
+                let common = &mut self.common;
+                if let Some(enc) = common.encrypted.as_mut() {
+                    new_size -= enc
+                        .flush_pending_with_writer(&mut common.packet_writer, channel_num)?
+                        as u32;
                 }
                 if let Some(chan) = self.channels.get(&channel_num) {
                     chan.window_size().update(new_size).await;
@@ -887,7 +890,7 @@ impl Session {
                 } => {
                     debug!("sending ssh-userauth service requset");
                     if !*sent {
-                        self.common.packet_writer.packet(|w| {
+                        self.common.packet_writer.write_packet(|w| {
                             msg::SERVICE_REQUEST.encode(w)?;
                             "ssh-userauth".encode(w)?;
                             Ok(())

russh/src/client/kex.rs

@@ -115,9 +115,7 @@ impl ClientKex {
 
                 let names = {
                     // read algorithms from packet.
-                    self.exchange
-                        .server_kex_init
-                        .extend_from_slice(&input.buffer);
+                    self.exchange.server_kex_init = input.buffer.clone().into();
                     negotiation::Client::read_kex(
                         &input.buffer,
                         &self.config.preferred,
@@ -147,7 +145,7 @@ impl ClientKex {
                         self.cause.session_id(),
                     )?;
 
-                    output.packet(|w| {
+                    output.write_packet(|w| {
                         msg::NEWKEYS.encode(w)?;
                         Ok(())
                     })?;
@@ -159,14 +157,14 @@ impl ClientKex {
                 }
 
                 if kex.is_dh_gex() {
-                    output.packet(|w| {
+                    output.write_packet(|w| {
                         kex.client_dh_gex_init(&self.config.gex, w)?;
                         Ok(())
                     })?;
 
                     self.state = ClientKexState::WaitingForGexReply { names, kex };
                 } else {
-                    output.packet(|w| {
+                    output.write_packet(|w| {
                         kex.client_dh(&mut self.exchange.client_ephemeral, w)?;
                         Ok(())
                     })?;
@@ -217,7 +215,7 @@ impl ClientKex {
                 let exchange = &mut self.exchange;
                 exchange.gex = Some((self.config.gex.clone(), group.clone()));
                 kex.dh_gex_set_group(group)?;
-                output.packet(|w| {
+                output.write_packet(|w| {
                     kex.client_dh(&mut exchange.client_ephemeral, w)?;
                     Ok(())
                 })?;
@@ -306,7 +304,7 @@ impl ClientKex {
                     self.cause.session_id(),
                 )?;
 
-                output.packet(|w| {
+                output.write_packet(|w| {
                     msg::NEWKEYS.encode(w)?;
                     Ok(())
                 })?;