1
0
Fork 0
mirror of https://github.com/vbatts/merkle.git synced 2024-12-12 06:58:04 +00:00
merkle/stream.go

202 lines
5.7 KiB
Go
Raw Normal View History

package merkle
2015-03-25 21:28:17 +00:00
import (
"hash"
"log"
2015-03-30 19:02:53 +00:00
"runtime"
2015-03-25 21:28:17 +00:00
)
// NewHash provides a hash.Hash to generate a merkle.Tree checksum, given a
// HashMaker for the checksums of the blocks written and the blockSize of each
// block per node in the tree.
2015-03-30 19:19:46 +00:00
func NewHash(hm HashMaker, merkleBlockLength int) HashTreeer {
2015-03-31 14:55:25 +00:00
return newMerkleHash(hm, merkleBlockLength)
}
func newMerkleHash(hm HashMaker, merkleBlockLength int) *merkleHash {
mh := new(merkleHash)
2015-03-25 21:28:17 +00:00
mh.blockSize = merkleBlockLength
mh.hm = hm
2015-03-25 21:28:17 +00:00
mh.tree = &Tree{Nodes: []*Node{}, BlockLength: merkleBlockLength}
2015-03-26 20:46:02 +00:00
mh.lastBlock = make([]byte, merkleBlockLength)
return mh
}
2015-03-30 19:19:46 +00:00
// Treeer (Tree-er) provides access to the Merkle tree internals
type Treeer interface {
Nodes() []*Node
Root() *Node
}
// HashTreeer can be used as a hash.Hash but also provide access to the Merkle
// tree internals
type HashTreeer interface {
hash.Hash
Treeer
}
// TODO make a similar hash.Hash, that accepts an argument of a merkle.Tree,
// that will validate nodes as the new bytes are written. If a new written
// block fails checksum, then return an error on the io.Writer
type merkleHash struct {
2015-03-26 22:24:34 +00:00
blockSize int
tree *Tree
hm HashMaker
lastBlock []byte // as needed, for Sum()
lastBlockLen int
partialLastNode bool // true when Sum() has appended a Node for a partial block
}
2015-03-31 14:55:25 +00:00
func (mh *merkleHash) Reset() {
mh1 := newMerkleHash(mh.hm, mh.blockSize)
*mh = *mh1
}
2015-03-30 19:19:46 +00:00
func (mh merkleHash) Nodes() []*Node {
return mh.tree.Nodes
}
func (mh merkleHash) Root() *Node {
return mh.tree.Root()
}
// XXX this will be tricky, as the last block can be less than the BlockSize.
// if they get the sum, it will be mh.tree.Root().Checksum() at that point.
//
// But if they continue writing, it would mean a continuation of the bytes in
// the last block. So popping the last node, and having a buffer for the bytes
// in that last partial block.
//
// if that last block was complete, then no worries. start the next node.
func (mh *merkleHash) Sum(b []byte) []byte {
2015-03-30 19:02:53 +00:00
var (
2015-03-30 19:20:50 +00:00
curBlock = []byte{}
offset int
2015-03-30 19:02:53 +00:00
)
if mh.partialLastNode {
// if this is true, then we need to pop the last node
mh.tree.Nodes = mh.tree.Nodes[:len(mh.tree.Nodes)-1]
mh.partialLastNode = false
2015-03-26 22:24:34 +00:00
}
2015-03-30 19:02:53 +00:00
if mh.lastBlockLen > 0 {
curBlock = append(curBlock[:], mh.lastBlock[:mh.lastBlockLen]...)
mh.lastBlockLen = 0
}
if b != nil && len(b) > 0 {
curBlock = append(curBlock, b...)
}
2015-03-31 14:55:25 +00:00
// incase we're at a new or reset state
if len(mh.tree.Nodes) == 0 && len(curBlock) == 0 {
return nil
}
2015-03-30 19:02:53 +00:00
for i := 0; i < len(curBlock)/mh.blockSize; i++ {
n, err := NewNodeHashBlock(mh.hm, curBlock[offset:(offset+mh.blockSize)])
if err != nil {
// XXX i hate to swallow an error here, but the `Sum() []byte` signature
// :-\
sBuf := make([]byte, 1024)
runtime.Stack(sBuf, false)
log.Printf("[ERROR]: %s %q", err, string(sBuf))
return nil
}
mh.tree.Nodes = append(mh.tree.Nodes, n)
offset = offset + mh.blockSize
}
// If there is remainder, we'll need to make a partial node and stash it
if m := (len(curBlock) % mh.blockSize); m != 0 {
mh.lastBlockLen = copy(mh.lastBlock, curBlock[offset:])
n, err := NewNodeHashBlock(mh.hm, curBlock[offset:])
if err != nil {
sBuf := make([]byte, 1024)
runtime.Stack(sBuf, false)
log.Printf("[ERROR]: %s %q", err, string(sBuf))
return nil
}
mh.tree.Nodes = append(mh.tree.Nodes, n)
mh.partialLastNode = true
2015-03-26 22:24:34 +00:00
}
2015-03-25 21:28:17 +00:00
sum, err := mh.tree.Root().Checksum()
if err != nil {
2015-03-30 19:02:53 +00:00
// XXX i hate to swallow an error here, but the `Sum() []byte` signature
// :-\
sBuf := make([]byte, 1024)
runtime.Stack(sBuf, false)
log.Printf("[ERROR]: %s %q", err, string(sBuf))
return nil
2015-03-25 21:28:17 +00:00
}
return sum
}
2015-03-25 21:28:17 +00:00
func (mh *merkleHash) Write(b []byte) (int, error) {
2015-03-25 21:28:17 +00:00
// basically we need to:
2015-03-26 20:46:02 +00:00
// * include prior partial lastBlock, if any
2015-03-25 21:28:17 +00:00
// * chunk these writes into blockSize
// * create Node of the sum
// * add the Node to the tree
// * stash remainder in the mh.lastBlock
var (
2015-03-30 19:20:50 +00:00
curBlock = make([]byte, mh.blockSize)
numBytes int
2015-03-25 21:28:17 +00:00
numWritten int
2015-03-30 19:20:50 +00:00
offset int
2015-03-25 21:28:17 +00:00
)
2015-03-26 20:46:02 +00:00
if mh.lastBlock != nil && mh.lastBlockLen > 0 {
if (mh.lastBlockLen + len(b)) < mh.blockSize {
mh.lastBlockLen += copy(mh.lastBlock[mh.lastBlockLen:], b[:])
return len(b), nil
}
2015-03-26 22:24:34 +00:00
// XXX off by one?
numBytes = copy(curBlock[:], mh.lastBlock[:mh.lastBlockLen])
2015-03-25 21:28:17 +00:00
// not adding to numWritten, since these blocks were accounted for in a
// prior Write()
2015-03-26 20:46:02 +00:00
// then we'll chunk the front of the incoming bytes
end := mh.blockSize - numBytes
if end > len(b) {
end = len(b)
}
offset = copy(curBlock[numBytes:], b[:end])
2015-03-26 20:46:02 +00:00
n, err := NewNodeHashBlock(mh.hm, curBlock)
if err != nil {
2015-03-26 22:24:34 +00:00
// XXX might need to stash again the prior lastBlock and first little chunk
2015-03-26 20:46:02 +00:00
return numWritten, err
}
mh.tree.Nodes = append(mh.tree.Nodes, n)
numWritten += offset
2015-03-25 21:28:17 +00:00
}
2015-03-26 20:46:02 +00:00
numBytes = (len(b) - offset)
2015-03-25 21:28:17 +00:00
for i := 0; i < numBytes/mh.blockSize; i++ {
2015-03-26 22:24:34 +00:00
//fmt.Printf("%s", b[offset:offset+mh.blockSize])
numWritten += copy(curBlock, b[offset:offset+mh.blockSize])
n, err := NewNodeHashBlock(mh.hm, curBlock)
if err != nil {
// XXX might need to stash again the prior lastBlock and first little chunk
return numWritten, err
}
mh.tree.Nodes = append(mh.tree.Nodes, n)
offset = offset + mh.blockSize
2015-03-25 21:28:17 +00:00
}
2015-03-26 22:24:34 +00:00
mh.lastBlockLen = numBytes % mh.blockSize
// XXX off by one?
numWritten += copy(mh.lastBlock[:], b[(len(b)-mh.lastBlockLen):])
2015-03-25 21:28:17 +00:00
return numWritten, nil
}
2015-03-25 21:28:17 +00:00
// likely not the best to pass this through and not use our own node block
// size, but let's revisit this.
func (mh *merkleHash) BlockSize() int { return mh.hm().BlockSize() }
func (mh *merkleHash) Size() int { return mh.hm().Size() }