trace: Reduce regexp memory use in tracestate (#4664)

trace/tracestate.go

@@ -28,9 +28,9 @@ const (
 
 	// based on the W3C Trace Context specification, see
 	// https://www.w3.org/TR/trace-context-1/#tracestate-header
-	noTenantKeyFormat   = `[a-z][_0-9a-z\-\*\/]{0,255}`
-	withTenantKeyFormat = `[a-z0-9][_0-9a-z\-\*\/]{0,240}@[a-z][_0-9a-z\-\*\/]{0,13}`
-	valueFormat         = `[\x20-\x2b\x2d-\x3c\x3e-\x7e]{0,255}[\x21-\x2b\x2d-\x3c\x3e-\x7e]`
+	noTenantKeyFormat   = `[a-z][_0-9a-z\-\*\/]*`
+	withTenantKeyFormat = `[a-z0-9][_0-9a-z\-\*\/]*@[a-z][_0-9a-z\-\*\/]*`
+	valueFormat         = `[\x20-\x2b\x2d-\x3c\x3e-\x7e]*[\x21-\x2b\x2d-\x3c\x3e-\x7e]`
 
 	errInvalidKey    errorConst = "invalid tracestate key"
 	errInvalidValue  errorConst = "invalid tracestate value"
@@ -40,9 +40,10 @@ const (
 )
 
 var (
-	keyRe    = regexp.MustCompile(`^((` + noTenantKeyFormat + `)|(` + withTenantKeyFormat + `))$`)
-	valueRe  = regexp.MustCompile(`^(` + valueFormat + `)$`)
-	memberRe = regexp.MustCompile(`^\s*((` + noTenantKeyFormat + `)|(` + withTenantKeyFormat + `))=(` + valueFormat + `)\s*$`)
+	noTenantKeyRe   = regexp.MustCompile(`^` + noTenantKeyFormat + `$`)
+	withTenantKeyRe = regexp.MustCompile(`^` + withTenantKeyFormat + `$`)
+	valueRe         = regexp.MustCompile(`^` + valueFormat + `$`)
+	memberRe        = regexp.MustCompile(`^\s*((?:` + noTenantKeyFormat + `)|(?:` + withTenantKeyFormat + `))=(` + valueFormat + `)\s*$`)
 )
 
 type member struct {
@@ -51,25 +52,34 @@ type member struct {
 }
 
 func newMember(key, value string) (member, error) {
-	if !keyRe.MatchString(key) {
+	if len(key) > 256 {
 		return member{}, fmt.Errorf("%w: %s", errInvalidKey, key)
 	}
-	if !valueRe.MatchString(value) {
+	if !noTenantKeyRe.MatchString(key) {
+		if !withTenantKeyRe.MatchString(key) {
+			return member{}, fmt.Errorf("%w: %s", errInvalidKey, key)
+		}
+		atIndex := strings.LastIndex(key, "@")
+		if atIndex > 241 || len(key)-1-atIndex > 14 {
+			return member{}, fmt.Errorf("%w: %s", errInvalidKey, key)
+		}
+	}
+	if len(value) > 256 || !valueRe.MatchString(value) {
 		return member{}, fmt.Errorf("%w: %s", errInvalidValue, value)
 	}
 	return member{Key: key, Value: value}, nil
 }
 
 func parseMember(m string) (member, error) {
 	matches := memberRe.FindStringSubmatch(m)
-	if len(matches) != 5 {
+	if len(matches) != 3 {
 		return member{}, fmt.Errorf("%w: %s", errInvalidMember, m)
 	}
-
-	return member{
-		Key:   matches[1],
-		Value: matches[4],
-	}, nil
+	result, e := newMember(matches[1], matches[2])
+	if e != nil {
+		return member{}, fmt.Errorf("%w: %s", errInvalidMember, m)
+	}
+	return result, nil
 }
 
 // String encodes member into a string compliant with the W3C Trace Context

trace/tracestate_test.go

@@ -552,3 +552,37 @@ func TestTraceStateImmutable(t *testing.T) {
 	assert.Equal(t, v0, ts2.Get(k0))
 	assert.Equal(t, "", ts3.Get(k0))
 }
+
+func BenchmarkParseTraceState(b *testing.B) {
+	benches := []struct {
+		name string
+		in   string
+	}{
+		{
+			name: "single key",
+			in:   "somewhatRealisticKeyLength=someValueAbcdefgh1234567890",
+		},
+		{
+			name: "tenant single key",
+			in:   "somewhatRealisticKeyLength@someTenant=someValueAbcdefgh1234567890",
+		},
+		{
+			name: "three keys",
+			in:   "someKeyName.One=someValue1,someKeyName.Two=someValue2,someKeyName.Three=someValue3",
+		},
+		{
+			name: "tenant three keys",
+			in:   "someKeyName.One@tenant=someValue1,someKeyName.Two@tenant=someValue2,someKeyName.Three@tenant=someValue3",
+		},
+	}
+	for _, bench := range benches {
+		b.Run(bench.name, func(b *testing.B) {
+			b.ReportAllocs()
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				_, _ = ParseTraceState(bench.in)
+			}
+		})
+	}
+}