PRW2: add labelRef validation

Signed-off-by: SungJin1212 <tjdwls1201@gmail.com>

Author: SungJin1212

CHANGELOG.md

@@ -85,6 +85,7 @@
 * [ENHANCEMENT] Ingester: Add `cortex_ingester_tsdb_head_stale_series` metric to keep track of number of stale series on head. #7071
 * [ENHANCEMENT] Expose more Go runtime metrics. #7070
 * [ENHANCEMENT] Distributor: Filter out label with empty value. #7069
+* [ENHANCEMENT] Distributor: Add a validation for label references. #7074
 * [BUGFIX] Ingester: Avoid error or early throttling when READONLY ingesters are present in the ring #6517
 * [BUGFIX] Ingester: Fix labelset data race condition. #6573
 * [BUGFIX] Compactor: Cleaner should not put deletion marker for blocks with no-compact marker. #6576

pkg/cortexpb/compatv2.go

@@ -1,22 +1,34 @@
 package cortexpb
 
-import "github.com/prometheus/prometheus/model/labels"
+import (
+	"fmt"
 
-func (e *ExemplarV2) ToLabels(b *labels.ScratchBuilder, symbols []string) labels.Labels {
+	"github.com/prometheus/prometheus/model/labels"
+)
+
+func (e *ExemplarV2) ToLabels(b *labels.ScratchBuilder, symbols []string) (labels.Labels, error) {
 	return desymbolizeLabels(b, e.GetLabelsRefs(), symbols)
 }
 
-func (t *TimeSeriesV2) ToLabels(b *labels.ScratchBuilder, symbols []string) labels.Labels {
+func (t *TimeSeriesV2) ToLabels(b *labels.ScratchBuilder, symbols []string) (labels.Labels, error) {
 	return desymbolizeLabels(b, t.GetLabelsRefs(), symbols)
 }
 
 // desymbolizeLabels decodes label references, with given symbols to labels.
-// Copied from the Prometheus: https://github.com/prometheus/prometheus/blob/v3.5.0/prompb/io/prometheus/write/v2/symbols.go#L76
-func desymbolizeLabels(b *labels.ScratchBuilder, labelRefs []uint32, symbols []string) labels.Labels {
+// Copied from the Prometheus: https://github.com/prometheus/prometheus/blob/v3.7.2/prompb/io/prometheus/write/v2/symbols.go#L83
+func desymbolizeLabels(b *labels.ScratchBuilder, labelRefs []uint32, symbols []string) (labels.Labels, error) {
+	if len(labelRefs)%2 != 0 {
+		return labels.EmptyLabels(), fmt.Errorf("invalid labelRefs length %d", len(labelRefs))
+	}
+
 	b.Reset()
 	for i := 0; i < len(labelRefs); i += 2 {
-		b.Add(symbols[labelRefs[i]], symbols[labelRefs[i+1]])
+		nameRef, valueRef := labelRefs[i], labelRefs[i+1]
+		if int(nameRef) >= len(symbols) || int(valueRef) >= len(symbols) {
+			return labels.EmptyLabels(), fmt.Errorf("labelRefs %d (name) = %d (value) outside of symbols table (size %d)", nameRef, valueRef, len(symbols))
+		}
+		b.Add(symbols[nameRef], symbols[valueRef])
 	}
 	b.Sort()
-	return b.Labels()
+	return b.Labels(), nil
 }

pkg/cortexpb/compatv2_test.go

@@ -0,0 +1,112 @@
+package cortexpb
+
+import (
+	"testing"
+
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_ExemplarV2ToLabels(t *testing.T) {
+	symbols := []string{"", "foo", "bar"}
+	b := &labels.ScratchBuilder{}
+
+	tests := []struct {
+		desc           string
+		e              *ExemplarV2
+		isErr          bool
+		expectedLabels labels.Labels
+	}{
+		{
+			desc: "Success",
+			e: &ExemplarV2{
+				LabelsRefs: []uint32{1, 2},
+			},
+			expectedLabels: labels.FromStrings("foo", "bar"),
+		},
+		{
+			desc: "Odd length of the label refs",
+			e: &ExemplarV2{
+				LabelsRefs: []uint32{1},
+			},
+			isErr: true,
+		},
+		{
+			desc: "Label name refs out of ranges",
+			e: &ExemplarV2{
+				LabelsRefs: []uint32{10, 2},
+			},
+			isErr: true,
+		},
+		{
+			desc: "Label value refs out of ranges",
+			e: &ExemplarV2{
+				LabelsRefs: []uint32{1, 10},
+			},
+			isErr: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			lbs, err := test.e.ToLabels(b, symbols)
+			if test.isErr {
+				require.Error(t, err)
+			} else {
+				require.Equal(t, test.expectedLabels.String(), lbs.String())
+			}
+		})
+	}
+}
+
+func Test_TimeSeriesV2ToLabels(t *testing.T) {
+	symbols := []string{"", "__name__", "test_metric", "job", "prometheus", "instance", "server-1"}
+	b := &labels.ScratchBuilder{}
+
+	tests := []struct {
+		desc           string
+		ts             *TimeSeriesV2
+		isErr          bool
+		expectedLabels labels.Labels
+	}{
+		{
+			desc: "Success",
+			ts: &TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 3, 4, 5, 6},
+			},
+			expectedLabels: labels.FromStrings("__name__", "test_metric", "job", "prometheus", "instance", "server-1"),
+		},
+		{
+			desc: "Odd length of the label refs",
+			ts: &TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 3},
+			},
+			isErr: true,
+		},
+		{
+			desc: "Label name refs out of ranges",
+			ts: &TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 10, 4},
+			},
+			isErr: true,
+		},
+		{
+			desc: "Label value refs out of ranges",
+			ts: &TimeSeriesV2{
+				LabelsRefs: []uint32{1, 2, 3, 10},
+			},
+			isErr: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			lbs, err := test.ts.ToLabels(b, symbols)
+			if test.isErr {
+				require.Error(t, err)
+			} else {
+				require.Equal(t, test.expectedLabels.String(), lbs.String())
+			}
+		})
+	}
+}

pkg/util/push/push.go

@@ -177,12 +177,19 @@ func convertV2RequestToV1(req *cortexpb.PreallocWriteRequestV2) (cortexpb.Preall
 	b := labels.NewScratchBuilder(0)
 	symbols := req.Symbols
 	for _, v2Ts := range req.Timeseries {
-		lbs := v2Ts.ToLabels(&b, symbols)
+		lbs, err := v2Ts.ToLabels(&b, symbols)
+		if err != nil {
+			return v1Req, err
+		}
+		exemplars, err := convertV2ToV1Exemplars(&b, symbols, v2Ts.Exemplars)
+		if err != nil {
+			return v1Req, err
+		}
 		v1Timeseries = append(v1Timeseries, cortexpb.PreallocTimeseries{
 			TimeSeries: &cortexpb.TimeSeries{
 				Labels:     cortexpb.FromLabelsToLabelAdapters(lbs),
 				Samples:    v2Ts.Samples,
-				Exemplars:  convertV2ToV1Exemplars(&b, symbols, v2Ts.Exemplars),
+				Exemplars:  exemplars,
 				Histograms: v2Ts.Histograms,
 			},
 		})
@@ -234,14 +241,18 @@ func convertV2ToV1Metadata(name string, symbols []string, metadata cortexpb.Meta
 	}
 }
 
-func convertV2ToV1Exemplars(b *labels.ScratchBuilder, symbols []string, v2Exemplars []cortexpb.ExemplarV2) []cortexpb.Exemplar {
+func convertV2ToV1Exemplars(b *labels.ScratchBuilder, symbols []string, v2Exemplars []cortexpb.ExemplarV2) ([]cortexpb.Exemplar, error) {
 	v1Exemplars := make([]cortexpb.Exemplar, 0, len(v2Exemplars))
 	for _, e := range v2Exemplars {
+		lbs, err := e.ToLabels(b, symbols)
+		if err != nil {
+			return nil, err
+		}
 		v1Exemplars = append(v1Exemplars, cortexpb.Exemplar{
-			Labels:      cortexpb.FromLabelsToLabelAdapters(e.ToLabels(b, symbols)),
+			Labels:      cortexpb.FromLabelsToLabelAdapters(lbs),
 			Value:       e.Value,
 			TimestampMs: e.Timestamp,
 		})
 	}
-	return v1Exemplars
+	return v1Exemplars, nil
 }