Skip to content

Commit

Permalink
fix(protobuf): Correctly decode multi-messages streams
Browse files Browse the repository at this point in the history
Signed-off-by: Sven Rebhan <[email protected]>
  • Loading branch information
srebhan committed Apr 3, 2024
1 parent 3dec13c commit 78585cf
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 3 deletions.
6 changes: 3 additions & 3 deletions expfmt/decode.go
Original file line number Diff line number Diff line change
Expand Up @@ -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())) {
Expand Down
27 changes: 27 additions & 0 deletions expfmt/decode_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,12 @@ package expfmt

import (
"bufio"
"bytes"
"errors"
"io"
"math"
"net/http"
"os"
"reflect"
"sort"
"strings"
Expand Down Expand Up @@ -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
Expand Down
Binary file added expfmt/testdata/protobuf-multimessage
Binary file not shown.

0 comments on commit 78585cf

Please sign in to comment.