Repository: aquasecurity/starboard
Pull request #970 (2022-02-17 17:53:27)

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


Patch:
diff --git a/pkg/configauditreport/io.go b/pkg/configauditreport/io.go
index cb173ce..a3265a8 100644
--- a/pkg/configauditreport/io.go
+++ b/pkg/configauditreport/io.go
@@ -128,7 +128,7 @@ func (r *readWriter) FindReportByOwnerInHierarchy(ctx context.Context, owner kub
 
 	// no reports found for provided owner, look for reports in related replicaset
 	if report == nil && (owner.Kind == kube.KindDeployment || owner.Kind == kube.KindPod) {
-		rsName, err := r.GetRelatedReplicasetName(ctx, owner)
+		rsName, err := r.RelatedReplicaSetName(ctx, owner)
 		if err != nil {
 			return nil, fmt.Errorf("getting replicaset related to %s/%s: %w", owner.Kind, owner.Name, err)
 		}
diff --git a/pkg/kube/object.go b/pkg/kube/object.go
index 6ae9c9d..99263ac 100644
--- a/pkg/kube/object.go
+++ b/pkg/kube/object.go
@@ -18,7 +18,6 @@ import (
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
-	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/validation"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/apiutil"
@@ -191,7 +190,7 @@ func KindForObject(object metav1.Object, scheme *runtime.Scheme) (string, error)
 	return gvk.Kind, nil
 }
 
-func ObjectRefFromKindAndNamespacedName(kind Kind, name types.NamespacedName) ObjectRef {
+func ObjectRefFromKindAndObjectKey(kind Kind, name client.ObjectKey) ObjectRef {
 	return ObjectRef{
 		Kind:      kind,
 		Name:      name.Name,
@@ -299,7 +298,10 @@ func (o *ObjectResolver) ObjectFromObjectRef(ctx context.Context, ref ObjectRef)
 	default:
 		return nil, fmt.Errorf("unknown kind: %s", ref.Kind)
 	}
-	err := o.Client.Get(ctx, types.NamespacedName{Name: ref.Name, Namespace: ref.Namespace}, obj)
+	err := o.Client.Get(ctx, client.ObjectKey{
+		Name:      ref.Name,
+		Namespace: ref.Namespace,
+	}, obj)
 	if err != nil {
 		return nil, err
 	}
@@ -348,49 +350,86 @@ func (o *ObjectResolver) ReportOwner(ctx context.Context, obj client.Object) (cl
 	}
 }
 
+// ReplicaSetByDeploymentRef returns the current revision of the specified
+// Deployment reference. If the current revision cannot be found the
+// ErrReplicaSetNotFound error is returned.
+func (o *ObjectResolver) ReplicaSetByDeploymentRef(ctx context.Context, deploymentRef ObjectRef) (*appsv1.ReplicaSet, error) {
+	deployment := &appsv1.Deployment{}
+	err := o.Client.Get(ctx, client.ObjectKey{
+		Namespace: deploymentRef.Namespace,
+		Name:      deploymentRef.Name,
+	}, deployment)
+	if err != nil {
+		return nil, fmt.Errorf("getting deployment %q: %w", deploymentRef.Namespace+"/"+deploymentRef.Name, err)
+	}
+	return o.ReplicaSetByDeployment(ctx, deployment)
+}
+
 // ReplicaSetByDeployment returns the current revision of the specified
 // Deployment. If the current revision cannot be found the ErrReplicaSetNotFound
-// error is returned.
-func (o *ObjectResolver) ReplicaSetByDeployment(ctx context.Context, deploy *appsv1.Deployment) (*appsv1.ReplicaSet, error) {
+//error is returned.
+func (o *ObjectResolver) ReplicaSetByDeployment(ctx context.Context, deployment *appsv1.Deployment) (*appsv1.ReplicaSet, error) {
 	var rsList appsv1.ReplicaSetList
 	err := o.Client.List(ctx, &rsList,
-		client.InNamespace(deploy.Namespace),
+		client.InNamespace(deployment.Namespace),
 		client.MatchingLabelsSelector{
-			Selector: labels.SelectorFromSet(deploy.Spec.Selector.MatchLabels),
+			Selector: labels.SelectorFromSet(deployment.Spec.Selector.MatchLabels),
 		})
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("listing replicasets for deployment %q: %w", deployment.Namespace+"/"+deployment.Name, err)
 	}
+
 	if len(rsList.Items) == 0 {
 		return nil, ErrReplicaSetNotFound
 	}
+
 	for _, rs := range rsList.Items {
-		if deploy.Annotations["deployment.kubernetes.io/revision"] !=
-			rs.Annotations["deployment.kubernetes.io/revision"] {
+		if deployment.Annotations[deploymentAnnotation] !=
+			rs.Annotations[deploymentAnnotation] {
 			continue
 		}
-		obj, err := o.ensureGVK(rs.DeepCopy())
-		return obj.(*appsv1.ReplicaSet), err
+		rsCopy := rs.DeepCopy()
+		_, err = o.ensureGVK(rsCopy)
+		return rsCopy, err
 	}
+
 	return nil, ErrReplicaSetNotFound
 }
 
+// ReplicaSetByPodRef returns the controller ReplicaSet of the specified Pod
+// reference.
+func (o *ObjectResolver) ReplicaSetByPodRef(ctx context.Context, object ObjectRef) (*appsv1.ReplicaSet, error) {
+	pod := &corev1.Pod{}
+	err := o.Client.Get(ctx, client.ObjectKey{
+		Namespace: object.Namespace,
+		Name:      object.Name,
+	}, pod)
+	if err != nil {
+		return nil, err
+	}
+	return o.ReplicaSetByPod(ctx, pod)
+}
+
 // ReplicaSetByPod returns the controller ReplicaSet of the specified Pod.
 func (o *ObjectResolver) ReplicaSetByPod(ctx context.Context, pod *corev1.Pod) (*appsv1.ReplicaSet, error) {
 	controller := metav1.GetControllerOf(pod)
 	if controller == nil {
-		return nil, fmt.Errorf("did not find a controller for pod %q", pod.Name)
+		return nil, fmt.Errorf("did not find a controller for pod %q", pod.Namespace+"/"+pod.Name)
 	}
 	if controller.Kind != "ReplicaSet" {
 		return nil, fmt.Errorf("pod %q is controlled by a %q, want replicaset", pod.Name, controller.Kind)
 	}
 	rs := &appsv1.ReplicaSet{}
-	err := o.Client.Get(ctx, client.ObjectKey{Namespace: pod.Namespace, Name: controller.Name}, rs)
+	err := o.Client.Get(ctx, client.ObjectKey{
+		Namespace: pod.Namespace,
+		Name:      controller.Name,
+	}, rs)
 	if err != nil {
 		return nil, err
 	}
-	obj, err := o.ensureGVK(rs)
-	return obj.(*appsv1.ReplicaSet), err
+	rsCopy := rs.DeepCopy()
+	_, err = o.ensureGVK(rsCopy)
+	return rsCopy, err
 }
 
 func (o *ObjectResolver) CronJobByJob(ctx context.Context, job *batchv1.Job) (*batchv1.CronJob, error) {
@@ -440,12 +479,20 @@ func (o *ObjectResolver) ensureGVK(obj client.Object) (client.Object, error) {
 // the given owner. If the owner is a Deployment, it will look for a ReplicaSet
 // that is controlled by the Deployment. If the owner is a Pod, it will look for
 // the ReplicaSet that owns the Pod.
-func (o *ObjectResolver) GetRelatedReplicasetName(ctx context.Context, object ObjectRef) (string, error) {
+func (o *ObjectResolver) RelatedReplicaSetName(ctx context.Context, object ObjectRef) (string, error) {
 	switch object.Kind {
 	case KindDeployment:
-		return o.getActiveReplicaSetByDeployment(ctx, object)
+		rs, err := o.ReplicaSetByDeploymentRef(ctx, object)
+		if err != nil {
+			return "", err
+		}
+		return rs.Name, nil
 	case KindPod:
-		return o.getReplicaSetByPod(ctx, object)
+		rs, err := o.ReplicaSetByPodRef(ctx, object)
+		if err != nil {
+			return "", err
+		}
+		return rs.Name, nil
 	}
 	return "", fmt.Errorf("can only get related ReplicaSet for Deployment or Pod, not %q", string(object.Kind))
 }
@@ -506,49 +553,6 @@ func (o *ObjectResolver) GetNodeName(ctx context.Context, obj client.Object) (st
 	}
 }
 
-func (o *ObjectResolver) getActiveReplicaSetByDeployment(ctx context.Context, object ObjectRef) (string, error) {
-	deploy := &appsv1.Deployment{}
-	err := o.Client.Get(ctx, types.NamespacedName{Namespace: object.Namespace, Name: object.Name}, deploy)
-	if err != nil {
-		return "", fmt.Errorf("getting deployment %q: %w", object.Namespace+"/"+object.Name, err)
-	}
-	var rsList appsv1.ReplicaSetList
-	err = o.Client.List(ctx, &rsList,
-		client.InNamespace(deploy.Namespace),
-		client.MatchingLabelsSelector{
-			Selector: labels.SelectorFromSet(deploy.Spec.Selector.MatchLabels),
-		})
-	if err != nil {
-		return "", fmt.Errorf("listing replicasets for deployment %q: %w", object.Name, err)
-	}
-	if len(rsList.Items) == 0 {
-		return "", fmt.Errorf("no replicasets associated with deployment %q", object.Name)
-	}
-	for _, rs := range rsList.Items {
-		if deploy.Annotations[deploymentAnnotation] != rs.Annotations[deploymentAnnotation] {
-			continue
-		}
-		return rs.Name, nil
-	}
-	return "", fmt.Errorf("did not find an active replicaset associated with deployment %q", object.Name)
-}
-
-func (o *ObjectResolver) getReplicaSetByPod(ctx context.Context, object ObjectRef) (string, error) {
-	pod := &corev1.Pod{}
-	err := o.Client.Get(ctx, types.NamespacedName{Namespace: object.Namespace, Name: object.Name}, pod)
-	if err != nil {
-		return "", err
-	}
-	controller := metav1.GetControllerOf(pod)
-	if controller == nil {
-		return "", fmt.Errorf("did not find a controller for pod %q", object.Name)
-	}
-	if controller.Kind != "ReplicaSet" {
-		return "", fmt.Errorf("pod %q is controlled by a %q, want replicaset", object.Name, controller.Kind)
-	}
-	return controller.Name, nil
-}
-
 func (o *ObjectResolver) IsActiveReplicaSet(ctx context.Context, workloadObj client.Object, controller *metav1.OwnerReference) (bool, error) {
 	if controller != nil && controller.Kind == string(KindDeployment) {
 		deploymentObject := &appsv1.Deployment{}
diff --git a/pkg/operator/controller/configauditreport.go b/pkg/operator/controller/configauditreport.go
index f60c279..8cc3364 100644
--- a/pkg/operator/controller/configauditreport.go
+++ b/pkg/operator/controller/configauditreport.go
@@ -132,10 +132,10 @@ func (r *ConfigAuditReportReconciler) reconcileResource(resourceKind kube.Kind)
 	return func(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
 		log := r.Logger.WithValues("kind", resourceKind, "name", req.NamespacedName)
 
-		resourcePartial := kube.ObjectRefFromKindAndNamespacedName(resourceKind, req.NamespacedName)
+		resourceRef := kube.ObjectRefFromKindAndObjectKey(resourceKind, req.NamespacedName)
 
 		log.V(1).Info("Getting resource from cache")
-		resource, err := r.ObjectFromObjectRef(ctx, resourcePartial)
+		resource, err := r.ObjectFromObjectRef(ctx, resourceRef)
 		if err != nil {
 			if errors.IsNotFound(err) {
 				log.V(1).Info("Ignoring cached resource that must have been deleted")
@@ -191,7 +191,7 @@ func (r *ConfigAuditReportReconciler) reconcileResource(resourceKind kube.Kind)
 		log = log.WithValues("resourceSpecHash", resourceSpecHash, "pluginConfigHash", pluginConfigHash)
 
 		log.V(1).Info("Checking whether configuration audit report exists")
-		hasReport, err := r.hasReport(ctx, resourcePartial, resourceSpecHash, pluginConfigHash)
+		hasReport, err := r.hasReport(ctx, resourceRef, resourceSpecHash, pluginConfigHash)
 		if err != nil {
 			return ctrl.Result{}, err
 		}
diff --git a/pkg/operator/controller/vulnerabilityreport.go b/pkg/operator/controller/vulnerabilityreport.go
index 6c14d2b..2ff14a7 100644
--- a/pkg/operator/controller/vulnerabilityreport.go
+++ b/pkg/operator/controller/vulnerabilityreport.go
@@ -87,10 +87,10 @@ func (r *VulnerabilityReportReconciler) reconcileWorkload(workloadKind kube.Kind
 	return func(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
 		log := r.Logger.WithValues("kind", workloadKind, "name", req.NamespacedName)
 
-		workloadPartial := kube.ObjectRefFromKindAndNamespacedName(workloadKind, req.NamespacedName)
+		workloadRef := kube.ObjectRefFromKindAndObjectKey(workloadKind, req.NamespacedName)
 
 		log.V(1).Info("Getting workload from cache")
-		workloadObj, err := r.ObjectFromObjectRef(ctx, workloadPartial)
+		workloadObj, err := r.ObjectFromObjectRef(ctx, workloadRef)
 		if err != nil {
 			if k8sapierror.IsNotFound(err) {
 				log.V(1).Info("Ignoring cached workload that must have been deleted")
@@ -140,7 +140,7 @@ func (r *VulnerabilityReportReconciler) reconcileWorkload(workloadKind kube.Kind
 		log = log.WithValues("podSpecHash", hash)
 
 		// Check if containers of the Pod have corresponding VulnerabilityReports.
-		hasReports, err := r.hasReports(ctx, workloadPartial, hash, containerImages)
+		hasReports, err := r.hasReports(ctx, workloadRef, hash, containerImages)
 		if err != nil {
 			return ctrl.Result{}, fmt.Errorf("getting vulnerability reports: %w", err)
 		}
@@ -150,7 +150,7 @@ func (r *VulnerabilityReportReconciler) reconcileWorkload(workloadKind kube.Kind
 			return ctrl.Result{}, nil
 		}
 
-		_, job, err := r.hasActiveScanJob(ctx, workloadPartial, hash)
+		_, job, err := r.hasActiveScanJob(ctx, workloadRef, hash)
 		if err != nil {
 			return ctrl.Result{}, fmt.Errorf("checking scan job: %w", err)
 		}
diff --git a/pkg/vulnerabilityreport/io.go b/pkg/vulnerabilityreport/io.go
index 68175f4..6f54cfa 100644
--- a/pkg/vulnerabilityreport/io.go
+++ b/pkg/vulnerabilityreport/io.go
@@ -104,7 +104,7 @@ func (r *readWriter) FindByOwnerInHierarchy(ctx context.Context, owner kube.Obje
 
 	// no reports found for provided owner, look for reports in related replicaset
 	if len(reports) == 0 && (owner.Kind == kube.KindDeployment || owner.Kind == kube.KindPod) {
-		rsName, err := r.GetRelatedReplicasetName(ctx, owner)
+		rsName, err := r.RelatedReplicaSetName(ctx, owner)
 		if err != nil {
 			return nil, fmt.Errorf("getting replicaset related to %s/%s: %w", owner.Kind, owner.Name, err)
 		}


Test patch:
diff --git a/pkg/kube/object_test.go b/pkg/kube/object_test.go
index d64c77b..6a77df2 100644
--- a/pkg/kube/object_test.go
+++ b/pkg/kube/object_test.go
@@ -13,7 +13,6 @@ import (
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/utils/pointer"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
@@ -241,8 +240,8 @@ func TestContainerImages_AsJSON_And_FromJSON(t *testing.T) {
 	assert.Equal(t, containerImages, newContainerImages)
 }
 
-func TestObjectRefFromKindAndNamespacedName(t *testing.T) {
-	partial := kube.ObjectRefFromKindAndNamespacedName(kube.KindReplicaSet, types.NamespacedName{
+func TestObjectRefFromKindAndObjectKey(t *testing.T) {
+	partial := kube.ObjectRefFromKindAndObjectKey(kube.KindReplicaSet, client.ObjectKey{
 		Namespace: "prod",
 		Name:      "wordpress",
 	})
@@ -481,7 +480,7 @@ func TestGetPodSpec(t *testing.T) {
 	}
 }
 
-func TestObjectResolver_GetRelatedReplicasetName(t *testing.T) {
+func TestObjectResolver_RelatedReplicaSetName(t *testing.T) {
 
 	instance := &kube.ObjectResolver{Client: fake.NewClientBuilder().WithScheme(starboard.NewScheme()).WithObjects(
 		&appsv1.Deployment{
@@ -567,7 +566,7 @@ func TestObjectResolver_GetRelatedReplicasetName(t *testing.T) {
 	).Build()}
 
 	t.Run("Should return error for unsupported kind", func(t *testing.T) {
-		_, err := instance.GetRelatedReplicasetName(context.Background(), kube.ObjectRef{
+		_, err := instance.RelatedReplicaSetName(context.Background(), kube.ObjectRef{
 			Kind:      kube.KindStatefulSet,
 			Name:      "statefulapp",
 			Namespace: corev1.NamespaceDefault,
@@ -576,7 +575,7 @@ func TestObjectResolver_GetRelatedReplicasetName(t *testing.T) {
 	})
 
 	t.Run("Should return ReplicaSet name for the specified Deployment", func(t *testing.T) {
-		name, err := instance.GetRelatedReplicasetName(context.Background(), kube.ObjectRef{
+		name, err := instance.RelatedReplicaSetName(context.Background(), kube.ObjectRef{
 			Kind:      kube.KindDeployment,
 			Name:      "nginx",
 			Namespace: corev1.NamespaceDefault,
@@ -586,7 +585,7 @@ func TestObjectResolver_GetRelatedReplicasetName(t *testing.T) {
 	})
 
 	t.Run("Should return ReplicaSet name for the specified Deployment", func(t *testing.T) {
-		name, err := instance.GetRelatedReplicasetName(context.Background(), kube.ObjectRef{
+		name, err := instance.RelatedReplicaSetName(context.Background(), kube.ObjectRef{
 			Kind:      kube.KindPod,
 			Name:      "nginx-549f5fcb58-7cr5b",
 			Namespace: corev1.NamespaceDefault,