Adds a user-agent identifier for preflight calls (#779)

Author: mcncl

cmd/preflight/preflight.go

@@ -8,12 +8,14 @@ import (
 	"os"
 	"os/signal"
 	"strconv"
+	"strings"
 	"syscall"
 	"time"
 
 	"github.com/alecthomas/kong"
 	"github.com/google/uuid"
 
+	"github.com/buildkite/cli/v3/cmd/version"
 	buildstate "github.com/buildkite/cli/v3/internal/build/state"
 	"github.com/buildkite/cli/v3/internal/build/watch"
 	"github.com/buildkite/cli/v3/internal/cli"
@@ -39,13 +41,28 @@ var (
 	newFactory    = factory.New
 )
 
+func preflightUserAgentSuffix() string {
+	major := strings.TrimPrefix(version.Version, "v")
+	if i := strings.IndexByte(major, '.'); i >= 0 {
+		major = major[:i]
+	}
+	if major == "" || major == "DEV" {
+		major = "DEV"
+	}
+	return "buildkite-cli-preflight/" + major + ".x"
+}
+
 func (c *PreflightCmd) Help() string {
 	return `Snapshots your working tree (uncommitted, staged, and untracked changes) and pushes it to a bk/preflight/<id> branch. If there are no local changes, pushes HEAD directly.`
 }
 
 func (c *PreflightCmd) Run(kongCtx *kong.Context, globals cli.GlobalFlags) error {
 	rlTransport := bkhttp.NewRateLimitTransport(http.DefaultTransport)
-	f, err := newFactory(factory.WithDebug(globals.EnableDebug()), factory.WithTransport(rlTransport))
+	f, err := newFactory(
+		factory.WithDebug(globals.EnableDebug()),
+		factory.WithTransport(rlTransport),
+		factory.WithUserAgentSuffix(preflightUserAgentSuffix()),
+	)
 	if err != nil {
 		return bkErrors.NewInternalError(err, "failed to initialize CLI", "This is likely a bug", "Report to Buildkite")
 	}

cmd/preflight/preflight_test.go

@@ -78,8 +78,10 @@ func TestPreflightCmd_Run(t *testing.T) {
 		t.Setenv("BUILDKITE_EXPERIMENTS", "preflight")
 
 		var gotReq buildkite.CreateBuild
+		var gotUserAgent string
 		s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			if r.Method == "POST" && strings.Contains(r.URL.Path, "/builds") {
+				gotUserAgent = r.Header.Get("User-Agent")
 				json.NewDecoder(r.Body).Decode(&gotReq)
 				w.Header().Set("Content-Type", "application/json")
 				json.NewEncoder(w).Encode(buildkite.Build{
@@ -128,6 +130,12 @@ func TestPreflightCmd_Run(t *testing.T) {
 		if gotReq.Env["BUILDKITE_PREFLIGHT"] != "true" {
 			t.Errorf("expected BUILDKITE_PREFLIGHT=true (deprecated), got %#v", gotReq.Env)
 		}
+		if !strings.Contains(gotUserAgent, buildkite.DefaultUserAgent) {
+			t.Errorf("expected User-Agent to contain %q, got %q", buildkite.DefaultUserAgent, gotUserAgent)
+		}
+		if !strings.Contains(gotUserAgent, "buildkite-cli-preflight/") {
+			t.Errorf("expected User-Agent to contain preflight token, got %q", gotUserAgent)
+		}
 	})
 
 	t.Run("falls back to git cli when factory cannot open repository", func(t *testing.T) {

pkg/cmd/factory/factory.go

@@ -16,7 +16,7 @@ import (
 	git "github.com/go-git/go-git/v5"
 )
 
-var userAgent string
+var baseUserAgent string
 
 type Factory struct {
 	Config        *config.Config
@@ -35,9 +35,10 @@ type Factory struct {
 type FactoryOpt func(*factoryConfig)
 
 type factoryConfig struct {
-	debug       bool
-	orgOverride string
-	transport   http.RoundTripper
+	debug           bool
+	orgOverride     string
+	transport       http.RoundTripper
+	userAgentSuffix string
 }
 
 // WithDebug enables debug output for REST API calls
@@ -64,6 +65,13 @@ func WithTransport(t http.RoundTripper) FactoryOpt {
 	}
 }
 
+// WithUserAgentSuffix appends an extra product token to the default user agent.
+func WithUserAgentSuffix(suffix string) FactoryOpt {
+	return func(c *factoryConfig) {
+		c.userAgentSuffix = suffix
+	}
+}
+
 // debugTransport wraps an http.RoundTripper and logs requests/responses with sensitive headers redacted
 type debugTransport struct {
 	transport http.RoundTripper
@@ -129,17 +137,25 @@ func redactHeaders(headers http.Header) {
 }
 
 type gqlHTTPClient struct {
-	client *http.Client
-	token  string
+	client    *http.Client
+	token     string
+	userAgent string
 }
 
 func init() {
-	userAgent = fmt.Sprintf("%s buildkite-cli/%s", buildkite.DefaultUserAgent, version.Version)
+	baseUserAgent = fmt.Sprintf("%s buildkite-cli/%s", buildkite.DefaultUserAgent, version.Version)
+}
+
+func buildUserAgent(suffix string) string {
+	if suffix == "" {
+		return baseUserAgent
+	}
+	return fmt.Sprintf("%s %s", baseUserAgent, suffix)
 }
 
 func (a *gqlHTTPClient) Do(req *http.Request) (*http.Response, error) {
 	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", a.token))
-	req.Header.Set("User-Agent", userAgent)
+	req.Header.Set("User-Agent", a.userAgent)
 	return a.client.Do(req)
 }
 
@@ -165,6 +181,8 @@ func New(opts ...FactoryOpt) (*Factory, error) {
 		}
 	}
 
+	userAgent := buildUserAgent(cfg.userAgentSuffix)
+
 	// Build client options
 	clientOpts := []buildkite.ClientOpt{
 		buildkite.WithBaseURL(conf.RESTAPIEndpoint()),
@@ -188,7 +206,7 @@ func New(opts ...FactoryOpt) (*Factory, error) {
 		return nil, fmt.Errorf("creating buildkite client: %w", err)
 	}
 
-	graphqlHTTPClient := &gqlHTTPClient{client: http.DefaultClient, token: token}
+	graphqlHTTPClient := &gqlHTTPClient{client: http.DefaultClient, token: token, userAgent: userAgent}
 
 	return &Factory{
 		Config:        conf,

pkg/cmd/factory/factory_test.go

@@ -6,6 +6,8 @@ import (
 	"net/http/httptest"
 	"strings"
 	"testing"
+
+	buildkite "github.com/buildkite/go-buildkite/v4"
 )
 
 func TestRedactHeaders(t *testing.T) {
@@ -152,3 +154,49 @@ func TestDebugTransportHandlesNilBody(t *testing.T) {
 		t.Errorf("expected status 200, got %d", resp.StatusCode)
 	}
 }
+
+func TestBuildUserAgent(t *testing.T) {
+	t.Run("default user agent has no preflight suffix", func(t *testing.T) {
+		got := buildUserAgent("")
+		if !strings.Contains(got, buildkite.DefaultUserAgent) {
+			t.Fatalf("expected default user agent %q in %q", buildkite.DefaultUserAgent, got)
+		}
+		if strings.Contains(got, "buildkite-cli-preflight/") {
+			t.Fatalf("expected no preflight suffix in %q", got)
+		}
+	})
+
+	t.Run("preflight suffix is appended when requested", func(t *testing.T) {
+		got := buildUserAgent("buildkite-cli-preflight/3.x")
+		if !strings.Contains(got, buildkite.DefaultUserAgent) {
+			t.Fatalf("expected default user agent %q in %q", buildkite.DefaultUserAgent, got)
+		}
+		if !strings.Contains(got, "buildkite-cli-preflight/3.x") {
+			t.Fatalf("expected preflight suffix in %q", got)
+		}
+	})
+}
+
+func TestNewUserAgent(t *testing.T) {
+	t.Chdir(t.TempDir())
+
+	t.Run("non-preflight factory does not set preflight suffix", func(t *testing.T) {
+		f, err := New()
+		if err != nil {
+			t.Fatalf("New() error = %v", err)
+		}
+		if strings.Contains(f.RestAPIClient.UserAgent, "buildkite-cli-preflight/") {
+			t.Fatalf("expected no preflight suffix in %q", f.RestAPIClient.UserAgent)
+		}
+	})
+
+	t.Run("factory can opt in to preflight suffix", func(t *testing.T) {
+		f, err := New(WithUserAgentSuffix("buildkite-cli-preflight/3.x"))
+		if err != nil {
+			t.Fatalf("New() error = %v", err)
+		}
+		if !strings.Contains(f.RestAPIClient.UserAgent, "buildkite-cli-preflight/3.x") {
+			t.Fatalf("expected preflight suffix in %q", f.RestAPIClient.UserAgent)
+		}
+	})
+}