[ETCM-102] Fix race condition in frame encoder/decoder

Author: KonradStaniec

src/main/scala/io/iohk/ethereum/network/PeerActor.scala

@@ -1,22 +1,22 @@
 package io.iohk.ethereum.network
 
-import java.net.{ InetSocketAddress, URI }
+import java.net.{InetSocketAddress, URI}
 
 import akka.actor.SupervisorStrategy.Escalate
 import akka.actor._
 import akka.util.ByteString
 import io.iohk.ethereum.network.PeerActor.Status._
-import io.iohk.ethereum.network.PeerEventBusActor.PeerEvent.{ MessageFromPeer, PeerHandshakeSuccessful }
+import io.iohk.ethereum.network.PeerEventBusActor.PeerEvent.{MessageFromPeer, PeerHandshakeSuccessful}
 import io.iohk.ethereum.network.PeerEventBusActor.Publish
 import io.iohk.ethereum.network.PeerManagerActor.PeerConfiguration
 import io.iohk.ethereum.network.handshaker.Handshaker
-import io.iohk.ethereum.network.handshaker.Handshaker.HandshakeComplete.{ HandshakeFailure, HandshakeSuccess }
-import io.iohk.ethereum.network.handshaker.Handshaker.{ HandshakeResult, NextMessage }
+import io.iohk.ethereum.network.handshaker.Handshaker.HandshakeComplete.{HandshakeFailure, HandshakeSuccess}
+import io.iohk.ethereum.network.handshaker.Handshaker.{HandshakeResult, NextMessage}
 import io.iohk.ethereum.network.p2p._
 import io.iohk.ethereum.network.p2p.messages.Versions
 import io.iohk.ethereum.network.p2p.messages.WireProtocol._
 import io.iohk.ethereum.network.rlpx.RLPxConnectionHandler.RLPxConfiguration
-import io.iohk.ethereum.network.rlpx.{ AuthHandshaker, RLPxConnectionHandler }
+import io.iohk.ethereum.network.rlpx.{AuthHandshaker, RLPxConnectionHandler}
 import org.bouncycastle.util.encoders.Hex
 
 
@@ -28,18 +28,18 @@ import org.bouncycastle.util.encoders.Hex
   * Once that's done it can send/receive messages with peer (HandshakedHandler.receive).
   */
 class PeerActor[R <: HandshakeResult](
-    peerAddress: InetSocketAddress,
-    rlpxConnectionFactory: ActorContext => ActorRef,
-    val peerConfiguration: PeerConfiguration,
-    peerEventBus: ActorRef,
-    knownNodesManager: ActorRef,
-    incomingConnection: Boolean,
-    externalSchedulerOpt: Option[Scheduler] = None,
-    initHandshaker: Handshaker[R])
+                                       peerAddress: InetSocketAddress,
+                                       rlpxConnectionFactory: ActorContext => ActorRef,
+                                       val peerConfiguration: PeerConfiguration,
+                                       peerEventBus: ActorRef,
+                                       knownNodesManager: ActorRef,
+                                       incomingConnection: Boolean,
+                                       externalSchedulerOpt: Option[Scheduler] = None,
+                                       initHandshaker: Handshaker[R])
   extends Actor with ActorLogging with Stash {
 
   import PeerActor._
-  import context.{ dispatcher, system }
+  import context.{dispatcher, system}
 
   override val supervisorStrategy: OneForOneStrategy =
     OneForOneStrategy() {
@@ -73,6 +73,7 @@ class PeerActor[R <: HandshakeResult](
     context watch ref
     RLPxConnection(ref, remoteAddress, uriOpt)
   }
+
   private def modifyOutGoingUri(remoteNodeId: ByteString, rlpxConnection: RLPxConnection, uri: URI): URI = {
     val host = getHostName(rlpxConnection.remoteAddress.getAddress)
     val port = rlpxConnection.remoteAddress.getPort
@@ -106,23 +107,14 @@ class PeerActor[R <: HandshakeResult](
 
   def processingHandshaking(handshaker: Handshaker[R], rlpxConnection: RLPxConnection,
                             timeout: Cancellable, numRetries: Int): Receive =
-      handleTerminated(rlpxConnection, numRetries, Handshaking(numRetries)) orElse
+    handleTerminated(rlpxConnection, numRetries, Handshaking(numRetries)) orElse
       handleDisconnectMsg(rlpxConnection, Handshaking(numRetries)) orElse
       handlePingMsg(rlpxConnection) orElse stashMessages orElse {
 
       case RLPxConnectionHandler.MessageReceived(msg) =>
-
-        // We need to determine p2p version just after hello message as next messages in handshake
-        // can be compressed.
-        msg match {
-          case Hello(p2pVersion, _, _, _, _) =>
-            rlpxConnection.ref ! PeerP2pVersion(p2pVersion)
-          case _ => ()
-        }
-
         // Processes the received message, cancels the timeout and processes a new message but only if the handshaker
         // handles the received message
-        handshaker.applyMessage(msg).foreach{ newHandshaker =>
+        handshaker.applyMessage(msg).foreach { newHandshaker =>
           timeout.cancel()
           processHandshakerNextMessage(newHandshaker, rlpxConnection, numRetries)
         }
@@ -143,7 +135,7 @@ class PeerActor[R <: HandshakeResult](
     *
     * @param handshaker
     * @param rlpxConnection
-    * @param numRetries, number of connection retries done during RLPxConnection establishment
+    * @param numRetries , number of connection retries done during RLPxConnection establishment
     */
   private def processHandshakerNextMessage(handshaker: Handshaker[R],
                                            rlpxConnection: RLPxConnection,
@@ -155,7 +147,7 @@ class PeerActor[R <: HandshakeResult](
         context become processingHandshaking(handshaker, rlpxConnection, newTimeout, numRetries)
 
       case Left(HandshakeSuccess(handshakeResult)) =>
-        rlpxConnection.uriOpt.foreach { uri =>knownNodesManager ! KnownNodesManager.AddKnownNode(uri) }
+        rlpxConnection.uriOpt.foreach { uri => knownNodesManager ! KnownNodesManager.AddKnownNode(uri) }
         context become new HandshakedPeer(rlpxConnection, handshakeResult).receive
         unstashAll()
 
@@ -252,8 +244,8 @@ class PeerActor[R <: HandshakeResult](
       */
     def receive: Receive =
       handlePingMsg(rlpxConnection) orElse
-      handleDisconnectMsg(rlpxConnection, Handshaked) orElse
-      handleTerminated(rlpxConnection, 0, Handshaked) orElse {
+        handleDisconnectMsg(rlpxConnection, Handshaked) orElse
+        handleTerminated(rlpxConnection, 0, Handshaked) orElse {
 
         case RLPxConnectionHandler.MessageReceived(message) =>
           log.debug(s"Received message: {} from $peerId", message)
@@ -268,7 +260,7 @@ class PeerActor[R <: HandshakeResult](
         case GetStatus =>
           sender() ! StatusResponse(Handshaked)
 
-    }
+      }
   }
 
 }
@@ -304,8 +296,6 @@ object PeerActor {
     }
   }
 
-  case class PeerP2pVersion(p2pVersion: Long)
-
   case class HandleConnection(connection: ActorRef, remoteAddress: InetSocketAddress)
 
   case class IncomingConnectionHandshakeSuccess(peer: Peer)
@@ -321,16 +311,25 @@ object PeerActor {
   private case object ResponseTimeout
 
   case object GetStatus
+
   case class StatusResponse(status: Status)
 
   case class DisconnectPeer(reason: Int)
 
   sealed trait Status
+
   object Status {
+
     case object Idle extends Status
+
     case object Connecting extends Status
+
     case class Handshaking(numRetries: Int) extends Status
+
     case object Handshaked extends Status
+
     case object Disconnected extends Status
+
   }
+
 }

src/main/scala/io/iohk/ethereum/network/rlpx/MessageCodec.scala

@@ -4,8 +4,10 @@ import java.util.concurrent.atomic.AtomicInteger
 
 import akka.util.ByteString
 import io.iohk.ethereum.network.handshaker.EtcHelloExchangeState
+import io.iohk.ethereum.network.p2p.messages.WireProtocol.Hello
 import io.iohk.ethereum.network.p2p.{Message, MessageDecoder, MessageSerializable}
 import org.xerial.snappy.Snappy
+
 import scala.util.{Failure, Success, Try}
 
 class MessageCodec(frameCodec: FrameCodec, messageDecoder: MessageDecoder, protocolVersion: Message.Version) {
@@ -17,19 +19,37 @@ class MessageCodec(frameCodec: FrameCodec, messageDecoder: MessageDecoder, proto
   // 16Mb in base 2
   val maxDecompressedLength = 16777216
 
-  def readMessages(data: ByteString, p2pVersion: Option[Long]): Seq[Try[Message]] = {
+  // MessageCodec is only used from actor context so it can be var
+  @volatile
+  private var remotePeerP2pVersion: Option[Long] = None
+
+  private def setRemoteVersionBasedOnHelloMessage(m: Message): Unit = {
+    if (remotePeerP2pVersion.isEmpty) {
+      m match {
+        case hello: Hello =>
+          remotePeerP2pVersion = Some(hello.p2pVersion)
+        case _ =>
+      }
+    }
+  }
+
+  def readMessages(data: ByteString): Seq[Try[Message]] = {
     val frames = frameCodec.readFrames(data)
 
     frames map { frame =>
       val frameData = frame.payload.toArray
       val payloadTry =
-        if (p2pVersion.contains(EtcHelloExchangeState.P2pVersion)){
+        if (remotePeerP2pVersion.exists(version => version >= EtcHelloExchangeState.P2pVersion)) {
           decompressData(frameData)
         } else {
           Success(frameData)
         }
 
-      payloadTry.map(payload => messageDecoder.fromBytes(frame.`type`, payload, protocolVersion))
+      payloadTry.map { payload =>
+        val m = messageDecoder.fromBytes(frame.`type`, payload, protocolVersion)
+        setRemoteVersionBasedOnHelloMessage(m)
+        m
+      }
     }
   }
 
@@ -42,14 +62,16 @@ class MessageCodec(frameCodec: FrameCodec, messageDecoder: MessageDecoder, proto
     }
   }
 
-  def encodeMessage(serializable: MessageSerializable, p2pVersion: Option[Long]): ByteString = {
+  def encodeMessage(serializable: MessageSerializable): ByteString = {
     val encoded: Array[Byte] = serializable.toBytes
     val numFrames = Math.ceil(encoded.length / MaxFramePayloadSize.toDouble).toInt
     val contextId = contextIdCounter.incrementAndGet()
     val frames = (0 until numFrames) map { frameNo =>
       val framedPayload = encoded.drop(frameNo * MaxFramePayloadSize).take(MaxFramePayloadSize)
       val payload =
-        if (p2pVersion.contains(EtcHelloExchangeState.P2pVersion)){
+        if (
+          remotePeerP2pVersion.exists(version => version >= EtcHelloExchangeState.P2pVersion) && serializable.code != Hello.code
+        ) {
           Snappy.compress(framedPayload)
         } else {
           framedPayload

src/main/scala/io/iohk/ethereum/network/rlpx/RLPxConnectionHandler.scala

@@ -6,7 +6,6 @@ import akka.actor._
 import akka.io.Tcp._
 import akka.io.{IO, Tcp}
 import akka.util.ByteString
-import io.iohk.ethereum.network.PeerActor.PeerP2pVersion
 import io.iohk.ethereum.network.p2p.{Message, MessageDecoder, MessageSerializable}
 import io.iohk.ethereum.network.rlpx.RLPxConnectionHandler.RLPxConfiguration
 import io.iohk.ethereum.utils.ByteUtils
@@ -24,15 +23,15 @@ import scala.util.{Failure, Success, Try}
   * 1. when created it waits for initial command (either handle incoming connection or connect usin g uri)
   * 2. when new connection is requested the actor waits for the result (waitingForConnectionResult)
   * 3. once underlying connection is established it either waits for handshake init message or for response message
-  *    (depending on who initiated the connection)
+  * (depending on who initiated the connection)
   * 4. once handshake is done (and secure connection established) actor can send/receive messages (`handshaked` state)
   */
 class RLPxConnectionHandler(
-    messageDecoder: MessageDecoder,
-    protocolVersion: Message.Version,
-    authHandshaker: AuthHandshaker,
-    messageCodecFactory: (Secrets, MessageDecoder, Message.Version) => MessageCodec,
-    rlpxConfiguration: RLPxConfiguration)
+                             messageDecoder: MessageDecoder,
+                             protocolVersion: Message.Version,
+                             authHandshaker: AuthHandshaker,
+                             messageCodecFactory: (Secrets, MessageDecoder, Message.Version) => MessageCodec,
+                             rlpxConfiguration: RLPxConfiguration)
   extends Actor with ActorLogging {
 
   import AuthHandshaker.{InitiatePacketLength, ResponsePacketLength}
@@ -126,10 +125,10 @@ class RLPxConnectionHandler(
     /**
       * Decode V4 packet
       *
-      * @param data, includes both the V4 packet with bytes from next messages
+      * @param data , includes both the V4 packet with bytes from next messages
       * @return data of the packet and the remaining data
       */
-    private def decodeV4Packet(data: ByteString): (ByteString, ByteString) =  {
+    private def decodeV4Packet(data: ByteString): (ByteString, ByteString) = {
       val encryptedPayloadSize = ByteUtils.bigEndianToShort(data.take(2).toArray)
       val (packetData, remainingData) = data.splitAt(encryptedPayloadSize + 2)
       packetData -> remainingData
@@ -148,7 +147,7 @@ class RLPxConnectionHandler(
           log.debug(s"Auth handshake succeeded for peer $peerId")
           context.parent ! ConnectionEstablished(remotePubKey)
           val messageCodec = messageCodecFactory(secrets, messageDecoder, protocolVersion)
-          val messagesSoFar = messageCodec.readMessages(remainingData ,None)
+          val messagesSoFar = messageCodec.readMessages(remainingData)
           messagesSoFar foreach processMessage
           context become handshaked(messageCodec)
 
@@ -170,59 +169,54 @@ class RLPxConnectionHandler(
       * Handles sending and receiving messages from the Akka TCP connection, while also handling acknowledgement of
       * messages sent. Messages are only sent when all Ack from previous messages were received.
       *
-      * @param messageCodec, for encoding the messages sent
-      * @param messagesNotSent, messages not yet sent
-      * @param cancellableAckTimeout, timeout for the message sent for which we are awaiting an acknowledgement (if there is one)
-      * @param seqNumber, sequence number for the next message to be sent
+      * @param messageCodec          , for encoding the messages sent
+      * @param messagesNotSent       , messages not yet sent
+      * @param cancellableAckTimeout , timeout for the message sent for which we are awaiting an acknowledgement (if there is one)
+      * @param seqNumber             , sequence number for the next message to be sent
       */
     def handshaked(messageCodec: MessageCodec,
                    messagesNotSent: Queue[MessageSerializable] = Queue.empty,
                    cancellableAckTimeout: Option[CancellableAckTimeout] = None,
-                   seqNumber: Int = 0,
-                   p2pVersion: Option[Long] = None): Receive =
+                   seqNumber: Int = 0): Receive =
       handleWriteFailed orElse handleConnectionClosed orElse {
         case sm: SendMessage =>
-          if(cancellableAckTimeout.isEmpty)
-            sendMessage(messageCodec, sm.serializable, seqNumber, messagesNotSent, p2pVersion)
+          if (cancellableAckTimeout.isEmpty)
+            sendMessage(messageCodec, sm.serializable, seqNumber, messagesNotSent)
           else
-            context become handshaked(messageCodec, messagesNotSent :+ sm.serializable, cancellableAckTimeout, seqNumber, p2pVersion)
+            context become handshaked(messageCodec, messagesNotSent :+ sm.serializable, cancellableAckTimeout, seqNumber)
 
         case Received(data) =>
-          val messages = messageCodec.readMessages(data, p2pVersion)
+          val messages = messageCodec.readMessages(data)
           messages foreach processMessage
 
         case Ack if cancellableAckTimeout.nonEmpty =>
           //Cancel pending message timeout
           cancellableAckTimeout.foreach(_.cancellable.cancel())
 
           //Send next message if there is one
-          if(messagesNotSent.nonEmpty)
-            sendMessage(messageCodec, messagesNotSent.head, seqNumber, messagesNotSent.tail, p2pVersion)
+          if (messagesNotSent.nonEmpty)
+            sendMessage(messageCodec, messagesNotSent.head, seqNumber, messagesNotSent.tail)
           else
-            context become handshaked(messageCodec, Queue.empty, None, seqNumber, p2pVersion)
+            context become handshaked(messageCodec, Queue.empty, None, seqNumber)
 
         case AckTimeout(ackSeqNumber) if cancellableAckTimeout.exists(_.seqNumber == ackSeqNumber) =>
           cancellableAckTimeout.foreach(_.cancellable.cancel())
           log.debug(s"[Stopping Connection] Write to $peerId failed")
           context stop self
-
-        case PeerP2pVersion(p2pVer) =>
-          // We have peer p2p version based on hello message, if version is >= 5 next messages will be compressed.
-          context.become(handshaked(messageCodec, messagesNotSent, cancellableAckTimeout, seqNumber, Some(p2pVer)))
       }
 
     /**
       * Sends an encoded message through the TCP connection, an Ack will be received when the message was
       * successfully queued for delivery. A cancellable timeout is created for the Ack message.
       *
-      * @param messageCodec, for encoding the messages sent
-      * @param messageToSend, message to be sent
-      * @param seqNumber, sequence number for the message to be sent
-      * @param remainingMsgsToSend, messages not yet sent
+      * @param messageCodec        , for encoding the messages sent
+      * @param messageToSend       , message to be sent
+      * @param seqNumber           , sequence number for the message to be sent
+      * @param remainingMsgsToSend , messages not yet sent
       */
     private def sendMessage(messageCodec: MessageCodec, messageToSend: MessageSerializable,
-                            seqNumber: Int, remainingMsgsToSend: Queue[MessageSerializable], p2pVersion: Option[Long]): Unit = {
-      val out = messageCodec.encodeMessage(messageToSend, p2pVersion)
+                            seqNumber: Int, remainingMsgsToSend: Queue[MessageSerializable]): Unit = {
+      val out = messageCodec.encodeMessage(messageToSend)
       connection ! Write(out, Ack)
       log.debug(s"Sent message: $messageToSend from $peerId")
 
@@ -231,15 +225,14 @@ class RLPxConnectionHandler(
         messageCodec = messageCodec,
         messagesNotSent = remainingMsgsToSend,
         cancellableAckTimeout = Some(CancellableAckTimeout(seqNumber, timeout)),
-        seqNumber = increaseSeqNumber(seqNumber),
-        p2pVersion = p2pVersion
+        seqNumber = increaseSeqNumber(seqNumber)
       )
     }
 
     /**
       * Given a sequence number for the AckTimeouts, the next seq number is returned
       *
-      * @param seqNumber, the current sequence number
+      * @param seqNumber , the current sequence number
       * @return the sequence number for the next message sent
       */
     private def increaseSeqNumber(seqNumber: Int): Int = seqNumber match {
@@ -258,13 +251,14 @@ class RLPxConnectionHandler(
         if (msg.isPeerClosed) {
           log.debug(s"[Stopping Connection] Connection with $peerId closed by peer")
         }
-        if(msg.isErrorClosed){
+        if (msg.isErrorClosed) {
           log.debug(s"[Stopping Connection] Connection with $peerId closed because of error ${msg.getErrorCause}")
         }
 
         context stop self
     }
   }
+
 }
 
 object RLPxConnectionHandler {