feat(image): Set User-Agent header for Trivy container registry reque…

…sts (#6868) Signed-off-by: m.nabokikh <[email protected]>
aquasecurity · Jun 10, 2024 · 9b31697 · 9b31697
1 parent 089b953
commit 9b31697
Show file tree

Hide file tree

Showing 11 changed files with 100 additions and 29 deletions.
diff --git a/goreleaser-canary.yml b/goreleaser-canary.yml
@@ -6,7 +6,7 @@ builds:
     ldflags:
       - -s -w
       - "-extldflags '-static'"
-      - -X github.com/aquasecurity/trivy/pkg/version.ver={{.Version}}
+      - -X github.com/aquasecurity/trivy/pkg/version/app.ver={{.Version}}
     env:
       - CGO_ENABLED=0
     goos:

diff --git a/goreleaser.yml b/goreleaser.yml
@@ -6,7 +6,7 @@ builds:
     ldflags:
       - -s -w
       - "-extldflags '-static'"
-      - -X github.com/aquasecurity/trivy/pkg/version.ver={{.Version}}
+      - -X github.com/aquasecurity/trivy/pkg/version/app.ver={{.Version}}
     env:
       - CGO_ENABLED=0
     goos:
@@ -26,7 +26,7 @@ builds:
     ldflags:
       - -s -w
       - "-extldflags '-static'"
-      - -X github.com/aquasecurity/trivy/pkg/version.ver={{.Version}}
+      - -X github.com/aquasecurity/trivy/pkg/version/app.ver={{.Version}}
     env:
       - CGO_ENABLED=0
     goos:
@@ -41,7 +41,7 @@ builds:
     ldflags:
       - -s -w
       - "-extldflags '-static'"
-      - -X github.com/aquasecurity/trivy/pkg/version.ver={{.Version}}
+      - -X github.com/aquasecurity/trivy/pkg/version/app.ver={{.Version}}
     env:
       - CGO_ENABLED=0
     goos:
@@ -57,7 +57,7 @@ builds:
     ldflags:
       - -s -w
       - "-extldflags '-static'"
-      - -X github.com/aquasecurity/trivy/pkg/version.ver={{.Version}}
+      - -X github.com/aquasecurity/trivy/pkg/version/app.ver={{.Version}}
     env:
       - CGO_ENABLED=0
     goos:

diff --git a/magefiles/magefile.go b/magefiles/magefile.go
@@ -48,7 +48,7 @@ func buildLdflags() (string, error) {
 	if err != nil {
 		return "", err
 	}
-	return fmt.Sprintf("-s -w -X=github.com/aquasecurity/trivy/pkg/version.ver=%s", ver), nil
+	return fmt.Sprintf("-s -w -X=github.com/aquasecurity/trivy/pkg/version/app.ver=%s", ver), nil
 }
 
 type Tool mg.Namespace

diff --git a/pkg/commands/app.go b/pkg/commands/app.go
@@ -28,6 +28,7 @@ import (
 	"github.com/aquasecurity/trivy/pkg/plugin"
 	"github.com/aquasecurity/trivy/pkg/types"
 	"github.com/aquasecurity/trivy/pkg/version"
+	"github.com/aquasecurity/trivy/pkg/version/app"
 	xstrings "github.com/aquasecurity/trivy/pkg/x/strings"
 )
 
@@ -178,7 +179,7 @@ func NewRootCommand(globalFlags *flag.GlobalFlagGroup) *cobra.Command {
 		Args: cobra.NoArgs,
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
 			// Set the Trivy version here so that we can override version printer.
-			cmd.Version = version.AppVersion()
+			cmd.Version = app.Version()
 
 			// viper.BindPFlag cannot be called in init().
 			// cf. https://github.com/spf13/cobra/issues/875

diff --git a/pkg/dependency/parser/golang/binary/parse.go b/pkg/dependency/parser/golang/binary/parse.go
@@ -217,7 +217,7 @@ func isValidSemVer(ver string) bool {
 
 // versionPrefix returns version prefix from `-ldflags` flag key
 // e.g.
-//   - `github.com/aquasecurity/trivy/pkg/version.version` => `version`
+//   - `github.com/aquasecurity/trivy/pkg/version/app.ver` => `version`
 //   - `github.com/google/go-containerregistry/cmd/crane/common.ver` => `common`
 func versionPrefix(s string) string {
 	// Trim module part.

diff --git a/pkg/dependency/parser/golang/binary/parse_test.go b/pkg/dependency/parser/golang/binary/parse_test.go
@@ -168,7 +168,7 @@ func TestParser_ParseLDFlags(t *testing.T) {
 					"-s",
 					"-w",
 					"-X=foo=bar",
-					"-X='github.com/aquasecurity/trivy/pkg/version.version=v0.50.1'",
+					"-X='github.com/aquasecurity/trivy/pkg/version/app.ver=v0.50.1'",
 				},
 			},
 			want: "v0.50.1",
@@ -194,7 +194,7 @@ func TestParser_ParseLDFlags(t *testing.T) {
 					"-s",
 					"-w",
 					"-X=foo=bar",
-					"-X='github.com/aquasecurity/trivy/pkg/version.ver=v0.50.1'",
+					"-X='github.com/aquasecurity/trivy/pkg/version/app.ver=v0.50.1'",
 				},
 			},
 			want: "v0.50.1",

diff --git a/pkg/flag/options.go b/pkg/flag/options.go
@@ -23,7 +23,7 @@ import (
 	"github.com/aquasecurity/trivy/pkg/plugin"
 	"github.com/aquasecurity/trivy/pkg/result"
 	"github.com/aquasecurity/trivy/pkg/types"
-	"github.com/aquasecurity/trivy/pkg/version"
+	"github.com/aquasecurity/trivy/pkg/version/app"
 )
 
 type FlagType interface {
@@ -602,7 +602,7 @@ func (f *Flags) Bind(cmd *cobra.Command) error {
 func (f *Flags) ToOptions(args []string) (Options, error) {
 	var err error
 	opts := Options{
-		AppVersion: version.AppVersion(),
+		AppVersion: app.Version(),
 	}
 
 	if f.GlobalFlagGroup != nil {

diff --git a/pkg/remote/remote.go b/pkg/remote/remote.go
@@ -3,6 +3,7 @@ package remote
 import (
 	"context"
 	"crypto/tls"
+	"fmt"
 	"net"
 	"net/http"
 	"time"
@@ -11,6 +12,7 @@ import (
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/remote"
+	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
 	v1types "github.com/google/go-containerregistry/pkg/v1/types"
 	"github.com/hashicorp/go-multierror"
 	"github.com/samber/lo"
@@ -19,14 +21,15 @@ import (
 	"github.com/aquasecurity/trivy/pkg/fanal/image/registry"
 	"github.com/aquasecurity/trivy/pkg/fanal/types"
 	"github.com/aquasecurity/trivy/pkg/log"
+	"github.com/aquasecurity/trivy/pkg/version/app"
 )
 
 type Descriptor = remote.Descriptor
 
 // Get is a wrapper of google/go-containerregistry/pkg/v1/remote.Get
 // so that it can try multiple authentication methods.
 func Get(ctx context.Context, ref name.Reference, option types.RegistryOptions) (*Descriptor, error) {
-	transport, err := httpTransport(option)
+	tr, err := httpTransport(option)
 	if err != nil {
 		return nil, xerrors.Errorf("failed to create http transport: %w", err)
 	}
@@ -35,7 +38,7 @@ func Get(ctx context.Context, ref name.Reference, option types.RegistryOptions)
 	// Try each authentication method until it succeeds
 	for _, authOpt := range authOptions(ctx, ref, option) {
 		remoteOpts := []remote.Option{
-			remote.WithTransport(transport),
+			remote.WithTransport(tr),
 			authOpt,
 		}
 
@@ -71,7 +74,7 @@ func Get(ctx context.Context, ref name.Reference, option types.RegistryOptions)
 // Image is a wrapper of google/go-containerregistry/pkg/v1/remote.Image
 // so that it can try multiple authentication methods.
 func Image(ctx context.Context, ref name.Reference, option types.RegistryOptions) (v1.Image, error) {
-	transport, err := httpTransport(option)
+	tr, err := httpTransport(option)
 	if err != nil {
 		return nil, xerrors.Errorf("failed to create http transport: %w", err)
 	}
@@ -80,7 +83,7 @@ func Image(ctx context.Context, ref name.Reference, option types.RegistryOptions
 	// Try each authentication method until it succeeds
 	for _, authOpt := range authOptions(ctx, ref, option) {
 		remoteOpts := []remote.Option{
-			remote.WithTransport(transport),
+			remote.WithTransport(tr),
 			authOpt,
 		}
 		index, err := remote.Image(ref, remoteOpts...)
@@ -98,7 +101,7 @@ func Image(ctx context.Context, ref name.Reference, option types.RegistryOptions
 // Referrers is a wrapper of google/go-containerregistry/pkg/v1/remote.Referrers
 // so that it can try multiple authentication methods.
 func Referrers(ctx context.Context, d name.Digest, option types.RegistryOptions) (v1.ImageIndex, error) {
-	transport, err := httpTransport(option)
+	tr, err := httpTransport(option)
 	if err != nil {
 		return nil, xerrors.Errorf("failed to create http transport: %w", err)
 	}
@@ -107,7 +110,7 @@ func Referrers(ctx context.Context, d name.Digest, option types.RegistryOptions)
 	// Try each authentication method until it succeeds
 	for _, authOpt := range authOptions(ctx, d, option) {
 		remoteOpts := []remote.Option{
-			remote.WithTransport(transport),
+			remote.WithTransport(tr),
 			authOpt,
 		}
 		index, err := remote.Referrers(d, remoteOpts...)
@@ -122,7 +125,7 @@ func Referrers(ctx context.Context, d name.Digest, option types.RegistryOptions)
 	return nil, errs
 }
 
-func httpTransport(option types.RegistryOptions) (*http.Transport, error) {
+func httpTransport(option types.RegistryOptions) (http.RoundTripper, error) {
 	d := &net.Dialer{
 		Timeout: 10 * time.Minute,
 	}
@@ -138,7 +141,8 @@ func httpTransport(option types.RegistryOptions) (*http.Transport, error) {
 		tr.TLSClientConfig.Certificates = []tls.Certificate{cert}
 	}
 
-	return tr, nil
+	tripper := transport.NewUserAgent(tr, fmt.Sprintf("trivy/%s", app.Version()))
+	return tripper, nil
 }
 
 func authOptions(ctx context.Context, ref name.Reference, option types.RegistryOptions) []remote.Option {

diff --git a/pkg/remote/remote_test.go b/pkg/remote/remote_test.go
@@ -4,9 +4,11 @@ import (
 	"context"
 	"encoding/base64"
 	"fmt"
+	"net/http"
 	"net/http/httptest"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
 
 	"github.com/google/go-containerregistry/pkg/name"
@@ -17,6 +19,7 @@ import (
 	"github.com/aquasecurity/testdocker/auth"
 	"github.com/aquasecurity/testdocker/registry"
 	"github.com/aquasecurity/trivy/pkg/fanal/types"
+	"github.com/aquasecurity/trivy/pkg/version/app"
 )
 
 func setupPrivateRegistry() *httptest.Server {
@@ -32,6 +35,7 @@ func setupPrivateRegistry() *httptest.Server {
 		},
 	})
 
+	tr.Config.Handler = newUserAgentsTrackingHandler(tr.Config.Handler)
 	return tr
 }
 
@@ -206,3 +210,63 @@ func TestGet(t *testing.T) {
 		})
 	}
 }
+
+type userAgentsTrackingHandler struct {
+	hr http.Handler
+
+	mu     sync.Mutex
+	agents map[string]struct{}
+}
+
+func newUserAgentsTrackingHandler(hr http.Handler) *userAgentsTrackingHandler {
+	return &userAgentsTrackingHandler{hr: hr, agents: make(map[string]struct{})}
+}
+
+func (uh *userAgentsTrackingHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+	for _, agent := range r.Header["User-Agent"] {
+		// Skip test framework user agent
+		if agent != "Go-http-client/1.1" {
+			uh.agents[agent] = struct{}{}
+		}
+	}
+	uh.hr.ServeHTTP(rw, r)
+}
+
+func setupAgentTrackingRegistry() (*httptest.Server, *userAgentsTrackingHandler) {
+	imagePaths := map[string]string{
+		"v2/library/alpine:3.10": "../fanal/test/testdata/alpine-310.tar.gz",
+	}
+	tr := registry.NewDockerRegistry(registry.Option{
+		Images: imagePaths,
+	})
+
+	tracker := newUserAgentsTrackingHandler(tr.Config.Handler)
+	tr.Config.Handler = tracker
+
+	return tr, tracker
+}
+
+func TestUserAgents(t *testing.T) {
+	tr, tracker := setupAgentTrackingRegistry()
+	defer tr.Close()
+
+	serverAddr := tr.Listener.Addr().String()
+
+	n, err := name.ParseReference(fmt.Sprintf("%s/library/alpine:3.10", serverAddr))
+	require.NoError(t, err)
+
+	_, err = Get(context.Background(), n, types.RegistryOptions{
+		Credentials: []types.Credential{
+			{
+				Username: "test",
+				Password: "testpass",
+			},
+		},
+		Insecure: true,
+	})
+	require.NoError(t, err)
+
+	require.Len(t, tracker.agents, 1)
+	_, ok := tracker.agents[fmt.Sprintf("trivy/%s go-containerregistry", app.Version())]
+	require.True(t, ok, `user-agent header equals to "trivy/dev go-containerregistry"`)
+}
diff --git a/pkg/version/app/version.go b/pkg/version/app/version.go
@@ -0,0 +1,9 @@
+package app
+
+var (
+	ver = "dev"
+)
+
+func Version() string {
+	return ver
+}
diff --git a/pkg/version/version.go b/pkg/version/version.go
@@ -8,16 +8,9 @@ import (
 	javadb "github.com/aquasecurity/trivy-java-db/pkg/db"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/policy"
+	"github.com/aquasecurity/trivy/pkg/version/app"
 )
 
-var (
-	ver = "dev"
-)
-
-func AppVersion() string {
-	return ver
-}
-
 type VersionInfo struct {
 	Version         string             `json:",omitempty"`
 	VulnerabilityDB *metadata.Metadata `json:",omitempty"`
@@ -99,7 +92,7 @@ func NewVersionInfo(cacheDir string) VersionInfo {
 	}
 
 	return VersionInfo{
-		Version:         ver,
+		Version:         app.Version(),
 		VulnerabilityDB: dbMeta,
 		JavaDB:          javadbMeta,
 		CheckBundle:     pbMeta,