yamux fixes

libp2p/security/noise/io.py

@@ -22,6 +22,9 @@
 MAX_NOISE_MESSAGE_LEN = 2 ** (8 * SIZE_NOISE_MESSAGE_LEN) - 1
 SIZE_NOISE_MESSAGE_BODY_LEN = 2
 MAX_NOISE_MESSAGE_BODY_LEN = MAX_NOISE_MESSAGE_LEN - SIZE_NOISE_MESSAGE_BODY_LEN
+# Max plaintext per Noise message: 65535 - 16 bytes Poly1305 MAC overhead.
+# Matches go-libp2p's MaxPlaintextLength in p2p/security/noise/rw.go.
+MAX_PLAINTEXT_LENGTH = MAX_NOISE_MESSAGE_LEN - 16
 BYTE_ORDER = "big"
 
 # |                         Noise packet                            |
@@ -53,14 +56,26 @@ def __init__(self, conn: IRawConnection, noise_state: NoiseState) -> None:
         self.noise_state = noise_state
 
     async def write_msg(self, msg: bytes, prefix_encoded: bool = False) -> None:
-        logger.debug(f"Noise write_msg: encrypting {len(msg)} bytes")
-        data_encrypted = self.encrypt(msg)
-        if prefix_encoded:
-            # Manually add the prefix if needed
-            data_encrypted = self.prefix + data_encrypted
-        logger.debug(f"Noise write_msg: writing {len(data_encrypted)} encrypted bytes")
-        await self.read_writer.write_msg(data_encrypted)
-        logger.debug("Noise write_msg: write completed successfully")
+        # Chunk large messages to stay within the Noise 65535-byte transport
+        # message limit, matching go-libp2p's noise/rw.go Write() approach.
+        if len(msg) <= MAX_PLAINTEXT_LENGTH:
+            # Fast path: single message (covers handshake and small writes)
+            data_encrypted = self.encrypt(msg)
+            if prefix_encoded:
+                data_encrypted = self.prefix + data_encrypted
+            await self.read_writer.write_msg(data_encrypted)
+        else:
+            # Slow path: chunk into multiple Noise messages
+            total = len(msg)
+            written = 0
+            while written < total:
+                end = min(written + MAX_PLAINTEXT_LENGTH, total)
+                chunk = msg[written:end]
+                data_encrypted = self.encrypt(chunk)
+                if prefix_encoded and written == 0:
+                    data_encrypted = self.prefix + data_encrypted
+                await self.read_writer.write_msg(data_encrypted)
+                written = end
 
     async def read_msg(self, prefix_encoded: bool = False) -> bytes:
         logger.debug("Noise read_msg: reading encrypted message")

libp2p/security/secure_session.py

@@ -94,24 +94,42 @@ async def read(self, n: int | None = None) -> bytes:
             return b""
 
         data_from_buffer = self._drain(n)
-        if len(data_from_buffer) > 0:
+        if n is None and len(data_from_buffer) > 0:
             return data_from_buffer
 
-        msg = await self.conn.read_msg()
+        if n is None:
+            msg = await self.conn.read_msg()
 
-        # If underlying connection returned empty bytes, treat as closed
-        # and raise to signal that reads after close are invalid.
-        if msg == b"":
-            raise Exception("Connection closed")
+            # If underlying connection returned empty bytes, treat as closed
+            # and raise to signal that reads after close are invalid.
+            if msg == b"":
+                raise Exception("Connection closed")
 
-        if n is None:
             return msg
 
-        if n < len(msg):
-            self._fill(msg)
-            return self._drain(n)
-        else:
-            return msg
+        if len(data_from_buffer) == n:
+            return data_from_buffer
+
+        result = bytearray(data_from_buffer)
+        while len(result) < n:
+            msg = await self.conn.read_msg()
+
+            # If the connection closes after a partial read, return the bytes
+            # we already assembled. This preserves the stream-read behavior
+            # expected by higher layers.
+            if msg == b"":
+                if result:
+                    return bytes(result)
+                raise Exception("Connection closed")
+
+            remaining = n - len(result)
+            if len(msg) <= remaining:
+                result.extend(msg)
+            else:
+                result.extend(msg[:remaining])
+                self._fill(msg[remaining:])
+
+        return bytes(result)
 
     async def write(self, data: bytes) -> None:
         await self.conn.write_msg(data)