diff --git a/merkle_tree.go b/merkle_tree.go index 7a4156b..8a0c505 100644 --- a/merkle_tree.go +++ b/merkle_tree.go @@ -38,6 +38,10 @@ func NewMerkleTree(data [][]byte) *MerkleTree { } nodes = newLevel + + if len(newLevel) == 1 { + break + } } mTree := MerkleTree{&nodes[0]} diff --git a/merkle_tree_test.go b/merkle_tree_test.go index acff5ff..e3001c8 100644 --- a/merkle_tree_test.go +++ b/merkle_tree_test.go @@ -49,7 +49,7 @@ func TestNewMerkleNode(t *testing.T) { ) } -func TestNewMerkleTree(t *testing.T) { +func TestNewMerkleTreeOddNodes(t *testing.T) { data := [][]byte{ []byte("node1"), []byte("node2"), @@ -73,3 +73,45 @@ func TestNewMerkleTree(t *testing.T) { assert.Equal(t, rootHash, fmt.Sprintf("%x", mTree.RootNode.Data), "Merkle tree root hash is correct") } + +func TestNewMerkleTreeEvenNodes(t *testing.T) { + + data := [][]byte{ + []byte("node1"), + []byte("node2"), + []byte("node3"), + []byte("node4"), + []byte("node5"), + []byte("node6"), + []byte("node7"), + []byte("node8"), + } + + // Level 1 + n1 := NewMerkleNode(nil, nil, data[0]) + n2 := NewMerkleNode(nil, nil, data[1]) + n3 := NewMerkleNode(nil, nil, data[2]) + n4 := NewMerkleNode(nil, nil, data[3]) + n5 := NewMerkleNode(nil, nil, data[4]) + n6 := NewMerkleNode(nil, nil, data[5]) + n7 := NewMerkleNode(nil, nil, data[6]) + n8 := NewMerkleNode(nil, nil, data[7]) + + // Level 2 + n9 := NewMerkleNode(n1, n2, nil) + n10 := NewMerkleNode(n3, n4, nil) + n11 := NewMerkleNode(n5, n6, nil) + n12 := NewMerkleNode(n7, n8, nil) + + // Level 3 + n13 := NewMerkleNode(n9, n10, nil) + n14 := NewMerkleNode(n11, n12, nil) + + // Level 4 + n15 := NewMerkleNode(n13, n14, nil) + + rootHash := fmt.Sprintf("%x", n15.Data) + mTree := NewMerkleTree(data) + + assert.Equal(t, rootHash, fmt.Sprintf("%x", mTree.RootNode.Data), "Merkle tree root hash is correct") +} \ No newline at end of file