[WIP] Init tree with an empty node as root node

merkletree/merkletree.go

@@ -29,15 +29,14 @@ const (
 // nonce.
 type MerkleTree struct {
 	nonce []byte
-	root  *interiorNode
+	root  merkleNode
 	hash  []byte
 }
 
 // NewMerkleTree returns an empty Merkle prefix tree
-// with a secure random nonce. The tree root is an interior node
-// and its children are two empty leaf nodes.
+// with a secure random nonce. The tree root is an empty node.
 func NewMerkleTree() (*MerkleTree, error) {
-	root := newInteriorNode(nil, 0, []bool{})
+	root := newEmptyNode(nil, 0, []bool{})
 	nonce, err := crypto.MakeRand()
 	if err != nil {
 		return nil, err
@@ -151,14 +150,22 @@ func (m *MerkleTree) insertNode(index []byte, toAdd *userLeafNode) {
 insertLoop:
 	for {
 		switch nodePointer.(type) {
+		case *emptyNode:
+			// assert that this is an empty tree.
+			// then we just replace the root node
+			// with a user leaf node and return.
+			if nodePointer.(*emptyNode).parent != nil {
+				panic(ErrInvalidTree)
+			}
+			toAdd.parent = nil
+			toAdd.level = 0
+			m.root = toAdd
+			return
 		case *userLeafNode:
 			// reached a "bottom" of the tree.
 			// add a new interior node and push the previous leaf down
 			// then continue insertion
 			currentNodeUL := nodePointer.(*userLeafNode)
-			if currentNodeUL.parent == nil {
-				panic(ErrInvalidTree)
-			}
 
 			if bytes.Equal(currentNodeUL.index, toAdd.index) {
 				// replace the value
@@ -178,6 +185,12 @@ insertLoop:
 			}
 			currentNodeUL.level = depth + 1
 			currentNodeUL.parent = newInteriorNode
+
+			if newInteriorNode.parent == nil {
+				m.root = newInteriorNode
+				nodePointer = newInteriorNode
+				break
+			}
 			if newInteriorNode.parent.(*interiorNode).leftChild == nodePointer {
 				newInteriorNode.parent.(*interiorNode).leftChild = newInteriorNode
 			} else {
@@ -249,7 +262,7 @@ func (m *MerkleTree) recomputeHash() {
 func (m *MerkleTree) Clone() *MerkleTree {
 	return &MerkleTree{
 		nonce: m.nonce,
-		root:  m.root.clone(nil).(*interiorNode),
+		root:  m.root.clone(nil),
 		hash:  append([]byte{}, m.hash...),
 	}
 }

merkletree/merkletree_test.go

@@ -5,8 +5,6 @@ import (
 	"testing"
 
 	"github.com/coniks-sys/coniks-go/crypto/vrf"
-	"github.com/coniks-sys/coniks-go/utils"
-	"golang.org/x/crypto/sha3"
 )
 
 var vrfPrivKey1, _ = vrf.GenerateKey(bytes.NewReader(
@@ -15,74 +13,137 @@ var vrfPrivKey1, _ = vrf.GenerateKey(bytes.NewReader(
 var vrfPrivKey2, _ = vrf.GenerateKey(bytes.NewReader(
 	[]byte("deterministic tests need 32 byte")))
 
-func TestOneEntry(t *testing.T) {
+func TestEmptyTree(t *testing.T) {
 	m, err := NewMerkleTree()
 	if err != nil {
 		t.Fatal(err)
 	}
+	m.recomputeHash()
 
-	var commit [32]byte
-	var expect [32]byte
+	newTree := m.Clone()
+	newTree.recomputeHash()
 
-	key := "key"
-	val := []byte("value")
-	index := vrfPrivKey1.Compute([]byte(key))
-	if err := m.Set(index, key, val); err != nil {
-		t.Fatal(err)
+	// assert tree clone.
+	if newTree.root.(*emptyNode).parent != nil {
+		t.Error("Expect a nil parent")
+	}
+	if !bytes.Equal(m.hash, newTree.hash) {
+		t.Error("Expect 2 identical trees, got different tree hashes")
 	}
-	m.recomputeHash()
 
-	// Check empty node hash
-	h := sha3.NewShake128()
-	h.Write([]byte{EmptyBranchIdentifier})
-	h.Write(m.nonce)
-	h.Write(utils.ToBytes([]bool{true}))
-	h.Write(utils.UInt32ToBytes(1))
-	h.Read(expect[:])
-	if !bytes.Equal(m.root.rightHash, expect[:]) {
-		t.Error("Wrong righ hash!",
-			"expected", expect,
-			"get", m.root.rightHash)
-	}
-
-	r := m.Get(index)
-	if r.Leaf.Value == nil {
-		t.Error("Cannot find value of key:", key)
-		return
+	key1 := "key1"
+	index1 := vrfPrivKey1.Compute([]byte(key1))
+	val1 := []byte("value1")
+	key2 := "key2"
+	index2 := vrfPrivKey1.Compute([]byte(key2))
+
+	// assert Get from empty tree
+	proof := newTree.Get(index1)
+	if proof.Leaf.Value != nil {
+		t.Fatal("Expect returned leaf's value is nil")
 	}
-	v := r.Leaf.Value
-	if !bytes.Equal(v, val) {
-		t.Errorf("Value mismatch %v / %v", v, val)
+	if bytes.Equal(proof.LookupIndex, proof.Leaf.Index) ||
+		!bytes.Equal(index1, proof.LookupIndex) {
+		t.Fatal("Expect a proof of absence")
 	}
 
-	// Check leaf node hash
-	h.Reset()
-	h.Write(r.Leaf.Commitment.Salt)
-	h.Write([]byte(key))
-	h.Write(val)
-	h.Read(commit[:])
+	// assert Set into empty tree
+	if nil != m.Set(index1, key1, val1) {
+		t.Fatal(err)
+	}
 
-	h.Reset()
-	h.Write([]byte{LeafIdentifier})
-	h.Write(m.nonce)
-	h.Write(index)
-	h.Write(utils.UInt32ToBytes(1))
-	h.Write(commit[:])
-	h.Read(expect[:])
+	if nil != newTree.Set(index1, key1, val1) {
+		t.Fatal("WTF")
+	}
 
-	if !bytes.Equal(m.root.leftHash, expect[:]) {
-		t.Error("Wrong left hash!",
-			"expected", expect,
-			"get", m.root.leftHash)
+	proof = newTree.Get(index1)
+	if proof.Leaf.Value == nil {
+		t.Fatal("Expect returned leaf's value is not nil")
+	}
+	if !bytes.Equal(proof.LookupIndex, proof.Leaf.Index) ||
+		!bytes.Equal(index1, proof.LookupIndex) {
+		t.Fatal("Expect a proof of inclusion")
 	}
 
-	r = m.Get([]byte("abc"))
-	if r.Leaf.Value != nil {
-		t.Error("Invalid look-up operation:", key)
-		return
+	proof = newTree.Get(index2)
+	if proof.Leaf.Value != nil {
+		t.Fatal("Expect returned leaf's value is nil")
+	}
+	if bytes.Equal(proof.LookupIndex, proof.Leaf.Index) ||
+		!bytes.Equal(index2, proof.LookupIndex) {
+		t.Fatal("Expect a proof of absence")
 	}
 }
 
+// func TestOneEntry(t *testing.T) {
+// 	m, err := NewMerkleTree()
+// 	if err != nil {
+// 		t.Fatal(err)
+// 	}
+
+// 	var commit [32]byte
+// 	var expect [32]byte
+
+// 	key := "key"
+// 	val := []byte("value")
+// 	index := vrfPrivKey1.Compute([]byte(key))
+// 	if err := m.Set(index, key, val); err != nil {
+// 		t.Fatal(err)
+// 	}
+// 	m.recomputeHash()
+
+// 	// Check empty node hash
+// 	h := sha3.NewShake128()
+// 	h.Write([]byte{EmptyBranchIdentifier})
+// 	h.Write(m.nonce)
+// 	h.Write(utils.ToBytes([]bool{true}))
+// 	h.Write(utils.UInt32ToBytes(1))
+// 	h.Read(expect[:])
+// 	root := m.root.(*interiorNode)
+// 	if !bytes.Equal(root.rightHash, expect[:]) {
+// 		t.Error("Wrong righ hash!",
+// 			"expected", expect,
+// 			"get", root.rightHash)
+// 	}
+
+// 	r := m.Get(index)
+// 	if r.Leaf.Value == nil {
+// 		t.Error("Cannot find value of key:", key)
+// 		return
+// 	}
+// 	v := r.Leaf.Value
+// 	if !bytes.Equal(v, val) {
+// 		t.Errorf("Value mismatch %v / %v", v, val)
+// 	}
+
+// 	// Check leaf node hash
+// 	h.Reset()
+// 	h.Write(r.Leaf.Commitment.Salt)
+// 	h.Write([]byte(key))
+// 	h.Write(val)
+// 	h.Read(commit[:])
+
+// 	h.Reset()
+// 	h.Write([]byte{LeafIdentifier})
+// 	h.Write(m.nonce)
+// 	h.Write(index)
+// 	h.Write(utils.UInt32ToBytes(1))
+// 	h.Write(commit[:])
+// 	h.Read(expect[:])
+// 	root = m.root.(*interiorNode)
+// 	if !bytes.Equal(root.leftHash, expect[:]) {
+// 		t.Error("Wrong left hash!",
+// 			"expected", expect,
+// 			"get", root.leftHash)
+// 	}
+
+// 	r = m.Get([]byte("abc"))
+// 	if r.Leaf.Value != nil {
+// 		t.Error("Invalid look-up operation:", key)
+// 		return
+// 	}
+// }
+
 func TestTwoEntries(t *testing.T) {
 	m, err := NewMerkleTree()
 	if err != nil {

merkletree/node.go

@@ -31,6 +31,16 @@ type emptyNode struct {
 	index []byte
 }
 
+func newEmptyNode(parent merkleNode, level uint32, prefixBits []bool) *emptyNode {
+	return &emptyNode{
+		node: node{
+			parent: parent,
+			level:  level,
+		},
+		index: utils.ToBytes(prefixBits),
+	}
+}
+
 func newInteriorNode(parent merkleNode, level uint32, prefixBits []bool) *interiorNode {
 	prefixLeft := append([]bool(nil), prefixBits...)
 	prefixLeft = append(prefixLeft, false)
@@ -68,7 +78,7 @@ func newInteriorNode(parent merkleNode, level uint32, prefixBits []bool) *interi
 type merkleNode interface {
 	isEmpty() bool
 	hash(*MerkleTree) []byte
-	clone(*interiorNode) merkleNode
+	clone(merkleNode) merkleNode
 }
 
 var _ merkleNode = (*userLeafNode)(nil)
@@ -104,7 +114,7 @@ func (n *emptyNode) hash(m *MerkleTree) []byte {
 	)
 }
 
-func (n *interiorNode) clone(parent *interiorNode) merkleNode {
+func (n *interiorNode) clone(parent merkleNode) merkleNode {
 	newNode := &interiorNode{
 		node: node{
 			parent: parent,
@@ -122,7 +132,7 @@ func (n *interiorNode) clone(parent *interiorNode) merkleNode {
 	return newNode
 }
 
-func (n *userLeafNode) clone(parent *interiorNode) merkleNode {
+func (n *userLeafNode) clone(parent merkleNode) merkleNode {
 	return &userLeafNode{
 		node: node{
 			parent: parent,
@@ -135,7 +145,7 @@ func (n *userLeafNode) clone(parent *interiorNode) merkleNode {
 	}
 }
 
-func (n *emptyNode) clone(parent *interiorNode) merkleNode {
+func (n *emptyNode) clone(parent merkleNode) merkleNode {
 	return &emptyNode{
 		node: node{
 			parent: parent,

merkletree/proof_test.go

@@ -184,10 +184,10 @@ func TestProofVerificationErrors(t *testing.T) {
 	if err := proof2.Verify([]byte("a"), nil, m.hash); err != ErrBindingsDiffer {
 		t.Error("Expect", ErrBindingsDiffer, "got", err)
 	}
-	// - ErrIndicesMismatch
+	// - ErrUnequalTreeHashes
 	proof2.Leaf.Value = nil
-	proof2.Leaf.Index[0] &= 0x01
-	if err := proof2.Verify([]byte("a"), nil, m.hash); err != ErrIndicesMismatch {
-		t.Error("Expect", ErrIndicesMismatch, "got", err)
+	proof2.Leaf.Index[0] ^= 0x01
+	if err := proof2.Verify([]byte("a"), nil, m.hash); err != ErrUnequalTreeHashes {
+		t.Error("Expect", ErrUnequalTreeHashes, "got", err)
 	}
 }