Fix some embarrassing bugsv0.3.1

1 files changed, 49 insertions, 26 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: