diff --git a/cmd/profilecli/client.go b/cmd/profilecli/client.go index 9aa835cdbf..cffcf6c17c 100644 --- a/cmd/profilecli/client.go +++ b/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) } diff --git a/cmd/profilecli/client_test.go b/cmd/profilecli/client_test.go new file mode 100644 index 0000000000..4073b0e6a8 --- /dev/null +++ b/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")) + }) + } +}