feat: Set ownerReference for vulnerabilities resource (aquasecurity#129)

Resolves: aquasecurity#26

Signed-off-by: Daniel Pacak <pacak.daniel@gmail.com>

cmd/starboard/main_test.go

@@ -0,0 +1 @@
+package main_test

pkg/cmd/commands.go

@@ -5,6 +5,10 @@ import (
 	"strings"
 	"time"
 
+	"k8s.io/client-go/kubernetes/scheme"
+
+	"k8s.io/apimachinery/pkg/runtime"
+
 	"github.com/aquasecurity/starboard/pkg/kube"
 	"github.com/spf13/cobra"
 	"k8s.io/cli-runtime/pkg/genericclioptions"
@@ -48,6 +52,10 @@ func WorkloadFromArgs(namespace string, args []string) (workload kube.Object, er
 	return
 }
 
+func GetScheme() *runtime.Scheme {
+	return scheme.Scheme
+}
+
 const (
 	scanJobTimeoutFlagName = "scan-job-timeout"
 	deleteScanJobFlagName  = "delete-scan-job"

pkg/cmd/find.go

@@ -10,7 +10,7 @@ func NewFindCmd(executable string, cf *genericclioptions.ConfigFlags) *cobra.Com
 		Use:   "find",
 		Short: "Manage security scanners",
 	}
-	findCmd.AddCommand(GetVulnerabilitiesCmd(executable, cf))
+	findCmd.AddCommand(NewFindVulnerabilitiesCmd(executable, cf))
 
 	return findCmd
 }

pkg/cmd/find_test.go

@@ -0,0 +1 @@
+package cmd_test

pkg/cmd/find_vulnerabilities.go

@@ -12,7 +12,7 @@ import (
 	"k8s.io/client-go/kubernetes"
 )
 
-func GetVulnerabilitiesCmd(executable string, cf *genericclioptions.ConfigFlags) *cobra.Command {
+func NewFindVulnerabilitiesCmd(executable string, cf *genericclioptions.ConfigFlags) *cobra.Command {
 	cmd := &cobra.Command{
 		Aliases: []string{"vulns", "vuln"},
 		Use:     "vulnerabilities (NAME | TYPE/NAME)",
@@ -70,15 +70,15 @@ NAME is the name of a particular Kubernetes workload.
 			if err != nil {
 				return
 			}
-			reports, err := trivy.NewScanner(opts, kubernetesClientset).Scan(ctx, workload)
+			reports, owner, err := trivy.NewScanner(opts, kubernetesClientset).Scan(ctx, workload)
 			if err != nil {
 				return
 			}
 			starboardClientset, err := starboardapi.NewForConfig(config)
 			if err != nil {
 				return
 			}
-			err = crd.NewReadWriter(starboardClientset).Write(ctx, workload, reports)
+			err = crd.NewReadWriter(GetScheme(), starboardClientset).Write(ctx, reports, owner)
 			return
 		},
 	}

pkg/cmd/find_vulnerabilities_test.go

@@ -0,0 +1 @@
+package cmd_test

pkg/cmd/get_report.go

@@ -42,7 +42,7 @@ NAME is the name of a particular Kubernetes workload.
 			}
 
 			caReader := configAuditCrd.NewReadWriter(starboardClientset)
-			vulnsReader := vulnsCrd.NewReadWriter(starboardClientset)
+			vulnsReader := vulnsCrd.NewReadWriter(GetScheme(), starboardClientset)
 
 			reporter := report.NewHTMLReporter(caReader, vulnsReader, workload)
 			err = reporter.GenerateReport(os.Stdout)

pkg/cmd/kube_bench.go

@@ -45,7 +45,7 @@ func NewKubeBenchCmd(cf *genericclioptions.ConfigFlags) *cobra.Command {
 				return
 			}
 			scanner := kubebench.NewScanner(opts, kubernetesClientset)
-			writer := crd.NewReadWriter(starboardClientset)
+			writer := crd.NewReadWriter(GetScheme(), starboardClientset)
 
 			var wg sync.WaitGroup
 

pkg/find/vulnerabilities/crd/writer.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"k8s.io/apimachinery/pkg/labels"
+	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/klog"
 
 	"github.com/aquasecurity/starboard/pkg/kube"
@@ -13,37 +14,44 @@ import (
 	"github.com/aquasecurity/starboard/pkg/find/vulnerabilities"
 	clientset "github.com/aquasecurity/starboard/pkg/generated/clientset/versioned"
 	"github.com/google/uuid"
-	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 type ReadWriter struct {
+	scheme *runtime.Scheme
 	client clientset.Interface
 }
 
-func NewReadWriter(client clientset.Interface) *ReadWriter {
+func NewReadWriter(scheme *runtime.Scheme, client clientset.Interface) *ReadWriter {
 	return &ReadWriter{
+		scheme: scheme,
 		client: client,
 	}
 }
 
-func (s *ReadWriter) Write(ctx context.Context, workload kube.Object, reports vulnerabilities.WorkloadVulnerabilities) (err error) {
+func (s *ReadWriter) Write(ctx context.Context, reports vulnerabilities.WorkloadVulnerabilities, owner metav1.Object) (err error) {
 	for container, report := range reports {
-		err = s.createVulnerability(ctx, workload, container, report)
+		err = s.createVulnerability(ctx, container, report, owner)
 		if err != nil {
 			return
 		}
 	}
 	return
 }
 
-func (s *ReadWriter) createVulnerability(ctx context.Context, workload kube.Object, container string, report starboard.VulnerabilityReport) (err error) {
-	namespace := workload.Namespace
+func (s *ReadWriter) createVulnerability(ctx context.Context, container string, report starboard.VulnerabilityReport, owner metav1.Object) (err error) {
+	namespace := owner.GetNamespace()
+	name := owner.GetName()
+	kind, err := kube.KindForObject(owner, s.scheme)
+	if err != nil {
+		return err
+	}
 
 	// Trying to find previously generated vulnerability report for this specific container
-	vulnsSearch, err := s.client.AquasecurityV1alpha1().Vulnerabilities(namespace).List(ctx, meta.ListOptions{
+	vulnsSearch, err := s.client.AquasecurityV1alpha1().Vulnerabilities(namespace).List(ctx, metav1.ListOptions{
 		LabelSelector: labels.Set{
-			kube.LabelResourceKind:  string(workload.Kind),
-			kube.LabelResourceName:  workload.Name,
+			kube.LabelResourceKind:  kind,
+			kube.LabelResourceName:  name,
 			kube.LabelContainerName: container,
 		}.String(),
 	})
@@ -53,32 +61,38 @@ func (s *ReadWriter) createVulnerability(ctx context.Context, workload kube.Obje
 
 	if len(vulnsSearch.Items) > 0 {
 		existingCR := vulnsSearch.Items[0]
-		klog.V(3).Infof("Updating vulnerability report: %s/%s", namespace, workload.Name)
+		klog.V(3).Infof("Updating vulnerability report: %s/%s", namespace, name)
 		deepCopy := existingCR.DeepCopy()
 		deepCopy.Report = report
-		_, err = s.client.AquasecurityV1alpha1().Vulnerabilities(namespace).Update(ctx, deepCopy, meta.UpdateOptions{})
+		_, err = s.client.AquasecurityV1alpha1().Vulnerabilities(namespace).Update(ctx, deepCopy, metav1.UpdateOptions{})
 	} else {
-		klog.V(3).Infof("Creating vulnerability report: %s/%s", namespace, workload.Name)
-		_, err = s.client.AquasecurityV1alpha1().Vulnerabilities(namespace).
-			Create(ctx, &starboard.Vulnerability{
-				ObjectMeta: meta.ObjectMeta{
-					Name: fmt.Sprintf(uuid.New().String()),
-					Labels: map[string]string{
-						kube.LabelResourceKind:      string(workload.Kind),
-						kube.LabelResourceName:      workload.Name,
-						kube.LabelResourceNamespace: workload.Namespace,
-						kube.LabelContainerName:     container,
-					},
+		klog.V(3).Infof("Creating vulnerability report: %s/%s", namespace, name)
+		report := &starboard.Vulnerability{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      fmt.Sprintf(uuid.New().String()),
+				Namespace: namespace,
+				Labels: map[string]string{
+					kube.LabelResourceKind:      kind,
+					kube.LabelResourceName:      name,
+					kube.LabelResourceNamespace: namespace,
+					kube.LabelContainerName:     container,
 				},
-				Report: report,
-			}, meta.CreateOptions{})
+			},
+			Report: report,
+		}
+		err = kube.SetOwnerReference(owner, report, s.scheme)
+		if err != nil {
+			return err
+		}
+		_, err = s.client.AquasecurityV1alpha1().Vulnerabilities(namespace).
+			Create(ctx, report, metav1.CreateOptions{})
 	}
 
 	return err
 }
 
 func (s *ReadWriter) Read(ctx context.Context, workload kube.Object) (vulnerabilities.WorkloadVulnerabilities, error) {
-	list, err := s.client.AquasecurityV1alpha1().Vulnerabilities(workload.Namespace).List(ctx, meta.ListOptions{
+	list, err := s.client.AquasecurityV1alpha1().Vulnerabilities(workload.Namespace).List(ctx, metav1.ListOptions{
 		LabelSelector: labels.Set{
 			kube.LabelResourceKind:      string(workload.Kind),
 			kube.LabelResourceName:      workload.Name,

pkg/find/vulnerabilities/crd/writer_test.go

@@ -1,9 +1,12 @@
-package crd
+package crd_test
 
 import (
 	"context"
+	"github.com/aquasecurity/starboard/pkg/find/vulnerabilities/crd"
 	"testing"
 
+	"github.com/aquasecurity/starboard/pkg/cmd"
+
 	"github.com/aquasecurity/starboard/pkg/find/vulnerabilities"
 
 	"github.com/aquasecurity/starboard/pkg/apis/aquasecurity/v1alpha1"
@@ -78,7 +81,7 @@ func TestReadWriter_Read(t *testing.T) {
 		Report: v1alpha1.VulnerabilityReport{},
 	})
 
-	reports, err := NewReadWriter(clientset).Read(context.TODO(), kube.Object{
+	reports, err := crd.NewReadWriter(cmd.GetScheme(), clientset).Read(context.TODO(), kube.Object{
 		Kind:      kube.KindDeployment,
 		Name:      "my-deploy",
 		Namespace: "my-namespace",

pkg/find/vulnerabilities/trivy/scanner.go

@@ -52,9 +52,9 @@ type Scanner struct {
 	scanners.Base
 }
 
-func (s *Scanner) Scan(ctx context.Context, workload kube.Object) (reports vulnerabilities.WorkloadVulnerabilities, err error) {
+func (s *Scanner) Scan(ctx context.Context, workload kube.Object) (reports vulnerabilities.WorkloadVulnerabilities, owner meta.Object, err error) {
 	klog.V(3).Infof("Getting Pod template for workload: %v", workload)
-	podSpec, err := s.pods.GetPodSpecByWorkload(ctx, workload)
+	podSpec, owner, err := s.pods.GetPodSpecByWorkload(ctx, workload)
 	if err != nil {
 		err = fmt.Errorf("getting Pod template: %w", err)
 		return

pkg/find/vulnerabilities/writer.go

@@ -3,11 +3,13 @@ package vulnerabilities
 import (
 	"context"
 
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
 	"github.com/aquasecurity/starboard/pkg/kube"
 )
 
 type Writer interface {
-	Write(ctx context.Context, workload kube.Object, reports WorkloadVulnerabilities) error
+	Write(ctx context.Context, reports WorkloadVulnerabilities, owner metav1.Object) error
 }
 
 type Reader interface {

pkg/kube/object.go

@@ -4,6 +4,10 @@ import (
 	"encoding/json"
 	"fmt"
 
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/schema"
+
 	"k8s.io/apimachinery/pkg/labels"
 )
 
@@ -84,3 +88,113 @@ func (ci ContainerImages) AsJSON() (string, error) {
 func (ci ContainerImages) FromJSON(value string) error {
 	return json.Unmarshal([]byte(value), &ci)
 }
+
+func GVKForObject(obj runtime.Object, scheme *runtime.Scheme) (schema.GroupVersionKind, error) {
+	gvks, isUnversioned, err := scheme.ObjectKinds(obj)
+	if err != nil {
+		return schema.GroupVersionKind{}, err
+	}
+	if isUnversioned {
+		return schema.GroupVersionKind{}, fmt.Errorf("cannot create group-version-kind for unversioned type %T", obj)
+	}
+
+	if len(gvks) < 1 {
+		return schema.GroupVersionKind{}, fmt.Errorf("no group-version-kinds associated with type %T", obj)
+	}
+	if len(gvks) > 1 {
+		// this should only trigger for things like metav1.XYZ --
+		// normal versioned types should be fine
+		return schema.GroupVersionKind{}, fmt.Errorf(
+			"multiple group-version-kinds associated with type %T, refusing to guess at one", obj)
+	}
+	return gvks[0], nil
+}
+
+func KindForObject(object metav1.Object, scheme *runtime.Scheme) (string, error) {
+	ro, ok := object.(runtime.Object)
+	if !ok {
+		return "", fmt.Errorf("%T is not a runtime.Object", object)
+	}
+	gvk, err := GVKForObject(ro, scheme)
+	if err != nil {
+		return "", err
+	}
+	return gvk.Kind, nil
+}
+
+func SetOwnerReference(owner, object metav1.Object, scheme *runtime.Scheme) error {
+	// Validate the owner.
+	ro, ok := owner.(runtime.Object)
+	if !ok {
+		return fmt.Errorf("%T is not a runtime.Object, cannot call SetOwnerReference", owner)
+	}
+	if err := validateOwner(owner, object); err != nil {
+		return err
+	}
+
+	// Create a new owner ref.
+	gvk, err := GVKForObject(ro, scheme)
+	if err != nil {
+		return err
+	}
+	ref := metav1.OwnerReference{
+		APIVersion: gvk.GroupVersion().String(),
+		Kind:       gvk.Kind,
+		UID:        owner.GetUID(),
+		Name:       owner.GetName(),
+	}
+
+	// Update owner references and return.
+	upsertOwnerRef(ref, object)
+	return nil
+}
+
+func validateOwner(owner, object metav1.Object) error {
+	ownerNs := owner.GetNamespace()
+	if ownerNs != "" {
+		objNs := object.GetNamespace()
+		if objNs == "" {
+			return fmt.Errorf("cluster-scoped resource must not have a namespace-scoped owner, owner's namespace %s", ownerNs)
+		}
+		if ownerNs != objNs {
+			return fmt.Errorf("cross-namespace owner references are disallowed, owner's namespace %s, obj's namespace %s", owner.GetNamespace(), object.GetNamespace())
+		}
+	}
+	return nil
+}
+
+func upsertOwnerRef(ref metav1.OwnerReference, object metav1.Object) {
+	owners := object.GetOwnerReferences()
+	idx := indexOwnerRef(owners, ref)
+	if idx == -1 {
+		owners = append(owners, ref)
+	} else {
+		owners[idx] = ref
+	}
+	object.SetOwnerReferences(owners)
+}
+
+// indexOwnerRef returns the index of the owner reference in the slice if found, or -1.
+func indexOwnerRef(ownerReferences []metav1.OwnerReference, ref metav1.OwnerReference) int {
+	for index, r := range ownerReferences {
+		if referSameObject(r, ref) {
+			return index
+		}
+	}
+	return -1
+}
+
+// Returns true if a and b point to the same object
+func referSameObject(a, b metav1.OwnerReference) bool {
+	aGV, err := schema.ParseGroupVersion(a.APIVersion)
+	if err != nil {
+		return false
+	}
+
+	bGV, err := schema.ParseGroupVersion(b.APIVersion)
+	if err != nil {
+		return false
+	}
+
+	return aGV.Group == bGV.Group && a.Kind == b.Kind && a.Name == b.Name
+}