Allow big protocol messages

crypto/protocol/src/protocol_transport/noise.rs

@@ -3,11 +3,12 @@
 //! We use the XK handshake pattern.
 //! This means the initiator has a static keypair, and the responder has a pre-shared static keypair
 //! That is, we already know the public key of the remote party we are connecting to before the
-//! handshake starts)
+//! handshake starts.
 //!
 //! See: <https://noiseexplorer.com/patterns/XK>
 use entropy_shared::X25519PublicKey;
 use snow::{params::NoiseParams, Builder, HandshakeState};
+use std::cmp::min;
 
 use super::{errors::EncryptedConnectionErr, WsConnection};
 
@@ -17,6 +18,12 @@ const NOISE_PARAMS: &str = "Noise_XK_25519_ChaChaPoly_BLAKE2s";
 /// This is used in the handshake as context
 const NOISE_PROLOGUE: &[u8; 24] = b"Entropy signing protocol";
 
+/// The maxiumum message size for the noise protocol
+const MAX_NOISE_MESSAGE_SIZE: usize = 65535;
+
+/// The size of the authentication data added to each encrypted noise message
+const NOISE_PAYLOAD_AUTHENTICATION_SIZE: usize = 16;
+
 /// Handshake as an initiator
 pub async fn noise_handshake_initiator<T: WsConnection>(
     mut ws_connection: T,
@@ -27,7 +34,7 @@ pub async fn noise_handshake_initiator<T: WsConnection>(
     let mut noise = setup_noise(local_private_key, Some(remote_public_key)).await?;
 
     // Used to hold handshake messages
-    let mut buf = vec![0u8; 65535];
+    let mut buf = vec![0u8; MAX_NOISE_MESSAGE_SIZE];
 
     // Initiator sends first message
     let len = noise.write_message(&[], &mut buf)?;
@@ -50,7 +57,7 @@ pub async fn noise_handshake_responder<T: WsConnection>(
     let mut noise = setup_noise(local_private_key, None).await?;
 
     // Used to hold handshake messages
-    let mut buf = vec![0u8; 65535];
+    let mut buf = vec![0u8; MAX_NOISE_MESSAGE_SIZE];
 
     // Responder reads first message
     noise.read_message(&ws_connection.recv().await?, &mut buf)?;
@@ -94,16 +101,38 @@ pub struct EncryptedWsConnection<T: WsConnection> {
 
 impl<T: WsConnection> EncryptedWsConnection<T> {
     /// Receive and decrypt the next message
+    /// This splits the incoming message into chunks of the maximum message size allowed
+    /// by the noise protocol, decrypts them individually and concatenates the results into
+    /// a single message
     pub async fn recv(&mut self) -> Result<Vec<u8>, EncryptedConnectionErr> {
         let ciphertext = self.ws_connection.recv().await?;
-        let len = self.noise_transport.read_message(&ciphertext, &mut self.buf)?;
-        Ok(self.buf[..len].to_vec())
+
+        let mut full_message = Vec::new();
+        let mut i = 0;
+        while i < ciphertext.len() {
+            let read_to = min(i + MAX_NOISE_MESSAGE_SIZE, ciphertext.len());
+            let len = self.noise_transport.read_message(&ciphertext[i..read_to], &mut self.buf)?;
+            full_message.extend_from_slice(&self.buf[..len]);
+            i += MAX_NOISE_MESSAGE_SIZE;
+        }
+
+        Ok(full_message)
     }
 
     /// Encrypt and send a message
+    /// This splits the outgoing message into chunks of the maximum size allowed by the noise
+    /// protocol, encrypts them individually and concatenates the results into a single message
     pub async fn send(&mut self, msg: Vec<u8>) -> Result<(), EncryptedConnectionErr> {
-        let len = self.noise_transport.write_message(&msg, &mut self.buf)?;
-        self.ws_connection.send(self.buf[..len].to_vec()).await?;
+        let mut messages = Vec::new();
+        let mut i = 0;
+        while i < msg.len() {
+            let read_to =
+                min(i + MAX_NOISE_MESSAGE_SIZE - NOISE_PAYLOAD_AUTHENTICATION_SIZE, msg.len());
+            let len = self.noise_transport.write_message(&msg[i..read_to], &mut self.buf)?;
+            messages.extend_from_slice(&self.buf[..len]);
+            i += MAX_NOISE_MESSAGE_SIZE - NOISE_PAYLOAD_AUTHENTICATION_SIZE;
+        }
+        self.ws_connection.send(messages).await?;
         Ok(())
     }
 
@@ -177,4 +206,40 @@ mod tests {
         assert_eq!(bob_connection.recv().await.unwrap(), b"hello bob".to_vec());
         assert_eq!(alice_connection.recv().await.unwrap(), b"hello alice".to_vec());
     }
+
+    #[tokio::test]
+    async fn test_encrypted_connection_with_big_message() {
+        let alice_sk = x25519_dalek::StaticSecret::new(rand_core::OsRng);
+
+        let bob_sk = x25519_dalek::StaticSecret::new(rand_core::OsRng);
+        let bob_pk = x25519_dalek::PublicKey::from(&bob_sk);
+
+        let (alice_tx, alice_rx) = mpsc::channel(100);
+        let (bob_tx, bob_rx) = mpsc::channel(100);
+
+        let (alice_connection_result, bob_connection_result) = futures::future::join(
+            noise_handshake_initiator(
+                MockWsConnection::new(alice_tx, bob_rx),
+                &alice_sk,
+                bob_pk.as_bytes().clone(),
+                Vec::new(),
+            ),
+            noise_handshake_responder(MockWsConnection::new(bob_tx, alice_rx), &bob_sk),
+        )
+        .await;
+
+        let mut alice_connection = alice_connection_result.unwrap();
+        let (mut bob_connection, _) = bob_connection_result.unwrap();
+
+        let big_message_for_bob: [u8; MAX_NOISE_MESSAGE_SIZE + 100] =
+            [1; MAX_NOISE_MESSAGE_SIZE + 100];
+        let big_message_for_alice: [u8; MAX_NOISE_MESSAGE_SIZE * 3] =
+            [2; MAX_NOISE_MESSAGE_SIZE * 3];
+
+        alice_connection.send(big_message_for_bob.to_vec()).await.unwrap();
+        bob_connection.send(big_message_for_alice.to_vec()).await.unwrap();
+
+        assert_eq!(bob_connection.recv().await.unwrap(), &big_message_for_bob);
+        assert_eq!(alice_connection.recv().await.unwrap(), &big_message_for_alice);
+    }
 }