Decoder improvements

cmd/decode/main.go

@@ -23,4 +23,8 @@ func main() {
 	if _, err := io.Copy(dec, os.Stdin); err != nil && err != cobs.EOD {
 		panic(err)
 	}
+
+	if dec.NeedsMoreData() {
+		panic(cobs.ErrIncompleteFrame)
+	}
 }

cobs.go

@@ -19,6 +19,9 @@ var EOD = errors.New("EOD")
 // ErrUnexpectedEOD means that a delimiter was encountered in a malformed frame.
 var ErrUnexpectedEOD = errors.New("unexpected EOD")
 
+// ErrIncompleteData means a decoder was closed with an incomplete frame.
+var ErrIncompleteFrame = errors.New("frame incomplete")
+
 // An Encoder implements the io.Writer and io.ByteWriter interfaces. Data
 // written will we be encoded into groups and forwarded.
 type Encoder struct {
@@ -170,12 +173,23 @@ func (d *Decoder) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+// NeedsMoreData returns true if the decoder needs more data for a full frame.
+func (d *Decoder) NeedsMoreData() bool {
+	return d.codeIndex != 0
+}
+
 // Decode decodes and returns a byte slice.
 func Decode(data []byte) ([]byte, error) {
 	buf := bytes.NewBuffer(make([]byte, 0, len(data)))
 	d := NewDecoder(buf)
 
-	_, err := d.Write(data)
+	if _, err := d.Write(data); err != nil {
+		return buf.Bytes(), err
+	}
+
+	if d.NeedsMoreData() {
+		return buf.Bytes(), ErrIncompleteFrame
+	}
 
-	return buf.Bytes(), err
+	return buf.Bytes(), nil
 }

cobs_test.go

@@ -166,7 +166,8 @@ func TestWriter(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			buf := bytes.NewBuffer(make([]byte, 0, len(tc.enc)))
-			e := NewEncoder(NewDecoder(buf))
+			d := NewDecoder(buf)
+			e := NewEncoder(d)
 
 			n, err := e.Write(tc.dec)
 			if err != nil {
@@ -176,6 +177,9 @@ func TestWriter(t *testing.T) {
 			if err != nil {
 				t.Errorf("writer close error: %v", err)
 			}
+			if d.NeedsMoreData() {
+				t.Error("writer incomplete decode data")
+			}
 			if n != len(tc.dec) {
 				t.Errorf("writer length got %d, want %d", n, len(tc.dec))
 			}
@@ -220,6 +224,10 @@ func TestStream(t *testing.T) {
 			t.Error("stream decode EOD missing")
 		}
 
+		if d.NeedsMoreData() {
+			t.Error("stream decode frame incomplete")
+		}
+
 		if !bytes.Equal(buf.Bytes(), tc.dec) {
 			t.Errorf("stream decode got %v, want %v", buf.Bytes(), tc.dec)
 		}
@@ -259,7 +267,8 @@ func FuzzChainWriter(f *testing.F) {
 	}
 	f.Fuzz(func(t *testing.T, a []byte) {
 		var buf bytes.Buffer
-		e := NewEncoder(NewDecoder(&buf))
+		d := NewDecoder(&buf)
+		e := NewEncoder(d)
 
 		n, err := e.Write(a)
 		if err != nil {
@@ -273,12 +282,65 @@ func FuzzChainWriter(f *testing.F) {
 		if err != nil {
 			t.Errorf("fuzz chain close error: %v", err)
 		}
+		if d.NeedsMoreData() {
+			t.Error("fuzz chain incomplete decode data")
+		}
 		if !bytes.Equal(buf.Bytes(), a) {
 			t.Errorf("fuzz chain got %v want %v", buf.Bytes(), a)
 		}
 	})
 }
 
+var incompleteFrame = []struct {
+	name string
+	data []byte
+}{
+	{
+		name: "Missing single byte",
+		data: []byte{0x02},
+	},
+	{
+		name: "2 zeroes and missing end",
+		data: []byte{0x01, 0x01, 0x05},
+	},
+}
+
+func TestDecodeIncomplete(t *testing.T) {
+	for _, tc := range incompleteFrame {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := Decode(tc.data)
+			if err != ErrIncompleteFrame {
+				t.Errorf("Unexpected decode incomplete error: %v", err)
+			}
+		})
+	}
+}
+
+var unexpectedDelimiter = []struct {
+	name string
+	data []byte
+}{
+	{
+		name: "Missing byte before delimiter",
+		data: []byte{0x02, 0x00},
+	},
+	{
+		name: "Unexpected embedded zero",
+		data: []byte("\x061234\x005\x056789"),
+	},
+}
+
+func TestDecodeUnexpectedEOD(t *testing.T) {
+	for _, tc := range unexpectedDelimiter {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := Decode(tc.data)
+			if err != ErrUnexpectedEOD {
+				t.Errorf("Unexpected decode EOD error: %v", err)
+			}
+		})
+	}
+}
+
 // https://github.com/golang/go/issues/54111
 type LimitedWriter struct {
 	W   io.Writer // underlying writer