Merge pull request #616 from srebhan/fix_protobuf_decode

fix(protobuf): Correctly decode multi-messages streams

expfmt/decode.go

@@ -75,22 +75,22 @@ func ResponseFormat(h http.Header) Format {
 func NewDecoder(r io.Reader, format Format) Decoder {
 	switch format.FormatType() {
 	case TypeProtoDelim:
-		return &protoDecoder{r: r}
+		return &protoDecoder{r: bufio.NewReader(r)}
 	}
 	return &textDecoder{r: r}
 }
 
 // protoDecoder implements the Decoder interface for protocol buffers.
 type protoDecoder struct {
-	r io.Reader
+	r protodelim.Reader
 }
 
 // Decode implements the Decoder interface.
 func (d *protoDecoder) Decode(v *dto.MetricFamily) error {
 	opts := protodelim.UnmarshalOptions{
 		MaxSize: -1,
 	}
-	if err := opts.UnmarshalFrom(bufio.NewReader(d.r), v); err != nil {
+	if err := opts.UnmarshalFrom(d.r, v); err != nil {
 		return err
 	}
 	if !model.IsValidMetricName(model.LabelValue(v.GetName())) {

expfmt/decode_test.go

@@ -15,10 +15,12 @@ package expfmt
 
 import (
 	"bufio"
+	"bytes"
 	"errors"
 	"io"
 	"math"
 	"net/http"
+	"os"
 	"reflect"
 	"sort"
 	"strings"
@@ -414,6 +416,31 @@ func TestProtoDecoder(t *testing.T) {
 	}
 }
 
+func TestProtoMultiMessageDecoder(t *testing.T) {
+	data, err := os.ReadFile("testdata/protobuf-multimessage")
+	if err != nil {
+		t.Fatalf("Reading file failed: %v", err)
+	}
+
+	buf := bytes.NewReader(data)
+	decoder := NewDecoder(buf, fmtProtoDelim)
+	var metrics []*dto.MetricFamily
+	for {
+		var mf dto.MetricFamily
+		if err := decoder.Decode(&mf); err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			t.Fatalf("Unmarshalling failed: %v", err)
+		}
+		metrics = append(metrics, &mf)
+	}
+
+	if len(metrics) != 6 {
+		t.Fatalf("Expected %d metrics but got %d!", 6, len(metrics))
+	}
+}
+
 func testDiscriminatorHTTPHeader(t testing.TB) {
 	scenarios := []struct {
 		input  map[string]string