Pre and post-serialization steps for send procs

eth/p2p/p2p_backends_helpers.nim

@@ -30,6 +30,14 @@ template networkState*(connection: Peer, Protocol: type): untyped =
 
 proc initProtocolState*[T](state: T, x: Peer|EthereumNode) {.gcsafe.} = discard
 
+proc initProtocolStates(peer: Peer, protocols: openarray[ProtocolInfo]) =
+  # Initialize all the active protocol states
+  newSeq(peer.protocolStates, allProtocols.len)
+  for protocol in protocols:
+    let peerStateInit = protocol.peerStateInitializer
+    if peerStateInit != nil:
+      peer.protocolStates[protocol.index] = peerStateInit(peer)
+
 proc resolveFuture[MsgType](msg: pointer, future: FutureBase) {.gcsafe.} =
   var f = Future[MsgType](future)
   doAssert(not f.finished())
@@ -91,8 +99,8 @@ proc handshakeImpl[T](peer: Peer,
   doAssert timeout.milliseconds > 0
   yield responseFut or sleepAsync(timeout)
   if not responseFut.finished:
-    discard disconnectAndRaise(peer, HandshakeTimeout,
-                               "Protocol handshake was not received in time.")
+    await disconnectAndRaise(peer, HandshakeTimeout,
+                             "Protocol handshake was not received in time.")
   elif responseFut.failed:
     raise responseFut.error
   else:

eth/p2p/p2p_protocol_dsl.nim

@@ -500,7 +500,8 @@ proc setBody*(sendProc: SendProc, body: NimNode) =
     msg.protocol.outSendProcs.add sendProc.extraDefs
 
 proc useStandardBody*(sendProc: SendProc,
-                      preludeGenerator: proc(stream: NimNode): NimNode,
+                      preSerializationStep: proc(stream: NimNode): NimNode,
+                      postSerializationStep: proc(stream: NimNode): NimNode,
                       sendCallGenerator: proc (peer, bytes: NimNode): NimNode) =
   let
     msg = sendProc.msg
@@ -514,8 +515,12 @@ proc useStandardBody*(sendProc: SendProc,
     msgRecName = msg.recIdent
     Format = msg.protocol.backend.SerializationFormat
 
-    prelude = if preludeGenerator.isNil: newStmtList()
-              else: preludeGenerator(outputStream)
+    preSerialization = if preSerializationStep.isNil: newStmtList()
+                       else: preSerializationStep(outputStream)
+
+    postSerialization = if postSerializationStep.isNil: newStmtList()
+                        else: postSerializationStep(outputStream)
+
     appendParams = newStmtList()
 
     initResultFuture = if msg.kind != msgRequest: newStmtList()
@@ -537,12 +542,13 @@ proc useStandardBody*(sendProc: SendProc,
 
     `initResultFuture`
     var `outputStream` = init OutputStream
-    `prelude`
+    `preSerialization`
     var `writer` = init(WriterType(`Format`), `outputStream`)
     var recordStartMemo = beginRecord(`writer`, `msgRecName`)
     `appendParams`
     `tracing`
     endRecord(`writer`, recordStartMemo)
+    `postSerialization`
     let `msgBytes` = getOutput(`outputStream`)
     `sendCall`
 

eth/p2p/rlpx.nim

@@ -862,15 +862,8 @@ proc initPeerState*(peer: Peer, capabilities: openarray[Capability]) =
   # Similarly, we need a bit of book-keeping data to keep track
   # of the potentially concurrent calls to `nextMsg`.
   peer.awaitedMessages.newSeq(peer.dispatcher.messages.len)
-
   peer.lastReqId = 0
-
-  # Initialize all the active protocol states
-  newSeq(peer.protocolStates, allProtocols.len)
-  for protocol in peer.dispatcher.activeProtocols:
-    let peerStateInit = protocol.peerStateInitializer
-    if peerStateInit != nil:
-      peer.protocolStates[protocol.index] = peerStateInit(peer)
+  peer.initProtocolStates peer.dispatcher.activeProtocols
 
 proc postHelloSteps(peer: Peer, h: devp2p.hello) {.async.} =
   initPeerState(peer, h.capabilities)