From 3f03127314bd768efd0bef57915320545afcdd78 Mon Sep 17 00:00:00 2001
From: Zeger-Jan van de Weg <git@zjvandeweg.nl>
Date: Wed, 14 Oct 2020 16:56:49 +0200
Subject: client: Allow User-Agent header to be overridden

The user agent for requests to the internal API endpoints used the
default Go provided user agent. This change updates that to always set
something else, by default `GitLab-Shell`.

Than for others importing the package, there's a new API to set it to
something else. This has been done with new method, a setter, to
maintain backwards compatibility in the API.
---
 client/gitlabnet.go       | 16 ++++++++++++++--
 client/httpclient_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/client/gitlabnet.go b/client/gitlabnet.go
index b908d04..fcefb24 100644
--- a/client/gitlabnet.go
+++ b/client/gitlabnet.go
@@ -19,6 +19,7 @@ import (
 const (
 	internalApiPath  = "/api/v4/internal"
 	secretHeaderName = "Gitlab-Shared-Secret"
+	defaultUserAgent = "GitLab-Shell"
 )
 
 type ErrorResponse struct {
@@ -26,8 +27,11 @@ type ErrorResponse struct {
 }
 
 type GitlabNetClient struct {
-	httpClient             *HttpClient
-	user, password, secret string
+	httpClient *HttpClient
+	user       string
+	password   string
+	secret     string
+	userAgent  string
 }
 
 func NewGitlabNetClient(
@@ -46,9 +50,16 @@ func NewGitlabNetClient(
 		user:       user,
 		password:   password,
 		secret:     secret,
+		userAgent:  defaultUserAgent,
 	}, nil
 }
 
+// SetUserAgent overrides the default user agent for the User-Agent header field
+// for subsequent requests for the GitlabNetClient
+func (c *GitlabNetClient) SetUserAgent(ua string) {
+	c.userAgent = ua
+}
+
 func normalizePath(path string) string {
 	if !strings.HasPrefix(path, "/") {
 		path = "/" + path
@@ -119,6 +130,7 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da
 	request.Header.Set(secretHeaderName, encodedSecret)
 
 	request.Header.Add("Content-Type", "application/json")
+	request.Header.Add("User-Agent", c.userAgent)
 	request.Close = true
 
 	start := time.Now()
diff --git a/client/httpclient_test.go b/client/httpclient_test.go
index 97e1384..59dabc4 100644
--- a/client/httpclient_test.go
+++ b/client/httpclient_test.go
@@ -94,6 +94,35 @@ func TestEmptyBasicAuthSettings(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestRequestWithUserAgent(t *testing.T) {
+	const gitalyUserAgent = "gitaly/13.5.0"
+	requests := []testserver.TestRequestHandler{
+		{
+			Path: "/api/v4/internal/default_user_agent",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				assert.Equal(t, defaultUserAgent, r.UserAgent())
+			},
+		},
+		{
+			Path: "/api/v4/internal/override_user_agent",
+			Handler: func(w http.ResponseWriter, r *http.Request) {
+				assert.Equal(t, gitalyUserAgent, r.UserAgent())
+			},
+		},
+	}
+
+	client, cleanup := setup(t, "", "", requests)
+	defer cleanup()
+
+	_, err := client.Get(context.Background(), "/default_user_agent")
+	require.NoError(t, err)
+
+	client.SetUserAgent(gitalyUserAgent)
+	_, err = client.Get(context.Background(), "/override_user_agent")
+	require.NoError(t, err)
+
+}
+
 func setup(t *testing.T, username, password string, requests []testserver.TestRequestHandler) (*GitlabNetClient, func()) {
 	url, cleanup := testserver.StartHttpServer(t, requests)
 
-- 
cgit v1.2.1