chore: sanitize label names

pkg/validation/validate.go

@@ -123,13 +123,18 @@ func ValidateLabels(limits LabelValidationLimits, tenantID string, ls []*typesv1
 	for _, l := range ls {
 		if len(l.Name) > limits.MaxLabelNameLength(tenantID) {
 			return NewErrorf(LabelNameTooLong, LabelNameTooLongErrorMsg, phlaremodel.LabelPairsString(ls), l.Name)
-		} else if len(l.Value) > limits.MaxLabelValueLength(tenantID) {
+		}
+		if len(l.Value) > limits.MaxLabelValueLength(tenantID) {
 			return NewErrorf(LabelValueTooLong, LabelValueTooLongErrorMsg, phlaremodel.LabelPairsString(ls), l.Value)
-		} else if !model.LabelName(l.Name).IsValid() {
+		}
+		var ok bool
+		if l.Name, ok = SanitizeLabelName(l.Name); !ok {
 			return NewErrorf(InvalidLabels, InvalidLabelsErrorMsg, phlaremodel.LabelPairsString(ls), "invalid label name '"+l.Name+"'")
-		} else if !model.LabelValue(l.Value).IsValid() {
+		}
+		if !model.LabelValue(l.Value).IsValid() {
 			return NewErrorf(InvalidLabels, InvalidLabelsErrorMsg, phlaremodel.LabelPairsString(ls), "invalid label value '"+l.Value+"'")
-		} else if cmp := strings.Compare(lastLabelName, l.Name); cmp == 0 {
+		}
+		if cmp := strings.Compare(lastLabelName, l.Name); cmp == 0 {
 			return NewErrorf(DuplicateLabelNames, DuplicateLabelNamesErrorMsg, phlaremodel.LabelPairsString(ls), l.Name)
 		}
 		lastLabelName = l.Name
@@ -138,6 +143,37 @@ func ValidateLabels(limits LabelValidationLimits, tenantID string, ls []*typesv1
 	return nil
 }
 
+/*func init() {
+	model.NameValidationScheme = model.UTF8Validation
+}*/
+
+// SanitizeLabelName replaces disallowed characters in the label name
+// with underscore character.
+//
+// The function only returns false if the input string is empty.
+func SanitizeLabelName(ln string) (string, bool) {
+	if len(ln) == 0 {
+		return ln, false
+	}
+	valid := true
+	for i, b := range ln {
+		if !((b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') || b == '_' || (b >= '0' && b <= '9' && i > 0)) {
+			valid = false
+		}
+	}
+	if valid {
+		return ln, true
+	}
+	// Replace invalid characters with underscore.
+	r := []rune(ln)
+	for i, b := range r {
+		if !((b >= 'a' && b <= 'z') || (b >= 'A' && b <= 'Z') || b == '_' || (b >= '0' && b <= '9' && i > 0)) {
+			r[i] = '_'
+		}
+	}
+	return string(r), true
+}
+
 type ProfileValidationLimits interface {
 	MaxProfileSizeBytes(tenantID string) int
 	MaxProfileStacktraceSamples(tenantID string) int

pkg/validation/validate_test.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
@@ -71,16 +72,6 @@ func TestValidateLabels(t *testing.T) {
 			expectedErr:    "invalid labels '{__name__=\"qux\", foo=\"\\xc5\", service_name=\"svc\"}' with error: invalid label value '\xc5'",
 			expectedReason: InvalidLabels,
 		},
-		{
-			name: "invalid label name",
-			lbs: []*typesv1.LabelPair{
-				{Name: phlaremodel.LabelNameServiceName, Value: "svc"},
-				{Name: model.MetricNameLabel, Value: "qux"},
-				{Name: "\xc5", Value: "foo"},
-			},
-			expectedErr:    "invalid labels '{__name__=\"qux\", service_name=\"svc\", \xc5=\"foo\"}' with error: invalid label name '\xc5'",
-			expectedReason: InvalidLabels,
-		},
 		{
 			name: "name too long",
 			lbs: []*typesv1.LabelPair{
@@ -400,3 +391,36 @@ func TestValidateFlamegraphMaxNodes(t *testing.T) {
 		})
 	}
 }
+
+func Test_SanitizeLabelName(t *testing.T) {
+	for _, tc := range []struct {
+		input    string
+		expected string
+		valid    bool
+	}{
+		{"", "", false},
+		{".", "_", true},
+		{".a", "_a", true},
+		{"a.", "a_", true},
+		{"..", "__", true},
+		{"..a", "__a", true},
+		{"a..", "a__", true},
+		{"a.a", "a_a", true},
+		{".a.", "_a_", true},
+		{"..a..", "__a__", true},
+		{"世界", "__", true},
+		{"界世_a", "___a", true},
+		{"界世__a", "____a", true},
+		{"a_世界", "a___", true},
+		{"_", "_", true},
+		{"__a", "__a", true},
+		{"__a__", "__a__", true},
+	} {
+		tc := tc
+		t.Run("", func(t *testing.T) {
+			actual, valid := SanitizeLabelName(tc.input)
+			assert.Equal(t, tc.expected, actual)
+			assert.Equal(t, tc.valid, valid)
+		})
+	}
+}