1  {.used.}
  2  
  3  import
  4    chronos,
  5    std/[sequtils, strformat, tables],
  6    testutils/unittests,
  7    ../dnsdisc/builder,
  8    ./test_utils
  9  
 10  procSuite "Test DNS Discovery: Merkle Tree builder":
 11    setup:
 12      let
 13        exampleRecords = initExampleRecords()
 14        exampleRoot = parseRootEntry(RootTxt).get()
 15        exampleLink = parseSubtreeEntry(LinkTxt).get()
 16        exampleBranch = parseSubtreeEntry(BranchTxt).get()
 17        exampleEnr1 = parseSubtreeEntry(Enr1Txt).get()
 18        exampleEnr2 = parseSubtreeEntry(Enr2Txt).get()
 19        exampleEnr3 = parseSubtreeEntry(Enr3Txt).get()
 20        exampleTree = Tree(rootEntry: exampleRoot,
 21                            entries: @[exampleLink,
 22                                       exampleBranch,
 23                                       exampleEnr1,
 24                                       exampleEnr2,
 25                                       exampleEnr3])
 26  
 27    # Test tree entries:
 28  
 29    asyncTest "Create TXT records":
 30      check:
 31        # Can convert (back) to TXT record
 32        exampleRoot.toTXTRecord().get() == RootTxt
 33        exampleEnr1.toTXTRecord().get() == Enr1Txt
 34        exampleLink.toTXTRecord().get() == LinkTxt
 35        exampleBranch.toTXTRecord().get() == BranchTxt
 36  
 37    asyncTest "Determine subdomain":
 38      check:
 39        # Successfully compute subdomain hash for each entry
 40        exampleEnr1.subdomain().get() == Enr1Subdomain
 41        exampleLink.subdomain().get() == LinkSubdomain
 42        exampleBranch.subdomain().get() == BranchSubdomain
 43  
 44    asyncTest "Build TXT records":
 45      check:
 46        # Successfully build all TXT records from example tree
 47        exampleTree.buildTXT("nodes.example.org").get() == exampleRecords
 48  
 49    asyncTest "Build subtree entries":
 50      # Test subtree from single leaf entry
 51      let subtreeLeaf = buildSubtree(@[parseSubtreeEntry(Enr1Txt).get()])
 52      check:
 53        subtreeLeaf.isOk()
 54        subtreeLeaf[].subtreeRoot.kind == Enr
 55        subtreeLeaf[].subtreeEntries.len == 0
 56  
 57      # Test subtree with single branch and leafs
 58      let subtreeSingle = buildSubtree(@[parseSubtreeEntry(Enr1Txt).get(),
 59                                         parseSubtreeEntry(Enr2Txt).get(),
 60                                         parseSubtreeEntry(Enr3Txt).get()])
 61      check:
 62        # Successfully build ENR subtree
 63        subtreeSingle.isOk()
 64        subtreeSingle[].subtreeRoot.kind == Branch
 65        subtreeSingle[].subtreeRoot.branchEntry.children == @[Enr1Subdomain, Enr2Subdomain, Enr3Subdomain]
 66        subtreeSingle[].subtreeEntries.len == 3
 67  
 68      # Test subtree with multiple branches-of-branches
 69      var entries: seq[SubtreeEntry]
 70      for i in 1..(MaxChildren*MaxChildren):
 71        # We need at least MaxChildren branch entries holding MaxChildren leaf nodes
 72        entries.add(parseSubtreeEntry(Enr1Txt).get())
 73  
 74      let
 75        subtreeMultiple = buildSubtree(entries)
 76        expectedLeafCount = MaxChildren*MaxChildren
 77        expectedBranchCount = MaxChildren # Excluding subtree root branch
 78      check:
 79        subtreeMultiple.isOk()
 80        subtreeMultiple[].subtreeRoot.kind == Branch
 81        subtreeMultiple[].subtreeEntries.len == expectedLeafCount + expectedBranchCount
 82        subtreeMultiple[].subtreeEntries.filterIt(it.kind == Enr).len == expectedLeafCount
 83        subtreeMultiple[].subtreeEntries.filterIt(it.kind == Branch).len == expectedBranchCount
 84  
 85    asyncTest "Build complete tree":
 86      var
 87        enr1: Record
 88        enr2: Record
 89        enr3: Record
 90        link = parseLinkEntry("enrtree://AM5FCQLWIZX2QFPNJAP7VUERCCRNGRHWZG3YYHIUV7BVDQ5FDPRT2@morenodes.example.org").get()
 91        seqNo = 1.uint32
 92  
 93      check:
 94        enr1.fromBase64("-HW4QOFzoVLaFJnNhbgMoDXPnOvcdVuj7pDpqRvh6BRDO68aVi5ZcjB3vzQRZH2IcLBGHzo8uUN3snqmgTiE56CH3AMBgmlkgnY0iXNlY3AyNTZrMaECC2_24YYkYHEgdzxlSNKQEnHhuNAbNlMlWJxrJxbAFvA")
 95        enr2.fromBase64("-HW4QAggRauloj2SDLtIHN1XBkvhFZ1vtf1raYQp9TBW2RD5EEawDzbtSmlXUfnaHcvwOizhVYLtr7e6vw7NAf6mTuoCgmlkgnY0iXNlY3AyNTZrMaECjrXI8TLNXU0f8cthpAMxEshUyQlK-AM0PW2wfrnacNI")
 96        enr3.fromBase64("-HW4QLAYqmrwllBEnzWWs7I5Ev2IAs7x_dZlbYdRdMUx5EyKHDXp7AV5CkuPGUPdvbv1_Ms1CPfhcGCvSElSosZmyoqAgmlkgnY0iXNlY3AyNTZrMaECriawHKWdDRk2xeZkrOXBQ0dfMFLHY4eENZwdufn1S1o")
 97  
 98      let tree = buildTree(seqNo, @[enr1, enr2, enr3], @[link])
 99  
100      check:
101        # Successfully build example tree
102        tree.isOk()
103        # Root entry
104        tree[].rootEntry.seqNo == 1
105        tree[].rootEntry.eroot == exampleRoot.eroot
106        tree[].rootEntry.lroot == exampleRoot.lroot
107        # Subtree entries
108        tree[].entries.len == 5
109        tree[].entries.filterIt(it.kind == Branch).len == 1
110        tree[].getNodes().len == 3
111        tree[].getLinks().len == 1
112  
113    asyncTest "Sign tree":
114      let
115        secKeyHex = "58d23b55bc9cdce1f18c2500f40ff4ab7245df9a89505e9b1fa4851f623d241d"
116        expectedSig = "8BGq3_ZasQUPvDyU7cqBpsHVjn4CP5GsFf7Xf9M1bJARCXII3SrD7e_I8Q6qw9oLItHapevgFlfPfwhqPXFRrAA"
117        expectedTXT = fmt"{RootPrefix} e={exampleTree.rootEntry.eroot} l={exampleTree.rootEntry.lroot} seq={exampleTree.rootEntry.seqNo} sig={expectedSig}"
118        testSecKey = PrivateKey.fromHex(secKeyHex).get()
119        testPubKey = testSecKey.toPublicKey()
120  
121      var tree = exampleTree
122  
123      let
124        signRes = signTree(tree, testSecKey)
125      check:
126        signRes.isOk()
127        tree.rootEntry.signature.len == SkRawRecoverableSignatureSize # 65 bytes
128  
129      # Verify signature
130      let
131        sigHash = hashableContent(tree.rootEntry)
132        sig = SignatureNR.fromRaw(tree.rootEntry.signature)
133  
134      check:
135        sig.isOk()
136        verify(sig = sig[],
137               msg = sigHash,
138               key = testPubKey)
139  
140      # Check TXT record
141      check:
142        tree.rootEntry.toTXTRecord().get() == expectedTXT