Skip to content

Commit

Permalink
feat(image): Set User-Agent header for Trivy container registry reque…
Browse files Browse the repository at this point in the history
…sts (#6868)

Signed-off-by: m.nabokikh <[email protected]>
  • Loading branch information
nabokihms authored Jun 10, 2024
1 parent 089b953 commit 9b31697
Show file tree
Hide file tree
Showing 11 changed files with 100 additions and 29 deletions.
2 changes: 1 addition & 1 deletion goreleaser-canary.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
8 changes: 4 additions & 4 deletions goreleaser.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion magefiles/magefile.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion pkg/commands/app.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pkg/dependency/parser/golang/binary/parse.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 2 additions & 2 deletions pkg/dependency/parser/golang/binary/parse_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand Down
4 changes: 2 additions & 2 deletions pkg/flag/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down
20 changes: 12 additions & 8 deletions pkg/remote/remote.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package remote
import (
"context"
"crypto/tls"
"fmt"
"net"
"net/http"
"time"
Expand All @@ -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"
Expand All @@ -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)
}
Expand All @@ -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,
}

Expand Down Expand Up @@ -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)
}
Expand All @@ -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...)
Expand All @@ -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)
}
Expand All @@ -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...)
Expand All @@ -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,
}
Expand All @@ -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 {
Expand Down
64 changes: 64 additions & 0 deletions pkg/remote/remote_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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 {
Expand All @@ -32,6 +35,7 @@ func setupPrivateRegistry() *httptest.Server {
},
})

tr.Config.Handler = newUserAgentsTrackingHandler(tr.Config.Handler)
return tr
}

Expand Down Expand Up @@ -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"`)
}
9 changes: 9 additions & 0 deletions pkg/version/app/version.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
package app

var (
ver = "dev"
)

func Version() string {
return ver
}
11 changes: 2 additions & 9 deletions pkg/version/version.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down Expand Up @@ -99,7 +92,7 @@ func NewVersionInfo(cacheDir string) VersionInfo {
}

return VersionInfo{
Version: ver,
Version: app.Version(),
VulnerabilityDB: dbMeta,
JavaDB: javadbMeta,
CheckBundle: pbMeta,
Expand Down

0 comments on commit 9b31697

Please sign in to comment.