refactor: remove duplicate artifact downloading

Signed-off-by: nikpivkin <[email protected]>
aquasecurity · Sep 30, 2024 · b08004a · b08004a
1 parent d75a584
commit b08004a
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 14 deletions.
diff --git a/pkg/db/db.go b/pkg/db/db.go
@@ -199,11 +199,8 @@ func (c *Client) updateDownloadedAt(ctx context.Context, dbDir string) error {
 }
 
 func (c *Client) initOCIArtifact(repository name.Reference, opt types.RegistryOptions) (*oci.Artifact, error) {
-	if c.artifact != nil {
-		return c.artifact, nil
-	}
-
 	art, err := oci.NewArtifact(repository.String(), c.quiet, opt)
+	// TODO: NewArtifact never returns an error
 	if err != nil {
 		var terr *transport.Error
 		if errors.As(err, &terr) {
@@ -221,28 +218,40 @@ func (c *Client) initOCIArtifact(repository name.Reference, opt types.RegistryOp
 	return art, nil
 }
 
-func (c *Client) downloadDB(ctx context.Context, opt types.RegistryOptions, dst string) error {
-	downloadOpt := oci.DownloadOption{MediaType: dbMediaType}
+func (c *Client) initArtifacts(opt types.RegistryOptions) ([]*oci.Artifact, error) {
 	if c.artifact != nil {
-		return c.artifact.Download(ctx, dst, downloadOpt)
+		return []*oci.Artifact{c.artifact}, nil
 	}
 
-	for i, repo := range c.dbRepositories {
+	artifacts := make([]*oci.Artifact, 0, len(c.dbRepositories))
+
+	for _, repo := range c.dbRepositories {
 		a, err := c.initOCIArtifact(repo, opt)
 		if err != nil {
-			return err
+			return nil, err
 		}
+		artifacts = append(artifacts, a)
+	}
+	return artifacts, nil
+}
+
+func (c *Client) downloadDB(ctx context.Context, opt types.RegistryOptions, dst string) error {
+	arts, err := c.initArtifacts(opt)
+	if err != nil {
+		return err
+	}
 
-		log.Info("Downloading vulnerability DB...", log.String("repo", repo.String()))
-		if err := a.Download(ctx, dst, downloadOpt); err != nil {
-			log.Error("Failed to download DB", log.String("repo", repo.String()), log.Err(err))
+	for i, art := range arts {
+		log.Info("Downloading vulnerability DB...", log.String("repo", art.Repository()))
+		if err := art.Download(ctx, dst, oci.DownloadOption{MediaType: dbMediaType}); err != nil {
+			log.Error("Failed to download DB", log.String("repo", art.Repository()), log.Err(err))
 			if i < len(c.dbRepositories)-1 {
 				log.Info("Trying to download DB from other repository...")
 			}
 			continue
 		}
 
-		log.Info("DB successfully downloaded", log.String("repo", repo.String()))
+		log.Info("DB successfully downloaded", log.String("repo", art.Repository()))
 		return nil
 	}
 

diff --git a/pkg/db/db_test.go b/pkg/db/db_test.go
@@ -159,7 +159,7 @@ func TestClient_Download(t *testing.T) {
 		{
 			name:    "invalid gzip",
 			input:   "testdata/trivy.db",
-			wantErr: "unexpected EOF",
+			wantErr: "OCI artifact error: failed to download vulnerability DB from any source",
 		},
 	}
 

diff --git a/pkg/oci/artifact.go b/pkg/oci/artifact.go
@@ -207,3 +207,7 @@ func (a *Artifact) Digest(ctx context.Context) (string, error) {
 	}
 	return digest.String(), nil
 }
+
+func (a *Artifact) Repository() string {
+	return a.repository
+}