Introduce EncryptingReader which will in the future handle encryption

bjorn3 · bjorn3 · commit 3fc5c5b09074 · 2026-01-05T16:40:14.000+01:00
diff --git a/src/key_exchange.rs b/src/key_exchange.rs
@@ -7,7 +7,6 @@ use aws_lc_rs::{
     rand::{self, SystemRandom},
     signature::KeyPair,
 };
-use tokio::io::AsyncWriteExt;
 use tracing::{debug, error, warn};
 
 use crate::{
@@ -104,16 +103,11 @@ impl EcdhKeyExchange {
             },
         };
 
-        conn.write_buf.clear();
-        let Ok(packet) = Packet::builder(&mut conn.write_buf)
-            .with_payload(&key_exchange_reply)
-            .without_mac()
-        else {
-            error!(addr = %conn.addr, "failed to build key exchange init packet");
-            return Err(());
-        };
-
-        if let Err(error) = conn.stream_write.write_all(packet).await {
+        if let Err(error) = conn
+            .stream_write
+            .write_packet(&key_exchange_reply, |_| {})
+            .await
+        {
             warn!(addr = %conn.addr, %error, "failed to send version exchange");
             return Err(());
         }
@@ -135,16 +129,11 @@ impl EcdhKeyExchange {
             return Err(());
         }
 
-        conn.write_buf.clear();
-        let Ok(packet) = Packet::builder(&mut conn.write_buf)
-            .with_payload(&MessageType::NewKeys)
-            .without_mac()
-        else {
-            error!(addr = %conn.addr, "failed to build newkeys packet");
-            return Err(());
-        };
-
-        if let Err(error) = conn.stream_write.write_all(packet).await {
+        if let Err(error) = conn
+            .stream_write
+            .write_packet(&MessageType::NewKeys, |_| {})
+            .await
+        {
             warn!(addr = %conn.addr, %error, "failed to send newkeys packet");
             return Err(());
         }
@@ -316,21 +305,15 @@ impl KeyExchange {
             }
         };
 
-        conn.write_buf.clear();
-        let builder = Packet::builder(&mut conn.write_buf).with_payload(&key_exchange_init);
-
-        if let Ok(kex_init_payload) = builder.payload() {
-            exchange.update(&(kex_init_payload.len() as u32).to_be_bytes());
-            exchange.update(kex_init_payload);
-        };
-
-        let Ok(packet) = builder.without_mac() else {
-            error!(addr = %conn.addr, "failed to build key exchange init packet");
-            return Err(());
-        };
-
-        if let Err(error) = conn.stream_write.write_all(packet).await {
-            warn!(addr = %conn.addr, %error, "failed to send version exchange");
+        if let Err(error) = conn
+            .stream_write
+            .write_packet(&key_exchange_init, |kex_init_payload| {
+                exchange.update(&(kex_init_payload.len() as u32).to_be_bytes());
+                exchange.update(kex_init_payload);
+            })
+            .await
+        {
+            warn!(addr = %conn.addr, %error, "failed to send key exchange init packet");
             return Err(());
         }
 
diff --git a/src/lib.rs b/src/lib.rs
@@ -4,25 +4,22 @@ use std::{io, str, sync::Arc};
 use aws_lc_rs::{digest, signature::Ed25519KeyPair};
 use thiserror::Error;
 use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
+    io::AsyncReadExt,
     net::{tcp, TcpStream},
 };
 use tracing::{debug, warn};
 
 mod key_exchange;
 use key_exchange::KeyExchange;
 mod proto;
-use proto::{DecryptingReader, Encode};
-
-use crate::proto::Packet;
+use proto::{DecryptingReader, Encode, EncryptingWriter, Packet};
 
 /// A single SSH connection
 pub struct Connection {
     stream_read: DecryptingReader<tcp::OwnedReadHalf>,
-    stream_write: tcp::OwnedWriteHalf,
+    stream_write: EncryptingWriter<tcp::OwnedWriteHalf>,
     addr: SocketAddr,
     host_key: Arc<Ed25519KeyPair>,
-    write_buf: Vec<u8>,
 }
 
 impl Connection {
@@ -38,10 +35,9 @@ impl Connection {
 
         Ok(Self {
             stream_read: DecryptingReader::new(stream_read),
-            stream_write,
+            stream_write: EncryptingWriter::new(stream_write),
             addr,
             host_key,
-            write_buf: Vec::with_capacity(16_384),
         })
     }
 
@@ -119,14 +115,18 @@ impl VersionExchange {
         exchange.update(v_c);
 
         let ident = Identification::outgoing();
-        ident.encode(&mut conn.write_buf);
-        if let Err(error) = conn.stream_write.write_all(&conn.write_buf).await {
+        let server_ident_bytes = ident.encode();
+        if let Err(error) = conn
+            .stream_write
+            .write_raw_cleartext(&server_ident_bytes)
+            .await
+        {
             warn!(addr = %conn.addr, %error, "failed to send version exchange");
             return Err(());
         }
 
-        let v_s_len = conn.write_buf.len() - 2;
-        if let Some(v_s) = conn.write_buf.get(..v_s_len) {
+        let v_s_len = server_ident_bytes.len() - 2;
+        if let Some(v_s) = server_ident_bytes.get(..v_s_len) {
             exchange.update(&(v_s.len() as u32).to_be_bytes());
             exchange.update(v_s);
         }
@@ -197,10 +197,9 @@ impl<'a> Identification<'a> {
             comments,
         })
     }
-}
 
-impl Encode for Identification<'_> {
-    fn encode(&self, buf: &mut Vec<u8>) {
+    fn encode(&self) -> Vec<u8> {
+        let mut buf = vec![];
         buf.extend_from_slice(b"SSH-");
         buf.extend_from_slice(self.protocol.as_bytes());
         buf.push(b'-');
@@ -210,6 +209,7 @@ impl Encode for Identification<'_> {
             buf.extend_from_slice(self.comments.as_bytes());
         }
         buf.extend_from_slice(b"\r\n");
+        buf
     }
 }
 
diff --git a/src/proto.rs b/src/proto.rs
@@ -1,8 +1,11 @@
 use core::iter;
 use std::io;
 
-use aws_lc_rs::{cipher::StreamingDecryptingKey, constant_time, hmac, rand};
-use tokio::io::AsyncReadExt;
+use aws_lc_rs::{
+    cipher::{StreamingDecryptingKey, StreamingEncryptingKey},
+    constant_time, hmac, rand,
+};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tracing::debug;
 
 use crate::Error;
@@ -257,6 +260,59 @@ impl<R: AsyncReadExt + Unpin> DecryptingReader<R> {
     }
 }
 
+pub(crate) struct EncryptingWriter<W: AsyncWriteExt + Unpin> {
+    stream: W,
+    buf: Vec<u8>,
+    encrypted_buf: Vec<u8>,
+
+    packet_number: u32,
+    encryption_key: Option<(StreamingEncryptingKey, hmac::Key)>,
+}
+
+impl<W: AsyncWriteExt + Unpin> EncryptingWriter<W> {
+    pub(crate) fn new(stream: W) -> Self {
+        Self {
+            stream,
+            buf: Vec::with_capacity(16_384),
+            encrypted_buf: Vec::with_capacity(16_384),
+            packet_number: 0,
+            encryption_key: None,
+        }
+    }
+
+    /// Write raw bytes without packet structure or encryption.
+    ///
+    /// This should only be used for writing the identification string.
+    pub(crate) async fn write_raw_cleartext(&mut self, bytes: &[u8]) -> Result<(), Error> {
+        assert!(self.encryption_key.is_none());
+        self.stream.write_all(bytes).await?;
+        Ok(())
+    }
+
+    /// Write a packet. Returns written [`Packet`].
+    pub(crate) async fn write_packet(
+        &mut self,
+        payload: &impl Encode,
+        update_exchange_hash: impl FnOnce(&[u8]),
+    ) -> Result<(), Error> {
+        self.buf.clear();
+        self.encrypted_buf.clear();
+
+        let packet_number = self.packet_number;
+        self.packet_number = self.packet_number.wrapping_add(1);
+
+        if let Some((encryption_key, integrity_key)) = &mut self.encryption_key {
+            todo!()
+        } else {
+            let packet = Packet::builder(&mut self.buf).with_payload(payload);
+            update_exchange_hash(packet.payload()?);
+            self.stream.write_all(packet.without_mac()?).await?;
+        };
+
+        Ok(())
+    }
+}
+
 pub(crate) struct Packet<'a> {
     pub(crate) payload: &'a [u8],
 }
