implement XOF

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 {
+// 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(
-			&o.input_chaining_value,
-			&o.block_words,
-			output_block_counter,
-			o.block_len,
-			o.flags|ROOT,
+				&or.o.input_chaining_value,
+				&or.o.block_words,
+				or.blocks_output,
+				or.o.block_len,
+				or.o.flags|ROOT,
 			)
-		var wordsBytes [16 * 4]byte
 			for i, w := range words {
-			binary.LittleEndian.PutUint32(wordsBytes[i*4:], w)
+				binary.LittleEndian.PutUint32(or.block[i*4:], w)
 			}
-		n := copy(out_slice, wordsBytes[:])
-		out_slice = out_slice[n:]
-		output_block_counter++
+			or.remaining = BLOCK_LEN
+			or.blocks_output++
 		}
+
+		// 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

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))
+}