Update timestamp and ackNr on packet resend (#444)

* Update timestamp and ackNr on packet resend

eth/utp/packets.nim

@@ -9,7 +9,7 @@
 import
   std/[monotimes],
   faststreams,
-  stew/[endians2, results, objects], bearssl,
+  stew/[endians2, results, objects, arrayops], bearssl,
   ../p2p/discoveryv5/random2
 
 export results
@@ -56,9 +56,11 @@ type
 # https://github.com/bittorrent/libutp/blob/master/utp_utils.cpp, to check all the
 # timing assumptions on different platforms
 proc getMonoTimeTimeStamp*(): uint32 = 
+  # TODO
   # this value is equivalent of:
   # uint32((Moment.now() - Moment.init(0, Microseconds)).microseconds())
   # on macOs
+  # so we we can use moment here and return (Moment, uint32) tuple from this function
   let time = getMonoTime()
   cast[uint32](time.ticks() div 1000)
 
@@ -136,6 +138,14 @@ proc decodePacket*(bytes: openArray[byte]): Result[Packet, string] =
 
     ok(Packet(header: header, payload: payload))
 
+proc modifyTimeStampAndAckNr*(packetBytes: var seq[byte], newTimestamp: uint32, newAckNr: uint16) =
+  ## Modifies timestamp and ack nr of already encoded packets. Those fields should be
+  ## filled right before sending, so when re-sending the packet we would like to update
+  ## it without decoding and re-encoding the packet once again
+  doAssert(len(packetBytes) >= minimalHeaderSize)
+  packetBytes[4..7] = toBytesBE(newTimestamp)
+  packetBytes[18..19] = toBytesBE(newAckNr)
+
 # connectionId - should be random not already used number
 # bufferSize - should be pre configured initial buffer size for socket
 # SYN packets are special, and should have the receive ID in the connid field,

eth/utp/utp_socket.nim

@@ -357,9 +357,16 @@ proc sendAck(socket: UtpSocket): Future[void] =
 
 # Should be called before sending packet
 proc setSend(s: UtpSocket, p: var OutgoingPacket): seq[byte] =
+  let currentMoment = Moment.now()
+  let currentTimeStamp = getMonoTimeTimeStamp()
+
   inc p.transmissions
   p.needResend = false
-  p.timeSent = Moment.now()
+  p.timeSent = currentMoment
+  # all bytearrays in outgoing buffer should be properly encoded utp packets
+  # so it is safe to directly modify fields
+  modifyTimeStampAndAckNr(p.packetBytes, currentTimeStamp, s.ackNr)
+
   return p.packetBytes
 
 proc flushPackets(socket: UtpSocket) {.async.} =

tests/utp/all_utp_tests.nim

@@ -12,4 +12,5 @@ import
   ./test_discv5_protocol,
   ./test_buffer,
   ./test_utp_socket,
-  ./test_utp_router
+  ./test_utp_router,
+  ./test_clock_drift_calculator

tests/utp/test_packets.nim

@@ -46,3 +46,19 @@ suite "Utp packets encoding/decoding":
       packet.header.wndSize == 1048576
       packet.header.seqNr == 16807
       packet.header.ackNr == 1
+
+  test "Modify timestamp of encoded packet":
+    let synPacket = synPacket(5, 10, 20)
+    let initialTimestamp = synPacket.header.timestamp
+    let initialAckNr = synPacket.header.ackNr
+    let modifiedTimeStamp = initialTimestamp + 120324
+    let modifiedAckNr = initialAckNr + 20
+    var encoded = encodePacket(synPacket)
+    modifyTimeStampAndAckNr(encoded, modifiedTimeStamp, modifiedAckNr)
+
+    let decoded = decodePacket(encoded)
+
+    check:
+      decoded.isOk()
+      decoded.get().header.timestamp == modifiedTimeStamp
+      decoded.get().header.ackNr == modifiedAckNr

tests/utp/test_utp_socket.nim

@@ -1214,3 +1214,51 @@ procSuite "Utp socket unit test":
       socketAck.header.timestampDiff > 0
 
     await outgoingSocket.destroyWait()
+
+  asyncTest "Re-sent packet should have updated timestamps and ack numbers":
+    let q = newAsyncQueue[Packet]()   
+    let initalRemoteSeqNr = 10'u16
+
+    let (outgoingSocket, initialPacket) = connectOutGoingSocket(initalRemoteSeqNr, q)
+
+    let writeResult = await outgoingSocket.write(@[1'u8])
+
+    check:
+      writeResult.isOk()
+
+    let firstSend = await q.get()
+
+    let secondSend = await q.get()
+
+    check:
+      # there was sometime between resend but no packet from remote
+      # so timestamp should be updated but not ackNr
+      secondSend.header.timestamp > firstSend.header.timestamp
+      firstSend.header.ackNr == secondSend.header.ackNr
+
+    let dataP1 = 
+      dataPacket(
+        initalRemoteSeqNr,
+        initialPacket.header.connectionId,
+        initialPacket.header.seqNr,
+        testBufferSize,
+        @[1'u8],
+        0
+      )
+    
+    await outgoingSocket.processPacket(dataP1)
+
+    let ack = await q.get()
+
+    check:
+      ack.header.pType == ST_STATE
+
+    let thirdSend = await q.get()
+
+    check:
+      # as there was some incoming data between resend, both timestamp and ackNr
+      # should be updated
+      thirdSend.header.timestamp > secondSend.header.timestamp
+      thirdSend.header.ackNr > secondSend.header.ackNr
+
+    await outgoingSocket.destroyWait()