From 1550bea1cc787b3ed1f5fea4a32896e68ebc7a22 Mon Sep 17 00:00:00 2001
From: Giovanni Petrantoni <giovanni@fragcolor.xyz>
Date: Wed, 25 Mar 2020 16:10:11 +0900
Subject: [PATCH] Support arbitrary stream write sizes in noise

---
 libp2p/protocols/secure/noise.nim | 31 +++++++++++++++--------
 tests/testnoise.nim               | 42 +++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/libp2p/protocols/secure/noise.nim b/libp2p/protocols/secure/noise.nim
index 8e4dd27..357abe0 100644
--- a/libp2p/protocols/secure/noise.nim
+++ b/libp2p/protocols/secure/noise.nim
@@ -35,6 +35,8 @@ const
   # Empty is a special value which indicates k has not yet been initialized.
   EmptyKey: ChaChaPolyKey = [0.byte, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
   NonceMax = uint64.high - 1 # max is reserved
+  NoiseSize = 32
+  MaxPlainSize = int(uint16.high - NoiseSize - ChaChaPolyTag.len)
 
 type
   KeyPair = object
@@ -282,7 +284,7 @@ proc sendHSMessage(sconn: Connection; buf: seq[byte]) {.async.} =
 
 proc packNoisePayload(payload: openarray[byte]): seq[byte] =
   let
-    noiselen = rand(2..31)
+    noiselen = rand(2..<NoiseSize)
     plen = payload.len.uint16
 
   var
@@ -427,17 +429,24 @@ method readMessage(sconn: NoiseConnection): Future[seq[byte]] {.async.} =
 
 method writeMessage(sconn: NoiseConnection, message: seq[byte]): Future[void] {.async.} =
   try:
-    let
-      packed = packNoisePayload(message)
-      cipher = sconn.writeCs.encryptWithAd([], packed)
     var
-      lesize = cipher.len.uint16
-      besize = lesize.toBytesBE
-      outbuf = newSeqOfCap[byte](cipher.len + 2)
-    trace "sendEncryptedMessage", size = lesize, peer = $sconn.peerInfo
-    outbuf &= besize
-    outbuf &= cipher
-    await sconn.write(outbuf)
+      left = message.len
+      offset = 0
+    while left > 0:
+      let
+        chunkSize = if left > MaxPlainSize: MaxPlainSize else: left
+        packed = packNoisePayload(message.toOpenArray(offset, offset + chunkSize - 1))
+        cipher = sconn.writeCs.encryptWithAd([], packed)
+      left = left - chunkSize
+      offset = offset + chunkSize
+      var
+        lesize = cipher.len.uint16
+        besize = lesize.toBytesBE
+        outbuf = newSeqOfCap[byte](cipher.len + 2)
+      trace "sendEncryptedMessage", size = lesize, peer = $sconn.peerInfo, left, offset
+      outbuf &= besize
+      outbuf &= cipher
+      await sconn.write(outbuf)
   except AsyncStreamWriteError:
     trace "Could not write to connection"
 
diff --git a/tests/testnoise.nim b/tests/testnoise.nim
index 0b96e9f..cb5b5e8 100644
--- a/tests/testnoise.nim
+++ b/tests/testnoise.nim
@@ -10,6 +10,7 @@
 import unittest, tables
 import chronos
 import chronicles
+import nimcrypto/sysrand
 import ../libp2p/crypto/crypto
 import ../libp2p/[switch,
                   multistream,
@@ -134,6 +135,47 @@ suite "Noise":
     check:
       waitFor(testListenerDialer()) == true
 
+  test "e2e: handle read + noise fragmented":
+    proc testListenerDialer(): Future[bool] {.async.} =
+      let
+        server: MultiAddress = Multiaddress.init("/ip4/0.0.0.0/tcp/0")
+        serverInfo = PeerInfo.init(PrivateKey.random(RSA), [server])
+        serverNoise = newNoise(serverInfo.privateKey, outgoing = false)
+        readTask = newFuture[void]()
+
+      var hugePayload = newSeq[byte](0xFFFFF)
+      check randomBytes(hugePayload) == hugePayload.len
+      trace "Sending huge payload", size = hugePayload.len
+
+      proc connHandler(conn: Connection) {.async, gcsafe.} =
+        let sconn = await serverNoise.secure(conn)
+        defer:
+          await sconn.close()
+        let msg = await sconn.readLp()
+        check msg == hugePayload
+        readTask.complete()
+
+      let
+        transport1: TcpTransport = newTransport(TcpTransport)
+      asyncCheck await transport1.listen(server, connHandler)
+
+      let
+        transport2: TcpTransport = newTransport(TcpTransport)
+        clientInfo = PeerInfo.init(PrivateKey.random(RSA), [transport1.ma])
+        clientNoise = newNoise(clientInfo.privateKey, outgoing = true)
+        conn = await transport2.dial(transport1.ma)
+        sconn = await clientNoise.secure(conn)
+
+      await sconn.writeLp(hugePayload)
+      await readTask
+      await sconn.close()
+      await transport1.close()
+
+      result = true
+
+    check:
+      waitFor(testListenerDialer()) == true
+
   test "e2e use switch dial proto string":
     proc testSwitch(): Future[bool] {.async, gcsafe.} =
       let ma1: MultiAddress = Multiaddress.init("/ip4/0.0.0.0/tcp/0")