github.com/avahowell/sia@v0.5.1-beta.0.20160524050156-83dcc3d37c94/crypto/merkle.go (about)

     1  package crypto
     2  
     3  import (
     4  	"bytes"
     5  
     6  	"github.com/NebulousLabs/Sia/encoding"
     7  
     8  	"github.com/NebulousLabs/merkletree"
     9  )
    10  
    11  const (
    12  	// SegmentSize is the chunk size that is used when taking the Merkle root
    13  	// of a file. 64 is chosen because bandwidth is scarce and it optimizes for
    14  	// the smallest possible storage proofs. Using a larger base, even 256
    15  	// bytes, would result in substantially faster hashing, but the bandwidth
    16  	// tradeoff was deemed to be more important, as blockchain space is scarce.
    17  	SegmentSize = 64
    18  )
    19  
    20  // MerkleTree wraps merkletree.Tree, changing some of the function definitions
    21  // to assume sia-specific constants and return sia-specific types.
    22  type MerkleTree struct {
    23  	merkletree.Tree
    24  }
    25  
    26  // NewTree returns a MerkleTree, which can be used for getting Merkle roots and
    27  // Merkle proofs on data. See merkletree.Tree for more details.
    28  func NewTree() *MerkleTree {
    29  	return &MerkleTree{*merkletree.New(NewHash())}
    30  }
    31  
    32  // PushObject encodes and adds the hash of the encoded object to the tree as a
    33  // leaf.
    34  func (t *MerkleTree) PushObject(obj interface{}) {
    35  	t.Push(encoding.Marshal(obj))
    36  }
    37  
    38  // Root is a redefinition of merkletree.Tree.Root, returning a Hash instead of
    39  // a []byte.
    40  func (t *MerkleTree) Root() (h Hash) {
    41  	copy(h[:], t.Tree.Root())
    42  	return
    43  }
    44  
    45  // CachedMerkleTree wraps merkletree.CachedTree, changing some of the function
    46  // definitions to assume sia-specific constants and return sia-specific types.
    47  type CachedMerkleTree struct {
    48  	merkletree.CachedTree
    49  }
    50  
    51  // NewCachedTree returns a CachedMerkleTree, which can be used for getting
    52  // Merkle roots and proofs from data that has cached subroots. See
    53  // merkletree.CachedTree for more details.
    54  func NewCachedTree(height uint64) *CachedMerkleTree {
    55  	return &CachedMerkleTree{*merkletree.NewCachedTree(NewHash(), height)}
    56  }
    57  
    58  // Prove is a redefinition of merkletree.CachedTree.Prove, so that Sia-specific
    59  // types are used instead of the generic types used by the parent package. The
    60  // base is not a return value because the base is used as input.
    61  func (ct *CachedMerkleTree) Prove(base []byte, cachedHashSet []Hash) []Hash {
    62  	// Turn the input in to a proof set that will be recognized by the high
    63  	// level tree.
    64  	cachedProofSet := make([][]byte, len(cachedHashSet)+1)
    65  	cachedProofSet[0] = base
    66  	for i := range cachedHashSet {
    67  		cachedProofSet[i+1] = cachedHashSet[i][:]
    68  	}
    69  	_, proofSet, _, _ := ct.CachedTree.Prove(cachedProofSet)
    70  
    71  	// convert proofSet to base and hashSet
    72  	hashSet := make([]Hash, len(proofSet)-1)
    73  	for i, proof := range proofSet[1:] {
    74  		copy(hashSet[i][:], proof)
    75  	}
    76  	return hashSet
    77  }
    78  
    79  // Push is a redefinition of merkletree.CachedTree.Push, with the added type
    80  // safety of only accepting a hash.
    81  func (ct *CachedMerkleTree) Push(h Hash) {
    82  	ct.CachedTree.Push(h[:])
    83  }
    84  
    85  // Root is a redefinition of merkletree.CachedTree.Root, returning a Hash
    86  // instead of a []byte.
    87  func (ct *CachedMerkleTree) Root() (h Hash) {
    88  	copy(h[:], ct.CachedTree.Root())
    89  	return
    90  }
    91  
    92  // CalculateLeaves calculates the number of leaves that would be pushed from
    93  // data of size 'dataSize'.
    94  func CalculateLeaves(dataSize uint64) uint64 {
    95  	numSegments := dataSize / SegmentSize
    96  	if dataSize == 0 || dataSize%SegmentSize != 0 {
    97  		numSegments++
    98  	}
    99  	return numSegments
   100  }
   101  
   102  // MerkleRoot returns the Merkle root of the input data.
   103  func MerkleRoot(b []byte) Hash {
   104  	t := NewTree()
   105  	buf := bytes.NewBuffer(b)
   106  	for buf.Len() > 0 {
   107  		t.Push(buf.Next(SegmentSize))
   108  	}
   109  	return t.Root()
   110  }
   111  
   112  // MerkleProof builds a Merkle proof that the data at segment 'proofIndex' is a
   113  // part of the Merkle root formed by 'b'.
   114  func MerkleProof(b []byte, proofIndex uint64) (base []byte, hashSet []Hash) {
   115  	// Create the tree.
   116  	t := NewTree()
   117  	t.SetIndex(proofIndex)
   118  
   119  	// Fill the tree.
   120  	buf := bytes.NewBuffer(b)
   121  	for buf.Len() > 0 {
   122  		t.Push(buf.Next(SegmentSize))
   123  	}
   124  
   125  	// Get the proof and convert it to a base + hash set.
   126  	_, proof, _, _ := t.Prove()
   127  	base = proof[0]
   128  	hashSet = make([]Hash, len(proof)-1)
   129  	for i, p := range proof[1:] {
   130  		copy(hashSet[i][:], p)
   131  	}
   132  	return base, hashSet
   133  }
   134  
   135  // VerifySegment will verify that a segment, given the proof, is a part of a
   136  // Merkle root.
   137  func VerifySegment(base []byte, hashSet []Hash, numSegments, proofIndex uint64, root Hash) bool {
   138  	// convert base and hashSet to proofSet
   139  	proofSet := make([][]byte, len(hashSet)+1)
   140  	proofSet[0] = base
   141  	for i := range hashSet {
   142  		proofSet[i+1] = hashSet[i][:]
   143  	}
   144  	return merkletree.VerifyProof(NewHash(), root[:], proofSet, proofIndex, numSegments)
   145  }