github.com/nebulouslabs/sia@v1.3.7/crypto/merkle_test.go (about)

     1  package crypto
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/NebulousLabs/fastrand"
     7  )
     8  
     9  // TestTreeBuilder builds a tree and gets the merkle root.
    10  func TestTreeBuilder(t *testing.T) {
    11  	tree := NewTree()
    12  	tree.PushObject("a")
    13  	tree.PushObject("b")
    14  	_ = tree.Root()
    15  
    16  	// Correctness is assumed, as it's tested by the merkletree package. This
    17  	// function is really for code coverage.
    18  }
    19  
    20  // TestCalculateLeaves probes the CalculateLeaves function.
    21  func TestCalculateLeaves(t *testing.T) {
    22  	tests := []struct {
    23  		size, expSegs uint64
    24  	}{
    25  		{0, 1},
    26  		{63, 1},
    27  		{64, 1},
    28  		{65, 2},
    29  		{127, 2},
    30  		{128, 2},
    31  		{129, 3},
    32  	}
    33  
    34  	for i, test := range tests {
    35  		if segs := CalculateLeaves(test.size); segs != test.expSegs {
    36  			t.Errorf("miscalculation for test %v: expected %v, got %v", i, test.expSegs, segs)
    37  		}
    38  	}
    39  }
    40  
    41  // TestStorageProof builds a storage proof and checks that it verifies
    42  // correctly.
    43  func TestStorageProof(t *testing.T) {
    44  	// Generate proof data.
    45  	numSegments := uint64(7)
    46  	data := fastrand.Bytes(int(numSegments * SegmentSize))
    47  	rootHash := MerkleRoot(data)
    48  
    49  	// Create and verify proofs for all indices.
    50  	for i := uint64(0); i < numSegments; i++ {
    51  		baseSegment, hashSet := MerkleProof(data, i)
    52  		if !VerifySegment(baseSegment, hashSet, numSegments, i, rootHash) {
    53  			t.Error("Proof", i, "did not pass verification")
    54  		}
    55  	}
    56  
    57  	// Try an incorrect proof.
    58  	baseSegment, hashSet := MerkleProof(data, 3)
    59  	if VerifySegment(baseSegment, hashSet, numSegments, 4, rootHash) {
    60  		t.Error("Verified a bad proof")
    61  	}
    62  }
    63  
    64  // TestNonMultipleNumberOfSegmentsStorageProof builds a storage proof that has
    65  // a last leaf of size less than SegmentSize.
    66  func TestNonMultipleLeafSizeStorageProof(t *testing.T) {
    67  	// Generate proof data.
    68  	data := fastrand.Bytes((2 * SegmentSize) + 10)
    69  	rootHash := MerkleRoot(data)
    70  
    71  	// Create and verify a proof for the last index.
    72  	baseSegment, hashSet := MerkleProof(data, 2)
    73  	if !VerifySegment(baseSegment, hashSet, 3, 2, rootHash) {
    74  		t.Error("padded segment proof failed")
    75  	}
    76  }
    77  
    78  // TestCachedTree tests the cached tree functions of the package.
    79  func TestCachedTree(t *testing.T) {
    80  	if testing.Short() {
    81  		t.SkipNow()
    82  	}
    83  
    84  	// Build a cached tree out of 4 subtrees, each subtree of height 2 (4
    85  	// elements).
    86  	tree1Bytes := fastrand.Bytes(SegmentSize * 4)
    87  	tree2Bytes := fastrand.Bytes(SegmentSize * 4)
    88  	tree3Bytes := fastrand.Bytes(SegmentSize * 4)
    89  	tree4Bytes := fastrand.Bytes(SegmentSize * 4)
    90  	tree1Root := MerkleRoot(tree1Bytes)
    91  	tree2Root := MerkleRoot(tree2Bytes)
    92  	tree3Root := MerkleRoot(tree3Bytes)
    93  	tree4Root := MerkleRoot(tree4Bytes)
    94  	fullRoot := MerkleRoot(append(tree1Bytes, append(tree2Bytes, append(tree3Bytes, tree4Bytes...)...)...))
    95  
    96  	// Get a cached proof for index 0.
    97  	base, cachedHashSet := MerkleProof(tree1Bytes, 0)
    98  	if !VerifySegment(base, cachedHashSet, 4, 0, tree1Root) {
    99  		t.Fatal("the proof for the subtree was invalid")
   100  	}
   101  	ct := NewCachedTree(2)
   102  	ct.SetIndex(0)
   103  	ct.Push(tree1Root)
   104  	ct.Push(tree2Root)
   105  	ct.Push(tree3Root)
   106  	ct.Push(tree4Root)
   107  	hashSet := ct.Prove(base, cachedHashSet)
   108  	if !VerifySegment(base, hashSet, 4*4, 0, fullRoot) {
   109  		t.Fatal("cached proof construction appears unsuccessful")
   110  	}
   111  	if ct.Root() != fullRoot {
   112  		t.Fatal("cached Merkle root is not matching the full Merkle root")
   113  	}
   114  
   115  	// Get a cached proof for index 6.
   116  	base, cachedHashSet = MerkleProof(tree2Bytes, 2)
   117  	if !VerifySegment(base, cachedHashSet, 4, 2, tree2Root) {
   118  		t.Fatal("the proof for the subtree was invalid")
   119  	}
   120  	ct = NewCachedTree(2)
   121  	ct.SetIndex(6)
   122  	ct.Push(tree1Root)
   123  	ct.Push(tree2Root)
   124  	ct.Push(tree3Root)
   125  	ct.Push(tree4Root)
   126  	hashSet = ct.Prove(base, cachedHashSet)
   127  	if !VerifySegment(base, hashSet, 4*4, 6, fullRoot) {
   128  		t.Fatal("cached proof construction appears unsuccessful")
   129  	}
   130  	if ct.Root() != fullRoot {
   131  		t.Fatal("cached Merkle root is not matching the full Merkle root")
   132  	}
   133  }
   134  
   135  // TestMerkleTreeOddDataSize checks that MerkleRoot and MerkleProof still
   136  // function correctly if you provide data which does not have a size evenly
   137  // divisible by SegmentSize.
   138  func TestOddDataSize(t *testing.T) {
   139  	if testing.Short() {
   140  		t.SkipNow()
   141  	}
   142  
   143  	// Create some random data that's not evenly padded.
   144  	for i := 0; i < 25; i++ {
   145  		randFullSegments := fastrand.Intn(65)
   146  		randOverflow := fastrand.Intn(63) + 1
   147  		randProofIndex := fastrand.Intn(randFullSegments + 1)
   148  		data := fastrand.Bytes(SegmentSize*randFullSegments + randOverflow)
   149  		root := MerkleRoot(data)
   150  		base, hashSet := MerkleProof(data, uint64(randProofIndex))
   151  		if !VerifySegment(base, hashSet, uint64(randFullSegments)+1, uint64(randProofIndex), root) {
   152  			t.Error("Padded data proof failed for", randFullSegments, randOverflow, randProofIndex)
   153  		}
   154  	}
   155  }