Improve error handlingHEADmain

2 files changed, 113 insertions, 13 deletions

diff --git a/oae2.go b/oae2.go
index ff59b76..c1ee627 100644
--- a/oae2.go
+++ b/oae2.go
@@ -33,6 +33,11 @@ import (
 	"math"
 )
 
+// ErrInvalid is returned from methods on Reader when the stream failed to
+// authenticate. This can indicate that a wrong key was used, or that the data
+// was tampered with.
+var ErrInvalid = errors.New("invalid stream")
+
 const (
 	nonceSize    = 12
 	aeadOverhead = 16
@@ -107,7 +112,11 @@ func (e *segmentEncrypter) decryptSegment(segment []byte, lastSegment bool) ([]b
 	if err != nil {
 		return nil, err
 	}
-	return e.aead.Open(segment[:0], nonce, segment, ad)
+	plainText, err := e.aead.Open(segment[:0], nonce, segment, ad)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %s", ErrInvalid, err)
+	}
+	return plainText, nil
 }
 
 // A Writer wraps an io.Writer and encrypts the data in segments. Make sure to
@@ -296,8 +305,8 @@ func (r *Reader) initialize() error {
 	r.initialized = true
 	buf := make([]byte, saltSize)
 	if _, r.err = io.ReadFull(&r.r, buf); r.err != nil {
-		if r.err == io.EOF {
-			r.err = io.ErrUnexpectedEOF
+		if r.err == io.EOF || r.err == io.ErrUnexpectedEOF {
+			r.err = fmt.Errorf("%w: too short", ErrInvalid)
 		}
 		return r.err
 	}
@@ -308,7 +317,7 @@ func (r *Reader) initialize() error {
 	// additional data
 	if err := r.fillBuf(); err != nil {
 		if err == io.EOF {
-			r.err = io.ErrUnexpectedEOF
+			r.err = fmt.Errorf("%w: too short", ErrInvalid)
 			return r.err
 		}
 		return err
@@ -349,7 +358,8 @@ func (r *Reader) fillBuf() error {
 	return nil
 }
 
-// Read implements io.Reader.
+// Read implements io.Reader. Read returns an error wrapping [ErrInvalid] if
+// the stream fails authentication.
 func (r *Reader) Read(buf []byte) (int, error) {
 	if err := r.init(); err != nil {
 		return 0, err
@@ -367,7 +377,8 @@ func (r *Reader) Read(buf []byte) (int, error) {
 	return n, nil
 }
 
-// WriteTo implements io.WriterTo.
+// WriteTo implements io.WriterTo. WriteTo returns an error wrapping
+// [ErrInvalid] if the stream fails authentication.
 func (r *Reader) WriteTo(w io.Writer) (int64, error) {
 	if err := r.init(); err != nil {
 		return 0, err
@@ -427,7 +438,9 @@ func (r *Reader) seek(offset int64) (int64, error) {
 }
 
 // Seek implements io.Seeker. If the underlying reader does not implement
-// io.Seeker, Seek returns [errors.ErrUnsupported].
+// io.Seeker, Seek returns an error wrapping [errors.ErrUnsupported]. Seek may
+// also return an error wrapping [ErrInvalid] if the seeked-to segment
+// fails authentication.
 func (r *Reader) Seek(offset int64, whence int) (int64, error) {
 	if _, ok := r.r.r.(io.Seeker); !ok {
 		return 0, fmt.Errorf("%w", errors.ErrUnsupported)
diff --git a/oae2_test.go b/oae2_test.go
index f0ed37f..5ca26e6 100644
--- a/oae2_test.go
+++ b/oae2_test.go
@@ -3,6 +3,7 @@ package oae2
 import (
 	"bytes"
 	"crypto/rand"
+	"errors"
 	"io"
 	"strings"
 	"testing"
@@ -66,8 +67,94 @@ func TestInvalid(t *testing.T) {
 		key = "asdf"
 	)
 	_, err := io.ReadAll(NewReader(strings.NewReader(msg), []byte(key), 1, nil))
-	if err == nil {
-		t.Errorf("Message %q passed validation, want error", msg)
+	if !errors.Is(err, ErrInvalid) {
+		t.Errorf("Reader.Read: decrypt message %q returned error %v, want ErrInvalid", msg, err)
+	}
+}
+
+func TestWrongKeyRead(t *testing.T) {
+	t.Parallel()
+
+	const (
+		key            = "asdf"
+		additionalData = "additional data"
+	)
+
+	encryptedBuf := new(bytes.Buffer)
+	if _, err := NewWriter(encryptedBuf, []byte(key), 4*1024*1024, []byte(additionalData)).Write([]byte("hello world")); err != nil {
+		t.Fatalf("Failed to encrypt test data: %s", err)
+	}
+
+	for _, tc := range []struct {
+		desc           string
+		key            string
+		additionalData string
+	}{{
+		desc: "WrongKey",
+		key:  "wrong key",
+	}, {
+		desc:           "WrongData",
+		additionalData: "wrong additional data",
+	}} {
+		t.Run(tc.desc, func(t *testing.T) {
+			t.Parallel()
+
+			k := tc.key
+			if k == "" {
+				k = key
+			}
+			ad := tc.additionalData
+			if ad == "" {
+				ad = additionalData
+			}
+			_, err := io.ReadAll(NewReader(bytes.NewReader(encryptedBuf.Bytes()), []byte(key), 1, nil))
+			if !errors.Is(err, ErrInvalid) {
+				t.Errorf("Reader.Read: message returned error %v with key %q and data %q, want ErrInvalid", err, k, ad)
+			}
+		})
+	}
+}
+
+func TestWrongKeyWriteTo(t *testing.T) {
+	t.Parallel()
+
+	const (
+		key            = "asdf"
+		additionalData = "additional data"
+	)
+
+	encryptedBuf := new(bytes.Buffer)
+	if _, err := NewWriter(encryptedBuf, []byte(key), 4*1024*1024, []byte(additionalData)).Write([]byte("hello world")); err != nil {
+		t.Fatalf("Failed to encrypt test data: %s", err)
+	}
+
+	for _, tc := range []struct {
+		desc           string
+		key            string
+		additionalData string
+	}{{
+		desc: "WrongKey",
+		key:  "wrong key",
+	}, {
+		desc:           "WrongData",
+		additionalData: "wrong additional data",
+	}} {
+		t.Run(tc.desc, func(t *testing.T) {
+			t.Parallel()
+
+			k := tc.key
+			if k == "" {
+				k = key
+			}
+			ad := tc.additionalData
+			if ad == "" {
+				ad = additionalData
+			}
+			_, err := NewReader(bytes.NewReader(encryptedBuf.Bytes()), []byte(key), 1, nil).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)
+			}
+		})
 	}
 }
 
@@ -220,8 +307,8 @@ func FuzzReadInvalid(f *testing.F) {
 		password := make([]byte, 32)
 		// Get a random password so that we can be sure msg is invalid
 		rand.Read(password)
-		if _, err := io.ReadAll(NewReader(strings.NewReader(msg), password, segmentSize, []byte(additionalData))); err == nil {
-			t.Errorf("Reader.Read: message passed validation")
+		if _, err := io.ReadAll(NewReader(strings.NewReader(msg), password, segmentSize, []byte(additionalData))); !errors.Is(err, ErrInvalid) {
+			t.Errorf("Reader.Read: invalid message got error %v, want ErrInvalid", err)
 		}
 	})
 }
@@ -235,8 +322,8 @@ func FuzzWriteToInvalid(f *testing.F) {
 		}
 		password := make([]byte, 32)
 		rand.Read(password)
-		if _, err := NewReader(strings.NewReader(msg), password, segmentSize, []byte(additionalData)).WriteTo(io.Discard); err == nil {
-			t.Errorf("Reader.WriteTo: message passed validation")
+		if _, err := NewReader(strings.NewReader(msg), password, segmentSize, []byte(additionalData)).WriteTo(io.Discard); !errors.Is(err, ErrInvalid) {
+			t.Errorf("Reader.WriteTo: invalid message got error %v, want ErrInvalid", err)
 		}
 	})
 }