feat: Implement utf8 label name client capability

pkg/featureflags/client_capability.go

@@ -0,0 +1,117 @@
+package featureflags
+
+import (
+	"context"
+	"mime"
+	"net/http"
+	"strings"
+
+	"connectrpc.com/connect"
+	"github.com/go-kit/log/level"
+	"github.com/grafana/dskit/middleware"
+	"github.com/grafana/pyroscope/pkg/util"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	// Capability names - update parseClientCapabilities below when new capabilities added
+	allowUtf8LabelNamesCapabilityName string = "allow-utf8-labelnames"
+)
+
+// Define a custom context key type to avoid collisions
+type contextKey struct{}
+
+type ClientCapabilities struct {
+	AllowUtf8LabelNames bool
+}
+
+func WithClientCapabilities(ctx context.Context, clientCapabilities ClientCapabilities) context.Context {
+	return context.WithValue(ctx, contextKey{}, clientCapabilities)
+}
+
+func GetClientCapabilities(ctx context.Context) (ClientCapabilities, bool) {
+	value, ok := ctx.Value(contextKey{}).(ClientCapabilities)
+	return value, ok
+}
+
+func ClientCapabilitiesGRPCMiddleware() grpc.UnaryServerInterceptor {
+	return func(
+		ctx context.Context,
+		req interface{},
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		// Extract metadata from context
+		md, ok := metadata.FromIncomingContext(ctx)
+		if !ok {
+			return handler(ctx, req)
+		}
+
+		// Convert metadata to http.Header for reuse of existing parsing logic
+		httpHeader := make(http.Header)
+		for key, values := range md {
+			// gRPC metadata keys are lowercase, HTTP headers are case-insensitive
+			httpHeader[http.CanonicalHeaderKey(key)] = values
+		}
+
+		// Reuse existing HTTP header parsing
+		clientCapabilities, err := parseClientCapabilities(httpHeader)
+		if err != nil {
+			return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		}
+
+		enhancedCtx := WithClientCapabilities(ctx, clientCapabilities)
+		return handler(enhancedCtx, req)
+	}
+}
+
+// ClientCapabilitiesHttpMiddleware creates middleware that extracts and parses the
+// `Accept` header for capabilities the client supports
+func ClientCapabilitiesHttpMiddleware() middleware.Interface {
+	return middleware.Func(func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			clientCapabilities, err := parseClientCapabilities(r.Header)
+			if err != nil {
+				http.Error(w, "Invalid header format: "+err.Error(), http.StatusBadRequest)
+				return
+			}
+
+			ctx := WithClientCapabilities(r.Context(), clientCapabilities)
+			next.ServeHTTP(w, r.WithContext(ctx))
+		})
+	})
+}
+
+func parseClientCapabilities(header http.Header) (ClientCapabilities, error) {
+	acceptHeaderValues := header.Values("Accept")
+
+	var capabilities ClientCapabilities
+
+	for _, acceptHeaderValue := range acceptHeaderValues {
+		if acceptHeaderValue != "" {
+			accepts := strings.Split(acceptHeaderValue, ",")
+
+			for _, accept := range accepts {
+				if _, params, err := mime.ParseMediaType(accept); err != nil {
+					return capabilities, err
+				} else {
+					for k, v := range params {
+						switch k {
+						case allowUtf8LabelNamesCapabilityName:
+							if v == "true" {
+								capabilities.AllowUtf8LabelNames = true
+							}
+						default:
+							level.Debug(util.Logger).Log(
+								"msg", "unknown capability parsed from Accept header",
+								"acceptHeaderKey", k,
+								"acceptHeaderValue", v)
+						}
+					}
+				}
+			}
+		}
+	}
+	return capabilities, nil
+}

pkg/featureflags/client_capability_test.go

@@ -0,0 +1,239 @@
+package featureflags
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_parseClientCapabilities(t *testing.T) {
+	tests := []struct {
+		Name         string
+		Header       http.Header
+		Want         ClientCapabilities
+		WantError    bool
+		ErrorMessage string
+	}{
+		{
+			Name:   "empty header returns default capabilities",
+			Header: http.Header{},
+			Want:   ClientCapabilities{AllowUtf8LabelNames: false},
+		},
+		{
+			Name: "no Accept header returns default capabilities",
+			Header: http.Header{
+				"Content-Type": []string{"application/json"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: false},
+		},
+		{
+			Name: "empty Accept header value returns default capabilities",
+			Header: http.Header{
+				"Accept": []string{""},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: false},
+		},
+		{
+			Name: "simple Accept header without capabilities",
+			Header: http.Header{
+				"Accept": []string{"application/json"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: false},
+		},
+		{
+			Name: "Accept header with utf8 label names capability true",
+			Header: http.Header{
+				"Accept": []string{"*/*; allow-utf8-labelnames=true"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "Accept header with utf8 label names capability false",
+			Header: http.Header{
+				"Accept": []string{"*/*; allow-utf8-labelnames=false"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: false},
+		},
+		{
+			Name: "Accept header with utf8 label names capability invalid value",
+			Header: http.Header{
+				"Accept": []string{"*/*; allow-utf8-labelnames=invalid"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: false},
+		},
+		{
+			Name: "Accept header with unknown capability",
+			Header: http.Header{
+				"Accept": []string{"*/*; unknown-capability=true"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: false},
+		},
+		{
+			Name: "Accept header with multiple capabilities",
+			Header: http.Header{
+				"Accept": []string{"*/*; allow-utf8-labelnames=true; unknown-capability=false"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "multiple Accept header values",
+			Header: http.Header{
+				"Accept": []string{"application/json", "*/*; allow-utf8-labelnames=true"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "multiple Accept header values with different capabilities",
+			Header: http.Header{
+				"Accept": []string{
+					"application/json; allow-utf8-labelnames=false",
+					"*/*; allow-utf8-labelnames=true",
+				},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "Accept header with quality values",
+			Header: http.Header{
+				"Accept": []string{"text/html; q=0.9; allow-utf8-labelnames=true"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "complex Accept header",
+			Header: http.Header{
+				"Accept": []string{
+					"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8;allow-utf8-labelnames=true",
+				},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "multiple Accept header entries",
+			Header: http.Header{
+				"Accept": []string{
+					"application/json",
+					"text/plain; allow-utf8-labelnames=true",
+					"*/*; q=0.1",
+				},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "invalid mime type in Accept header",
+			Header: http.Header{
+				"Accept": []string{"invalid/mime/type/format"},
+			},
+			Want:         ClientCapabilities{},
+			WantError:    true,
+			ErrorMessage: "mime: unexpected content after media subtype",
+		},
+		{
+			Name: "Accept header with invalid syntax",
+			Header: http.Header{
+				"Accept": []string{"text/html; invalid-parameter-syntax"},
+			},
+			Want:         ClientCapabilities{},
+			WantError:    true,
+			ErrorMessage: "mime: invalid media parameter",
+		},
+		{
+			Name: "mixed valid and invalid Accept header values",
+			Header: http.Header{
+				"Accept": []string{
+					"application/json",
+					"invalid/mime/type/format",
+				},
+			},
+			Want:         ClientCapabilities{},
+			WantError:    true,
+			ErrorMessage: "mime: unexpected content after media subtype",
+		},
+		{
+			// Parameter names are case-insensitive in mime.ParseMediaType
+			Name: "case sensitivity test for capability name",
+			Header: http.Header{
+				"Accept": []string{"*/*; Allow-Utf8-Labelnames=true"},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "whitespace handling in Accept header",
+			Header: http.Header{
+				"Accept": []string{" application/json ; allow-utf8-labelnames=true "},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.Name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := parseClientCapabilities(tt.Header)
+
+			if tt.WantError {
+				require.Error(t, err)
+				if tt.ErrorMessage != "" {
+					require.Contains(t, err.Error(), tt.ErrorMessage)
+				}
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tt.Want, got)
+		})
+	}
+}
+
+func Test_parseClientCapabilities_MultipleCapabilities(t *testing.T) {
+	// This test specifically checks that when the same capability appears
+	// multiple times with different values, the last "true" value wins
+	tests := []struct {
+		Name   string
+		Header http.Header
+		Want   ClientCapabilities
+	}{
+		{
+			Name: "capability appears multiple times - last true wins",
+			Header: http.Header{
+				"Accept": []string{
+					"application/json; allow-utf8-labelnames=false",
+					"text/plain; allow-utf8-labelnames=true",
+				},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "capability appears multiple times - last false loses to earlier true",
+			Header: http.Header{
+				"Accept": []string{
+					"application/json; allow-utf8-labelnames=true",
+					"text/plain; allow-utf8-labelnames=false",
+				},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: true},
+		},
+		{
+			Name: "capability appears multiple times - all false",
+			Header: http.Header{
+				"Accept": []string{
+					"application/json; allow-utf8-labelnames=false",
+					"text/plain; allow-utf8-labelnames=false",
+				},
+			},
+			Want: ClientCapabilities{AllowUtf8LabelNames: false},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.Name, func(t *testing.T) {
+			t.Parallel()
+
+			got, err := parseClientCapabilities(tt.Header)
+			require.NoError(t, err)
+			require.Equal(t, tt.Want, got)
+		})
+	}
+}