feat: Use allow-utf8-labelnames client capability in profilecli (#4490)

Author: jake-kramer

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
@@ -49,7 +80,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"))
+		})
+	}
+}