Skip to content

Commit

Permalink
Add RemainingBytes utility to streaming decoder
Browse files Browse the repository at this point in the history
Signed-off-by: Michael Andersen <michael.andersen@antimatter.io>
  • Loading branch information
immesys committed Apr 26, 2023
1 parent 6dfffb1 commit 3bfccb1
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 0 deletions.
11 changes: 11 additions & 0 deletions stream.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
package cbor

import (
"bytes"
"errors"
"io"
"reflect"
Expand Down Expand Up @@ -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.
Expand Down
53 changes: 53 additions & 0 deletions stream_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

0 comments on commit 3bfccb1

Please sign in to comment.