1package merkletree
  2
  3import (
  4	"encoding/hex"
  5	"fmt"
  6	"strconv"
  7	"testing"
  8	"time"
  9
 10	"github.com/dchest/uniuri"
 11	"github.com/stretchr/testify/assert"
 12	"github.com/syndtr/goleveldb/leveldb"
 13)
 14
 15type testBase struct {
 16	Length    [4]byte
 17	Namespace Hash
 18	Type      Hash
 19	Version   uint32
 20}
 21type testLeaf struct {
 22	testBase
 23	extraIndex struct {
 24		Data []byte
 25	}
 26}
 27
 28func parseTestLeafBytes(b []byte) testLeaf {
 29	var c testLeaf
 30	copy(c.testBase.Length[:], b[0:4])
 31	copy(c.testBase.Namespace[:], b[4:36])
 32	copy(c.testBase.Type[:], b[36:68])
 33	versionBytes := b[68:72]
 34	c.testBase.Version = BytesToUint32(versionBytes)
 35	c.extraIndex.Data = b[72:]
 36	return c
 37}
 38func (c testLeaf) Bytes() (b []byte) {
 39	b = append(b, c.testBase.Length[:]...)
 40	b = append(b, c.testBase.Namespace[:]...)
 41	b = append(b, c.testBase.Type[:]...)
 42	versionBytes := Uint32ToBytes(c.testBase.Version)
 43	b = append(b, versionBytes[:]...)
 44	b = append(b, c.extraIndex.Data[:]...)
 45	return b
 46}
 47func (c testLeaf) IndexLength() uint32 {
 48	return uint32(len(c.Bytes()))
 49}
 50func (c testLeaf) hi() Hash {
 51	h := HashBytes(c.Bytes())
 52	return h
 53}
 54func newTestLeaf(namespaceStr, typeStr string, data []byte) testLeaf {
 55	var c testLeaf
 56	c.testBase.Length = [4]byte{0x00, 0x00, 0x00, 0x48}
 57	c.testBase.Namespace = HashBytes([]byte(namespaceStr))
 58	c.testBase.Type = HashBytes([]byte(typeStr))
 59	c.testBase.Version = 0
 60	c.extraIndex.Data = data
 61	return c
 62}
 63
 64type Fatalable interface {
 65	Fatal(args ...interface{})
 66}
 67
 68func newTestingMerkle(f Fatalable, numLevels int) *MerkleTree {
 69	tim := time.Now().Unix()
 70	tstr := fmt.Sprint(tim)
 71	s := uniuri.New()
 72	db, err := leveldb.OpenFile("tmp/db"+s+"-"+tstr, nil)
 73	if err != nil {
 74		f.Fatal(err)
 75		return nil
 76	}
 77	// defer db.Close()
 78	mt, err := New(db, numLevels)
 79	if err != nil {
 80		f.Fatal(err)
 81		return nil
 82	}
 83	return mt
 84}
 85
 86func TestNewMT(t *testing.T) {
 87
 88	//create a new MT
 89	mt := newTestingMerkle(t, 140)
 90	defer mt.storage.Close()
 91	assert.Equal(t, "0x0000000000000000000000000000000000000000000000000000000000000000", mt.Root().Hex())
 92}
 93
 94func TestAddLeaf(t *testing.T) {
 95
 96	mt := newTestingMerkle(t, 140)
 97	defer mt.storage.Close()
 98
 99	leaf := newTestLeaf("iden3.io", "typespec", []byte("c1"))
100	assert.Equal(t, "0x939862c94ca9772fc9e2621df47128b1d4041b514e19edc969a92d8f0dae558f", leaf.hi().Hex())
101
102	assert.Nil(t, mt.Add(leaf))
103	assert.Equal(t, "0x9d3c407ff02c813cd474c0a6366b4f7c58bf417a38268f7a0d73a8bca2490b9b", mt.Root().Hex())
104}
105
106func TestAddLeafs(t *testing.T) {
107
108	mt := newTestingMerkle(t, 140)
109	defer mt.storage.Close()
110
111	leaf := newTestLeaf("iden3.io", "typespec", []byte("c1"))
112	assert.Equal(t, "0x939862c94ca9772fc9e2621df47128b1d4041b514e19edc969a92d8f0dae558f", leaf.hi().Hex())
113	assert.Nil(t, mt.Add(leaf))
114
115	assert.Nil(t, mt.Add(newTestLeaf("iden3.io2", "typespec2", []byte("c2"))))
116	assert.Equal(t, "0xebae8fb483b48ba6c337136535198eb8bcf891daba40ac81e28958c09b9b229b", mt.Root().Hex())
117
118	mt.Add(newTestLeaf("iden3.io3", "typespec3", []byte("c3")))
119	mt.Add(newTestLeaf("iden3.io4", "typespec4", []byte("c4")))
120	assert.Equal(t, "0xb4b51aa0c77a8e5ed0a099d7c11c7d2a9219ef241da84f0689da1f40a5f6ac31", mt.Root().Hex())
121}
122
123func TestAddLeafsCollision(t *testing.T) {
124
125	mt := newTestingMerkle(t, 140)
126	defer mt.storage.Close()
127
128	leaf := newTestLeaf("iden3.io", "typespec", []byte("c1"))
129	assert.Nil(t, mt.Add(leaf))
130
131	root1 := mt.Root()
132	assert.EqualError(t, mt.Add(leaf), ErrNodeAlreadyExists.Error())
133
134	assert.Equal(t, root1.Hex(), mt.Root().Hex())
135}
136
137func TestAddLeafsDifferentOrders(t *testing.T) {
138
139	mt1 := newTestingMerkle(t, 140)
140	defer mt1.storage.Close()
141
142	mt1.Add(newTestLeaf("iden3.io", "typespec", []byte("c1")))
143	mt1.Add(newTestLeaf("iden3.io2", "typespec2", []byte("c2")))
144	mt1.Add(newTestLeaf("iden3.io3", "typespec3", []byte("c3")))
145	mt1.Add(newTestLeaf("iden3.io4", "typespec4", []byte("c4")))
146	mt1.Add(newTestLeaf("iden3.io5", "typespec5", []byte("c5")))
147
148	mt2 := newTestingMerkle(t, 140)
149	defer mt2.storage.Close()
150
151	mt2.Add(newTestLeaf("iden3.io3", "typespec3", []byte("c3")))
152	mt2.Add(newTestLeaf("iden3.io2", "typespec2", []byte("c2")))
153	mt2.Add(newTestLeaf("iden3.io", "typespec", []byte("c1")))
154	mt2.Add(newTestLeaf("iden3.io4", "typespec4", []byte("c4")))
155	mt2.Add(newTestLeaf("iden3.io5", "typespec5", []byte("c5")))
156
157	assert.Equal(t, mt1.Root().Hex(), mt2.Root().Hex())
158}
159func TestBenchmarkAddingLeafs(t *testing.T) {
160
161	mt := newTestingMerkle(t, 140)
162	defer mt.storage.Close()
163
164	start := time.Now()
165	numToAdd := 1000
166	for i := 0; i < numToAdd; i++ {
167		leaf := newTestLeaf("iden3.io"+strconv.Itoa(i), "typespec"+strconv.Itoa(i), []byte("c"+strconv.Itoa(i)))
168		mt.Add(leaf)
169	}
170	fmt.Print("time elapsed adding " + strconv.Itoa(numToAdd) + " leafs: ")
171	fmt.Println(time.Since(start))
172}
173func BenchmarkAddingLeafs(b *testing.B) {
174
175	mt := newTestingMerkle(b, 140)
176	defer mt.storage.Close()
177
178	b.ResetTimer()
179	for i := 0; i < b.N; i++ {
180		leaf := newTestLeaf("iden3.io"+strconv.Itoa(i), "typespec"+strconv.Itoa(i), []byte("c"+strconv.Itoa(i)))
181		if err := mt.Add(leaf); err != nil {
182			b.Fatal(err)
183		}
184	}
185}
186
187func TestGenerateProof(t *testing.T) {
188	mt := newTestingMerkle(t, 140)
189	defer mt.storage.Close()
190
191	mt.Add(newTestLeaf("iden3.io_3", "typespec_3", []byte("c3")))
192	mt.Add(newTestLeaf("iden3.io_2", "typespec_2", []byte("c2")))
193
194	leaf1 := newTestLeaf("iden3.io_1", "typespec_1", []byte("c1"))
195	assert.Nil(t, mt.Add(leaf1))
196
197	mp, err := mt.GenerateProof(parseTestLeafBytes(leaf1.Bytes()).hi())
198	assert.Nil(t, err)
199
200	mpHexExpected := "0000000000000000000000000000000000000000000000000000000000000002beb0fd6dcf18d37fe51cf34beacd4c524d9c039ef9da2a27ccd3e7edf662c39c"
201	assert.Equal(t, mpHexExpected, hex.EncodeToString(mp))
202}
203
204func TestCheckProof(t *testing.T) {
205	mt := newTestingMerkle(t, 140)
206	defer mt.storage.Close()
207
208	leaf1 := newTestLeaf("iden3.io_1", "typespec_1", []byte("c1"))
209	assert.Nil(t, mt.Add(leaf1))
210
211	leaf3 := newTestLeaf("iden3.io_3", "typespec_3", []byte("c3"))
212	assert.Nil(t, mt.Add(leaf3))
213
214	mp, err := mt.GenerateProof(parseTestLeafBytes(leaf1.Bytes()).hi())
215	assert.Nil(t, err)
216	verified := CheckProof(mt.Root(), mp, leaf1.hi(), HashBytes(leaf1.Bytes()), mt.NumLevels())
217	assert.True(t, verified)
218
219}
220
221func TestProofOfEmpty(t *testing.T) { // proof of a non revocated leaf, prove that is empty the hi position of the leaf.version+1
222	mt := newTestingMerkle(t, 140)
223	defer mt.storage.Close()
224
225	leaf1 := newTestLeaf("iden3.io_1", "typespec_1", []byte("c1"))
226	// proof when there is nothing in the tree
227	mp, err := mt.GenerateProof(leaf1.hi())
228	assert.Nil(t, err)
229	verified := CheckProof(mt.Root(), mp, leaf1.hi(), EmptyNodeValue, mt.NumLevels())
230	assert.True(t, verified)
231
232	// add the first leaf
233	assert.Nil(t, mt.Add(leaf1))
234
235	// proof when there is only one leaf in the tree
236	leaf2 := newTestLeaf("iden3.io_2", "typespec_2", []byte("c2"))
237	mp, err = mt.GenerateProof(leaf2.hi())
238	assert.Nil(t, err)
239	verified = CheckProof(mt.Root(), mp, leaf2.hi(), EmptyNodeValue, mt.NumLevels())
240	assert.True(t, verified)
241
242	// check that the value in Hi is Empty
243	valueInPos, err := mt.GetValueInPos(leaf2.hi())
244	assert.Nil(t, err)
245	assert.Equal(t, EmptyNodeValue.Bytes(), valueInPos)
246}
247
248func DifferentNonExistenceProofs(t *testing.T) {
249	mt1 := newTestingMerkle(t, 140)
250	defer mt1.storage.Close()
251
252	mt2 := newTestingMerkle(t, 140)
253	defer mt2.storage.Close()
254
255	leaf1 := newTestLeaf("iden3.io_1", "typespec_1", []byte("c1"))
256	leaf2 := newTestLeaf("iden3.io_1", "typespec_1", []byte("c2"))
257
258	assert.Nil(t, mt1.Add(leaf1))
259	assert.Nil(t, mt2.Add(leaf2))
260
261	leaf1.Version++
262	leaf2.Version++
263
264	np1, err := mt1.GenerateProof(leaf1.hi())
265	assert.Nil(t, err)
266	np2, err := mt2.GenerateProof(leaf2.hi())
267	assert.Nil(t, err)
268
269	assert.True(t, CheckProof(mt1.Root(), np1, leaf1.hi(), EmptyNodeValue, mt1.NumLevels()))
270	assert.True(t, CheckProof(mt2.Root(), np2, leaf2.hi(), EmptyNodeValue, mt2.NumLevels()))
271
272	assert.Equal(t, "0000000000000000000000000000000000000000000000000000000000000010a40617c8c3390736831d00b2003e2133353190f5d3b3a586cf829f0f2009aacc", hex.EncodeToString(np1))
273	assert.Equal(t, "0000000000000000000000000000000000000000000000000000000000000001b274a34a3bd95915fe982a0163e3e0a2f79a371b8307661341f8914e22b313e1", hex.EncodeToString(np2))
274
275}
276
277func TestGetLeafInPos(t *testing.T) {
278	mt := newTestingMerkle(t, 140)
279	defer mt.storage.Close()
280
281	for i := 0; i < 50; i++ {
282		leaf := newTestLeaf("iden3.io"+strconv.Itoa(i), "typespec"+strconv.Itoa(i), []byte("c"+strconv.Itoa(i)))
283		mt.Add(leaf)
284
285	}
286	leaf1 := newTestLeaf("iden3.io_x", "typespec_x", []byte("cx"))
287	assert.Nil(t, mt.Add(leaf1))
288
289	leaf := parseTestLeafBytes(leaf1.Bytes())
290	leafInPosBytes, err := mt.GetValueInPos(leaf.hi())
291	assert.Nil(t, err)
292	assert.Equal(t, leaf1.Bytes(), leafInPosBytes)
293
294	// empty value in position
295	leaf2 := newTestLeaf("iden3.io_y", "typespec_y", []byte("cy"))
296	leafInPosBytes, err = mt.GetValueInPos(leaf2.hi())
297	assert.Nil(t, err)
298	assert.Equal(t, EmptyNodeValue[:], leafInPosBytes)
299}
300
301type vt struct {
302	v      []byte
303	idxlen uint32
304}
305
306func (v vt) IndexLength() uint32 {
307	return v.idxlen
308}
309func (v vt) Bytes() []byte {
310	return v.v
311}
312
313func TestVector4(t *testing.T) {
314	mt := newTestingMerkle(t, 4)
315	defer mt.storage.Close()
316
317	zeros := make([]byte, 32, 32)
318	zeros[31] = 1 // to avoid adding Empty element
319	assert.Nil(t, mt.Add(vt{zeros, uint32(1)}))
320	v := vt{zeros, uint32(2)}
321	assert.Nil(t, mt.Add(v))
322	proof, _ := mt.GenerateProof(HashBytes(v.Bytes()[:v.IndexLength()]))
323	assert.True(t, CheckProof(mt.Root(), proof, HashBytes(v.Bytes()[:v.IndexLength()]), HashBytes(v.Bytes()), mt.NumLevels()))
324	assert.Equal(t, 4, mt.NumLevels())
325	assert.Equal(t, "0000000000000000000000000000000000000000000000000000000000000001", hex.EncodeToString(v.Bytes()))
326	assert.Equal(t, "0xc1b95ffbb999a6dd7a472a610a98891ffae95cc973d1d1e21acfdd68db830b51", mt.Root().Hex())
327	assert.Equal(t, "00000000000000000000000000000000000000000000000000000000000000023cf025e4b4fc3ebe57374bf0e0c78ceb0009bdc4466a45174d80e8f508d1a4e3", hex.EncodeToString(proof))
328}
329
330func TestVector140(t *testing.T) {
331	mt := newTestingMerkle(t, 140)
332	defer mt.storage.Close()
333
334	zeros := make([]byte, 32, 32)
335	zeros[31] = 1 // to avoid adding Empty element
336	for i := 1; i < len(zeros)-1; i++ {
337		v := vt{zeros, uint32(i)}
338		assert.Nil(t, mt.Add(v))
339		proof, err := mt.GenerateProof(HashBytes(v.Bytes()[:v.IndexLength()]))
340		assert.Nil(t, err)
341		assert.True(t, CheckProof(mt.Root(), proof, HashBytes(v.Bytes()[:v.IndexLength()]), HashBytes(v.Bytes()), mt.NumLevels()))
342		if i == len(zeros)-2 {
343			assert.Equal(t, 140, mt.NumLevels())
344			assert.Equal(t, uint32(30), v.IndexLength())
345			assert.Equal(t, "0000000000000000000000000000000000000000000000000000000000000001", hex.EncodeToString(v.Bytes()))
346			assert.Equal(t, "0x35f83288adf03bfb61d8d57fab9ed092da79833b58bbdbe9579b636753494ebd", mt.Root().Hex())
347			assert.Equal(t, "000000000000000000000000000000000000000000000000000000000000001f0d1f363115f3333197a009b6674f46bba791308af220ad71515567702b3b44a2b540c1abad0ff81386a78b77e8907a56b7268d24513928ae83497adf4ad93a55e380267ead8305202da0640c1518e144dee87717c732b738fa182c6ef458defd6baf50022b01e3222715d4fca4c198e94536101f6ac314b3d261d3aaa0684395c1db60626e01c39fe4f69418055c2ebd70e0c07b6d9db5c4aed0a11ed2b6a773", hex.EncodeToString(proof))
348		}
349	}
350}