Merge pull request #52 from Shopify/voucher-client-user-agent

voucher-client: send User-Agent

v2/client/client.go

@@ -24,8 +24,11 @@ type Client struct {
 	httpClient *http.Client
 	username   string
 	password   string
+	userAgent  string
 }
 
+const DefaultUserAgent = "voucher-client/2"
+
 // NewClient creates a new Client set to connect to the passed
 // hostname.
 func NewClient(voucherURL string) (*Client, error) {
@@ -54,6 +57,11 @@ func (c *Client) SetBasicAuth(username, password string) {
 	c.password = password
 }
 
+// SetUserAgent customizes the user agent used by the client
+func (c *Client) SetUserAgent(userAgent string) {
+	c.userAgent = userAgent
+}
+
 // CopyURL returns a copy of this client's URL
 func (c *Client) CopyURL() *url.URL {
 	urlCopy := (*c.url)
@@ -74,6 +82,9 @@ func (c *Client) newVoucherRequest(ctx context.Context, url string, image refere
 		return nil, err
 	}
 	req.Header.Set("Content-Type", "application/json")
+	if c.userAgent != "" {
+		req.Header.Set("User-Agent", c.userAgent)
+	}
 	if c.username != "" && c.password != "" {
 		req.SetBasicAuth(c.username, c.password)
 	}
@@ -122,6 +133,7 @@ func newClient(voucherURL string, httpClient *http.Client) (*Client, error) {
 	client := &Client{
 		url:        u,
 		httpClient: httpClient,
+		userAgent:  DefaultUserAgent,
 	}
 	return client, nil
 }

v2/client/client_test.go

@@ -38,7 +38,7 @@ func TestNewClient(t *testing.T) {
 		t.Run(label, func(t *testing.T) {
 			c, err := client.NewClient(tc.input)
 			if tc.errMsg != "" {
-				assert.Equal(t, tc.errMsg, err.Error())
+				assert.EqualError(t, err, tc.errMsg)
 			} else {
 				require.NoError(t, err)
 				assert.Equal(t, tc.hostname, c.CopyURL().String())
@@ -50,7 +50,7 @@ func TestNewClient(t *testing.T) {
 const image = "gcr.io/project/image@sha256:0000000000000000000000000000000000000000000000000000000000000000"
 
 func TestVoucher_Check(t *testing.T) {
-	v := &mockVoucher{}
+	v := &mockVoucher{t: t}
 	v.checks = append(v.checks, &voucher.Response{Image: image, Success: true})
 	srv := httptest.NewServer(v)
 	defer srv.Close()
@@ -62,8 +62,24 @@ func TestVoucher_Check(t *testing.T) {
 	assert.True(t, res.Success)
 }
 
+func TestVoucher_CustomUserAgent(t *testing.T) {
+	const customUserAgent = "my-awesome-voucher/1.0"
+	v := &mockVoucher{t: t}
+	v.ua = customUserAgent
+	v.checks = append(v.checks, &voucher.Response{Image: image, Success: true})
+	srv := httptest.NewServer(v)
+	defer srv.Close()
+
+	c, err := client.NewClient(srv.URL)
+	c.SetUserAgent(customUserAgent)
+	require.NoError(t, err)
+	res, err := c.Check(context.Background(), "diy", canonical(t, image))
+	require.NoError(t, err)
+	assert.True(t, res.Success)
+}
+
 func TestVoucher_Verify(t *testing.T) {
-	v := &mockVoucher{}
+	v := &mockVoucher{t: t}
 	v.verifications = append(v.verifications, &voucher.Response{Image: image, Success: true})
 	srv := httptest.NewServer(v)
 	defer srv.Close()
@@ -76,11 +92,21 @@ func TestVoucher_Verify(t *testing.T) {
 }
 
 type mockVoucher struct {
+	t             *testing.T
+	ua            string
 	checks        []*voucher.Response
 	verifications []*voucher.Response
 }
 
 func (v *mockVoucher) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	hdr := r.Header
+	assert.Equal(v.t, "application/json", hdr.Get("Content-Type"))
+	if v.ua != "" {
+		assert.Equal(v.t, v.ua, hdr.Get("User-Agent"))
+	} else {
+		assert.Equal(v.t, client.DefaultUserAgent, hdr.Get("User-Agent"))
+	}
+
 	var req voucher.Request
 	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
 		http.Error(w, "", http.StatusBadRequest)

v2/cmd/voucher_client/config.go

@@ -26,25 +26,31 @@ func getCheck() string {
 }
 
 func getVoucherClient() (voucher.Interface, error) {
+	var newClient *client.Client
+	var err error
 	switch strings.ToLower(defaultConfig.Auth) {
 	case "idtoken":
-		newClient, err := client.NewAuthClient(defaultConfig.Server)
-		return newClient, err
+		newClient, err = client.NewAuthClient(defaultConfig.Server)
+		if err != nil {
+			return nil, err
+		}
 	case "basic":
-		newClient, err := client.NewClient(defaultConfig.Server)
-		if err == nil {
-			newClient.SetBasicAuth(defaultConfig.Username, defaultConfig.Password)
+		newClient, err = client.NewClient(defaultConfig.Server)
+		if err != nil {
+			return nil, err
 		}
-		return newClient, err
+		newClient.SetBasicAuth(defaultConfig.Username, defaultConfig.Password)
 	case "default-access-token":
-		newClient, err := NewAuthClientWithToken(context.Background(), defaultConfig.Server)
+		newClient, err = NewAuthClientWithToken(context.Background(), defaultConfig.Server)
 		if err != nil {
 			return nil, err
 		}
-		return newClient, err
 	default:
 		return nil, fmt.Errorf("invalid auth value: %q", defaultConfig.Auth)
 	}
+
+	newClient.SetUserAgent(fmt.Sprintf("voucher-client/%s", version))
+	return newClient, nil
 }
 
 func newContext() (context.Context, context.CancelFunc) {