controller.go

package vulnerabilityreport

import (
	. "github.com/kube-security-manager/kube-security-manager/pkg/operator/predicate"

	"context"
	"errors"
	"fmt"
	"reflect"

	"github.com/go-logr/logr"
	"github.com/kube-security-manager/kube-security-manager/pkg/apis/aquasecurity/v1alpha1"
	"github.com/kube-security-manager/kube-security-manager/pkg/kube"
	"github.com/kube-security-manager/kube-security-manager/pkg/operator/controller"
	"github.com/kube-security-manager/kube-security-manager/pkg/operator/etc"
	"github.com/kube-security-manager/kube-security-manager/pkg/starboard"
	appsv1 "k8s.io/api/apps/v1"
	batchv1 "k8s.io/api/batch/v1"
	batchv1beta1 "k8s.io/api/batch/v1beta1"
	corev1 "k8s.io/api/core/v1"
	k8sapierror "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"
)

// WorkloadController watches Kubernetes workloads and generates
// v1alpha1.VulnerabilityReport instances using vulnerability scanner that that
// implements the Plugin interface.
type WorkloadController struct {
	logr.Logger
	etc.Config
	client.Client
	kube.ObjectResolver
	controller.LimitChecker
	kube.LogsReader
	kube.SecretsReader
	Plugin
	starboard.PluginContext
	ReadWriter
	starboard.ConfigData
}

func (r *WorkloadController) SetupWithManager(mgr ctrl.Manager) error {
	installModePredicate, err := InstallModePredicate(r.Config)
	if err != nil {
		return err
	}

	workloads := []struct {
		kind       kube.Kind
		forObject  client.Object
		ownsObject client.Object
	}{
		{kind: kube.KindPod, forObject: &corev1.Pod{}, ownsObject: &v1alpha1.VulnerabilityReport{}},
		{kind: kube.KindReplicaSet, forObject: &appsv1.ReplicaSet{}, ownsObject: &v1alpha1.VulnerabilityReport{}},
		{kind: kube.KindReplicationController, forObject: &corev1.ReplicationController{}, ownsObject: &v1alpha1.VulnerabilityReport{}},
		{kind: kube.KindStatefulSet, forObject: &appsv1.StatefulSet{}, ownsObject: &v1alpha1.VulnerabilityReport{}},
		{kind: kube.KindDaemonSet, forObject: &appsv1.DaemonSet{}, ownsObject: &v1alpha1.VulnerabilityReport{}},
		{kind: kube.KindCronJob, forObject: &batchv1beta1.CronJob{}, ownsObject: &v1alpha1.VulnerabilityReport{}},
		{kind: kube.KindJob, forObject: &batchv1.Job{}, ownsObject: &v1alpha1.VulnerabilityReport{}},
	}

	for _, workload := range workloads {
		err = ctrl.NewControllerManagedBy(mgr).
			For(workload.forObject, builder.WithPredicates(
				Not(ManagedByStarboardOperator),
				Not(IsBeingTerminated),
				installModePredicate,
			)).
			Owns(workload.ownsObject).
			Complete(r.reconcileWorkload(workload.kind))
		if err != nil {
			return err
		}
	}
	var predicates []predicate.Predicate
	if !r.ConfigData.VulnerabilityScanJobsInSameNamespace() {
		predicates = append(predicates, InNamespace(r.Config.Namespace))
	}
	predicates = append(predicates, ManagedByStarboardOperator, IsVulnerabilityReportScan, JobHasAnyCondition)
	return ctrl.NewControllerManagedBy(mgr).
		For(&batchv1.Job{}, builder.WithPredicates(predicates...)).
		Complete(r.reconcileJobs())
}

func (r *WorkloadController) reconcileWorkload(workloadKind kube.Kind) reconcile.Func {
	return func(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
		log := r.Logger.WithValues("kind", workloadKind, "name", req.NamespacedName)

		workloadRef := kube.ObjectRefFromKindAndObjectKey(workloadKind, req.NamespacedName)

		log.V(1).Info("Getting workload from cache")
		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")
				return ctrl.Result{}, nil
			}
			return ctrl.Result{}, fmt.Errorf("getting %s from cache: %w", workloadKind, err)
		}

		// Skip processing if it's a Pod controlled by a built-in K8s workload.
		if workloadKind == kube.KindPod {
			controller := metav1.GetControllerOf(workloadObj)
			if kube.IsBuiltInWorkload(controller) {
				log.V(1).Info("Ignoring managed pod", "controllerKind", controller.Kind, "controllerName", controller.Name)
				return ctrl.Result{}, nil
			}
		}

		if r.Config.VulnerabilityScannerScanOnlyCurrentRevisions && workloadKind == kube.KindReplicaSet {
			controller := metav1.GetControllerOf(workloadObj)
			activeReplicaSet, err := r.IsActiveReplicaSet(ctx, workloadObj, controller)
			if err != nil {
				return ctrl.Result{}, fmt.Errorf("failed checking current revision: %w", err)
			}
			if !activeReplicaSet {
				log.V(1).Info("Ignoring inactive ReplicaSet", "controllerKind", controller.Kind, "controllerName", controller.Name)
				return ctrl.Result{}, nil
			}
		}

		// Skip processing if it's a Job controlled by CronJob.
		if workloadKind == kube.KindJob {
			controller := metav1.GetControllerOf(workloadObj)
			if controller != nil && controller.Kind == string(kube.KindCronJob) {
				log.V(1).Info("Ignoring managed job", "controllerKind", controller.Kind, "controllerName", controller.Name)
				return ctrl.Result{}, nil
			}
		}

		podSpec, err := kube.GetPodSpec(workloadObj)
		if err != nil {
			return ctrl.Result{}, err
		}

		containerImages := kube.GetContainerImagesFromPodSpec(podSpec)
		hash := kube.ComputeHash(podSpec)

		log = log.WithValues("podSpecHash", hash)

		// Check if containers of the Pod have corresponding VulnerabilityReports.
		hasReports, err := r.hasReports(ctx, workloadRef, hash, containerImages)
		if err != nil {
			return ctrl.Result{}, fmt.Errorf("getting vulnerability reports: %w", err)
		}

		if hasReports {
			log.V(1).Info("VulnerabilityReports already exist")
			return ctrl.Result{}, nil
		}

		_, job, err := r.hasActiveScanJob(ctx, workloadRef, hash)
		if err != nil {
			return ctrl.Result{}, fmt.Errorf("checking scan job: %w", err)
		}

		if job != nil {
			log.V(1).Info("Scan job already exists",
				"job", fmt.Sprintf("%s/%s", job.Namespace, job.Name))
			return ctrl.Result{}, nil
		}

		limitExceeded, scanJobsCount, err := r.LimitChecker.Check(ctx)
		if err != nil {
			return ctrl.Result{}, err
		}
		log.V(1).Info("Checking scan jobs limit", "count", scanJobsCount, "limit", r.ConcurrentScanJobsLimit)

		if limitExceeded {
			log.V(1).Info("Pushing back scan job", "count", scanJobsCount, "retryAfter", r.ScanJobRetryAfter)
			return ctrl.Result{RequeueAfter: r.Config.ScanJobRetryAfter}, nil
		}

		return ctrl.Result{}, r.submitScanJob(ctx, workloadObj)
	}
}

func (r *WorkloadController) hasReports(ctx context.Context, owner kube.ObjectRef, hash string, images kube.ContainerImages) (bool, error) {
	// TODO FindByOwner should accept optional label selector to further narrow down search results
	list, err := r.FindByOwner(ctx, owner)
	if err != nil {
		return false, err
	}

	actual := map[string]bool{}
	for _, report := range list {
		if containerName, ok := report.Labels[starboard.LabelContainerName]; ok {
			if hash == report.Labels[starboard.LabelResourceSpecHash] {
				actual[containerName] = true
			}
		}
	}

	expected := map[string]bool{}
	for containerName := range images {
		expected[containerName] = true
	}

	return reflect.DeepEqual(actual, expected), nil
}

func (r *WorkloadController) hasActiveScanJob(ctx context.Context, owner kube.ObjectRef, hash string) (bool, *batchv1.Job, error) {
	jobName := fmt.Sprintf("scan-vulnerabilityreport-%s", kube.ComputeHash(owner))
	job := &batchv1.Job{}
	err := r.Get(ctx, client.ObjectKey{Namespace: r.Config.Namespace, Name: jobName}, job)
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			return false, nil, nil
		}
		return false, nil, fmt.Errorf("getting job from cache: %w", err)
	}
	if job.Labels[starboard.LabelResourceSpecHash] == hash {
		return true, job, nil
	}
	return false, nil, nil
}

func (r *WorkloadController) submitScanJob(ctx context.Context, owner client.Object) error {
	log := r.Logger.WithValues("kind", owner.GetObjectKind().GroupVersionKind().Kind,
		"name", owner.GetName(), "namespace", owner.GetNamespace())
	credentials, err := r.CredentialsByWorkload(ctx, owner)
	if err != nil {
		return err
	}

	scanJobTolerations, err := r.GetScanJobTolerations()
	if err != nil {
		return fmt.Errorf("getting scan job tolerations: %w", err)
	}

	scanJobAnnotations, err := r.GetScanJobAnnotations()
	if err != nil {
		return fmt.Errorf("getting scan job annotations: %w", err)
	}

	scanJobPodTemplateLabels, err := r.GetScanJobPodTemplateLabels()
	if err != nil {
		return fmt.Errorf("getting scan job template labels: %w", err)
	}

	scanJob, secrets, err := NewScanJobBuilder().
		WithPlugin(r.Plugin).
		WithPluginContext(r.PluginContext).
		WithTimeout(r.Config.ScanJobTimeout).
		WithObject(owner).
		WithTolerations(scanJobTolerations).
		WithAnnotations(scanJobAnnotations).
		WithPodTemplateLabels(scanJobPodTemplateLabels).
		WithCredentials(credentials).
		Get()

	if err != nil {
		if errors.Is(err, kube.ErrReplicaSetNotFound) || errors.Is(err, kube.ErrNoRunningPods) ||
			errors.Is(err, kube.ErrUnSupportedKind) {
			log.V(1).Info("ignoring vulnerability scan", "reason", err)
			return nil
		}
		return fmt.Errorf("constructing scan job: %w", err)
	}

	for _, secret := range secrets {
		err = r.Client.Create(ctx, secret)
		if err != nil {
			if k8sapierror.IsAlreadyExists(err) {
				return nil
			}
			return fmt.Errorf("creating secret used by scan job failed: %s: %w", secret.Namespace+"/"+secret.Name, err)
		}
	}

	err = r.Client.Create(ctx, scanJob)
	if err != nil {
		if k8sapierror.IsAlreadyExists(err) {
			// TODO Delete secrets that were created in the previous step. Alternatively we can delete them on schedule.
			return nil
		}
		return fmt.Errorf("creating scan job failed: %s: %w", scanJob.Namespace+"/"+scanJob.Name, err)
	}

	for _, secret := range secrets {
		err = controllerutil.SetOwnerReference(scanJob, secret, r.Client.Scheme())
		if err != nil {
			return fmt.Errorf("setting owner reference: %w", err)
		}
		err := r.Client.Update(ctx, secret)
		if err != nil {
			return fmt.Errorf("setting owner reference of secret used by scan job failed: %s: %w", secret.Namespace+"/"+secret.Name, err)
		}
	}

	return nil
}

func (r *WorkloadController) reconcileJobs() reconcile.Func {
	return func(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
		log := r.Logger.WithValues("job", req.NamespacedName)

		job := &batchv1.Job{}
		err := r.Client.Get(ctx, req.NamespacedName, job)
		if err != nil {
			if k8sapierror.IsNotFound(err) {
				log.V(1).Info("Ignoring cached job that must have been deleted")
				return ctrl.Result{}, nil
			}
			return ctrl.Result{}, fmt.Errorf("getting job from cache: %w", err)
		}

		if len(job.Status.Conditions) == 0 {
			log.V(1).Info("Ignoring Job without conditions")
			return ctrl.Result{}, nil
		}

		switch jobCondition := job.Status.Conditions[0].Type; jobCondition {
		case batchv1.JobComplete:
			err = r.processCompleteScanJob(ctx, job)
		case batchv1.JobFailed:
			err = r.processFailedScanJob(ctx, job)
		default:
			err = fmt.Errorf("unrecognized scan job condition: %v", jobCondition)
		}

		return ctrl.Result{}, err
	}

}

func (r *WorkloadController) processCompleteScanJob(ctx context.Context, job *batchv1.Job) error {
	log := r.Logger.WithValues("job", fmt.Sprintf("%s/%s", job.Namespace, job.Name))

	ownerRef, err := kube.ObjectRefFromObjectMeta(job.ObjectMeta)
	if err != nil {
		return fmt.Errorf("getting owner ref from scan job metadata: %w", err)
	}

	owner, err := r.ObjectFromObjectRef(ctx, ownerRef)
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			log.V(1).Info("Report owner must have been deleted", "owner", owner)
			return r.deleteJob(ctx, job)
		}
		return fmt.Errorf("getting object from object ref: %w", err)
	}

	containerImages, err := kube.GetContainerImagesFromJob(job)
	if err != nil {
		return fmt.Errorf("getting container images: %w", err)
	}

	podSpecHash, ok := job.Labels[starboard.LabelResourceSpecHash]
	if !ok {
		return fmt.Errorf("expected label %s not set", starboard.LabelResourceSpecHash)
	}

	hasReports, err := r.hasReports(ctx, ownerRef, podSpecHash, containerImages)
	if err != nil {
		return err
	}

	if hasReports {
		log.V(1).Info("VulnerabilityReports already exist", "owner", owner)
		log.V(1).Info("Deleting complete scan job", "owner", owner)
		return r.deleteJob(ctx, job)
	}

	var vulnerabilityReports []v1alpha1.VulnerabilityReport

	for containerName, containerImage := range containerImages {
		logsStream, err := r.LogsReader.GetLogsByJobAndContainerName(ctx, job, containerName)
		if err != nil {
			if k8sapierror.IsNotFound(err) {
				log.V(1).Info("Cached job must have been deleted")
				return nil
			}
			if kube.IsPodControlledByJobNotFound(err) {
				log.V(1).Info("Pod must have been deleted")
				return r.deleteJob(ctx, job)
			}
			return fmt.Errorf("getting logs for pod %q: %w", job.Namespace+"/"+job.Name, err)
		}
		reportData, err := r.Plugin.ParseVulnerabilityReportData(r.PluginContext, containerImage, logsStream)
		if err != nil {
			return err
		}
		_ = logsStream.Close()

		reportBuilder := NewReportBuilder(r.Client.Scheme()).
			Controller(owner).
			Container(containerName).
			Data(reportData).
			PodSpecHash(podSpecHash)

		if r.Config.VulnerabilityScannerReportTTL != nil {
			reportBuilder.ReportTTL(r.Config.VulnerabilityScannerReportTTL)
		}

		report, err := reportBuilder.Get()

		if err != nil {
			return err
		}

		vulnerabilityReports = append(vulnerabilityReports, report)
	}

	err = r.ReadWriter.Write(ctx, vulnerabilityReports)
	if err != nil {
		return err
	}

	log.V(1).Info("Deleting complete scan job", "owner", owner)
	return r.deleteJob(ctx, job)
}

func (r *WorkloadController) processFailedScanJob(ctx context.Context, scanJob *batchv1.Job) error {
	log := r.Logger.WithValues("job", fmt.Sprintf("%s/%s", scanJob.Namespace, scanJob.Name))

	statuses, err := r.GetTerminatedContainersStatusesByJob(ctx, scanJob)
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			log.V(1).Info("Cached job must have been deleted")
			return nil
		}
		if kube.IsPodControlledByJobNotFound(err) {
			log.V(1).Info("Pod must have been deleted")
			return r.deleteJob(ctx, scanJob)
		}
		return err
	}
	for container, status := range statuses {
		if status.ExitCode == 0 {
			continue
		}
		log.Error(nil, "Scan job container", "container", container, "status.reason", status.Reason, "status.message", status.Message)
	}
	log.V(1).Info("Deleting failed scan job")
	return r.deleteJob(ctx, scanJob)
}

func (r *WorkloadController) deleteJob(ctx context.Context, job *batchv1.Job) error {
	err := r.Client.Delete(ctx, job, client.PropagationPolicy(metav1.DeletePropagationBackground))
	if err != nil {
		if k8sapierror.IsNotFound(err) {
			return nil
		}
		return fmt.Errorf("deleting job: %w", err)
	}
	return nil
}