From 3bfccb10338fd6c1fb868369962d3c07941a84b3 Mon Sep 17 00:00:00 2001
From: Michael Andersen <michael@antimatter.io>
Date: Wed, 26 Apr 2023 12:54:41 -0700
Subject: [PATCH] Add RemainingBytes utility to streaming decoder

Signed-off-by: Michael Andersen <michael.andersen@antimatter.io>
---
 stream.go      | 11 +++++++++++
 stream_test.go | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/stream.go b/stream.go
index 4e8c1e87..adbcbbc8 100644
--- a/stream.go
+++ b/stream.go
@@ -4,6 +4,7 @@
 package cbor
 
 import (
+	"bytes"
 	"errors"
 	"io"
 	"reflect"
@@ -65,6 +66,16 @@ func (dec *Decoder) NumBytesRead() int {
 	return dec.bytesRead
 }
 
+// RemainingBytes returns a reader that contains the bytes that have not
+// yet been consumed by Decode() or Skip(). Note that undefined behavior will
+// occur if this decoder is used after the returned reader is read from.
+func (dec *Decoder) RemainingBytes() io.Reader {
+	if dec.off == len(dec.buf) {
+		return dec.r
+	}
+	return io.MultiReader(bytes.NewBuffer(dec.buf[dec.off:]), dec.r)
+}
+
 // readNext() reads next CBOR data item from Reader to buffer.
 // It returns the size of next CBOR data item.
 // It also returns validation error or read error if any.
diff --git a/stream_test.go b/stream_test.go
index bbd74ac5..dbed4e5f 100644
--- a/stream_test.go
+++ b/stream_test.go
@@ -992,3 +992,56 @@ func (r *recoverableReader) Read(b []byte) (int, error) {
 	}
 	return r.nBytesReader.Read(b)
 }
+
+func TestDecoderRemainingBytes(t *testing.T) {
+	var buf bytes.Buffer
+
+	// Prepare test cbor bytes
+	encoder := NewEncoder(&buf)
+	if err := encoder.Encode("Item1"); err != nil {
+		t.Fatalf("Encode returned error %v", err)
+	}
+	if err := encoder.Encode(2); err != nil {
+		t.Fatalf("Encode returned error %v", err)
+	}
+
+	// Now write some data that is not CBOR to the end of the buffer
+	trailer := []byte("customContent")
+	buf.Write(trailer)
+
+	// Now read off one object
+	var obj1 string
+	decoder1 := NewDecoder(&buf)
+	if err := decoder1.Decode(&obj1); err != nil {
+		t.Fatalf("Decode returned error %v", err)
+	}
+	if obj1 != "Item1" {
+		t.Fatalf("Decode gave %v, want 'Item1'", obj1)
+	}
+
+	// Obtain the "remaining reader" which should point to the start of the second object
+	remainingReader1 := decoder1.RemainingBytes()
+
+	// This reader should be valid to pass to another decoder to read the second cbor object
+	// (although this is not the common use case)
+	// This also tests that the invalid bytes after the object do not cause issues with
+	// decoding the first object
+	var obj2 int
+	decoder2 := NewDecoder(remainingReader1)
+	if err := decoder2.Decode(&obj2); err != nil {
+		t.Fatalf("Decode returned error %v", err)
+	}
+	if obj2 != 2 {
+		t.Fatalf("Decode gave %v, want int(2)", obj2)
+	}
+
+	// Obtain another remaining reader which should point to the start of the non-cbor trailing bytes
+	remainingReader2 := decoder2.RemainingBytes()
+	gotTrailer, err := io.ReadAll(remainingReader2)
+	if err != nil {
+		t.Fatalf("RemainingBytes().Read returned error %v", err)
+	}
+	if !bytes.Equal(gotTrailer, trailer) {
+		t.Fatalf("RemainingBytes got %v, expected %v", gotTrailer, trailer)
+	}
+}