Implement actual decryption

Author: bjorn3

src/key_exchange.rs

@@ -2,7 +2,8 @@ use std::str;
 
 use aws_lc_rs::{
     agreement::{self, EphemeralPrivateKey, UnparsedPublicKey, X25519},
-    digest,
+    cipher::{StreamingDecryptingKey, UnboundCipherKey, AES_128},
+    digest, hmac,
     rand::{self, SystemRandom},
     signature::KeyPair,
 };
@@ -117,7 +118,36 @@ impl EcdhKeyExchange {
             return Err(());
         }
 
-        // FIXME wait for and send newkey packet
+        let packet = match conn.stream_read.read_packet().await {
+            Ok(packet) => packet,
+            Err(error) => {
+                warn!(addr = %conn.addr, %error, "failed to read packet");
+                return Err(());
+            }
+        };
+        let Decoded {
+            value: r#type,
+            next: _,
+        } = MessageType::decode(packet.payload)
+            .map_err(|error| warn!(addr = %conn.addr, %error, "failed to read packet type"))?;
+        if r#type != MessageType::NewKeys {
+            warn!(addr = %conn.addr,  "unexpected message type {:?}", r#type);
+            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 {
+            warn!(addr = %conn.addr, %error, "failed to send newkeys packet");
+            return Err(());
+        }
 
         // The first exchange hash is used as session id.
         let session_id = self.session_id.as_ref().unwrap_or(&exchange_hash);
@@ -126,12 +156,31 @@ impl EcdhKeyExchange {
             exchange_hash,
             session_id,
         };
-        #[expect(clippy::unnecessary_operation)]
-        RawKeySet {
+        let raw_keys = RawKeySet {
             client_to_server: RawKeys::client_to_server(&derivation),
             server_to_client: RawKeys::server_to_client(&derivation),
         };
 
+        conn.stream_read.set_decryption_key(
+            StreamingDecryptingKey::ctr(
+                UnboundCipherKey::new(
+                    &AES_128,
+                    &raw_keys.client_to_server.encryption_key.as_ref()[..16],
+                )
+                .unwrap(),
+                aws_lc_rs::cipher::DecryptionContext::Iv128(
+                    raw_keys.client_to_server.initial_iv.as_ref()[..16]
+                        .try_into()
+                        .unwrap(),
+                ),
+            )
+            .unwrap(),
+            hmac::Key::new(
+                hmac::HMAC_SHA256,
+                &raw_keys.client_to_server.integrity_key.as_ref()[..32],
+            ),
+        );
+
         Ok(())
     }
 }
@@ -555,7 +604,6 @@ struct RawKeySet {
     server_to_client: RawKeys,
 }
 
-#[expect(dead_code)] // FIXME implement encryption/decryption and MAC
 struct RawKeys {
     initial_iv: digest::Digest,
     encryption_key: digest::Digest,

src/main.rs

@@ -59,7 +59,7 @@ async fn main() -> anyhow::Result<()> {
             Ok((stream, addr)) => {
                 debug!(%addr, "accepted connection");
                 let conn = Connection::new(stream, addr, host_key.clone())?;
-                tokio::spawn(conn.run());
+                conn.run().await; // FIXME(aws/aws-lc-rs#975) use tokio::spawn() once StreamingDecryptingKey is Send
             }
             Err(error) => {
                 warn!(%error, "failed to accept connection");

src/proto.rs

@@ -1,7 +1,7 @@
 use core::iter;
 use std::io;
 
-use aws_lc_rs::rand;
+use aws_lc_rs::{cipher::StreamingDecryptingKey, constant_time, hmac, rand};
 use tokio::io::AsyncReadExt;
 use tracing::debug;
 
@@ -75,25 +75,36 @@ impl From<u8> for MessageType {
 ///      v
 /// |read|unread and not yet decrypted|
 /// ```
-// FIXME implement actual decryption
 pub(crate) struct DecryptingReader<R: AsyncReadExt + Unpin> {
     stream: R,
     buf: Vec<u8>,
+    decrypted_buf: Vec<u8>,
     unread_start: usize,
+
+    packet_number: u32,
+    decryption_key: Option<(StreamingDecryptingKey, hmac::Key)>,
 }
 
 impl<R: AsyncReadExt + Unpin> DecryptingReader<R> {
     pub(crate) fn new(stream: R) -> Self {
         Self {
             stream,
             buf: Vec::with_capacity(16_384),
+            decrypted_buf: Vec::with_capacity(16_384),
             unread_start: 0,
+            packet_number: 0,
+            decryption_key: None,
         }
     }
 
-    async fn ensure_at_least(&mut self, n: u32) -> Result<(), Error> {
-        while self.buf.len() - self.unread_start < n as usize {
-            let read = self.stream.read_buf(&mut self.buf).await?;
+    async fn ensure_at_least(
+        stream: &mut R,
+        buf: &mut Vec<u8>,
+        unread_start: &mut usize,
+        n: u32,
+    ) -> Result<(), Error> {
+        while buf.len() - *unread_start < n as usize {
+            let read = stream.read_buf(buf).await?;
             debug!(bytes = read, "read from stream");
             if read == 0 {
                 return Err(Error::Io(io::Error::new(
@@ -109,49 +120,149 @@ impl<R: AsyncReadExt + Unpin> DecryptingReader<R> {
     ///
     /// This should only be used for reading the identification string.
     pub(crate) async fn read_u8_cleartext(&mut self) -> Result<u8, Error> {
-        self.ensure_at_least(1).await?;
+        assert!(self.decryption_key.is_none());
+
+        Self::ensure_at_least(&mut self.stream, &mut self.buf, &mut self.unread_start, 1).await?;
 
         let byte = self.buf[self.unread_start];
         self.unread_start += 1;
         Ok(byte)
     }
 
+    pub(crate) fn set_decryption_key(
+        &mut self,
+        decryption_key: StreamingDecryptingKey,
+        integrity_key: hmac::Key,
+    ) {
+        self.decrypted_buf.clear();
+        self.decryption_key = Some((decryption_key, integrity_key));
+    }
+
     pub(crate) async fn read_packet<'a>(&'a mut self) -> Result<Packet<'a>, Error> {
         // Compact the internal buffer
         if self.unread_start > 0 {
             self.buf.copy_within(self.unread_start.., 0);
         }
         self.buf.truncate(self.buf.len() - self.unread_start);
+        self.decrypted_buf.clear();
         self.unread_start = 0;
 
-        self.ensure_at_least(4).await?;
-        let Decoded {
-            value: packet_length,
-            next,
-        } = PacketLength::decode(&self.buf[self.unread_start..self.unread_start + 4])?;
-        assert!(next.is_empty());
-
-        self.ensure_at_least(4 + packet_length.inner).await?;
-        let Decoded {
-            value: packet,
-            next,
-        } = Packet::decode(
-            &self.buf[self.unread_start..self.unread_start + 4 + packet_length.inner as usize],
-        )?;
-        assert!(next.is_empty());
-
-        self.unread_start += 4 + packet_length.inner as usize;
-
-        Ok(packet)
+        let packet_number = self.packet_number;
+        self.packet_number = self.packet_number.wrapping_add(1);
+
+        if let Some((decrypting_key, integrity_key)) = &mut self.decryption_key {
+            let block_len = decrypting_key.algorithm().block_len();
+
+            Self::ensure_at_least(
+                &mut self.stream,
+                &mut self.buf,
+                &mut self.unread_start,
+                block_len as u32,
+            )
+            .await?;
+            self.decrypted_buf.resize(self.buf.len() + block_len, 0);
+
+            let update = decrypting_key
+                .update(
+                    &self.buf[self.unread_start..self.unread_start + block_len],
+                    &mut self.decrypted_buf[self.unread_start..self.unread_start + 2 * block_len],
+                )
+                .unwrap();
+            assert_eq!(update.remainder().len(), block_len);
+
+            let Decoded {
+                value: packet_length,
+                next,
+            } = PacketLength::decode(
+                &self.decrypted_buf[self.unread_start..self.unread_start + 4],
+            )?;
+            assert!(next.is_empty());
+
+            Self::ensure_at_least(
+                &mut self.stream,
+                &mut self.buf,
+                &mut self.unread_start,
+                4 + packet_length.inner
+                    + integrity_key.algorithm().digest_algorithm().output_len as u32,
+            )
+            .await?;
+
+            let update = decrypting_key
+                .update(
+                    &self.buf[self.unread_start + block_len
+                        ..self.unread_start + 4 + packet_length.inner as usize],
+                    &mut self.decrypted_buf[self.unread_start + block_len
+                        ..self.unread_start + 4 + packet_length.inner as usize + block_len],
+                )
+                .unwrap();
+            assert_eq!(update.remainder().len(), block_len);
+
+            let mut hmac_ctx = hmac::Context::with_key(integrity_key);
+            hmac_ctx.update(&packet_number.to_be_bytes());
+            hmac_ctx.update(
+                &self.decrypted_buf
+                    [self.unread_start..self.unread_start + 4 + packet_length.inner as usize],
+            );
+            let actual_mac = hmac_ctx.sign();
+            let expected_mac = &self.buf[self.unread_start + 4 + packet_length.inner as usize
+                ..self.unread_start
+                    + 4
+                    + packet_length.inner as usize
+                    + integrity_key.algorithm().digest_algorithm().output_len];
+            constant_time::verify_slices_are_equal(actual_mac.as_ref(), expected_mac).unwrap(); // FIXME report error
+
+            let Decoded {
+                value: packet,
+                next,
+            } = Packet::decode(
+                &self.decrypted_buf
+                    [self.unread_start..self.unread_start + 4 + packet_length.inner as usize],
+            )?;
+            assert!(next.is_empty());
+
+            self.unread_start += 4
+                + packet_length.inner as usize
+                + integrity_key.algorithm().digest_algorithm().output_len;
+
+            Ok(packet)
+        } else {
+            Self::ensure_at_least(&mut self.stream, &mut self.buf, &mut self.unread_start, 4)
+                .await?;
+            let Decoded {
+                value: packet_length,
+                next,
+            } = PacketLength::decode(&self.buf[self.unread_start..self.unread_start + 4])?;
+            assert!(next.is_empty());
+
+            Self::ensure_at_least(
+                &mut self.stream,
+                &mut self.buf,
+                &mut self.unread_start,
+                4 + packet_length.inner,
+            )
+            .await?;
+
+            let Decoded {
+                value: packet,
+                next,
+            } = Packet::decode(
+                &self.buf[self.unread_start..self.unread_start + 4 + packet_length.inner as usize],
+            )?;
+            assert!(next.is_empty());
+
+            self.unread_start += 4 + packet_length.inner as usize;
+
+            Ok(packet)
+        }
     }
 }
 
 pub(crate) struct Packet<'a> {
     pub(crate) payload: &'a [u8],
 }
 
-impl Packet<'_> {
-    pub(crate) fn builder(buf: &mut Vec<u8>) -> PacketBuilder<'_> {
+impl<'a> Packet<'a> {
+    pub(crate) fn builder(buf: &'a mut Vec<u8>) -> PacketBuilder<'a> {
         let start = buf.len();
         buf.extend_from_slice(&[0, 0, 0, 0]); // packet_length
         buf.push(0); // padding_length