1 files changed, 51 insertions, 6 deletions

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