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.

137 lines
3.5 KiB

10 years ago
  1. package merkle
  2. import (
  3. "bytes"
  4. "crypto/sha256"
  5. "github.com/tendermint/tendermint/binary"
  6. . "github.com/tendermint/tendermint/common"
  7. )
  8. type IAVLProof struct {
  9. Root []byte
  10. Branches []IAVLProofBranch
  11. Leaf IAVLProofLeaf
  12. }
  13. func (proof *IAVLProof) Verify() bool {
  14. hash := proof.Leaf.Hash()
  15. // fmt.Printf("leaf hash: %X\n", hash)
  16. for i := len(proof.Branches) - 1; 0 <= i; i-- {
  17. hash = proof.Branches[i].Hash(hash)
  18. // fmt.Printf("branch hash: %X\n", hash)
  19. }
  20. // fmt.Printf("root: %X, computed: %X\n", proof.Root, hash)
  21. return bytes.Equal(proof.Root, hash)
  22. }
  23. type IAVLProofBranch struct {
  24. Height uint8
  25. Size uint
  26. Left []byte
  27. Right []byte
  28. }
  29. func (branch IAVLProofBranch) Hash(childHash []byte) []byte {
  30. hasher := sha256.New()
  31. buf := new(bytes.Buffer)
  32. n, err := int64(0), error(nil)
  33. binary.WriteUint8(branch.Height, buf, &n, &err)
  34. binary.WriteUvarint(branch.Size, buf, &n, &err)
  35. if branch.Left == nil {
  36. binary.WriteByteSlice(childHash, buf, &n, &err)
  37. binary.WriteByteSlice(branch.Right, buf, &n, &err)
  38. } else {
  39. binary.WriteByteSlice(branch.Left, buf, &n, &err)
  40. binary.WriteByteSlice(childHash, buf, &n, &err)
  41. }
  42. if err != nil {
  43. panic(Fmt("Failed to hash IAVLProofBranch: %v", err))
  44. }
  45. // fmt.Printf("Branch hash bytes: %X\n", buf.Bytes())
  46. hasher.Write(buf.Bytes())
  47. return hasher.Sum(nil)
  48. }
  49. type IAVLProofLeaf struct {
  50. KeyBytes []byte
  51. ValueBytes []byte
  52. }
  53. func (leaf IAVLProofLeaf) Hash() []byte {
  54. hasher := sha256.New()
  55. buf := new(bytes.Buffer)
  56. n, err := int64(0), error(nil)
  57. binary.WriteUint8(0, buf, &n, &err)
  58. binary.WriteUvarint(1, buf, &n, &err)
  59. binary.WriteByteSlice(leaf.KeyBytes, buf, &n, &err)
  60. binary.WriteByteSlice(leaf.ValueBytes, buf, &n, &err)
  61. if err != nil {
  62. panic(Fmt("Failed to hash IAVLProofLeaf: %v", err))
  63. }
  64. // fmt.Printf("Leaf hash bytes: %X\n", buf.Bytes())
  65. hasher.Write(buf.Bytes())
  66. return hasher.Sum(nil)
  67. }
  68. func (node *IAVLNode) constructProof(t *IAVLTree, key interface{}, proof *IAVLProof) (exists bool) {
  69. if node.height == 0 {
  70. if t.keyCodec.Compare(node.key, key) == 0 {
  71. keyBuf, valueBuf := new(bytes.Buffer), new(bytes.Buffer)
  72. n, err := int64(0), error(nil)
  73. t.keyCodec.Encode(node.key, keyBuf, &n, &err)
  74. if err != nil {
  75. panic(Fmt("Failed to encode node.key: %v", err))
  76. }
  77. t.valueCodec.Encode(node.value, valueBuf, &n, &err)
  78. if err != nil {
  79. panic(Fmt("Failed to encode node.value: %v", err))
  80. }
  81. leaf := IAVLProofLeaf{
  82. KeyBytes: keyBuf.Bytes(),
  83. ValueBytes: valueBuf.Bytes(),
  84. }
  85. proof.Leaf = leaf
  86. return true
  87. } else {
  88. return false
  89. }
  90. } else {
  91. if t.keyCodec.Compare(key, node.key) < 0 {
  92. branch := IAVLProofBranch{
  93. Height: node.height,
  94. Size: node.size,
  95. Left: nil,
  96. Right: node.getRightNode(t).hash,
  97. }
  98. if node.getRightNode(t).hash == nil {
  99. // fmt.Println(node.getRightNode(t))
  100. panic("WTF")
  101. }
  102. proof.Branches = append(proof.Branches, branch)
  103. return node.getLeftNode(t).constructProof(t, key, proof)
  104. } else {
  105. branch := IAVLProofBranch{
  106. Height: node.height,
  107. Size: node.size,
  108. Left: node.getLeftNode(t).hash,
  109. Right: nil,
  110. }
  111. proof.Branches = append(proof.Branches, branch)
  112. return node.getRightNode(t).constructProof(t, key, proof)
  113. }
  114. }
  115. }
  116. // Returns nil if key is not in tree.
  117. func (t *IAVLTree) ConstructProof(key interface{}) *IAVLProof {
  118. if t.root == nil {
  119. return nil
  120. }
  121. t.root.hashWithCount(t) // Ensure that all hashes are calculated.
  122. proof := &IAVLProof{
  123. Root: t.root.hash,
  124. }
  125. t.root.constructProof(t, key, proof)
  126. return proof
  127. }