diff --git a/libp2p/protocols/mix/seqno_generator.nim b/libp2p/protocols/mix/seqno_generator.nim new file mode 100644 index 000000000..c035a593b --- /dev/null +++ b/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 diff --git a/libp2p/protocols/mix/tag_manager.nim b/libp2p/protocols/mix/tag_manager.nim new file mode 100644 index 000000000..eaa6cc1e2 --- /dev/null +++ b/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() diff --git a/tests/mix/testseqnogenerator.nim b/tests/mix/testseqnogenerator.nim new file mode 100644 index 000000000..89c803011 --- /dev/null +++ b/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 diff --git a/tests/mix/testtagmanager.nim b/tests/mix/testtagmanager.nim new file mode 100644 index 000000000..2ff132601 --- /dev/null +++ b/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) diff --git a/tests/testnative.nim b/tests/testnative.nim index 90888375b..fa9060555 100644 --- a/tests/testnative.nim +++ b/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]