feat(mix): sequence number generator and tag manager (#1688)

libp2p/protocols/mix/seqno_generator.nim

@@ -0,0 +1,29 @@
+import std/endians, times
+import ../../peerid
+import ./crypto
+import ../../utils/sequninit
+
+type SeqNo* = uint32
+
+proc init*(T: typedesc[SeqNo], data: seq[byte]): T =
+  var seqNo: SeqNo = 0
+  let hash = sha256_hash(data)
+  for i in 0 .. 3:
+    seqNo = seqNo or (uint32(hash[i]) shl (8 * (3 - i)))
+  return seqNo
+
+proc init*(T: typedesc[SeqNo], peerId: PeerId): T =
+  SeqNo.init(peerId.data)
+
+proc generate*(seqNo: var SeqNo, messageBytes: seq[byte]) =
+  let
+    currentTime = getTime().toUnix() * 1000
+    currentTimeBytes = newSeqUninit[byte](8)
+  bigEndian64(unsafeAddr currentTimeBytes[0], unsafeAddr currentTime)
+
+  let s = SeqNo.init(messageBytes & currentTimeBytes)
+  seqNo = (seqNo + s) mod high(uint32)
+
+proc inc*(seqNo: var SeqNo) =
+  seqNo = (seqNo + 1) mod high(uint32)
+  # TODO: Manage sequence no. overflow in a way that it does not affect re-assembly

libp2p/protocols/mix/tag_manager.nim

@@ -0,0 +1,28 @@
+import tables, locks
+import ./curve25519
+
+type TagManager* = ref object
+  lock: Lock
+  seenTags: Table[FieldElement, bool]
+
+proc new*(T: typedesc[TagManager]): T =
+  let tm = T()
+  tm.seenTags = initTable[FieldElement, bool]()
+  initLock(tm.lock)
+  return tm
+
+proc addTag*(tm: TagManager, tag: FieldElement) {.gcsafe.} =
+  withLock tm.lock:
+    tm.seenTags[tag] = true
+
+proc isTagSeen*(tm: TagManager, tag: FieldElement): bool {.gcsafe.} =
+  withLock tm.lock:
+    return tm.seenTags.contains(tag)
+
+proc removeTag*(tm: TagManager, tag: FieldElement) {.gcsafe.} =
+  withLock tm.lock:
+    tm.seenTags.del(tag)
+
+proc clearTags*(tm: TagManager) {.gcsafe.} =
+  withLock tm.lock:
+    tm.seenTags.clear()

tests/mix/testseqnogenerator.nim

@@ -0,0 +1,94 @@
+{.used.}
+
+import chronicles, sets, unittest
+import std/[os, times]
+import ../../libp2p/peerid
+include ../../libp2p/protocols/mix/seqno_generator
+
+const second = 1000
+
+suite "Sequence Number Generator":
+  test "init_seq_no_from_peer_id":
+    let
+      peerId =
+        PeerId.init("16Uiu2HAmFkwLVsVh6gGPmSm9R3X4scJ5thVdKfWYeJsKeVrbcgVC").get()
+      seqNo = SeqNo.init(peerId)
+
+    check seqNo != 0
+
+  test "generate_seq_nos_for_different_messages":
+    let
+      peerId =
+        PeerId.init("16Uiu2HAmFkwLVsVh6gGPmSm9R3X4scJ5thVdKfWYeJsKeVrbcgVC").get()
+      msg1 = @[byte 1, 2, 3]
+      msg2 = @[byte 4, 5, 6]
+
+    var seqNo = SeqNo.init(peerId)
+
+    seqNo.generate(msg1)
+    let seqNo1 = seqNo
+
+    seqNo.generate(msg2)
+    let seqNo2 = seqNo
+
+    check seqNo1 != seqNo2
+
+  test "generate_seq_nos_for_same_message":
+    let
+      peerId =
+        PeerId.init("16Uiu2HAmFkwLVsVh6gGPmSm9R3X4scJ5thVdKfWYeJsKeVrbcgVC").get()
+      msg = @[byte 1, 2, 3]
+    var seqNo = SeqNo.init(peerId)
+
+    seqNo.generate(msg)
+    let seqNo1 = seqNo
+
+    sleep(second)
+
+    seqNo.generate(msg)
+    let seqNo2 = seqNo
+
+    check seqNo1 != seqNo2
+
+  test "generate_seq_nos_for_different_peer_ids":
+    let
+      peerId1 =
+        PeerId.init("16Uiu2HAmFkwLVsVh6gGPmSm9R3X4scJ5thVdKfWYeJsKeVrbcgVC").get()
+      peerId2 =
+        PeerId.init("16Uiu2HAm6WNzw8AssyPscYYi8x1bY5wXyQrGTShRH75bh5dPCjBQ").get()
+
+    var
+      seqNo1 = SeqNo.init(peerId1)
+      seqNo2 = SeqNo.init(peerId2)
+
+    check seqNo1 != seqNo2
+
+  test "increment_seq_no":
+    let peerId =
+      PeerId.init("16Uiu2HAmFkwLVsVh6gGPmSm9R3X4scJ5thVdKfWYeJsKeVrbcgVC").get()
+    var seqNo: SeqNo = SeqNo.init(peerId)
+    let initialCounter = seqNo
+
+    seqNo.inc()
+
+    check seqNo == initialCounter + 1
+
+  test "seq_no_wraps_around_at_max_value":
+    var seqNo = high(uint32) - 1
+
+    seqNo.inc()
+
+    check seqNo == 0
+
+  test "generate_seq_no_uses_entire_uint32_range":
+    let peerId =
+      PeerId.init("16Uiu2HAmFkwLVsVh6gGPmSm9R3X4scJ5thVdKfWYeJsKeVrbcgVC").get()
+    var
+      seqNo = SeqNo.init(peerId)
+      seenValues = initHashSet[uint32]()
+
+    for i in 0 ..< 10000:
+      seqNo.generate(@[i.uint8])
+      seenValues.incl(seqNo)
+
+    check seenValues.len > 9000

tests/mix/testtagmanager.nim

@@ -0,0 +1,61 @@
+{.used.}
+
+import chronicles, results, unittest
+import ../../libp2p/protocols/mix/[curve25519, tag_manager]
+
+suite "tag_manager_tests":
+  var tm: TagManager
+
+  setup:
+    tm = TagManager.new()
+
+  teardown:
+    tm.clearTags()
+
+  test "add_and_check_tag":
+    let
+      tag = generateRandomFieldElement().expect("should generate FE")
+      nonexistentTag = generateRandomFieldElement().expect("should generate FE")
+
+    tm.addTag(tag)
+
+    check:
+      tm.isTagSeen(tag)
+      not tm.isTagSeen(nonexistentTag)
+
+  test "remove_tag":
+    let tag = generateRandomFieldElement().expect("should generate FE")
+
+    tm.addTag(tag)
+    check tm.isTagSeen(tag)
+
+    tm.removeTag(tag)
+    check not tm.isTagSeen(tag)
+
+  test "check_tag_presence":
+    let tag = generateRandomFieldElement().expect("should generate FE")
+    check not tm.isTagSeen(tag)
+
+    tm.addTag(tag)
+    check tm.isTagSeen(tag)
+
+    tm.removeTag(tag)
+    check not tm.isTagSeen(tag)
+
+  test "handle_multiple_tags":
+    let tag1 = generateRandomFieldElement().expect("should generate FE")
+    let tag2 = generateRandomFieldElement().expect("should generate FE")
+
+    tm.addTag(tag1)
+    tm.addTag(tag2)
+
+    check:
+      tm.isTagSeen(tag1)
+      tm.isTagSeen(tag2)
+
+    tm.removeTag(tag1)
+    tm.removeTag(tag2)
+
+    check:
+      not tm.isTagSeen(tag1)
+      not tm.isTagSeen(tag2)

tests/testnative.nim

@@ -42,4 +42,4 @@ import kademlia/[testencoding, testroutingtable, testfindnode, testputval]
 when defined(libp2p_autotls_support):
   import testautotls
 
-import mix/[testcrypto, testcurve25519]
+import mix/[testcrypto, testcurve25519, testtagmanager, testseqnogenerator]