Fix some embarrassing bugsv0.3.1

4 files changed, 108 insertions, 32 deletions

diff --git a/oae2.go b/oae2.go
index c1ee627..8e41061 100644
--- a/oae2.go
+++ b/oae2.go
@@ -31,6 +31,7 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"slices"
 )
 
 // ErrInvalid is returned from methods on Reader when the stream failed to
@@ -72,7 +73,7 @@ func (e *segmentEncrypter) firstSegment() bool {
 	for _, b := range n[:nonceSize-1] {
 		or |= b
 	}
-	return or != 0
+	return or == 0
 }
 
 func (e *segmentEncrypter) nextNonce(lastSegment bool) ([]byte, error) {
@@ -140,6 +141,9 @@ func NewWriter(w io.Writer, key []byte, segmentSize int, additionalData []byte)
 	if segmentSize <= 0 {
 		panic("oae2.NewWriter: segmentSize must be strictly greater than 0")
 	}
+	segmentSizeBytes := make([]byte, 8)
+	binary.LittleEndian.PutUint64(segmentSizeBytes, uint64(segmentSize))
+	additionalData = slices.Concat(additionalData, segmentSizeBytes)
 	return &Writer{
 		w:           w,
 		segmentSize: segmentSize,
@@ -293,6 +297,9 @@ func NewReader(r io.Reader, key []byte, segmentSize int, additionalData []byte)
 	if segmentSize <= 0 {
 		panic("oae2.NewReader: segmentSize must be strictly greater than 0")
 	}
+	segmentSizeBytes := make([]byte, 8)
+	binary.LittleEndian.PutUint64(segmentSizeBytes, uint64(segmentSize))
+	additionalData = slices.Concat(additionalData, segmentSizeBytes)
 	return &Reader{
 		r:           bufReader{r: r},
 		segmentSize: segmentSize,
@@ -404,7 +411,12 @@ func (r *Reader) WriteTo(w io.Writer) (int64, error) {
 
 func (r *Reader) encryptedToPlaintextSize(offset int64) int64 {
 	segments := (offset - saltSize) / int64(r.segmentSize+aeadOverhead)
-	if (offset-saltSize)%int64(r.segmentSize+aeadOverhead) != 0 {
+	segmentOffset := (offset - saltSize) % int64(r.segmentSize+aeadOverhead)
+	if segmentOffset > 0 {
+		if segmentOffset < aeadOverhead {
+			// The final block is too short
+			return -1
+		}
 		segments++
 	}
 	return offset - saltSize - segments*aeadOverhead
@@ -430,11 +442,34 @@ func (r *Reader) seek(offset int64) (int64, error) {
 	r.buf = r.buf[:0]
 	binary.LittleEndian.PutUint64(r.decrypter.nonce[:], uint64(segment))
 	clear(r.decrypter.nonce[8:])
-	if err := r.fillBuf(); err != nil && err != io.EOF {
-		return 0, err
+	if err := r.fillBuf(); err != nil {
+		if err != io.EOF {
+			return 0, err
+		}
+		// In case we seeked to EOF, we need to back up and validate
+		// the last segment to avoid truncation attacks.
+		plaintextOffset := r.encryptedToPlaintextSize(n)
+		prevSegment := plaintextOffset / int64(r.segmentSize)
+		segmentOffset := plaintextOffset % int64(r.segmentSize)
+		if segmentOffset == 0 && prevSegment > 0 {
+			prevSegment--
+		}
+		if _, err := r.r.seek(saltSize+prevSegment*int64(r.segmentSize+aeadOverhead), io.SeekStart); err != nil {
+			return 0, err
+		}
+		binary.LittleEndian.PutUint64(r.decrypter.nonce[:], uint64(prevSegment))
+		clear(r.decrypter.nonce[8:])
+		if err := r.fillBuf(); err != nil {
+			if err == io.EOF {
+				err = fmt.Errorf("%w", ErrInvalid)
+			}
+			return 0, err
+		}
+		r.buf = r.buf[:0]
 	}
 	r.nRead = min(int(segmentOffset), len(r.buf))
-	return r.encryptedToPlaintextSize(n) + int64(r.nRead), nil
+	segmentStart := r.encryptedToPlaintextSize(n)
+	return segmentStart + int64(r.nRead), nil
 }
 
 // Seek implements io.Seeker. If the underlying reader does not implement
@@ -452,35 +487,23 @@ func (r *Reader) Seek(offset int64, whence int) (int64, error) {
 	case io.SeekStart:
 		return r.seek(offset)
 	case io.SeekCurrent:
-		if r.decrypter.nonce[8]|r.decrypter.nonce[9]|r.decrypter.nonce[10] != 0 {
-			return 0, fmt.Errorf("oae2.Reader.Seek: internal error: current stream position would overflow int64")
-		}
-		currentSegment := binary.LittleEndian.Uint64(r.decrypter.nonce[:8])
-		if currentSegment > 0 {
-			currentSegment--
-		}
-		if currentSegment > math.MaxInt64/uint64(r.segmentSize) {
-			return 0, fmt.Errorf("oae2.Reader.Seek: internal error: current stream position would overflow int64")
+		readerOffset, err := r.r.seek(0, io.SeekCurrent)
+		if err != nil {
+			return 0, err
 		}
-		currentOffset := int64(currentSegment)*int64(r.segmentSize) + int64(r.nRead)
-		if currentOffset < 0 {
-			return 0, fmt.Errorf("oae2.Reader.Seek: internal error: current stream position would overflow int64")
+		segment := (readerOffset - saltSize) / int64(r.segmentSize+aeadOverhead)
+		if len(r.buf) >= r.segmentSize {
+			segment--
 		}
-		return r.seek(currentOffset + offset)
+		return r.seek(segment*int64(r.segmentSize) + int64(r.nRead) + offset)
 	case io.SeekEnd:
-		// Unfortunately we have to do an additional seek for SeekEnd
-		// in the general case to determine the total size of the
-		// underlying stream. Fortunately, the most common case for
-		// SeekEnd is with offset 0 and we can avoid the additional
-		// seek in that case.
 		encryptedSize, err := r.r.seek(0, io.SeekEnd)
 		if err != nil {
 			return 0, err
 		}
-		r.buf = r.buf[:0]
 		plaintextSize := r.encryptedToPlaintextSize(encryptedSize)
-		if offset == 0 {
-			return plaintextSize, nil
+		if plaintextSize < 0 {
+			return 0, fmt.Errorf("%w: invalid length", ErrInvalid)
 		}
 		return r.seek(plaintextSize + offset)
 	default:
diff --git a/oae2_test.go b/oae2_test.go
index 5ca26e6..0380245 100644
--- a/oae2_test.go
+++ b/oae2_test.go
@@ -78,10 +78,15 @@ func TestWrongKeyRead(t *testing.T) {
 	const (
 		key            = "asdf"
 		additionalData = "additional data"
+		segmentSize    = 4 * 1024 * 1024
 	)
 
 	encryptedBuf := new(bytes.Buffer)
-	if _, err := NewWriter(encryptedBuf, []byte(key), 4*1024*1024, []byte(additionalData)).Write([]byte("hello world")); err != nil {
+	w := NewWriter(encryptedBuf, []byte(key), segmentSize, []byte(additionalData))
+	if _, err := w.Write([]byte("hello world")); err != nil {
+		t.Fatalf("Failed to encrypt test data: %s", err)
+	}
+	if err := w.Close(); err != nil {
 		t.Fatalf("Failed to encrypt test data: %s", err)
 	}
 
@@ -89,12 +94,16 @@ func TestWrongKeyRead(t *testing.T) {
 		desc           string
 		key            string
 		additionalData string
+		segmentSize    int
 	}{{
 		desc: "WrongKey",
 		key:  "wrong key",
 	}, {
 		desc:           "WrongData",
 		additionalData: "wrong additional data",
+	}, {
+		desc:        "WrongSegmentSize",
+		segmentSize: 8 * 1024 * 1024,
 	}} {
 		t.Run(tc.desc, func(t *testing.T) {
 			t.Parallel()
@@ -107,7 +116,11 @@ func TestWrongKeyRead(t *testing.T) {
 			if ad == "" {
 				ad = additionalData
 			}
-			_, err := io.ReadAll(NewReader(bytes.NewReader(encryptedBuf.Bytes()), []byte(key), 1, nil))
+			segSize := tc.segmentSize
+			if segSize == 0 {
+				segSize = segmentSize
+			}
+			_, err := io.ReadAll(NewReader(bytes.NewReader(encryptedBuf.Bytes()), []byte(k), segSize, []byte(ad)))
 			if !errors.Is(err, ErrInvalid) {
 				t.Errorf("Reader.Read: message returned error %v with key %q and data %q, want ErrInvalid", err, k, ad)
 			}
@@ -121,10 +134,15 @@ func TestWrongKeyWriteTo(t *testing.T) {
 	const (
 		key            = "asdf"
 		additionalData = "additional data"
+		segmentSize    = 4 * 1024 * 1024
 	)
 
 	encryptedBuf := new(bytes.Buffer)
-	if _, err := NewWriter(encryptedBuf, []byte(key), 4*1024*1024, []byte(additionalData)).Write([]byte("hello world")); err != nil {
+	w := NewWriter(encryptedBuf, []byte(key), segmentSize, []byte(additionalData))
+	if _, err := w.Write([]byte("hello world")); err != nil {
+		t.Fatalf("Failed to encrypt test data: %s", err)
+	}
+	if err := w.Close(); err != nil {
 		t.Fatalf("Failed to encrypt test data: %s", err)
 	}
 
@@ -132,12 +150,16 @@ func TestWrongKeyWriteTo(t *testing.T) {
 		desc           string
 		key            string
 		additionalData string
+		segmentSize    int
 	}{{
 		desc: "WrongKey",
 		key:  "wrong key",
 	}, {
 		desc:           "WrongData",
 		additionalData: "wrong additional data",
+	}, {
+		desc:        "WrongSegmentSize",
+		segmentSize: 8 * 1024 * 1024,
 	}} {
 		t.Run(tc.desc, func(t *testing.T) {
 			t.Parallel()
@@ -150,7 +172,11 @@ func TestWrongKeyWriteTo(t *testing.T) {
 			if ad == "" {
 				ad = additionalData
 			}
-			_, err := NewReader(bytes.NewReader(encryptedBuf.Bytes()), []byte(key), 1, nil).WriteTo(io.Discard)
+			segSize := tc.segmentSize
+			if segSize == 0 {
+				segSize = segmentSize
+			}
+			_, err := NewReader(bytes.NewReader(encryptedBuf.Bytes()), []byte(k), segSize, []byte(ad)).WriteTo(io.Discard)
 			if !errors.Is(err, ErrInvalid) {
 				t.Errorf("Reader.WriteTo: message returned error %v with key %q and data %q, want ErrInvalid", err, k, ad)
 			}
@@ -227,6 +253,25 @@ func TestReader_Seek(t *testing.T) {
 	}
 }
 
+func TestReader_Seek_TruncatedStream(t *testing.T) {
+	t.Parallel()
+
+	buf := new(bytes.Buffer)
+	w := NewWriter(buf, []byte("asdf"), 2, nil)
+	if _, err := io.WriteString(w, "Hello World!"); err != nil {
+		t.Fatalf("Failed to write test string: %s", err)
+	}
+	if err := w.Close(); err != nil {
+		t.Fatalf("Failed to write test string: %s", err)
+	}
+	truncated := buf.Bytes()[:53]
+	r := NewReader(bytes.NewReader(truncated), []byte("asdf"), 2, nil)
+	_, err := r.Seek(0, io.SeekEnd)
+	if !errors.Is(err, ErrInvalid) {
+		t.Fatalf("r.Seek(0, io.SeekEnd) returned error %v, want ErrInvalid", err)
+	}
+}
+
 func FuzzRoundTrip(f *testing.F) {
 	f.Add(1, []byte("Hello World!"))
 	f.Fuzz(func(t *testing.T, segmentSize int, msg []byte) {
@@ -403,14 +448,14 @@ func FuzzSeekEnd(f *testing.F) {
 			t.Fatalf("Seek(%d, io.SeekEnd) failed: %s", offset, err)
 		}
 		if n != int64(len(msg)+offset) {
-			t.Errorf("Seek(%d, io.SeekEnd) = %d, want %d", offset, n, offset)
+			t.Errorf("Seek(%d, io.SeekEnd) = %d, want %d", offset, n, len(msg)+offset)
 		}
 		got, err := io.ReadAll(r)
 		if err != nil {
 			t.Fatalf("Reader.Read failed: %s", err)
 		}
 		if !bytes.Equal(got, msg[len(msg)+offset:]) {
-			t.Errorf("Message after seek is %q, want %q", got, msg[offset:])
+			t.Errorf("Message after seek is %q, want %q", got, msg[len(msg)+offset:])
 		}
 	})
 }
diff --git a/testdata/fuzz/FuzzSeekEnd/0ae92f8de4f43da0 b/testdata/fuzz/FuzzSeekEnd/0ae92f8de4f43da0
new file mode 100644
index 0000000..f5aa34d
--- /dev/null
+++ b/testdata/fuzz/FuzzSeekEnd/0ae92f8de4f43da0
@@ -0,0 +1,4 @@
+go test fuzz v1
+int(54)
+[]byte("")
+int(0)
diff --git a/testdata/fuzz/FuzzSeekEnd/e169ff98bf9a0ad3 b/testdata/fuzz/FuzzSeekEnd/e169ff98bf9a0ad3
new file mode 100644
index 0000000..dd73082
--- /dev/null
+++ b/testdata/fuzz/FuzzSeekEnd/e169ff98bf9a0ad3
@@ -0,0 +1,4 @@
+go test fuzz v1
+int(2)
+[]byte("00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000")
+int(-191)