From 4833fabd0e432a2d66db573f31af9ea5aff9a9ad Mon Sep 17 00:00:00 2001
From: lukechampine <luke.champine@gmail.com>
Date: Fri, 10 Jan 2020 14:39:34 -0500
Subject: [PATCH] fix SeekEnd behavior

---
 blake3.go      | 67 +++++++++++++++++++++++++++-----------------------
 blake3_test.go | 13 +++++++++-
 2 files changed, 48 insertions(+), 32 deletions(-)

diff --git a/blake3.go b/blake3.go
index 094582d..4f08a60 100644
--- a/blake3.go
+++ b/blake3.go
@@ -9,6 +9,7 @@ import (
 	"errors"
 	"hash"
 	"io"
+	"math"
 	"math/bits"
 )
 
@@ -136,64 +137,68 @@ func wordsToBytes(words []uint32, bytes []byte) {
 	}
 }
 
-// An OutputReader produces an seekable stream of output. Up to 2^64 - 1 bytes
-// can be safely read from the stream.
+// An OutputReader produces an seekable stream of 2^64 - 1 output bytes.
 type OutputReader struct {
-	n      node
-	block  [blockLen]byte
-	unread int
+	n     node
+	block [blockLen]byte
+	off   uint64
 }
 
-// Read implements io.Reader. It always return len(p), nil.
+// Read implements io.Reader. Callers may assume that Read returns len(p), nil
+// unless the read would extend beyond the end of the stream.
 func (or *OutputReader) Read(p []byte) (int, error) {
+	if or.off == math.MaxUint64 {
+		return 0, io.EOF
+	} else if rem := math.MaxUint64 - or.off; uint64(len(p)) > rem {
+		p = p[:rem]
+	}
 	lenp := len(p)
 	for len(p) > 0 {
-		if or.unread == 0 {
+		if or.off%blockLen == 0 {
+			or.n.counter = or.off / blockLen
 			words := or.n.compress()
 			wordsToBytes(words[:], or.block[:])
-			or.unread = blockLen
-			or.n.counter++
 		}
 
-		// copy from output buffer
-		n := copy(p, or.block[blockLen-or.unread:])
-		or.unread -= n
+		n := copy(p, or.block[or.off%blockLen:])
 		p = p[n:]
+		or.off += uint64(n)
 	}
 	return lenp, nil
 }
 
-// Seek implements io.Seeker. SeekEnd is defined as 2^64 - 1 bytes, the maximum
-// safe output of a BLAKE3 stream.
+// Seek implements io.Seeker.
 func (or *OutputReader) Seek(offset int64, whence int) (int64, error) {
-	off := int64(or.n.counter*blockLen) + int64(blockLen-or.unread)
+	off := or.off
 	switch whence {
 	case io.SeekStart:
-		off = offset
+		if offset < 0 {
+			return 0, errors.New("seek position cannot be negative")
+		}
+		off = uint64(offset)
 	case io.SeekCurrent:
-		off += offset
+		if offset < 0 {
+			if uint64(-offset) > off {
+				return 0, errors.New("seek position cannot be negative")
+			}
+			off -= uint64(-offset)
+		} else {
+			off += uint64(offset)
+		}
 	case io.SeekEnd:
-		// BLAKE3 can safely output up to 2^64 - 1 bytes. Seeking to the "end"
-		// of this stream is kind of strange, but perhaps could be useful for
-		// testing overflow scenarios.
-		off = int64(^uint64(0) - uint64(offset))
+		off = uint64(offset) - 1
 	default:
 		panic("invalid whence")
 	}
-	if off < 0 {
-		return 0, errors.New("seek position cannot be negative")
-	}
+	or.off = off
 	or.n.counter = uint64(off) / blockLen
-	or.unread = blockLen - (int(off) % blockLen)
-
-	// If the new offset is not a block boundary, generate the block we are
-	// "inside."
-	if or.unread != 0 {
+	if or.off%blockLen != 0 {
 		words := or.n.compress()
 		wordsToBytes(words[:], or.block[:])
 	}
-
-	return off, nil
+	// NOTE: or.off >= 2^63 will result in a negative return value.
+	// Nothing we can do about this.
+	return int64(or.off), nil
 }
 
 type chunkState struct {
diff --git a/blake3_test.go b/blake3_test.go
index 269f2e0..f31d737 100644
--- a/blake3_test.go
+++ b/blake3_test.go
@@ -104,9 +104,20 @@ func TestXOF(t *testing.T) {
 			if !bytes.Equal(outRead[:n], xofRead[:n]) {
 				t.Errorf("XOF output did not match test vector at offset %v:\n\texpected: %x...\n\t     got: %x...", offset, outRead[:10], xofRead[:10])
 			}
-
 		}
 	}
+	// test behavior at end of stream
+	xof := blake3.New(0, nil).XOF()
+	buf := make([]byte, 1024)
+	xof.Seek(-1000, io.SeekEnd)
+	n, err := xof.Read(buf)
+	if n != 1000 || err != nil {
+		t.Errorf("expected (1000, nil) when reading near end of stream, got (%v, %v)", n, err)
+	}
+	n, err = xof.Read(buf)
+	if n != 0 || err != io.EOF {
+		t.Errorf("expected (0, EOF) when reading past end of stream, got (%v, %v)", n, err)
+	}
 }
 
 type nopReader struct{}