From cafc7f30638900ea6723f7ed127673f1c94db982 Mon Sep 17 00:00:00 2001
From: gagliardetto <gagliardetto@users.noreply.github.com>
Date: Sun, 28 Aug 2022 15:20:26 +0200
Subject: [PATCH] Improve readNBytes

---
 decoder.go      | 15 ++++++---------
 decoder_test.go | 20 ++++++++++++++++++++
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/decoder.go b/decoder.go
index e96250b..49fe2e2 100644
--- a/decoder.go
+++ b/decoder.go
@@ -274,22 +274,19 @@ type peekAbleByteReader interface {
 	Peek(n int) ([]byte, error)
 }
 
-func readNBytes(n int, reader peekAbleByteReader) ([]byte, error) {
+func readNBytes(n int, reader *Decoder) ([]byte, error) {
 	if n == 0 {
 		return make([]byte, 0), nil
 	}
 	if n < 0 || n > 0x7FFF_FFFF {
 		return nil, fmt.Errorf("invalid length n: %v", n)
 	}
-	buf := make([]byte, n)
-	for i := 0; i < n; i++ {
-		b, err := reader.ReadByte()
-		if err != nil {
-			return nil, err
-		}
-		buf[i] = b
+	if reader.pos+n > len(reader.data) {
+		return nil, fmt.Errorf("not enough data: %d bytes missing", reader.pos+n-len(reader.data))
 	}
-	return buf, nil
+	out := reader.data[reader.pos : reader.pos+n]
+	reader.pos += n
+	return out, nil
 }
 
 func discardNBytes(n int, reader *Decoder) error {
diff --git a/decoder_test.go b/decoder_test.go
index 595ecf3..43094ca 100644
--- a/decoder_test.go
+++ b/decoder_test.go
@@ -1174,6 +1174,26 @@ func Test_Decode_custom(t *testing.T) {
 	})
 }
 
+func Test_ReadNBytes(t *testing.T) {
+	{
+		b1 := []byte{123, 99, 88, 77, 66, 55, 44, 33, 22, 11}
+		b2 := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
+		buf := concatByteSlices(
+			b1,
+			b2,
+		)
+		decoder := NewBinDecoder(buf)
+
+		got, err := decoder.ReadNBytes(10)
+		require.NoError(t, err)
+		require.Equal(t, b1, got)
+
+		got, err = decoder.ReadNBytes(10)
+		require.NoError(t, err)
+		require.Equal(t, b2, got)
+	}
+}
+
 func Test_Decode_readArrayOfUint(t *testing.T) {
 	{
 		{