From 56392913f902424f92166cb0cf2a266f314844ac Mon Sep 17 00:00:00 2001
From: Giovanni Petrantoni <sinkingsugar@gmail.com>
Date: Thu, 26 Mar 2020 15:06:47 +0900
Subject: [PATCH] Allow to specify a read size in readLp (#117)

* Allow to specify a read size in readLp

Also fix secio fragments + test

* Revert readLp changes, use just `read` in test
---
 libp2p/protocols/secure/secio.nim |  2 +-
 tests/testswitch.nim              | 45 ++++++++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/libp2p/protocols/secure/secio.nim b/libp2p/protocols/secure/secio.nim
index 2cb765a41..4d57d80a0 100644
--- a/libp2p/protocols/secure/secio.nim
+++ b/libp2p/protocols/secure/secio.nim
@@ -208,7 +208,7 @@ method writeMessage(sconn: SecioConn, message: seq[byte]) {.async.} =
       offset = 0
     while left > 0:
       let
-        chunkSize = if left > SecioMaxMessageSize: SecioMaxMessageSize else: left
+        chunkSize = if left > SecioMaxMessageSize - 64: SecioMaxMessageSize - 64 else: left
       let macsize = sconn.writerMac.sizeDigest()
       var msg = newSeq[byte](chunkSize + 4 + macsize)
       sconn.writerCoder.encrypt(message.toOpenArray(offset, offset + chunkSize - 1),
diff --git a/tests/testswitch.nim b/tests/testswitch.nim
index b9d968c3e..eba8641f8 100644
--- a/tests/testswitch.nim
+++ b/tests/testswitch.nim
@@ -1,5 +1,7 @@
 import unittest, tables
 import chronos
+import chronicles
+import nimcrypto/sysrand
 import ../libp2p/[switch,
                   multistream,
                   protocols/identify,
@@ -113,4 +115,45 @@ suite "Switch":
       result = true
 
     check:
-      waitFor(testSwitch()) == true
\ No newline at end of file
+      waitFor(testSwitch()) == true
+
+  test "e2e: handle read + secio 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 = newSecio(serverInfo.privateKey)
+        readTask = newFuture[void]()
+
+      var hugePayload = newSeq[byte](0x1200000)
+      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.read(0x1200000)
+        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 = newSecio(clientInfo.privateKey)
+        conn = await transport2.dial(transport1.ma)
+        sconn = await clientNoise.secure(conn)
+
+      await sconn.write(hugePayload)
+      await readTask
+      await sconn.close()
+      await transport1.close()
+
+      result = true
+
+    check:
+      waitFor(testListenerDialer()) == true