Utp improvements (#489)

* Move connection finalization to separate function

* Do not process data unless in correct state

eth/utp/utp_socket.nim

@@ -1039,6 +1039,25 @@ proc sendAck(socket: UtpSocket) =
 
   socket.sendData(encodePacket(ackPacket))
 
+
+proc tryfinalizeConnection(socket: UtpSocket, p: Packet) =
+  # To avoid amplification attacks, server socket is in SynRecv state until
+  # it receices first data transfer
+  # https://www.usenix.org/system/files/conference/woot15/woot15-paper-adamsky.pdf
+  # Socket is in SynRecv state only when recv timeout is configured
+  if (socket.state == SynRecv and p.header.pType == ST_DATA):
+    socket.state = Connected
+
+  if (socket.state == SynSent and p.header.pType == ST_STATE):
+    socket.state = Connected
+    socket.ackNr = p.header.seqNr - 1
+
+    debug "Received Syn-Ack finalizing connection",
+      socketAckNr = socker.ackNr
+
+    if (not socket.connectionFuture.finished()):
+      socket.connectionFuture.complete()
+
 # TODO at socket level we should handle only FIN/DATA/ACK packets. Refactor to make
 # it enforcable by type system
 proc processPacketInternal(socket: UtpSocket, p: Packet) =
@@ -1227,6 +1246,8 @@ proc processPacketInternal(socket: UtpSocket, p: Packet) =
       resetZeroWindowTime = socket.zeroWindowTimer,
       currentPacketSize = currentPacketSize
 
+  socket.tryfinalizeConnection(p)
+
   # socket.curWindowPackets == acks means that this packet acked all remaining packets
   # including the sent fin packets
   if (socket.finSent and socket.curWindowPackets == acks):
@@ -1280,14 +1301,15 @@ proc processPacketInternal(socket: UtpSocket, p: Packet) =
   if (p.eack.isSome()):
     socket.selectiveAckPackets(pkAckNr, p.eack.unsafeGet(), timestampInfo.moment)
 
-  case p.header.pType
+  if p.header.pType == ST_DATA or p.header.pType == ST_FIN:
-    of ST_DATA, ST_FIN:
+    if socket.state != Connected:
-      # To avoid amplification attacks, server socket is in SynRecv state until
+      debug "Unexpected packet",
-      # it receices first data transfer
+        socketState = socket.state,
-      # https://www.usenix.org/system/files/conference/woot15/woot15-paper-adamsky.pdf
+        packetType = p.header.pType
-      # Socket is in SynRecv state only when recv timeout is configured
+
-      if (socket.state == SynRecv and p.header.pType == ST_DATA):
+      # we have received user generated packet (DATA or FIN), in not connected
-        socket.state = Connected
+      # state. Stop processing it.
+      return
 
     if (p.header.pType == ST_FIN and (not socket.gotFin)):
       debug "Received FIN packet",
@@ -1436,22 +1458,6 @@ proc processPacketInternal(socket: UtpSocket, p: Packet) =
           # generated
           socket.sendAck()
 
-    of ST_STATE:
-      if (socket.state == SynSent and (not socket.connectionFuture.finished())):
-        socket.state = Connected
-        # TODO reference implementation sets ackNr (p.header.seqNr - 1), although
-        # spec mention that it should be equal p.header.seqNr. For now follow the
-        # reference impl to be compatible with it. Later investigate trin compatibility.
-        socket.ackNr = p.header.seqNr - 1
-        # In case of SynSent complate the future as last thing to make sure user of libray will
-        # receive socket in correct state
-        socket.connectionFuture.complete()
-
-    of ST_RESET:
-      debug "Received ST_RESET on known socket, ignoring"
-    of ST_SYN:
-      debug "Received ST_SYN on known socket, ignoring"
-
 proc processPacket*(socket: UtpSocket, p: Packet): Future[void] = 
   socket.eventQueue.put(SocketEvent(kind: NewPacket, packet: p))
 

tests/utp/test_utp_socket.nim

@@ -1407,3 +1407,61 @@ procSuite "Utp socket unit test":
       resent3.header.seqNr == sent3.header.seqNr
 
     await outgoingSocket.destroyWait()
+
+  asyncTest "Socket should accept data only in connected state":
+    let q = newAsyncQueue[Packet]()
+    let initialRemoteSeq = 10'u16
+    let cfg = SocketConfig.init()
+    let remoteReciveBuffer = 1024'u32
+
+    let dataDropped = @[1'u8]
+    let dataRecived = @[2'u8]
+
+    let sock1 = newOutgoingSocket[TransportAddress](testAddress, initTestSnd(q), cfg, defaultRcvOutgoingId, rng[])
+
+    asyncSpawn sock1.startOutgoingSocket()
+
+    let initialPacket = await q.get()
+
+    check:
+      initialPacket.header.pType == ST_SYN
+
+    let dpDropped = dataPacket(
+      initialRemoteSeq,
+      initialPacket.header.connectionId,
+      initialPacket.header.seqNr,
+      testBufferSize,
+      dataDropped,
+      0
+    )
+
+    let dpReceived = dataPacket(
+      initialRemoteSeq,
+      initialPacket.header.connectionId,
+      initialPacket.header.seqNr,
+      testBufferSize,
+      dataRecived,
+      0
+    )
+
+    let responseAck =
+      ackPacket(
+        initialRemoteSeq,
+        initialPacket.header.connectionId,
+        initialPacket.header.seqNr,
+        remoteReciveBuffer,
+        0
+      )
+
+    # even though @[1'u8] is received first, it should be dropped as socket is not
+    # yet in connected state
+    await sock1.processPacket(dpDropped)
+    await sock1.processPacket(responseAck)
+    await sock1.processPacket(dpReceived)
+
+    let receivedData = await sock1.read(1)
+
+    check:
+      receivedData == dataRecived
+
+    await sock1.destroyWait()