protocol/bc/types/merkle.go

   1 package types
   2
   3 import (
   4         "container/list"
   5         "io"
   6         "math"
   7
   8         "gopkg.in/fatih/set.v0"
   9
  10         "github.com/bytom/bytom/crypto/sha3pool"
  11         "github.com/bytom/bytom/protocol/bc"
  12 )
  13
  14 // merkleFlag represent the type of merkle tree node, it's used to generate the structure of merkle tree
  15 // Bitcoin has only two flags, which zero means the hash of assist node. And one means the hash of the related
  16 // transaction node or it's parents, which distinguish them according to the height of the tree. But in the bytom,
  17 // the height of transaction node is not fixed, so we need three flags to distinguish these nodes.
  18 const (
  19         // FlagAssist represent assist node
  20         FlagAssist = iota
  21         // FlagTxParent represent the parent of transaction of node
  22         FlagTxParent
  23         // FlagTxLeaf represent transaction of node
  24         FlagTxLeaf
  25 )
  26
  27 var (
  28         leafPrefix     = []byte{0x00}
  29         interiorPrefix = []byte{0x01}
  30 )
  31
  32 type merkleNode interface {
  33         WriteTo(io.Writer) (int64, error)
  34 }
  35
  36 func merkleRoot(nodes []merkleNode) (root bc.Hash, err error) {
  37         switch {
  38         case len(nodes) == 0:
  39                 return bc.EmptyStringHash, nil
  40
  41         case len(nodes) == 1:
  42                 root = leafMerkleHash(nodes[0])
  43                 return root, nil
  44
  45         default:
  46                 k := prevPowerOfTwo(len(nodes))
  47                 left, err := merkleRoot(nodes[:k])
  48                 if err != nil {
  49                         return root, err
  50                 }
  51
  52                 right, err := merkleRoot(nodes[k:])
  53                 if err != nil {
  54                         return root, err
  55                 }
  56
  57                 root = interiorMerkleHash(&left, &right)
  58                 return root, nil
  59         }
  60 }
  61
  62 func interiorMerkleHash(left merkleNode, right merkleNode) (hash bc.Hash) {
  63         h := sha3pool.Get256()
  64         defer sha3pool.Put256(h)
  65         h.Write(interiorPrefix)
  66         left.WriteTo(h)
  67         right.WriteTo(h)
  68         hash.ReadFrom(h)
  69         return hash
  70 }
  71
  72 func leafMerkleHash(node merkleNode) (hash bc.Hash) {
  73         h := sha3pool.Get256()
  74         defer sha3pool.Put256(h)
  75         h.Write(leafPrefix)
  76         node.WriteTo(h)
  77         hash.ReadFrom(h)
  78         return hash
  79 }
  80
  81 type merkleTreeNode struct {
  82         hash  bc.Hash
  83         left  *merkleTreeNode
  84         right *merkleTreeNode
  85 }
  86
  87 // buildMerkleTree construct a merkle tree based on the provide node data
  88 func buildMerkleTree(rawDatas []merkleNode) *merkleTreeNode {
  89         switch len(rawDatas) {
  90         case 0:
  91                 return nil
  92         case 1:
  93                 rawData := rawDatas[0]
  94                 merkleHash := leafMerkleHash(rawData)
  95                 node := newMerkleTreeNode(merkleHash, nil, nil)
  96                 return node
  97         default:
  98                 k := prevPowerOfTwo(len(rawDatas))
  99                 left := buildMerkleTree(rawDatas[:k])
 100                 right := buildMerkleTree(rawDatas[k:])
 101                 merkleHash := interiorMerkleHash(&left.hash, &right.hash)
 102                 node := newMerkleTreeNode(merkleHash, left, right)
 103                 return node
 104         }
 105 }
 106
 107 func (node *merkleTreeNode) getMerkleTreeProof(merkleHashSet *set.Set) ([]*bc.Hash, []uint8) {
 108         var hashes []*bc.Hash
 109         var flags []uint8
 110
 111         if node.left == nil && node.right == nil {
 112                 if key := node.hash.String(); merkleHashSet.Has(key) {
 113                         hashes = append(hashes, &node.hash)
 114                         flags = append(flags, FlagTxLeaf)
 115                         return hashes, flags
 116                 }
 117                 return hashes, flags
 118         }
 119         var leftHashes, rightHashes []*bc.Hash
 120         var leftFlags, rightFlags []uint8
 121         if node.left != nil {
 122                 leftHashes, leftFlags = node.left.getMerkleTreeProof(merkleHashSet)
 123         }
 124         if node.right != nil {
 125                 rightHashes, rightFlags = node.right.getMerkleTreeProof(merkleHashSet)
 126         }
 127         leftFind, rightFind := len(leftHashes) > 0, len(rightHashes) > 0
 128
 129         if leftFind || rightFind {
 130                 flags = append(flags, FlagTxParent)
 131         } else {
 132                 return hashes, flags
 133         }
 134
 135         if leftFind {
 136                 hashes = append(hashes, leftHashes...)
 137                 flags = append(flags, leftFlags...)
 138         } else {
 139                 hashes = append(hashes, &node.left.hash)
 140                 flags = append(flags, FlagAssist)
 141         }
 142
 143         if rightFind {
 144                 hashes = append(hashes, rightHashes...)
 145                 flags = append(flags, rightFlags...)
 146         } else {
 147                 hashes = append(hashes, &node.right.hash)
 148                 flags = append(flags, FlagAssist)
 149         }
 150         return hashes, flags
 151 }
 152
 153 func getMerkleTreeProof(rawDatas []merkleNode, relatedRawDatas []merkleNode) ([]*bc.Hash, []uint8) {
 154         merkleTree := buildMerkleTree(rawDatas)
 155         if merkleTree == nil {
 156                 return []*bc.Hash{}, []uint8{}
 157         }
 158         merkleHashSet := set.New()
 159         for _, data := range relatedRawDatas {
 160                 merkleHash := leafMerkleHash(data)
 161                 merkleHashSet.Add(merkleHash.String())
 162         }
 163         if merkleHashSet.Size() == 0 {
 164                 return []*bc.Hash{&merkleTree.hash}, []uint8{FlagAssist}
 165         }
 166         return merkleTree.getMerkleTreeProof(merkleHashSet)
 167 }
 168
 169 func (node *merkleTreeNode) getMerkleTreeProofByFlags(flagList *list.List) []*bc.Hash {
 170         var hashes []*bc.Hash
 171
 172         if flagList.Len() == 0 {
 173                 return hashes
 174         }
 175         flagEle := flagList.Front()
 176         flag := flagEle.Value.(uint8)
 177         flagList.Remove(flagEle)
 178
 179         if flag == FlagTxLeaf || flag == FlagAssist {
 180                 hashes = append(hashes, &node.hash)
 181                 return hashes
 182         }
 183         if node.left != nil {
 184                 leftHashes := node.left.getMerkleTreeProofByFlags(flagList)
 185                 hashes = append(hashes, leftHashes...)
 186         }
 187         if node.right != nil {
 188                 rightHashes := node.right.getMerkleTreeProofByFlags(flagList)
 189                 hashes = append(hashes, rightHashes...)
 190         }
 191         return hashes
 192 }
 193
 194 func getMerkleTreeProofByFlags(rawDatas []merkleNode, flagList *list.List) []*bc.Hash {
 195         tree := buildMerkleTree(rawDatas)
 196         return tree.getMerkleTreeProofByFlags(flagList)
 197 }
 198
 199 // GetTxMerkleTreeProof return a proof of merkle tree, which used to proof the transaction does
 200 // exist in the merkle tree
 201 func GetTxMerkleTreeProof(txs []*Tx, relatedTxs []*Tx) ([]*bc.Hash, []uint8) {
 202         var rawDatas []merkleNode
 203         var relatedRawDatas []merkleNode
 204         for _, tx := range txs {
 205                 rawDatas = append(rawDatas, &tx.ID)
 206         }
 207         for _, relatedTx := range relatedTxs {
 208                 relatedRawDatas = append(relatedRawDatas, &relatedTx.ID)
 209         }
 210         return getMerkleTreeProof(rawDatas, relatedRawDatas)
 211 }
 212
 213 // getMerkleRootByProof caculate the merkle root hash according to the proof
 214 func getMerkleRootByProof(hashList *list.List, flagList *list.List, merkleHashes *list.List) bc.Hash {
 215         if flagList.Len() == 0 || hashList.Len() == 0 {
 216                 return bc.EmptyStringHash
 217         }
 218         flagEle := flagList.Front()
 219         flag := flagEle.Value.(uint8)
 220         flagList.Remove(flagEle)
 221         switch flag {
 222         case FlagAssist:
 223                 {
 224                         hash := hashList.Front()
 225                         hashList.Remove(hash)
 226                         return hash.Value.(bc.Hash)
 227                 }
 228         case FlagTxLeaf:
 229                 {
 230                         if merkleHashes.Len() == 0 {
 231                                 return bc.EmptyStringHash
 232                         }
 233                         hashEle := hashList.Front()
 234                         hash := hashEle.Value.(bc.Hash)
 235                         relatedHashEle := merkleHashes.Front()
 236                         relatedHash := relatedHashEle.Value.(bc.Hash)
 237                         if hash == relatedHash {
 238                                 hashList.Remove(hashEle)
 239                                 merkleHashes.Remove(relatedHashEle)
 240                                 return hash
 241                         }
 242                 }
 243         case FlagTxParent:
 244                 {
 245                         leftHash := getMerkleRootByProof(hashList, flagList, merkleHashes)
 246                         rightHash := getMerkleRootByProof(hashList, flagList, merkleHashes)
 247                         hash := interiorMerkleHash(&leftHash, &rightHash)
 248                         return hash
 249                 }
 250         }
 251         return bc.EmptyStringHash
 252 }
 253
 254 func newMerkleTreeNode(merkleHash bc.Hash, left *merkleTreeNode, right *merkleTreeNode) *merkleTreeNode {
 255         return &merkleTreeNode{
 256                 hash:  merkleHash,
 257                 left:  left,
 258                 right: right,
 259         }
 260 }
 261
 262 // ValidateMerkleTreeProof caculate the merkle root according to the hash of node and the flags
 263 // only if the merkle root by caculated equals to the specify merkle root, and the merkle tree
 264 // contains all of the related raw datas, the validate result will be true.
 265 func validateMerkleTreeProof(hashes []*bc.Hash, flags []uint8, relatedNodes []merkleNode, merkleRoot bc.Hash) bool {
 266         merkleHashes := list.New()
 267         for _, relatedNode := range relatedNodes {
 268                 merkleHashes.PushBack(leafMerkleHash(relatedNode))
 269         }
 270         hashList := list.New()
 271         for _, hash := range hashes {
 272                 hashList.PushBack(*hash)
 273         }
 274         flagList := list.New()
 275         for _, flag := range flags {
 276                 flagList.PushBack(flag)
 277         }
 278         root := getMerkleRootByProof(hashList, flagList, merkleHashes)
 279         return root == merkleRoot && merkleHashes.Len() == 0
 280 }
 281
 282 // ValidateTxMerkleTreeProof validate the merkle tree of transactions
 283 func ValidateTxMerkleTreeProof(hashes []*bc.Hash, flags []uint8, relatedHashes []*bc.Hash, merkleRoot bc.Hash) bool {
 284         var relatedNodes []merkleNode
 285         for _, hash := range relatedHashes {
 286                 relatedNodes = append(relatedNodes, hash)
 287         }
 288         return validateMerkleTreeProof(hashes, flags, relatedNodes, merkleRoot)
 289 }
 290
 291 // TxMerkleRoot creates a merkle tree from a slice of transactions
 292 // and returns the root hash of the tree.
 293 func TxMerkleRoot(transactions []*bc.Tx) (root bc.Hash, err error) {
 294         nodes := []merkleNode{}
 295         for _, tx := range transactions {
 296                 nodes = append(nodes, &tx.ID)
 297         }
 298         return merkleRoot(nodes)
 299 }
 300
 301 // prevPowerOfTwo returns the largest power of two that is smaller than a given number.
 302 // In other words, for some input n, the prevPowerOfTwo k is a power of two such that
 303 // k < n <= 2k. This is a helper function used during the calculation of a merkle tree.
 304 func prevPowerOfTwo(n int) int {
 305         // If the number is a power of two, divide it by 2 and return.
 306         if n&(n-1) == 0 {
 307                 return n / 2
 308         }
 309
 310         // Otherwise, find the previous PoT.
 311         exponent := uint(math.Log2(float64(n)))
 312         return 1 << exponent // 2^exponent
 313 }