[Updt] Mtp 0.3.0

iota-ipc/src/transport.rs

@@ -3,7 +3,10 @@ use serde::de::DeserializeOwned;
 use std::io::{Error, ErrorKind, Result};
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 
-const MAX_MESSAGE_SIZE: usize = 1024 * 1024;
+/// Maximum encoded payload size for a single IPC frame.
+///
+/// This is a wire-level contract shared by both sides of the connection.
+pub const MAX_MESSAGE_SIZE: usize = 1024 * 1024;
 
 /* Length-prefixing preserves message boundaries on a byte stream and bounds
  * allocations before JSON is deserialized. */
@@ -14,6 +17,12 @@ where
 {
     let payload =
         serde_json::to_vec(message).map_err(|error| Error::new(ErrorKind::InvalidData, error))?;
+    if payload.len() > MAX_MESSAGE_SIZE {
+        return Err(Error::new(
+            ErrorKind::InvalidData,
+            "IPC message exceeds limit",
+        ));
+    }
     let len = u32::try_from(payload.len())
         .map_err(|_| Error::new(ErrorKind::InvalidData, "IPC message is too large"))?;
     writer.write_u32(len).await?;
@@ -40,7 +49,7 @@ where
 
 #[cfg(test)]
 mod tests {
-    use super::{read_msg, write_msg};
+    use super::{MAX_MESSAGE_SIZE, read_msg, write_msg};
     use crate::protocol::{ClientMessage, LocalRequest, RequestEnvelope};
 
     #[tokio::test]
@@ -57,4 +66,31 @@ mod tests {
         let received: ClientMessage = read_msg(&mut reader).await.expect("read succeeds");
         assert!(matches!(received, ClientMessage::Request(req) if req.request_id == 4));
     }
+
+    #[tokio::test]
+    async fn write_rejects_message_above_frame_limit() {
+        let (mut writer, _reader) = tokio::io::duplex(MAX_MESSAGE_SIZE + 16);
+        let message = "x".repeat(MAX_MESSAGE_SIZE + 1);
+
+        let error = write_msg(&mut writer, &message)
+            .await
+            .expect_err("oversized payload must be rejected before framing");
+
+        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
+        assert!(error.to_string().contains("exceeds limit"));
+    }
+
+    #[tokio::test]
+    async fn read_rejects_frame_above_limit_before_allocating_payload() {
+        let (mut writer, mut reader) = tokio::io::duplex(16);
+        tokio::io::AsyncWriteExt::write_u32(&mut writer, (MAX_MESSAGE_SIZE + 1) as u32)
+            .await
+            .expect("length prefix write succeeds");
+
+        let error = read_msg::<_, ClientMessage>(&mut reader)
+            .await
+            .expect_err("oversized frame must be rejected");
+
+        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
+    }
 }