refactor: update api with new rpc

* add package to grpc api
* create testdata for nais-api
* change docker compose ports to not interfere with nais api
* change numbers for severity to better fit with sorting

Co-authored-by: ybelmekk <[email protected]>

Author: tommytroen

.env.sample

@@ -1,5 +1,5 @@
 V13S_LISTEN_ADDR=localhost:50051
-V13S_DATABASE_URL=postgres://v13s:[email protected]:3002/v13s?sslmode=disable
+V13S_DATABASE_URL=postgres://v13s:[email protected]:4002/v13s?sslmode=disable
 V13S_DEPENDENCYTRACK_URL=todo
 V13S_DEPENDENCYTRACK_TEAM=todo
 V13S_DEPENDENCYTRACK_USERNAME=todo

Makefile

@@ -33,8 +33,8 @@ generate: generate-proto generate_dp_track generate-sql generate-mocks
 
 generate-proto:
 	protoc \
-		-I pkg/api/vulnerabilities/schema/ \
-		./pkg/api/vulnerabilities/schema/*.proto \
+		-I pkg/api/vulnerabilities/schemas/ \
+		./pkg/api/vulnerabilities/schemas/*.proto \
 		--go_out=. \
 		--go-grpc_out=.
 

cmd/seed/main.go

@@ -23,33 +23,37 @@ func main() {
 		fmt.Println("No .env file found")
 	}
 
-	dpClient, err := dependencytrack.NewClient(
-		os.Getenv("V13S_DEPENDENCYTRACK_URL"),
-		os.Getenv("V13S_DEPENDENCYTRACK_TEAM"),
-		os.Getenv("V13S_DEPENDENCYTRACK_USERNAME"),
-		os.Getenv("V13S_DEPENDENCYTRACK_PASSWORD"),
-	)
+	log.Infof("initializing database")
+
+	pool, err := database.New(ctx, "postgres://v13s:[email protected]:4002/v13s?sslmode=disable", log.WithField("component", "database"))
 	if err != nil {
 		panic(err)
 	}
+	defer pool.Close()
 
-	projects, err := dpClient.GetProjectsByTag(ctx, "team:nais-system", 10, 0)
+	db := sql.New(pool)
+
+	log.Infof("reseting database")
+	err = db.ResetDatabase(ctx)
 	if err != nil {
 		panic(err)
 	}
 
-	log.Infof("initializing database")
+	createNaisApiTestdata(ctx, db)
+}
 
-	pool, err := database.New(ctx, "postgres://v13s:[email protected]:3002/v13s?sslmode=disable", log.WithField("component", "database"))
+func seedFromDependencyTrack(ctx context.Context, db sql.Querier) {
+	dpClient, err := dependencytrack.NewClient(
+		os.Getenv("V13S_DEPENDENCYTRACK_URL"),
+		os.Getenv("V13S_DEPENDENCYTRACK_TEAM"),
+		os.Getenv("V13S_DEPENDENCYTRACK_USERNAME"),
+		os.Getenv("V13S_DEPENDENCYTRACK_PASSWORD"),
+	)
 	if err != nil {
 		panic(err)
 	}
-	defer pool.Close()
-
-	db := sql.New(pool)
 
-	log.Infof("reseting database")
-	err = db.ResetDatabase(ctx)
+	projects, err := dpClient.GetProjectsByTag(ctx, "team:nais-system", 10, 0)
 	if err != nil {
 		panic(err)
 	}
@@ -96,6 +100,109 @@ func main() {
 	}
 }
 
+func createNaisApiTestdata(ctx context.Context, db sql.Querier) {
+	for chicken := 1; chicken < 9; chicken++ {
+		if err := db.CreateImage(ctx, sql.CreateImageParams{
+			Name:     "ghcr.io/nais/nais-deploy-chicken-" + fmt.Sprintf("%d", chicken),
+			Tag:      "1",
+			Metadata: map[string]string{},
+		}); err != nil {
+			panic(err)
+		}
+
+		_, err := db.CreateWorkload(ctx, sql.CreateWorkloadParams{
+			Cluster:      "dev",
+			Namespace:    "devteam",
+			WorkloadType: "app",
+			Name:         "nais-deploy-chicken-" + fmt.Sprintf("%d", chicken),
+			ImageName:    "ghcr.io/nais/nais-deploy-chicken-" + fmt.Sprintf("%d", chicken),
+			ImageTag:     "1",
+		})
+		if err != nil {
+			panic(err)
+		}
+
+		batch := generateVulnerabilities(chicken, "ghcr.io/nais/nais-deploy-chicken-"+fmt.Sprintf("%d", chicken), "1")
+
+		db.BatchUpsertCve(ctx, batch.cve).Exec(func(i int, err error) {
+			if err != nil {
+				panic(err)
+			}
+		})
+
+		db.BatchUpsertVulnerabilities(ctx, batch.vuln).Exec(func(i int, err error) {
+			if err != nil {
+				panic(err)
+			}
+		})
+
+		sumRow, err := db.GenerateVulnerabilitySummaryForImage(ctx, sql.GenerateVulnerabilitySummaryForImageParams{
+			ImageName: "ghcr.io/nais/nais-deploy-chicken-" + fmt.Sprintf("%d", chicken),
+			ImageTag:  "1",
+		})
+		if err != nil {
+			panic(err)
+		}
+
+		summary := sql.CreateVulnerabilitySummaryParams{
+			ImageName:  "ghcr.io/nais/nais-deploy-chicken-" + fmt.Sprintf("%d", chicken),
+			ImageTag:   "1",
+			Critical:   int32(sumRow.Critical),
+			High:       int32(sumRow.High),
+			Medium:     int32(sumRow.Medium),
+			Low:        int32(sumRow.Low),
+			Unassigned: int32(sumRow.Unassigned),
+			RiskScore:  sumRow.RiskScore,
+		}
+
+		_, err = db.CreateVulnerabilitySummary(ctx, summary)
+		if err != nil {
+			panic(fmt.Errorf("summary error: %v", err))
+		}
+	}
+}
+
+type BatchVulnerabilities struct {
+	vuln []sql.BatchUpsertVulnerabilitiesParams
+	cve  []sql.BatchUpsertCveParams
+}
+
+// generateVulnerabilities creates a different number of vulnerabilities per workload
+func generateVulnerabilities(chicken int, imageName string, imageTag string) BatchVulnerabilities {
+	vulns := make([]sql.BatchUpsertVulnerabilitiesParams, 0)
+	cves := make([]sql.BatchUpsertCveParams, 0)
+
+	for j := 1; j <= chicken; j++ {
+		for s := 1; s <= 5; s++ {
+			v, c := createVulnerability(s, fmt.Sprintf("CWE-%d-%d-%d", chicken, j, s), imageName, imageTag)
+			vulns = append(vulns, v)
+			cves = append(cves, c)
+		}
+	}
+
+	return BatchVulnerabilities{
+		vuln: vulns,
+		cve:  cves,
+	}
+}
+
+// createVulnerability generates a predictable vulnerability instance
+func createVulnerability(severity int, cveID string, imageName string, imageTag string) (sql.BatchUpsertVulnerabilitiesParams, sql.BatchUpsertCveParams) {
+	return sql.BatchUpsertVulnerabilitiesParams{
+			ImageName: imageName,
+			ImageTag:  imageTag,
+			Package:   fmt.Sprintf("package-%s", cveID),
+			Source:    "seed",
+			CveID:     cveID,
+		}, sql.BatchUpsertCveParams{
+			CveID:    cveID,
+			CveTitle: "Title for " + cveID,
+			CveDesc:  "description for " + cveID,
+			CveLink:  "https://example.com/" + cveID,
+			Severity: int32(severity),
+		}
+}
+
 func toCreateWorkloadParams(p client.Project, image sql.CreateImageParams) []*sql.CreateWorkloadParams {
 	workloads := make([]*sql.CreateWorkloadParams, 0)
 	for _, t := range p.Tags {

docker-compose.yaml

@@ -3,7 +3,7 @@ services:
     image: postgres:15-alpine
     command: ["postgres", "-c", "log_statement=all", "-c", "log_destination=stderr"]
     ports:
-      - "3002:5432"
+      - "4002:5432"
     environment:
       POSTGRES_USER: v13s
       POSTGRES_PASSWORD: v13s
@@ -16,7 +16,7 @@ services:
       - postgres
     image: adminer:latest
     ports:
-      - "3003:8080"
+      - "4003:8080"
     environment:
       ADMINER_DEFAULT_SERVER: postgres
 

internal/api/grpcvulnerabilities/server.go

@@ -2,14 +2,14 @@ package grpcvulnerabilities
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"github.com/jackc/pgx/v5"
 	"github.com/nais/v13s/internal/api/grpcpagination"
-	"github.com/nais/v13s/internal/database/sql"
-	"google.golang.org/protobuf/types/known/timestamppb"
-	"time"
-
 	"github.com/nais/v13s/internal/collections"
+	"github.com/nais/v13s/internal/database/sql"
 	"github.com/nais/v13s/pkg/api/vulnerabilities"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 var _ vulnerabilities.VulnerabilitiesServer = (*Server)(nil)
@@ -172,20 +172,76 @@ func (s *Server) GetVulnerabilitySummary(ctx context.Context, request *vulnerabi
 		return nil, err
 	}
 
-	var summary = vulnerabilities.Summary{
-		Critical:    sum.CriticalVulnerabilities,
-		High:        sum.HighVulnerabilities,
-		Medium:      sum.MediumVulnerabilities,
-		Low:         sum.LowVulnerabilities,
-		Unassigned:  sum.UnassignedVulnerabilities,
-		RiskScore:   sum.TotalRiskScore,
-		LastUpdated: timestamppb.New(time.Now()),
+	summary := &vulnerabilities.Summary{}
+	if sum != nil {
+		summary.Critical = sum.CriticalVulnerabilities
+		summary.High = sum.HighVulnerabilities
+		summary.Medium = sum.MediumVulnerabilities
+		summary.Low = sum.LowVulnerabilities
+		summary.Unassigned = sum.UnassignedVulnerabilities
+		summary.RiskScore = sum.TotalRiskScore
+		summary.HasSbom = true
 	}
 
 	response := &vulnerabilities.GetVulnerabilitySummaryResponse{
 		Filter:               request.Filter,
-		VulnerabilitySummary: &summary,
+		VulnerabilitySummary: summary,
 		WorkloadCount:        sum.WorkloadCount,
 	}
 	return response, nil
 }
+
+func (s *Server) GetVulnerabilitySummaryForImage(ctx context.Context, request *vulnerabilities.GetVulnerabilitySummaryForImageRequest) (*vulnerabilities.GetVulnerabilitySummaryForImageResponse, error) {
+	summary, err := s.db.GetVulnerabilitySummaryForImage(ctx, sql.GetVulnerabilitySummaryForImageParams{
+		ImageName: request.ImageName,
+		ImageTag:  request.ImageTag,
+	})
+	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return &vulnerabilities.GetVulnerabilitySummaryForImageResponse{
+				VulnerabilitySummary: &vulnerabilities.Summary{},
+				WorkloadRef:          make([]*vulnerabilities.Workload, 0),
+			}, nil
+		}
+
+		return nil, fmt.Errorf("failed to get vulnerability summary for image: %w", err)
+	}
+	workloads, err := s.db.ListWorkloadsByImage(ctx, sql.ListWorkloadsByImageParams{
+		ImageName: request.ImageName,
+		ImageTag:  request.ImageTag,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to list workloads by image: %w", err)
+	}
+
+	refs := make([]*vulnerabilities.Workload, 0)
+	for _, w := range workloads {
+		refs = append(refs, &vulnerabilities.Workload{
+			Cluster:   w.Cluster,
+			Namespace: w.Namespace,
+			Name:      w.Name,
+			Type:      w.WorkloadType,
+			ImageName: w.ImageName,
+			ImageTag:  w.ImageTag,
+		})
+	}
+
+	vulnSummary := &vulnerabilities.Summary{}
+	if summary != nil {
+		vulnSummary = &vulnerabilities.Summary{
+			Critical:    summary.Critical,
+			High:        summary.High,
+			Medium:      summary.Medium,
+			Low:         summary.Low,
+			Unassigned:  summary.Unassigned,
+			RiskScore:   summary.RiskScore,
+			LastUpdated: timestamppb.New(summary.UpdatedAt.Time),
+			HasSbom:     true,
+		}
+	}
+
+	return &vulnerabilities.GetVulnerabilitySummaryForImageResponse{
+		VulnerabilitySummary: vulnSummary,
+		WorkloadRef:          refs,
+	}, nil
+}

internal/database/queries/vulnerabilities.sql

@@ -143,3 +143,22 @@ WHERE image_name = @image_name
 ORDER BY updated_at DESC
 LIMIT sqlc.arg('limit') OFFSET sqlc.arg('offset')
 ;
+
+-- name: GenerateVulnerabilitySummaryForImage :one
+SELECT COUNT(*) AS total,
+       SUM(CASE WHEN c.severity = 5 THEN 1 ELSE 0 END) AS critical,
+       SUM(CASE WHEN c.severity = 4 THEN 1 ELSE 0 END) AS high,
+       SUM(CASE WHEN c.severity = 3 THEN 1 ELSE 0 END) AS medium,
+       SUM(CASE WHEN c.severity = 2 THEN 1 ELSE 0 END) AS low,
+       SUM(CASE WHEN c.severity = 1 THEN 1 ELSE 0 END) AS unassigned,
+       -- 10*critical + 5*high + 3*medium + 1*low + 5*unassigned
+         10 * SUM(CASE WHEN c.severity = 5 THEN 1 ELSE 0 END) +
+            5 * SUM(CASE WHEN c.severity = 4 THEN 1 ELSE 0 END) +
+            3 * SUM(CASE WHEN c.severity = 3 THEN 1 ELSE 0 END) +
+            1 * SUM(CASE WHEN c.severity = 2 THEN 1 ELSE 0 END) +
+            5 * SUM(CASE WHEN c.severity = 1 THEN 1 ELSE 0 END) AS risk_score
+
+FROM vulnerabilities v
+         JOIN cve c ON v.cve_id = c.cve_id
+WHERE v.image_name = @image_name
+    AND v.image_tag = @image_tag;

internal/database/queries/vulnerbility_summary.sql

@@ -93,6 +93,8 @@ WHERE
   AND (CASE WHEN sqlc.narg('namespace')::TEXT is not null THEN w.namespace = sqlc.narg('namespace')::TEXT ELSE TRUE END)
   AND (CASE WHEN sqlc.narg('workload_type')::TEXT is not null THEN w.workload_type = sqlc.narg('workload_type')::TEXT ELSE TRUE END)
   AND (CASE WHEN sqlc.narg('workload_name')::TEXT is not null THEN w.name = sqlc.narg('workload_name')::TEXT ELSE TRUE END)
+  AND (CASE WHEN sqlc.narg('image_name')::TEXT is not null THEN v.image_name = sqlc.narg('image_name')::TEXT ELSE TRUE END)
+  AND (CASE WHEN sqlc.narg('image_tag')::TEXT is not null THEN v.image_tag = sqlc.narg('image_tag')::TEXT ELSE TRUE END)
 ORDER BY w.updated_at DESC
 LIMIT
     sqlc.arg('limit')
@@ -117,3 +119,8 @@ WHERE
   AND (CASE WHEN sqlc.narg('namespace')::TEXT IS NOT NULL THEN w.namespace = sqlc.narg('namespace')::TEXT ELSE TRUE END)
   AND (CASE WHEN sqlc.narg('workload_type')::TEXT IS NOT NULL THEN w.workload_type = sqlc.narg('workload_type')::TEXT ELSE TRUE END)
   AND (CASE WHEN sqlc.narg('workload_name')::TEXT IS NOT NULL THEN w.name = sqlc.narg('workload_name')::TEXT ELSE TRUE END);
+
+-- name: GetVulnerabilitySummaryForImage :one
+SELECT * FROM vulnerability_summary
+WHERE image_name = @image_name
+  AND image_tag = @image_tag;

internal/database/queries/workloads.sql

@@ -43,4 +43,12 @@ WHERE
     name = @name
 RETURNING
     *
-;
+;
+
+-- name: ListWorkloadsByImage :many
+SELECT *
+FROM workloads
+WHERE image_name = @image_name
+  AND image_tag = @image_tag
+ORDER BY
+    (name, cluster, updated_at) DESC;