You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

475 lines
15 KiB

  1. package merkletree
  2. import (
  3. "bytes"
  4. "encoding/hex"
  5. "fmt"
  6. "math/big"
  7. "testing"
  8. "github.com/iden3/go-merkletree/db/memory"
  9. "github.com/stretchr/testify/assert"
  10. "github.com/stretchr/testify/require"
  11. )
  12. var debug = false
  13. type Fatalable interface {
  14. Fatal(args ...interface{})
  15. }
  16. func newTestingMerkle(f Fatalable, numLevels int) *MerkleTree {
  17. mt, err := NewMerkleTree(memory.NewMemoryStorage(), numLevels)
  18. if err != nil {
  19. f.Fatal(err)
  20. return nil
  21. }
  22. return mt
  23. }
  24. func TestHashParsers(t *testing.T) {
  25. h0 := NewHashFromBigInt(big.NewInt(0))
  26. assert.Equal(t, "0", h0.String())
  27. h1 := NewHashFromBigInt(big.NewInt(1))
  28. assert.Equal(t, "1", h1.String())
  29. h10 := NewHashFromBigInt(big.NewInt(10))
  30. assert.Equal(t, "10", h10.String())
  31. h7l := NewHashFromBigInt(big.NewInt(1234567))
  32. assert.Equal(t, "1234567", h7l.String())
  33. h8l := NewHashFromBigInt(big.NewInt(12345678))
  34. assert.Equal(t, "12345678...", h8l.String())
  35. b, ok := new(big.Int).SetString("4932297968297298434239270129193057052722409868268166443802652458940273154854", 10)
  36. assert.True(t, ok)
  37. h := NewHashFromBigInt(b)
  38. assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", h.BigInt().String())
  39. assert.Equal(t, "49322979...", h.String())
  40. assert.Equal(t, "0ae794eb9c3d8bbb9002e993fc2ed301dcbd2af5508ed072c375e861f1aa5b26", h.Hex())
  41. }
  42. func TestNewTree(t *testing.T) {
  43. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
  44. assert.Nil(t, err)
  45. assert.Equal(t, "0", mt.Root().String())
  46. // test vectors generated using https://github.com/iden3/circomlib smt.js
  47. err = mt.Add(big.NewInt(1), big.NewInt(2))
  48. assert.Nil(t, err)
  49. assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", mt.Root().BigInt().String())
  50. err = mt.Add(big.NewInt(33), big.NewInt(44))
  51. assert.Nil(t, err)
  52. assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
  53. err = mt.Add(big.NewInt(1234), big.NewInt(9876))
  54. assert.Nil(t, err)
  55. assert.Equal(t, "16970503620176669663662021947486532860010370357132361783766545149750777353066", mt.Root().BigInt().String())
  56. proof, err := mt.GenerateProof(big.NewInt(33), nil)
  57. assert.Nil(t, err)
  58. assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(44)))
  59. assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(33), big.NewInt(45)))
  60. }
  61. func TestAddDifferentOrder(t *testing.T) {
  62. mt1 := newTestingMerkle(t, 140)
  63. defer mt1.db.Close()
  64. for i := 0; i < 16; i++ {
  65. k := big.NewInt(int64(i))
  66. v := big.NewInt(0)
  67. if err := mt1.Add(k, v); err != nil {
  68. t.Fatal(err)
  69. }
  70. }
  71. mt2 := newTestingMerkle(t, 140)
  72. defer mt2.db.Close()
  73. for i := 16 - 1; i >= 0; i-- {
  74. k := big.NewInt(int64(i))
  75. v := big.NewInt(0)
  76. if err := mt2.Add(k, v); err != nil {
  77. t.Fatal(err)
  78. }
  79. }
  80. assert.Equal(t, mt1.Root().Hex(), mt2.Root().Hex())
  81. assert.Equal(t, "0967b777d660e54aa3a0f0f3405bb962504d3d69d6b930146cd212dff9913bee", mt1.Root().Hex())
  82. }
  83. func TestAddRepeatedIndex(t *testing.T) {
  84. mt := newTestingMerkle(t, 140)
  85. defer mt.db.Close()
  86. k := big.NewInt(int64(3))
  87. v := big.NewInt(int64(12))
  88. if err := mt.Add(k, v); err != nil {
  89. t.Fatal(err)
  90. }
  91. err := mt.Add(k, v)
  92. assert.NotNil(t, err)
  93. assert.Equal(t, err, ErrEntryIndexAlreadyExists)
  94. }
  95. func TestGenerateAndVerifyProof128(t *testing.T) {
  96. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
  97. require.Nil(t, err)
  98. defer mt.db.Close()
  99. for i := 0; i < 128; i++ {
  100. k := big.NewInt(int64(i))
  101. v := big.NewInt(0)
  102. if err := mt.Add(k, v); err != nil {
  103. t.Fatal(err)
  104. }
  105. }
  106. proof, err := mt.GenerateProof(big.NewInt(42), nil)
  107. assert.Nil(t, err)
  108. assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(42), big.NewInt(0)))
  109. }
  110. func TestTreeLimit(t *testing.T) {
  111. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 5)
  112. require.Nil(t, err)
  113. defer mt.db.Close()
  114. for i := 0; i < 16; i++ {
  115. err = mt.Add(big.NewInt(int64(i)), big.NewInt(int64(i)))
  116. assert.Nil(t, err)
  117. }
  118. // here the tree is full, should not allow to add more data as reaches the maximum number of levels
  119. err = mt.Add(big.NewInt(int64(16)), big.NewInt(int64(16)))
  120. assert.NotNil(t, err)
  121. assert.Equal(t, ErrReachedMaxLevel, err)
  122. }
  123. func TestSiblingsFromProof(t *testing.T) {
  124. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
  125. require.Nil(t, err)
  126. defer mt.db.Close()
  127. for i := 0; i < 64; i++ {
  128. k := big.NewInt(int64(i))
  129. v := big.NewInt(0)
  130. if err := mt.Add(k, v); err != nil {
  131. t.Fatal(err)
  132. }
  133. }
  134. proof, err := mt.GenerateProof(big.NewInt(4), nil)
  135. if err != nil {
  136. t.Fatal(err)
  137. }
  138. siblings := SiblingsFromProof(proof)
  139. assert.Equal(t, 6, len(siblings))
  140. assert.Equal(t, "23db1f6fb07af47d7715f18960548c215fc7a2e6d25cb4a7eb82c9d3cf69bc26", siblings[0].Hex())
  141. assert.Equal(t, "2156e64dedcb76719ec732414dd6a8aa4348dafb24c19351a68fbc4158bb7fba", siblings[1].Hex())
  142. assert.Equal(t, "04a8e9b34d5a8b55268ca96b0b8c7c5aaef4f606ec3437ec67e4152d9b323913", siblings[2].Hex())
  143. assert.Equal(t, "0ff484133e0d25deb4a7c0cb46d90432e00fcc280948c2fab6fed9476f1e26b2", siblings[3].Hex())
  144. assert.Equal(t, "015dff482e87eb2046b8f5323049afd05f8dd8554e2c9aa1ef28991cf205c9b6", siblings[4].Hex())
  145. assert.Equal(t, "1e4da486ad68b07acec1406bed5a60732de5ff72d63910f7afbb491f953a8769", siblings[5].Hex())
  146. }
  147. func TestVerifyProofCases(t *testing.T) {
  148. mt := newTestingMerkle(t, 140)
  149. defer mt.DB().Close()
  150. for i := 0; i < 8; i++ {
  151. if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
  152. t.Fatal(err)
  153. }
  154. }
  155. // Existence proof
  156. proof, err := mt.GenerateProof(big.NewInt(4), nil)
  157. if err != nil {
  158. t.Fatal(err)
  159. }
  160. assert.Equal(t, proof.Existence, true)
  161. assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(4), big.NewInt(0)))
  162. assert.Equal(t, "000300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f", hex.EncodeToString(proof.Bytes()))
  163. for i := 8; i < 32; i++ {
  164. proof, err = mt.GenerateProof(big.NewInt(int64(i)), nil)
  165. assert.Nil(t, err)
  166. if debug {
  167. fmt.Println(i, proof)
  168. }
  169. }
  170. // Non-existence proof, empty aux
  171. proof, err = mt.GenerateProof(big.NewInt(12), nil)
  172. if err != nil {
  173. t.Fatal(err)
  174. }
  175. assert.Equal(t, proof.Existence, false)
  176. // assert.True(t, proof.nodeAux == nil)
  177. assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(12), big.NewInt(0)))
  178. assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b16621630b9104d995843c7cffa685009a1b28dcd371022a3b27b3a4d6987f7c8b39b0f2fffc165330710754ca0fc24451bdd5d5f82a05f42f1427fbdf17879c0b84be60f04000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
  179. // Non-existence proof, diff. node aux
  180. proof, err = mt.GenerateProof(big.NewInt(10), nil)
  181. if err != nil {
  182. t.Fatal(err)
  183. }
  184. assert.Equal(t, proof.Existence, false)
  185. assert.True(t, proof.NodeAux != nil)
  186. assert.True(t, VerifyProof(mt.Root(), proof, big.NewInt(10), big.NewInt(0)))
  187. assert.Equal(t, "030300000000000000000000000000000000000000000000000000000000000728ea2b267d2a9436657f20b5827285175e030f58c07375535106903b1662163097fcf8f911b271df196e0a75667b8a4f3024ef39f87201ed2b7cda349ba202296b7aeba35dc19ab0d4f65e175536c9952a90b6de18c3205611c3cd4fb408f01602000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000", hex.EncodeToString(proof.Bytes()))
  188. }
  189. func TestVerifyProofFalse(t *testing.T) {
  190. mt := newTestingMerkle(t, 140)
  191. defer mt.DB().Close()
  192. for i := 0; i < 8; i++ {
  193. if err := mt.Add(big.NewInt(int64(i)), big.NewInt(0)); err != nil {
  194. t.Fatal(err)
  195. }
  196. }
  197. // Invalid existence proof (node used for verification doesn't
  198. // correspond to node in the proof)
  199. proof, err := mt.GenerateProof(big.NewInt(int64(4)), nil)
  200. if err != nil {
  201. t.Fatal(err)
  202. }
  203. assert.Equal(t, proof.Existence, true)
  204. assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(5)), big.NewInt(int64(5))))
  205. // Invalid non-existence proof (Non-existence proof, diff. node aux)
  206. proof, err = mt.GenerateProof(big.NewInt(int64(4)), nil)
  207. if err != nil {
  208. t.Fatal(err)
  209. }
  210. assert.Equal(t, proof.Existence, true)
  211. // Now we change the proof from existence to non-existence, and add e's
  212. // data as auxiliary node.
  213. proof.Existence = false
  214. proof.NodeAux = &NodeAux{Key: NewHashFromBigInt(big.NewInt(int64(4))), Value: NewHashFromBigInt(big.NewInt(4))}
  215. assert.True(t, !VerifyProof(mt.Root(), proof, big.NewInt(int64(4)), big.NewInt(0)))
  216. }
  217. func TestGraphViz(t *testing.T) {
  218. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
  219. assert.Nil(t, err)
  220. mt.Add(big.NewInt(1), big.NewInt(0))
  221. mt.Add(big.NewInt(2), big.NewInt(0))
  222. mt.Add(big.NewInt(3), big.NewInt(0))
  223. mt.Add(big.NewInt(4), big.NewInt(0))
  224. mt.Add(big.NewInt(5), big.NewInt(0))
  225. mt.Add(big.NewInt(100), big.NewInt(0))
  226. // mt.PrintGraphViz(nil)
  227. expected := `digraph hierarchy {
  228. node [fontname=Monospace,fontsize=10,shape=box]
  229. "60195538..." -> {"19759736..." "18893277..."}
  230. "19759736..." -> {"16152312..." "43945008..."}
  231. "16152312..." -> {"empty0" "13952255..."}
  232. "empty0" [style=dashed,label=0];
  233. "13952255..." -> {"61769925..." "empty1"}
  234. "empty1" [style=dashed,label=0];
  235. "61769925..." -> {"92723289..." "empty2"}
  236. "empty2" [style=dashed,label=0];
  237. "92723289..." -> {"21082735..." "82784818..."}
  238. "21082735..." [style=filled];
  239. "82784818..." [style=filled];
  240. "43945008..." [style=filled];
  241. "18893277..." -> {"19855703..." "17718670..."}
  242. "19855703..." -> {"11499909..." "15828714..."}
  243. "11499909..." [style=filled];
  244. "15828714..." [style=filled];
  245. "17718670..." [style=filled];
  246. }
  247. `
  248. w := bytes.NewBufferString("")
  249. mt.GraphViz(w, nil)
  250. assert.Equal(t, []byte(expected), w.Bytes())
  251. }
  252. func TestDelete(t *testing.T) {
  253. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
  254. assert.Nil(t, err)
  255. assert.Equal(t, "0", mt.Root().String())
  256. // test vectors generated using https://github.com/iden3/circomlib smt.js
  257. err = mt.Add(big.NewInt(1), big.NewInt(2))
  258. assert.Nil(t, err)
  259. assert.Equal(t, "4932297968297298434239270129193057052722409868268166443802652458940273154854", mt.Root().BigInt().String())
  260. err = mt.Add(big.NewInt(33), big.NewInt(44))
  261. assert.Nil(t, err)
  262. assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
  263. err = mt.Add(big.NewInt(1234), big.NewInt(9876))
  264. assert.Nil(t, err)
  265. assert.Equal(t, "16970503620176669663662021947486532860010370357132361783766545149750777353066", mt.Root().BigInt().String())
  266. // mt.PrintGraphViz(nil)
  267. err = mt.Delete(big.NewInt(33))
  268. // mt.PrintGraphViz(nil)
  269. assert.Nil(t, err)
  270. assert.Equal(t, "12820263606494630162816839760750120928463716794691735985748071431547370997091", mt.Root().BigInt().String())
  271. err = mt.Delete(big.NewInt(1234))
  272. assert.Nil(t, err)
  273. err = mt.Delete(big.NewInt(1))
  274. assert.Nil(t, err)
  275. assert.Equal(t, "0", mt.Root().String())
  276. }
  277. func TestDelete2(t *testing.T) {
  278. mt := newTestingMerkle(t, 140)
  279. defer mt.db.Close()
  280. for i := 0; i < 8; i++ {
  281. k := big.NewInt(int64(i))
  282. v := big.NewInt(0)
  283. if err := mt.Add(k, v); err != nil {
  284. t.Fatal(err)
  285. }
  286. }
  287. expectedRoot := mt.Root()
  288. k := big.NewInt(8)
  289. v := big.NewInt(0)
  290. err := mt.Add(k, v)
  291. require.Nil(t, err)
  292. err = mt.Delete(big.NewInt(8))
  293. assert.Nil(t, err)
  294. assert.Equal(t, expectedRoot, mt.Root())
  295. mt2 := newTestingMerkle(t, 140)
  296. defer mt2.db.Close()
  297. for i := 0; i < 8; i++ {
  298. k := big.NewInt(int64(i))
  299. v := big.NewInt(0)
  300. if err := mt2.Add(k, v); err != nil {
  301. t.Fatal(err)
  302. }
  303. }
  304. assert.Equal(t, mt2.Root(), mt.Root())
  305. }
  306. func TestDelete3(t *testing.T) {
  307. mt := newTestingMerkle(t, 140)
  308. defer mt.db.Close()
  309. err := mt.Add(big.NewInt(1), big.NewInt(1))
  310. assert.Nil(t, err)
  311. err = mt.Add(big.NewInt(2), big.NewInt(2))
  312. assert.Nil(t, err)
  313. assert.Equal(t, "2427629547967522489273866134471574861207714751136138191708011221765688788661", mt.Root().BigInt().String())
  314. err = mt.Delete(big.NewInt(1))
  315. assert.Nil(t, err)
  316. assert.Equal(t, "10822920717809411688334493481050035035708810950159417482558569847174767667301", mt.Root().BigInt().String())
  317. mt2 := newTestingMerkle(t, 140)
  318. defer mt2.db.Close()
  319. err = mt2.Add(big.NewInt(2), big.NewInt(2))
  320. assert.Nil(t, err)
  321. assert.Equal(t, mt2.Root(), mt.Root())
  322. }
  323. func TestDelete4(t *testing.T) {
  324. mt := newTestingMerkle(t, 140)
  325. defer mt.db.Close()
  326. err := mt.Add(big.NewInt(1), big.NewInt(1))
  327. assert.Nil(t, err)
  328. err = mt.Add(big.NewInt(2), big.NewInt(2))
  329. assert.Nil(t, err)
  330. err = mt.Add(big.NewInt(3), big.NewInt(3))
  331. assert.Nil(t, err)
  332. assert.Equal(t, "16614298246517994771186095530428786749320098419259206061045083278756632941513", mt.Root().BigInt().String())
  333. err = mt.Delete(big.NewInt(1))
  334. assert.Nil(t, err)
  335. assert.Equal(t, "6117330520107511783353383870014397665359816230889739699667943862706617498952", mt.Root().BigInt().String())
  336. mt2 := newTestingMerkle(t, 140)
  337. defer mt2.db.Close()
  338. err = mt2.Add(big.NewInt(2), big.NewInt(2))
  339. assert.Nil(t, err)
  340. err = mt2.Add(big.NewInt(3), big.NewInt(3))
  341. assert.Nil(t, err)
  342. assert.Equal(t, mt2.Root(), mt.Root())
  343. }
  344. func TestDelete5(t *testing.T) {
  345. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
  346. assert.Nil(t, err)
  347. err = mt.Add(big.NewInt(1), big.NewInt(2))
  348. assert.Nil(t, err)
  349. err = mt.Add(big.NewInt(33), big.NewInt(44))
  350. assert.Nil(t, err)
  351. assert.Equal(t, "13563340744765267202993741297198970774200042973817962221376874695587906013050", mt.Root().BigInt().String())
  352. err = mt.Delete(big.NewInt(1))
  353. assert.Nil(t, err)
  354. assert.Equal(t, "12075524681474630909546786277734445038384732558409197537058769521806571391765", mt.Root().BigInt().String())
  355. mt2 := newTestingMerkle(t, 140)
  356. defer mt2.db.Close()
  357. err = mt2.Add(big.NewInt(33), big.NewInt(44))
  358. assert.Nil(t, err)
  359. assert.Equal(t, mt2.Root(), mt.Root())
  360. }
  361. func TestDeleteNonExistingKeys(t *testing.T) {
  362. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 10)
  363. assert.Nil(t, err)
  364. err = mt.Add(big.NewInt(1), big.NewInt(2))
  365. assert.Nil(t, err)
  366. err = mt.Add(big.NewInt(33), big.NewInt(44))
  367. assert.Nil(t, err)
  368. err = mt.Delete(big.NewInt(33))
  369. assert.Nil(t, err)
  370. err = mt.Delete(big.NewInt(33))
  371. assert.Equal(t, ErrKeyNotFound, err)
  372. err = mt.Delete(big.NewInt(1))
  373. assert.Nil(t, err)
  374. assert.Equal(t, "0", mt.Root().String())
  375. err = mt.Delete(big.NewInt(33))
  376. assert.Equal(t, ErrKeyNotFound, err)
  377. }
  378. func TestDumpLeafsImportLeafs(t *testing.T) {
  379. mt, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
  380. require.Nil(t, err)
  381. defer mt.db.Close()
  382. for i := 0; i < 10; i++ {
  383. k := big.NewInt(int64(i))
  384. v := big.NewInt(0)
  385. err = mt.Add(k, v)
  386. require.Nil(t, err)
  387. }
  388. d, err := mt.DumpLeafs(nil)
  389. assert.Nil(t, err)
  390. mt2, err := NewMerkleTree(memory.NewMemoryStorage(), 140)
  391. require.Nil(t, err)
  392. defer mt2.db.Close()
  393. err = mt2.ImportDumpedLeafs(d)
  394. assert.Nil(t, err)
  395. assert.Equal(t, mt.Root(), mt2.Root())
  396. }