From 0b5de5f9020137edcc0533f6dd1aec298fed6555 Mon Sep 17 00:00:00 2001
From: Anton Tolchanov <1687799+knyar@users.noreply.github.com>
Date: Thu, 17 Aug 2023 19:41:40 +0100
Subject: [PATCH] tailscale: support a custom user-agent (#57)

This allows specifying a customer user-agent used for outgoing HTTP
requests.
---
 tailscale/client.go         | 28 ++++++++++++++++++++++------
 tailscale/client_test.go    | 22 ++++++++++++++++++++++
 tailscale/tailscale_test.go |  3 +++
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/tailscale/client.go b/tailscale/client.go
index 34fb51c..1538e7f 100644
--- a/tailscale/client.go
+++ b/tailscale/client.go
@@ -21,10 +21,11 @@ import (
 type (
 	// Client type is used to perform actions against the Tailscale API.
 	Client struct {
-		apiKey  string
-		http    *http.Client
-		baseURL *url.URL
-		tailnet string
+		apiKey    string
+		http      *http.Client
+		baseURL   *url.URL
+		tailnet   string
+		userAgent string // empty string means Go's default value.
 	}
 
 	// APIError type describes an error as returned by the Tailscale API.
@@ -47,6 +48,7 @@ type (
 const baseURL = "https://api.tailscale.com"
 const contentType = "application/json"
 const defaultHttpClientTimeout = time.Minute
+const defaultUserAgent = "tailscale-client-go"
 
 // NewClient returns a new instance of the Client type that will perform operations against a chosen tailnet and will
 // provide the apiKey for authorization. Additional options can be provided, see ClientOption for more details.
@@ -65,8 +67,9 @@ func NewClient(apiKey, tailnet string, options ...ClientOption) (*Client, error)
 	}
 
 	c := &Client{
-		baseURL: u,
-		tailnet: tailnet,
+		baseURL:   u,
+		tailnet:   tailnet,
+		userAgent: defaultUserAgent,
 	}
 
 	if apiKey != "" {
@@ -122,6 +125,15 @@ func WithOAuthClientCredentials(clientID, clientSecret string, scopes []string)
 	}
 }
 
+// WithUserAgent sets a custom User-Agent header in HTTP requests.
+// Passing an empty string will make the client use Go's default value.
+func WithUserAgent(ua string) ClientOption {
+	return func(c *Client) error {
+		c.userAgent = ua
+		return nil
+	}
+}
+
 // TODO: consider setting `headers` and `body` via opts to decrease the number of arguments.
 func (c *Client) buildRequest(ctx context.Context, method, uri string, headers map[string]string, body interface{}) (*http.Request, error) {
 	u, err := c.baseURL.Parse(uri)
@@ -142,6 +154,10 @@ func (c *Client) buildRequest(ctx context.Context, method, uri string, headers m
 		return nil, err
 	}
 
+	if c.userAgent != "" {
+		req.Header.Set("User-Agent", c.userAgent)
+	}
+
 	for k, v := range headers {
 		req.Header.Set(k, v)
 	}
diff --git a/tailscale/client_test.go b/tailscale/client_test.go
index 2126ccb..9b97d1c 100644
--- a/tailscale/client_test.go
+++ b/tailscale/client_test.go
@@ -952,3 +952,25 @@ func TestClient_ValidateACL(t *testing.T) {
 	assert.EqualValues(t, http.MethodPost, server.Method)
 	assert.EqualValues(t, "/api/v2/tailnet/example.com/acl/validate", server.Path)
 }
+
+func TestClient_UserAgent(t *testing.T) {
+	t.Parallel()
+	client, server := NewTestHarness(t)
+	server.ResponseCode = http.StatusOK
+
+	// Check the default user-agent.
+	assert.NoError(t, client.SetDeviceAuthorized(context.Background(), "test", true))
+	assert.Equal(t, "tailscale-client-go", server.Header.Get("User-Agent"))
+
+	// Check a custom user-agent.
+	client, err := tailscale.NewClient("fake key", "", tailscale.WithBaseURL(server.BaseURL), tailscale.WithUserAgent("custom-user-agent"))
+	assert.NoError(t, err)
+	assert.NoError(t, client.SetDeviceAuthorized(context.Background(), "test", true))
+	assert.Equal(t, "custom-user-agent", server.Header.Get("User-Agent"))
+
+	// Overriding with an empty string uses runtime's default value.
+	client, err = tailscale.NewClient("fake key", "", tailscale.WithBaseURL(server.BaseURL), tailscale.WithUserAgent(""))
+	assert.NoError(t, err)
+	assert.NoError(t, client.SetDeviceAuthorized(context.Background(), "test", true))
+	assert.Contains(t, server.Header.Get("User-Agent"), "Go-http-client")
+}
diff --git a/tailscale/tailscale_test.go b/tailscale/tailscale_test.go
index f869511..eb27831 100644
--- a/tailscale/tailscale_test.go
+++ b/tailscale/tailscale_test.go
@@ -17,6 +17,8 @@ import (
 type TestServer struct {
 	t *testing.T
 
+	BaseURL string
+
 	Method string
 	Path   string
 	Body   *bytes.Buffer
@@ -53,6 +55,7 @@ func NewTestHarness(t *testing.T) (*tailscale.Client, *TestServer) {
 	})
 
 	baseURL := fmt.Sprintf("http://localhost:%v", listener.Addr().(*net.TCPAddr).Port)
+	testServer.BaseURL = baseURL
 	client, err := tailscale.NewClient("not a real key", "example.com", tailscale.WithBaseURL(baseURL))
 	assert.NoError(t, err)