Merge pull request #21 from status-im/expectedLength2

Fix ack/auth message handling.

ethp2p/auth.nim

@@ -72,6 +72,7 @@ type
     remoteEPubkey*: PublicKey   ## remote host ephemeral public key
     initiatorNonce*: Nonce      ## initiator nonce
     responderNonce*: Nonce      ## responder nonce
+    expectedLength*: int        ## expected incoming message length
 
   ConnectionSecret* = object
     aesKey*: array[aes256.sizeKey, byte]
@@ -96,9 +97,11 @@ proc newHandshake*(flags: set[HandshakeFlag] = {Initiator},
   result.flags = flags
   result.ephemeral = newKeyPair()
   if Initiator in flags:
+    result.expectedLength = AckMessageV4Length
     if randomBytes(result.initiatorNonce) != len(result.initiatorNonce):
       raise newException(AuthException, "Could not obtain random data!")
   else:
+    result.expectedLength = AuthMessageV4Length
     if randomBytes(result.responderNonce) != len(result.responderNonce):
       raise newException(AuthException, "Could not obtain random data!")
 
@@ -334,17 +337,17 @@ proc decodeAuthMessageV4(h: var Handshake, m: openarray[byte]): AuthStatus =
   h.remoteHPubkey = pubkey
   result = Success
 
-proc expectedAuthMsgLenEip8*(input: openarray[byte]): uint16 {.inline.} =
-  bigEndian16(addr result, unsafeAddr input[0])
-
 proc decodeAuthMessageEip8(h: var Handshake, m: openarray[byte]): AuthStatus =
   ## Decodes EIP-8 AuthMessage.
   var
     pubkey: PublicKey
     nonce: Nonce
     secret: SharedSecret
-  let size = expectedAuthMsgLenEip8(m)
-  if 2 + int(size) > len(m):
+    size: uint16
+
+  bigEndian16(addr size, unsafeAddr m[0])
+  h.expectedLength = int(size) + 2
+  if h.expectedLength > len(m):
     return(IncompleteError)
   var buffer = newSeq[byte](eciesDecryptedLength(int(size)))
   if eciesDecrypt(toa(m, 2, int(size)), buffer, h.host.seckey,
@@ -388,9 +391,9 @@ proc decodeAuthMessageEip8(h: var Handshake, m: openarray[byte]): AuthStatus =
 proc decodeAckMessageEip8*(h: var Handshake, m: openarray[byte]): AuthStatus =
   ## Decodes EIP-8 AckMessage.
   var size: uint16
-  assert(len(m) > 2)
   bigEndian16(addr size, unsafeAddr m[0])
-  if 2 + int(size) > len(m):
+  h.expectedLength = 2 + int(size)
+  if h.expectedLength > len(m):
     return(IncompleteError)
   var buffer = newSeq[byte](eciesDecryptedLength(int(size)))
   if eciesDecrypt(toa(m, 2, int(size)), buffer, h.host.seckey,
@@ -456,9 +459,10 @@ proc decodeAckMessage*(h: var Handshake, input: openarray[byte]): AuthStatus =
   if len(input) < AckMessageV4Length:
     return(IncompleteError)
   elif len(input) == AckMessageV4Length:
-    let res = h.decodeAckMessageV4(input)
+    var res = h.decodeAckMessageV4(input)
     if res != Success:
-      if h.decodeAckMessageEip8(input) != Success:
+      res = h.decodeAckMessageEip8(input)
+      if res != Success:
         result = res
       else:
         h.flags.incl(EIP8)

ethp2p/rlpx.nim

@@ -552,15 +552,21 @@ proc rlpxConnect*(myKeys: KeyPair, listenPort: Port, remote: Node): Future[Peer]
   var authMsg: array[AuthMessageMaxEIP8, byte]
   var authMsgLen = 0
   check authMessage(handshake, remote.node.pubkey, authMsg, authMsgLen)
-
   await result.socket.send(addr authMsg[0], authMsgLen)
 
-  var ackMsg: array[AckMessageMaxEIP8, byte]
-  let ackMsgLen = handshake.ackSize()
-  await result.socket.fullRecvInto(addr ackMsg, ackMsgLen)
+  let initialSize = handshake.expectedLength
+  var ackMsg = newSeqOfCap[byte](1024)
+  ackMsg.setLen(initialSize)
+  await result.socket.fullRecvInto(ackMsg)
+  var ret = handshake.decodeAckMessage(ackMsg)
+  if ret == AuthStatus.IncompleteError:
+    ackMsg.setLen(handshake.expectedLength)
+    await result.socket.fullRecvInto(addr ackMsg[initialSize],
+                                     len(ackMsg) - initialSize)
+    ret = handshake.decodeAckMessage(ackMsg)
+  check ret
 
-  check handshake.decodeAckMessage(^ackMsg)
-  initSecretState(handshake, ^authMsg, ^ackMsg, result)
+  initSecretState(handshake, ^authMsg, ackMsg, result)
 
   if handshake.remoteHPubkey != remote.node.pubKey:
     raise newException(Exception, "Remote pubkey is wrong")
@@ -581,23 +587,22 @@ proc rlpxConnectIncoming*(myKeys: KeyPair, listenPort: Port, address: IpAddress,
   var handshake = newHandshake({Responder})
   handshake.host = myKeys
 
+  let initialSize = handshake.expectedLength
   var authMsg = newSeqOfCap[byte](1024)
-  authMsg.setLen(AuthMessageV4Length)
-
+  authMsg.setLen(initialSize)
   await s.fullRecvInto(authMsg)
   var ret = handshake.decodeAuthMessage(authMsg)
   if ret == AuthStatus.IncompleteError: # Eip8 auth message is likely
-    authMsg.setLen(expectedAuthMsgLenEip8(authMsg).int + 2)
-    await s.fullRecvInto(addr authMsg[AuthMessageV4Length], authMsg.len - AuthMessageV4Length)
+    authMsg.setLen(handshake.expectedLength)
+    await s.fullRecvInto(addr authMsg[initialSize], len(authMsg) - initialSize)
     ret = handshake.decodeAuthMessage(authMsg)
-
   check ret
 
   var ackMsg: array[AckMessageMaxEIP8, byte]
   var ackMsgLen: int
   check handshake.ackMessage(ackMsg, ackMsgLen)
-
   await s.send(addr ackMsg[0], ackMsgLen)
+
   initSecretState(handshake, authMsg, ^ackMsg, result)
 
   var response = await result.nextMsg(p2p.hello, discardOthers = true)