diff --git a/blake3.go b/blake3.go index bbba202..5f5c917 100644 --- a/blake3.go +++ b/blake3.go @@ -141,24 +141,40 @@ func (o *output) chaining_value() [8]uint32 { )) } -func (o *output) root_output_bytes(out_slice []byte) { - output_block_counter := uint64(0) - for len(out_slice) > 0 { - words := compress( - &o.input_chaining_value, - &o.block_words, - output_block_counter, - o.block_len, - o.flags|ROOT, - ) - var wordsBytes [16 * 4]byte - for i, w := range words { - binary.LittleEndian.PutUint32(wordsBytes[i*4:], w) +// An OutputReader produces an unbounded stream of output from its initial +// state. +type OutputReader struct { + o *output + block [BLOCK_LEN]byte + remaining int + blocks_output uint64 +} + +// Read implements io.Reader. Read always return len(p), nil. +func (or *OutputReader) Read(p []byte) (int, error) { + lenp := len(p) + for len(p) > 0 { + if or.remaining == 0 { + words := compress( + &or.o.input_chaining_value, + &or.o.block_words, + or.blocks_output, + or.o.block_len, + or.o.flags|ROOT, + ) + for i, w := range words { + binary.LittleEndian.PutUint32(or.block[i*4:], w) + } + or.remaining = BLOCK_LEN + or.blocks_output++ } - n := copy(out_slice, wordsBytes[:]) - out_slice = out_slice[n:] - output_block_counter++ + + // copy from output buffer + n := copy(p, or.block[BLOCK_LEN-or.remaining:]) + or.remaining -= n + p = p[n:] } + return lenp, nil } type chunkState struct { @@ -348,6 +364,13 @@ func (h *Hasher) Write(input []byte) (int, error) { // Sum implements hash.Hash. func (h *Hasher) Sum(out_slice []byte) []byte { + out := make([]byte, h.Size()) + h.XOF().Read(out) + return append(out_slice, out...) +} + +// XOF returns an OutputReader initialized with the current hash state. +func (h *Hasher) XOF() *OutputReader { // Starting with the output from the current chunk, compute all the // parent chaining values along the right edge of the tree, until we // have the root output. @@ -362,9 +385,9 @@ func (h *Hasher) Sum(out_slice []byte) []byte { h.flags, ) } - out := make([]byte, h.Size()) - output.root_output_bytes(out) - return append(out_slice, out...) + return &OutputReader{ + o: output, + } } // ensure that Hasher implements hash.Hash diff --git a/blake3_test.go b/blake3_test.go index 4d0f40b..343dbca 100644 --- a/blake3_test.go +++ b/blake3_test.go @@ -1,8 +1,10 @@ package blake3_test import ( + "bytes" "encoding/hex" "encoding/json" + "io" "io/ioutil" "testing" @@ -65,6 +67,14 @@ func TestVectors(t *testing.T) { if out := toHex(h.Sum(nil)); out != vec.DeriveKey { t.Errorf("output did not match test vector:\n\texpected: %v...\n\t got: %v...", vec.DeriveKey[:10], out[:10]) } + // XOF should produce identical results, even when outputting 7 bytes at a time + h = blake3.New(len(vec.Hash)/2, nil) + h.Write(in) + var xofBuf bytes.Buffer + io.CopyBuffer(&xofBuf, io.LimitReader(h.XOF(), int64(len(vec.Hash)/2)), make([]byte, 7)) + if out := toHex(xofBuf.Bytes()); out != vec.Hash { + t.Errorf("XOF output did not match test vector:\n\texpected: %v...\n\t got: %v...", vec.Hash[:10], out[:10]) + } } } @@ -77,12 +87,17 @@ func BenchmarkWrite(b *testing.B) { } } -func BenchmarkBlock(b *testing.B) { +func BenchmarkChunk(b *testing.B) { h := blake3.New(32, nil) - buf := make([]byte, h.BlockSize()) + buf := make([]byte, blake3.CHUNK_LEN) out := make([]byte, 32) for i := 0; i < b.N; i++ { h.Write(buf) h.Sum(out) } } + +func BenchmarkXOF(b *testing.B) { + b.SetBytes(1) + io.CopyN(ioutil.Discard, blake3.New(32, nil).XOF(), int64(b.N)) +}