Add token reuse

equinor · Apr 2, 2024 · 7758602 · 7758602
1 parent b07f0ac
commit 7758602
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 30 deletions.
diff --git a/main.go b/main.go
@@ -12,6 +12,7 @@ import (
 	"github.com/equinor/radix-vulnerability-scanner/pkg/dockercfg"
 	"github.com/equinor/radix-vulnerability-scanner/pkg/options"
 	"github.com/equinor/radix-vulnerability-scanner/pkg/scan"
+	"github.com/equinor/radix-vulnerability-scanner/pkg/tokenstore"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"k8s.io/client-go/kubernetes"
@@ -35,7 +36,7 @@ func main() {
 
 	logOptions(opts)
 
-	scanner, err := scan.New(ctx, &opts.Docker)
+	scanner, err := newSnykScanner(ctx, &opts.Docker)
 	if err != nil {
 		log.Fatal().Msg(err.Error())
 	}
@@ -123,3 +124,22 @@ func getKubernetesClients(opts *options.KubeOptions) (kubernetes.Interface, radi
 
 	return kubeClient, radixClient, nil
 }
+
+func newSnykScanner(ctx context.Context, opts *options.DockerOptions) (scan.Scanner, error) {
+	var dockerConfig dockercfg.DockerConfig
+	var err error
+
+	if opts.AuthsFile != "" {
+		dockerConfig, err = dockercfg.ReadDockerAuthConfigFromFile(opts.AuthsFile)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	tokenStore, err := tokenstore.NewTokenStore(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return scan.NewSnyk(dockerConfig, tokenStore), nil
+}
diff --git a/pkg/dockercfg/config.go b/pkg/dockercfg/config.go
@@ -7,7 +7,7 @@ import (
 	"os"
 
 	"github.com/containerd/containerd/reference/docker"
-	"github.com/equinor/radix-vulnerability-scanner/pkg/acr"
+	"github.com/equinor/radix-vulnerability-scanner/pkg/tokenstore"
 	"github.com/rs/zerolog/log"
 )
 
@@ -51,7 +51,7 @@ func ReadDockerAuthConfigFromBytes(contents []byte) (DockerConfig, error) {
 	return cfgJSON, nil
 }
 
-func (c DockerConfig) GetAuth(image string, tokenStore *acr.TokenStore) (DockerAuthConfigEntry, error) {
+func (c DockerConfig) GetAuth(image string, tokenStore *tokenstore.TokenStore) (DockerAuthConfigEntry, error) {
 	named, err := docker.ParseDockerRef(image)
 	if err != nil {
 		return DockerAuthConfigEntry{}, err

diff --git a/pkg/scan/scanner.go b/pkg/scan/scanner.go
@@ -3,32 +3,11 @@ package scan
 import (
 	"context"
 
-	"github.com/equinor/radix-vulnerability-scanner/pkg/acr"
 	"github.com/equinor/radix-vulnerability-scanner/pkg/dockercfg"
-	"github.com/equinor/radix-vulnerability-scanner/pkg/options"
 )
 
 // Scanner defines methods for scanning Docker images for vulnerabilities
 type Scanner interface {
 	// Scan scans a Docker image for vulnerabilities
 	Scan(ctx context.Context, image string, dockerConfig dockercfg.DockerConfig) (*ScanResult, error)
 }
-
-func New(ctx context.Context, opts *options.DockerOptions) (Scanner, error) {
-	var dockerAuth dockercfg.DockerConfig
-	var err error
-
-	if opts.AuthsFile != "" {
-		dockerAuth, err = dockercfg.ReadDockerAuthConfigFromFile(opts.AuthsFile)
-		if err != nil {
-			return nil, err
-		}
-	}
-
-	tokenStore, err := acr.NewTokenStore(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	return NewSnyk(dockerAuth, tokenStore), nil
-}
diff --git a/pkg/scan/snyk.go b/pkg/scan/snyk.go
@@ -8,8 +8,8 @@ import (
 	"io"
 	"os/exec"
 
-	"github.com/equinor/radix-vulnerability-scanner/pkg/acr"
 	"github.com/equinor/radix-vulnerability-scanner/pkg/dockercfg"
+	"github.com/equinor/radix-vulnerability-scanner/pkg/tokenstore"
 	"github.com/equinor/radix-vulnerability-scanner/pkg/utils/logwriter"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
@@ -38,11 +38,11 @@ func (commandExecutorImpl) Execute(ctx context.Context, command string, args []s
 type snykScanner struct {
 	commonDockerConfig dockercfg.DockerConfig
 	executor           commandExecutor
-	tokenStore         *acr.TokenStore
+	tokenStore         *tokenstore.TokenStore
 }
 
 // NewSnyk create a Scanner that use SNYK to scan for vulnerabilities
-func NewSnyk(commonAuths dockercfg.DockerConfig, tokenStore *acr.TokenStore) Scanner {
+func NewSnyk(commonAuths dockercfg.DockerConfig, tokenStore *tokenstore.TokenStore) Scanner {
 	return &snykScanner{commonDockerConfig: commonAuths, executor: commandExecutorImpl{}, tokenStore: tokenStore}
 }
 

diff --git a/pkg/acr/token_store.go → pkg/tokenstore/token_store.go b/pkg/acr/token_store.go → pkg/tokenstore/token_store.go
@@ -1,4 +1,4 @@
-package acr
+package tokenstore
 
 import (
 	"context"
@@ -58,7 +58,7 @@ func (t *TokenStore) GetToken(registryName string) (string, error) {
 
 	_, ok := t.tokens[registryName]
 	if !ok {
-		t.tokens[registryName] = t.getTokenSource(t.ctx, registryName)
+		t.tokens[registryName] = oauth2.ReuseTokenSource(nil, t.getACRTokenSource(t.ctx, registryName))
 	}
 
 	token, err := t.tokens[registryName].Token()
@@ -75,7 +75,7 @@ func (s TokenSourceFunc) Token() (*oauth2.Token, error) {
 	return s()
 }
 
-func (t *TokenStore) getTokenSource(ctx context.Context, registryName string) oauth2.TokenSource {
+func (t *TokenStore) getACRTokenSource(ctx context.Context, registryName string) oauth2.TokenSource {
 	return TokenSourceFunc(func() (*oauth2.Token, error) {
 		log.Ctx(ctx).Debug().Str("registry", registryName).Msg("Fetching new ACR token")