From 78585cf51e25f67c820ee6637cfe5560cdc5f55d Mon Sep 17 00:00:00 2001
From: Sven Rebhan <srebhan@influxdata.com>
Date: Wed, 3 Apr 2024 11:47:10 +0200
Subject: [PATCH] fix(protobuf): Correctly decode multi-messages streams

Signed-off-by: Sven Rebhan <srebhan@influxdata.com>
---
 expfmt/decode.go                      |   6 +++---
 expfmt/decode_test.go                 |  27 ++++++++++++++++++++++++++
 expfmt/testdata/protobuf-multimessage | Bin 0 -> 414 bytes
 3 files changed, 30 insertions(+), 3 deletions(-)
 create mode 100644 expfmt/testdata/protobuf-multimessage

diff --git a/expfmt/decode.go b/expfmt/decode.go
index b2b89b01..25cfaa21 100644
--- a/expfmt/decode.go
+++ b/expfmt/decode.go
@@ -75,14 +75,14 @@ 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.
@@ -90,7 +90,7 @@ 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())) {
diff --git a/expfmt/decode_test.go b/expfmt/decode_test.go
index e5e245d3..19560ffc 100644
--- a/expfmt/decode_test.go
+++ b/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
diff --git a/expfmt/testdata/protobuf-multimessage b/expfmt/testdata/protobuf-multimessage
new file mode 100644
index 0000000000000000000000000000000000000000..d9fa9fe14667b9f6e087b3c8997b960e6afa1463
GIT binary patch
literal 414
zcmZ?f;w&ytEQn7lN=+4#3`xyNO)pAJQ%KIw$w^HvNlj76O)V+POqO6&lH}syV#&xa
zE)inM&n?at;^brif(I&R9UZyY!DeLU;WL2&-2^F!2@DMOj!s-0U=#97OUN=pF3-`0
ziwkT-Nq$LU4n9M0`6OfJSx0ARXp|PGrjTO>gHo=eH<uvTAs{p23sQ@cQ}gh83Y&X!
MUu`+6I{Sw`0AKrh6951J

literal 0
HcmV?d00001