diff --git a/storage/remote/client.go b/storage/remote/client.go index 62218cfba9..715ea42bf6 100644 --- a/storage/remote/client.go +++ b/storage/remote/client.go @@ -63,8 +63,15 @@ const ( ) var ( - // UserAgent represents Prometheus version to use for user agent header. - UserAgent = fmt.Sprintf("Prometheus/%s", version.Version) + // internalUserAgent should not be modified as it is to allow tracking of the + // specific Prometheus version in use if UserAgent below is over-ridden. + // This value is used to set the X-Prometheus-User-Agent header. + internalUserAgent = fmt.Sprintf("Prometheus/%s", version.Version) + + // UserAgent represents Prometheus version to use for the User-Agent header. + // This is now modifiable. + // It defaults to the internalUserAgent value defined above. + UserAgent = internalUserAgent remoteWriteContentTypeHeaders = map[config.RemoteWriteProtoMsg]string{ config.RemoteWriteProtoMsgV1: appProtoContentType, // Also application/x-protobuf;proto=prometheus.WriteRequest but simplified for compatibility with 1.x spec. @@ -151,6 +158,11 @@ type ReadClient interface { Read(ctx context.Context, query *prompb.Query, sortSeries bool) (storage.SeriesSet, error) } +// SetUserAgent allows the User-Agent header value to be over-ridden. +func SetUserAgent(userAgent string) { + UserAgent = userAgent +} + // NewReadClient creates a new client for remote read. func NewReadClient(name string, conf *ClientConfig) (ReadClient, error) { httpClient, err := config_util.NewClientFromConfig(conf.HTTPClientConfig, "remote_storage_read_client") @@ -262,6 +274,7 @@ func (c *Client) Store(ctx context.Context, req []byte, attempt int) (WriteRespo httpReq.Header.Add("Content-Encoding", string(c.writeCompression)) httpReq.Header.Set("Content-Type", remoteWriteContentTypeHeaders[c.writeProtoMsg]) httpReq.Header.Set("User-Agent", UserAgent) + httpReq.Header.Set("X-Prometheus-User-Agent", internalUserAgent) if c.writeProtoMsg == config.RemoteWriteProtoMsgV1 { // Compatibility mode for 1.0. httpReq.Header.Set(RemoteWriteVersionHeader, RemoteWriteVersion1HeaderValue) @@ -363,6 +376,7 @@ func (c *Client) Read(ctx context.Context, query *prompb.Query, sortSeries bool) httpReq.Header.Add("Accept-Encoding", "snappy") httpReq.Header.Set("Content-Type", "application/x-protobuf") httpReq.Header.Set("User-Agent", UserAgent) + httpReq.Header.Set("X-Prometheus-User-Agent", internalUserAgent) httpReq.Header.Set("X-Prometheus-Remote-Read-Version", "0.1.0") ctx, cancel := context.WithTimeout(ctx, c.timeout) diff --git a/storage/remote/client_test.go b/storage/remote/client_test.go index c8b3d487e7..94510ce906 100644 --- a/storage/remote/client_test.go +++ b/storage/remote/client_test.go @@ -90,6 +90,146 @@ func TestStoreHTTPErrorHandling(t *testing.T) { } } +func TestReadClientUserAgent(t *testing.T) { + tests := []struct { + name string + userAgent string + }{ + { + name: "default-no-override", + userAgent: "", + }, + { + name: "overridden", + userAgent: "ArgleBargle/1.2.3", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var called bool + server := httptest.NewServer( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + receivedHeaders := r.Header + + // Check the X-Prometheus-User-Agent header. + require.Equal(t, []string{internalUserAgent}, receivedHeaders.Values("X-Prometheus-User-Agent"), + "expected X-Prometheus-User-Agent header to be default value of %q", internalUserAgent) + + if test.userAgent == "" { + // Expect original header value. + require.Equal(t, []string{internalUserAgent}, receivedHeaders.Values("User-Agent"), + "expected User-Agent header to be default value of %q", internalUserAgent) + } else { + // Expect over-ridden header value. + require.Equal(t, []string{test.userAgent}, receivedHeaders.Values("User-Agent"), + "expected User-Agent header to be over-ridden value of %q", test.userAgent) + } + w.Header().Set("Content-Type", "text/plain") + }), + ) + defer server.Close() + + u, err := url.Parse(server.URL) + require.NoError(t, err) + + conf := &ClientConfig{ + URL: &config_util.URL{URL: u}, + Timeout: model.Duration(5 * time.Second), + ChunkedReadLimit: config.DefaultChunkedReadLimit, + } + + // Set the User-Agent. + if test.userAgent == "" { + // Ensure it is set to the default. + SetUserAgent(internalUserAgent) + } else { + SetUserAgent(test.userAgent) + } + + c, err := NewReadClient("test", conf) + require.NoError(t, err) + + query := &prompb.Query{} + + _, err = c.Read(context.Background(), query, false) + + require.ErrorContains(t, err, "unsupported content type") + + require.True(t, called, "The remote server wasn't called") + }) + } +} + +func TestWriteClientUserAgent(t *testing.T) { + tests := []struct { + name string + userAgent string + }{ + { + name: "default-no-override", + userAgent: "", + }, + { + name: "overridden", + userAgent: "ArgleBargle/1.2.3", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var called bool + server := httptest.NewServer( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + receivedHeaders := r.Header + + // Check the X-Prometheus-User-Agent header. + require.Equal(t, []string{internalUserAgent}, receivedHeaders.Values("X-Prometheus-User-Agent"), + "expected X-Prometheus-User-Agent header to be default value of %q", internalUserAgent) + + if test.userAgent == "" { + // Expect original header value. + require.Equal(t, []string{internalUserAgent}, receivedHeaders.Values("User-Agent"), + "expected User-Agent header to be default value of %q", internalUserAgent) + } else { + // Expect over-ridden header value. + require.Equal(t, []string{test.userAgent}, receivedHeaders.Values("User-Agent"), + "expected User-Agent header to be over-ridden value of %q", test.userAgent) + } + // w.Header().Set("Content-Type", "text/plain") + }), + ) + defer server.Close() + + serverURL, err := url.Parse(server.URL) + require.NoError(t, err) + + conf := &ClientConfig{ + URL: &config_util.URL{URL: serverURL}, + Timeout: model.Duration(time.Second), + } + + // Set the User-Agent. + if test.userAgent == "" { + // Ensure it is set to the default. + SetUserAgent(internalUserAgent) + } else { + SetUserAgent(test.userAgent) + } + + hash, err := toHash(conf) + require.NoError(t, err) + c, err := NewWriteClient(hash, conf) + require.NoError(t, err) + + _, err = c.Store(context.Background(), []byte{}, 0) + require.NoError(t, err) + + require.True(t, called, "The remote server wasn't called") + }) + } +} + func TestClientRetryAfter(t *testing.T) { setupServer := func(statusCode int) *httptest.Server { return httptest.NewServer(