feat: Implement utf8 label name client capability

cmd/profilecli/client.go

@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"net/http"
+	"strings"
 
 	"connectrpc.com/connect"
 	"github.com/prometheus/common/version"
@@ -15,10 +16,40 @@ const (
 	protocolTypeConnect = "connect"
 	protocolTypeGRPC    = "grpc"
 	protocolTypeGRPCWeb = "grpc-web"
+
+	acceptHeaderMimeType = "*/*"
 )
 
+var acceptHeaderClientCapabilities = []string{
+	"allow-utf8-labelnames=true",
+}
+
 var userAgentHeader = fmt.Sprintf("pyroscope/%s", version.Version)
 
+func addClientCapabilitiesHeader(r *http.Request, mime string, clientCapabilities []string) {
+	missingClientCapabilities := make([]string, 0, len(clientCapabilities))
+	for _, capability := range clientCapabilities {
+		found := false
+		// Check if any header value already contains this capability
+		for _, value := range r.Header.Values("Accept") {
+			if strings.Contains(value, capability) {
+				found = true
+				break
+			}
+		}
+
+		if !found {
+			missingClientCapabilities = append(missingClientCapabilities, capability)
+		}
+	}
+
+	if len(missingClientCapabilities) > 0 {
+		acceptHeader := mime
+		acceptHeader += "; " + strings.Join(missingClientCapabilities, "; ")
+		r.Header.Add("Accept", acceptHeader)
+	}
+}
+
 type phlareClient struct {
 	TenantID  string
 	URL       string
@@ -46,7 +77,9 @@ func (a *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error)
 		}
 	}
 
+	addClientCapabilitiesHeader(req, acceptHeaderMimeType, acceptHeaderClientCapabilities)
 	req.Header.Set("User-Agent", userAgentHeader)
+
 	return a.next.RoundTrip(req)
 }
 

cmd/profilecli/client_test.go

@@ -0,0 +1,79 @@
+package main
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func Test_AcceptHeader(t *testing.T) {
+	tests := []struct {
+		Name               string
+		Header             http.Header
+		ClientCapabilities []string
+		Want               []string
+	}{
+		{
+			Name:   "empty header adds capability",
+			Header: http.Header{},
+			ClientCapabilities: []string{
+				"allow-utf8-labelnames=true",
+			},
+			Want: []string{"*/*; allow-utf8-labelnames=true"},
+		},
+		{
+			Name: "existing header appends capability",
+			Header: http.Header{
+				"Accept": []string{"application/json"},
+			},
+			ClientCapabilities: []string{
+				"allow-utf8-labelnames=true",
+			},
+			Want: []string{"application/json", "*/*; allow-utf8-labelnames=true"},
+		},
+		{
+			Name: "multiple existing values appends capability",
+			Header: http.Header{
+				"Accept": []string{"application/json", "text/plain"},
+			},
+			ClientCapabilities: []string{
+				"allow-utf8-labelnames=true",
+			},
+			Want: []string{"application/json", "text/plain", "*/*; allow-utf8-labelnames=true"},
+		},
+		{
+			Name: "existing capability is not duplicated",
+			Header: http.Header{
+				"Accept": []string{"*/*; allow-utf8-labelnames=true"},
+			},
+			ClientCapabilities: []string{
+				"allow-utf8-labelnames=true",
+			},
+			Want: []string{"*/*; allow-utf8-labelnames=true"},
+		},
+		{
+			Name: "multiple client capabilities appends capability",
+			Header: http.Header{
+				"Accept": []string{"*/*; allow-utf8-labelnames=true"},
+			},
+			ClientCapabilities: []string{
+				"allow-utf8-labelnames=true",
+				"capability2=false",
+			},
+			Want: []string{"*/*; allow-utf8-labelnames=true", "*/*; capability2=false"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.Name, func(t *testing.T) {
+			t.Parallel()
+			req, _ := http.NewRequest("GET", "example.com", nil)
+			req.Header = tt.Header
+			clientCapabilities := tt.ClientCapabilities
+
+			addClientCapabilitiesHeader(req, acceptHeaderMimeType, clientCapabilities)
+			require.Equal(t, tt.Want, req.Header.Values("Accept"))
+		})
+	}
+}

pkg/featureflags/client_capability.go

@@ -0,0 +1,119 @@
+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
+		// TODO add metrics = # requests like this and # clients [need
+		//  labels for requests and clients/tenet and user agent(?)]
+		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
+}