Improve error handlingHEADmain

1 files changed, 93 insertions, 6 deletions

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