From 74f0b705f0d1d529c9e247d9dbf4f0083b012b93 Mon Sep 17 00:00:00 2001
From: Yuriy Glukhov <yuriy.glukhov@gmail.com>
Date: Tue, 15 May 2018 21:31:56 +0300
Subject: [PATCH] Eip8 auth

---
 ethp2p/auth.nim | 11 +++++++----
 ethp2p/rlpx.nim | 17 +++++++++++++----
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/ethp2p/auth.nim b/ethp2p/auth.nim
index 5b4fe3f..00da52f 100644
--- a/ethp2p/auth.nim
+++ b/ethp2p/auth.nim
@@ -334,14 +334,16 @@ 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
-    size: uint16
     secret: SharedSecret
-  bigEndian16(addr size, unsafeAddr m[0])
+  let size = expectedAuthMsgLenEip8(m)
   if 2 + int(size) > len(m):
     return(IncompleteError)
   var buffer = newSeq[byte](eciesDecryptedLength(int(size)))
@@ -434,9 +436,10 @@ proc decodeAuthMessage*(h: var Handshake, input: openarray[byte]): AuthStatus =
   if len(input) < AuthMessageV4Length:
     result = IncompleteError
   elif len(input) == AuthMessageV4Length:
-    let res = h.decodeAuthMessageV4(input)
+    var res = h.decodeAuthMessageV4(input)
     if res != Success:
-      if h.decodeAuthMessageEip8(input) != Success:
+      res = h.decodeAuthMessageEip8(input)
+      if res != Success:
         result = res
       else:
         h.flags.incl(EIP8)
diff --git a/ethp2p/rlpx.nim b/ethp2p/rlpx.nim
index 40621c7..7aad096 100644
--- a/ethp2p/rlpx.nim
+++ b/ethp2p/rlpx.nim
@@ -358,7 +358,6 @@ macro rlpxProtocol*(protoIdentifier: untyped,
         error("The only type that can be defined inside a RLPx protocol is the protocol's State type.")
 
     of nnkProcDef:
-      inc nextId
       let
         msgIdent = n.name.ident
         msgName = $msgIdent
@@ -464,6 +463,7 @@ macro rlpxProtocol*(protoIdentifier: untyped,
                          newStrLitNode($n.name),
                          thunkName)
 
+      inc nextId
     else:
       error("illegal syntax in a RLPx protocol definition", n)
 
@@ -574,11 +574,20 @@ proc rlpxConnectIncoming*(myKeys: KeyPair, listenPort: Port, address: IpAddress,
   var handshake = newHandshake({Responder})
   handshake.host = myKeys
 
-  var authMsg: array[1024, byte]
+  var authMsg: array[4086, byte]
   var authMsgLen = AuthMessageV4Length
-  # TODO: Handle both auth methods
+
   await s.fullRecvInto(addr authMsg[0], authMsgLen)
-  check handshake.decodeAuthMessage(^authMsg)
+  var ret = handshake.decodeAuthMessage(^authMsg)
+  if ret == AuthStatus.IncompleteError: # Eip8 auth message is likely
+    let expectedLen = expectedAuthMsgLenEip8(authMsg).int + 2
+    if expectedLen > authMsg.len:
+      raise newException(Exception, "Auth message too big: " & $authMsgLen)
+    await s.fullRecvInto(addr authMsg[authMsgLen], expectedLen - authMsgLen)
+    authMsgLen = expectedLen
+    ret = handshake.decodeAuthMessage(^authMsg)
+
+  check ret
 
   var ackMsg: array[AckMessageMaxEIP8, byte]
   var ackMsgLen: int