Make ParseKey validate user input

go.mod

@@ -39,15 +39,16 @@ require (
 	github.com/onsi/gomega v1.12.0
 	github.com/pelletier/go-toml v1.8.1 // indirect
 	github.com/peterbourgon/ff/v3 v3.0.0
-	github.com/prometheus/client_golang v1.10.0
+	github.com/prometheus/client_golang v1.11.0
+	github.com/prometheus/common v0.29.0
 	github.com/pyroscope-io/dotnetdiag v1.2.1
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/shirou/gopsutil v3.21.4+incompatible
 	github.com/sirupsen/logrus v1.7.0
 	github.com/tklauser/go-sysconf v0.3.6 // indirect
 	github.com/twmb/murmur3 v1.1.5
 	github.com/wacul/ptr v1.0.0 // indirect
-	golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421
+	golang.org/x/oauth2 v0.0.0-20210514164344-f6687ab2804c
 	golang.org/x/sync v0.0.0-20201207232520-09787c993a3a
 	golang.org/x/sys v0.0.0-20210616094352-59db8d763f22
 	golang.org/x/tools v0.1.0

pkg/agent/session.go

@@ -109,7 +109,11 @@ func (ps *ProfileSession) createNames(tags map[string]string) error {
 		if err != nil {
 			return err
 		}
-		tagsCopy["__name__"] = appName + "." + string(t)
+		appName += "." + string(t)
+		if err = flameql.ValidateAppName(appName); err != nil {
+			return err
+		}
+		tagsCopy["__name__"] = appName
 		ps.names[t] = segment.NewKey(tagsCopy).Normalized()
 	}
 	return nil

pkg/flameql/error.go

@@ -33,5 +33,13 @@ func (e *Error) Error() string { return e.Inner.Error() + ": " + e.Expr }
 func (e *Error) Unwrap() error { return e.Inner }
 
 func newInvalidTagKeyRuneError(k string, r rune) *Error {
-	return newErr(ErrInvalidTagKey, fmt.Sprintf("%s: character is not allowed: %q", k, r))
+	return newInvalidRuneError(ErrInvalidTagKey, k, r)
+}
+
+func newInvalidAppNameRuneError(k string, r rune) *Error {
+	return newInvalidRuneError(ErrInvalidAppName, k, r)
+}
+
+func newInvalidRuneError(err error, k string, r rune) *Error {
+	return newErr(err, fmt.Sprintf("%s: character is not allowed: %q", k, r))
 }

pkg/flameql/flameql.go

@@ -31,6 +31,14 @@ const (
 	EQL_REGEX    // =~
 )
 
+const (
+	ReservedTagKeyName = "__name__"
+)
+
+var reservedTagKeys = []string{
+	ReservedTagKeyName,
+}
+
 // IsNegation reports whether the operator assumes negation.
 func (o Op) IsNegation() bool { return o < EQL }
 
@@ -55,3 +63,52 @@ func (m *TagMatcher) Match(v string) bool {
 		panic("invalid match operator")
 	}
 }
+
+// ValidateTagKey report an error if the given key k violates constraints.
+//
+// The function should be used to validate user input. The function returns
+// ErrTagKeyReserved if the key is valid but reserved for internal use.
+func ValidateTagKey(k string) error {
+	if len(k) == 0 {
+		return ErrTagKeyIsRequired
+	}
+	for _, r := range k {
+		if !IsTagKeyRuneAllowed(r) {
+			return newInvalidTagKeyRuneError(k, r)
+		}
+	}
+	if IsTagKeyReserved(k) {
+		return newErr(ErrTagKeyReserved, k)
+	}
+	return nil
+}
+
+// ValidateAppName report an error if the given app name n violates constraints.
+func ValidateAppName(n string) error {
+	if len(n) == 0 {
+		return ErrAppNameIsRequired
+	}
+	for _, r := range n {
+		if !IsAppNameRuneAllowed(r) {
+			return newInvalidAppNameRuneError(n, r)
+		}
+	}
+	return nil
+}
+
+func IsTagKeyRuneAllowed(r rune) bool {
+	return (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '_'
+}
+
+func IsAppNameRuneAllowed(r rune) bool {
+	return (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '_' || r == '-' || r == '.'
+}
+
+func IsTagKeyReserved(k string) bool {
+	for _, s := range reservedTagKeys {
+		if s == k {
+			return true
+		}
+	}
+	return false
+}

pkg/flameql/flameql_test.go

@@ -1,6 +1,8 @@
 package flameql
 
 import (
+	"errors"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
@@ -31,3 +33,54 @@ var _ = Describe("TagMatcher", func() {
 		}
 	})
 })
+
+var _ = Describe("ValidateTagKey", func() {
+	It("reports error if a key violates constraints", func() {
+		type testCase struct {
+			key string
+			err error
+		}
+
+		testCases := []testCase{
+			{"foo_BAR_12_baz_qux", nil},
+
+			{ReservedTagKeyName, ErrTagKeyReserved},
+			{"", ErrTagKeyIsRequired},
+			{"#", ErrInvalidTagKey},
+		}
+
+		for _, tc := range testCases {
+			err := ValidateTagKey(tc.key)
+			if tc.err != nil {
+				Expect(errors.Is(err, tc.err)).To(BeTrue())
+				continue
+			}
+			Expect(err).To(BeNil())
+		}
+	})
+})
+
+var _ = Describe("ValidateAppName", func() {
+	It("reports error if an app name violates constraints", func() {
+		type testCase struct {
+			appName string
+			err     error
+		}
+
+		testCases := []testCase{
+			{"foo.BAR-1.2_baz_qux", nil},
+
+			{"", ErrAppNameIsRequired},
+			{"#", ErrInvalidAppName},
+		}
+
+		for _, tc := range testCases {
+			err := ValidateAppName(tc.appName)
+			if tc.err != nil {
+				Expect(errors.Is(err, tc.err)).To(BeTrue())
+				continue
+			}
+			Expect(err).To(BeNil())
+		}
+	})
+})

pkg/flameql/parse.go

@@ -28,7 +28,7 @@ func ParseQuery(s string) (*Query, error) {
 			q.Matchers = m
 			return &q, nil
 		default:
-			if !IsTagKeyRuneAllowed(c) {
+			if !IsAppNameRuneAllowed(c) {
 				return nil, newErr(ErrInvalidAppName, s[:offset+1])
 			}
 		}

pkg/storage/query_test.go

@@ -142,21 +142,21 @@ var _ = Describe("Querying", func() {
 						dimension.Key("app.name{baz=xxx,foo=bar}"),
 						dimension.Key("app.name{baz=xxx,waldo=fred}"),
 					}},
-					{`app.name{non-existing-key!="bar"}`, []dimension.Key{
+					{`app.name{non_existing_key!="bar"}`, []dimension.Key{
 						dimension.Key("app.name{baz=qux,foo=bar}"),
 						dimension.Key("app.name{baz=xxx,foo=bar}"),
 						dimension.Key("app.name{baz=xxx,waldo=fred}"),
 					}},
-					{`app.name{non-existing-key!~"bar"}`, []dimension.Key{
+					{`app.name{non_existing_key!~"bar"}`, []dimension.Key{
 						dimension.Key("app.name{baz=qux,foo=bar}"),
 						dimension.Key("app.name{baz=xxx,foo=bar}"),
 						dimension.Key("app.name{baz=xxx,waldo=fred}"),
 					}},
 
 					{`app.name{foo="non-existing-value"}`, nil},
 					{`app.name{foo=~"non-existing-.*"}`, nil},
-					{`app.name{non-existing-key="bar"}`, nil},
-					{`app.name{non-existing-key=~"bar"}`, nil},
+					{`app.name{non_existing_key="bar"}`, nil},
+					{`app.name{non_existing_key=~"bar"}`, nil},
 
 					{`non-existing-app{}`, nil},
 				}

pkg/storage/segment/key.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/pyroscope-io/pyroscope/pkg/flameql"
 	"github.com/pyroscope-io/pyroscope/pkg/structs/sortedmap"
 )
 
@@ -28,26 +29,21 @@ const (
 
 func NewKey(labels map[string]string) *Key { return &Key{labels: labels} }
 
-// TODO: should rewrite this at some point to not rely on regular expressions & splits
 func ParseKey(name string) (*Key, error) {
-	k := &Key{
-		labels: make(map[string]string),
-	}
-
-	p := parser{
-		parserState: nameParserState,
-		key:         "",
-		value:       "",
-	}
-
+	k := &Key{labels: make(map[string]string)}
+	p := parser{parserState: nameParserState}
+	var err error
 	for _, r := range name + "{" {
 		switch p.parserState {
 		case nameParserState:
-			p.nameParserCase(r, k)
+			err = p.nameParserCase(r, k)
 		case tagKeyParserState:
 			p.tagKeyParserCase(r)
 		case tagValueParserState:
-			p.tagValueParserCase(r, k)
+			err = p.tagValueParserCase(r, k)
+		}
+		if err != nil {
+			return nil, err
 		}
 	}
 	return k, nil
@@ -60,14 +56,19 @@ type parser struct {
 }
 
 // ParseKey's nameParserState switch case
-func (p *parser) nameParserCase(r int32, k *Key) {
+func (p *parser) nameParserCase(r int32, k *Key) error {
 	switch r {
 	case '{':
 		p.parserState = tagKeyParserState
-		k.labels["__name__"] = strings.TrimSpace(p.value)
+		appName := strings.TrimSpace(p.value)
+		if err := flameql.ValidateAppName(appName); err != nil {
+			return err
+		}
+		k.labels["__name__"] = appName
 	default:
 		p.value += string(r)
 	}
+	return nil
 }
 
 // ParseKey's tagKeyParserState switch case
@@ -84,15 +85,22 @@ func (p *parser) tagKeyParserCase(r int32) {
 }
 
 // ParseKey's tagValueParserState switch case
-func (p *parser) tagValueParserCase(r int32, k *Key) {
+func (p *parser) tagValueParserCase(r int32, k *Key) error {
 	switch r {
 	case ',', '}':
 		p.parserState = tagKeyParserState
-		k.labels[strings.TrimSpace(p.key)] = strings.TrimSpace(p.value)
+		key := strings.TrimSpace(p.key)
+		if !flameql.IsTagKeyReserved(key) {
+			if err := flameql.ValidateTagKey(key); err != nil {
+				return err
+			}
+		}
+		k.labels[key] = strings.TrimSpace(p.value)
 		p.key = ""
 	default:
 		p.value += string(r)
 	}
+	return nil
 }
 
 func (k *Key) SegmentKey() string {