mirror of
https://github.com/PeernetOfficial/core.git
synced 2026-07-17 02:47:51 +01:00
Rename fragment to merkle #48
This commit is contained in:
Kleissner
306
merkle/Merkle Tree.go
Normal file
306
merkle/Merkle Tree.go
Normal file
@@ -0,0 +1,306 @@
|
||||
/*
|
||||
File Name: Merkle Tree.go
|
||||
Copyright: 2021 Peernet s.r.o.
|
||||
Author: Peter Kleissner
|
||||
|
||||
Generates the merkle tree based on input data.
|
||||
In case of uneven number of fragments, the last uneven fragment is moved up a level.
|
||||
*/
|
||||
|
||||
package merkle
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"io"
|
||||
|
||||
"lukechampine.com/blake3"
|
||||
)
|
||||
|
||||
// MerkleTree represents an entire merkle tree
|
||||
type MerkleTree struct {
|
||||
// information about the original file
|
||||
FileSize uint64
|
||||
FragmentSize uint64
|
||||
FragmentCount uint64
|
||||
|
||||
// list of hashes
|
||||
FragmentHashes [][]byte // List of hashes for each fragment
|
||||
MiddleHashes [][][]byte // All hashes in the middle, bottom up.
|
||||
RootHash []byte // Root hash.
|
||||
}
|
||||
|
||||
// NewMerkleTree creates a new merkle tree from the input
|
||||
func NewMerkleTree(fileSize, fragmentSize uint64, reader io.Reader) (tree *MerkleTree, err error) {
|
||||
if fragmentSize == 0 {
|
||||
return nil, errors.New("invalid fragment size")
|
||||
}
|
||||
|
||||
tree = &MerkleTree{
|
||||
FileSize: fileSize,
|
||||
FragmentSize: fragmentSize,
|
||||
FragmentCount: fileSizeToFragmentCount(fileSize, fragmentSize),
|
||||
}
|
||||
|
||||
// Special case: No fragments, in case of empty data.
|
||||
if tree.FragmentCount == 0 {
|
||||
hash := blake3.Sum256(nil)
|
||||
tree.RootHash = hash[:]
|
||||
|
||||
return tree, nil
|
||||
} else if tree.FragmentCount == 1 {
|
||||
// Special case: Single fragment.
|
||||
data := make([]byte, fileSize)
|
||||
if _, err := io.ReadAtLeast(reader, data, int(fileSize)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
hash := blake3.Sum256(data)
|
||||
tree.RootHash = hash[:]
|
||||
|
||||
return tree, nil
|
||||
}
|
||||
|
||||
// calculate the hash per fragment
|
||||
data := make([]byte, fragmentSize)
|
||||
remaining := fileSize
|
||||
|
||||
for n := uint64(0); n < tree.FragmentCount; n++ {
|
||||
if fragmentSize > remaining {
|
||||
fragmentSize = remaining
|
||||
}
|
||||
|
||||
if _, err := io.ReadAtLeast(reader, data, int(fragmentSize)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// hash the fragment
|
||||
hash := blake3.Sum256(data[:fragmentSize])
|
||||
|
||||
tree.FragmentHashes = append(tree.FragmentHashes, hash[:])
|
||||
|
||||
remaining -= fragmentSize
|
||||
}
|
||||
|
||||
// calculate the intermediate hashes
|
||||
tree.calculateMiddleHashes(0)
|
||||
|
||||
return tree, nil
|
||||
}
|
||||
|
||||
func fileSizeToFragmentCount(fileSize, fragmentSize uint64) (count uint64) {
|
||||
return (fileSize + fragmentSize - 1) / fragmentSize
|
||||
}
|
||||
|
||||
func (tree *MerkleTree) calculateMiddleHashes(level uint64) {
|
||||
if len(tree.FragmentHashes) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
var newHashes, inputHashes [][]byte
|
||||
|
||||
if level == 0 {
|
||||
inputHashes = tree.FragmentHashes
|
||||
} else {
|
||||
inputHashes = tree.MiddleHashes[level-1]
|
||||
}
|
||||
|
||||
for n := 0; n+1 <= len(inputHashes)-1; n += 2 {
|
||||
newHashes = append(newHashes, calculateMiddleHash(inputHashes[n], inputHashes[n+1]))
|
||||
}
|
||||
|
||||
// Uneven leafs? in this case the new hash is just a copy of the uneven one. No point in artifically recalcualting it with itself like Bitcoin does.
|
||||
// For other possible implementations see https://medium.com/coinmonks/merkle-trees-concepts-and-use-cases-5da873702318.
|
||||
if len(inputHashes)%2 != 0 {
|
||||
newHashes = append(newHashes, inputHashes[len(inputHashes)-1])
|
||||
}
|
||||
|
||||
if len(newHashes) == 1 {
|
||||
// Only one hash generated.
|
||||
tree.RootHash = newHashes[0]
|
||||
} else if len(newHashes) > 1 {
|
||||
tree.MiddleHashes = append(tree.MiddleHashes, newHashes)
|
||||
|
||||
tree.calculateMiddleHashes(level + 1)
|
||||
}
|
||||
}
|
||||
|
||||
func calculateMiddleHash(hash1 []byte, hash2 []byte) (newHash []byte) {
|
||||
var data []byte
|
||||
data = append(data, hash1...)
|
||||
data = append(data, hash2...)
|
||||
|
||||
hash := blake3.Sum256(data)
|
||||
|
||||
return hash[:]
|
||||
}
|
||||
|
||||
// CreateVerification returns the verification hashes for the given fragment number. The root hash itself is not included.
|
||||
// The result might be empty if there is no or a single fragment.
|
||||
// Each verification hash has a preceding left (= 0)/right (= 1) indicator that indicates where the verification is positioned.
|
||||
// This makes the algorithm future proof, in case uneven leafs will be handled differently.
|
||||
func (tree *MerkleTree) CreateVerification(fragment uint64) (verificationHashes [][]byte) {
|
||||
// 0 fragments: Empty data.
|
||||
// 1 fragment: The hash of the fragment is the root hash.
|
||||
if tree.FragmentCount <= 1 {
|
||||
return nil
|
||||
} else if fragment >= tree.FragmentCount {
|
||||
// invalid fragment index
|
||||
return nil
|
||||
}
|
||||
|
||||
// first hash it he neighbor fragment hash, if available
|
||||
if fragment == tree.FragmentCount-1 && fragment%2 == 0 {
|
||||
} else if fragment%2 == 0 {
|
||||
verificationHashes = append(verificationHashes, append([]byte{1}, tree.FragmentHashes[fragment+1]...))
|
||||
} else {
|
||||
verificationHashes = append(verificationHashes, append([]byte{0}, tree.FragmentHashes[fragment-1]...))
|
||||
}
|
||||
|
||||
// go through all middle hash levels
|
||||
for n := 0; n < len(tree.MiddleHashes); n++ {
|
||||
fragment = fragment / 2
|
||||
|
||||
if fragment == uint64(len(tree.MiddleHashes[n])-1) && fragment%2 == 0 {
|
||||
} else if fragment%2 == 0 {
|
||||
verificationHashes = append(verificationHashes, append([]byte{1}, tree.MiddleHashes[n][fragment+1]...))
|
||||
} else {
|
||||
verificationHashes = append(verificationHashes, append([]byte{0}, tree.MiddleHashes[n][fragment-1]...))
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// MerkleVerify validates the hashed data against the verification hashes and the known root hash.
|
||||
func MerkleVerify(rootHash []byte, dataHash []byte, verificationHashes [][]byte) (valid bool) {
|
||||
for _, verifyHash := range verificationHashes {
|
||||
if verifyHash[0] == 0 {
|
||||
dataHash = calculateMiddleHash(verifyHash[1:], dataHash)
|
||||
} else {
|
||||
dataHash = calculateMiddleHash(dataHash, verifyHash[1:])
|
||||
}
|
||||
}
|
||||
|
||||
return bytes.Equal(rootHash, dataHash)
|
||||
}
|
||||
|
||||
/*
|
||||
Export/Import of the merkle tree structure:
|
||||
|
||||
Offset Size Info
|
||||
0 8 File Size
|
||||
8 8 Fragment Size
|
||||
16 32 Merkle Root Hash
|
||||
48 32 * n Fragment Hashes
|
||||
? 32 * n Middle Hashes
|
||||
|
||||
*/
|
||||
|
||||
const merkleTreeFileHeaderSize = 8 + 8 + 32
|
||||
|
||||
// calculateTotalHashCount returns the total number of fragment and middle hashes needed for the given count of fragments
|
||||
func calculateTotalHashCount(fragmentCount uint64) (count uint64) {
|
||||
// Special case no or 1 fragment: None needed, since the fragment hash is directly stored as root hash.
|
||||
if fragmentCount <= 1 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// Equal count of fragment hashes needed
|
||||
count = fragmentCount
|
||||
|
||||
// Calculate middle hashes number
|
||||
for countHashesLast := fragmentCount; ; {
|
||||
countMiddleNew := (countHashesLast + 1) / 2 // round up
|
||||
if countMiddleNew <= 1 {
|
||||
break
|
||||
}
|
||||
|
||||
count += countMiddleNew
|
||||
countHashesLast = countMiddleNew
|
||||
}
|
||||
|
||||
return count
|
||||
}
|
||||
|
||||
// Export stores the tree as blob
|
||||
func (tree *MerkleTree) Export() (data []byte) {
|
||||
data = make([]byte, merkleTreeFileHeaderSize+calculateTotalHashCount(tree.FragmentCount)*32)
|
||||
|
||||
// header
|
||||
binary.LittleEndian.PutUint64(data[0:8], tree.FileSize)
|
||||
binary.LittleEndian.PutUint64(data[8:16], tree.FragmentSize)
|
||||
copy(data[16:16+32], tree.RootHash)
|
||||
|
||||
// fragment hashes
|
||||
offset := 48
|
||||
for _, hash := range tree.FragmentHashes {
|
||||
copy(data[offset:offset+32], hash)
|
||||
offset += 32
|
||||
}
|
||||
|
||||
// middle hashes
|
||||
for n := 0; n < len(tree.MiddleHashes); n++ {
|
||||
for _, hash := range tree.MiddleHashes[n] {
|
||||
copy(data[offset:offset+32], hash)
|
||||
offset += 32
|
||||
}
|
||||
}
|
||||
|
||||
return data[:offset]
|
||||
}
|
||||
|
||||
// Import reads the tree from the input data
|
||||
func ImportMerkleTree(data []byte) (tree *MerkleTree) {
|
||||
// Read the header. Enforce the minimum size.
|
||||
if len(data) < 8+8+32 {
|
||||
return nil
|
||||
}
|
||||
|
||||
tree = &MerkleTree{
|
||||
FileSize: binary.LittleEndian.Uint64(data[0:8]),
|
||||
FragmentSize: binary.LittleEndian.Uint64(data[8:16]),
|
||||
}
|
||||
tree.FragmentCount = fileSizeToFragmentCount(tree.FileSize, tree.FragmentSize)
|
||||
tree.RootHash = data[16 : 16+32]
|
||||
|
||||
if tree.FragmentCount <= 1 {
|
||||
return tree
|
||||
}
|
||||
|
||||
// verify size
|
||||
if uint64(len(data)) < merkleTreeFileHeaderSize+calculateTotalHashCount(tree.FragmentCount)*32 {
|
||||
return nil
|
||||
}
|
||||
|
||||
// fragment hashes
|
||||
offset := 48
|
||||
for n := 0; n < int(tree.FragmentCount); n++ {
|
||||
hash := data[offset : offset+32]
|
||||
tree.FragmentHashes = append(tree.FragmentHashes, hash)
|
||||
offset += 32
|
||||
}
|
||||
|
||||
// middle hashes
|
||||
n := tree.FragmentCount / 2
|
||||
if tree.FragmentCount > 2 && tree.FragmentCount%2 != 0 {
|
||||
n++
|
||||
}
|
||||
|
||||
for ; n > 1; n = n / 2 {
|
||||
var hashList [][]byte
|
||||
for m := uint64(0); m < n; m++ {
|
||||
hash := data[offset : offset+32]
|
||||
hashList = append(hashList, hash)
|
||||
offset += 32
|
||||
}
|
||||
|
||||
tree.MiddleHashes = append(tree.MiddleHashes, hashList)
|
||||
if len(hashList)%2 != 0 {
|
||||
n++
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
Reference in New Issue
Block a user