Support arbitrary stream write sizes in noise

This commit is contained in:
Giovanni Petrantoni 2020-03-25 16:10:11 +09:00
parent 4199508f17
commit 1550bea1cc
2 changed files with 62 additions and 11 deletions

View File

@ -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"

View File

@ -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")