Fix server controls workload logic

deepfence · Sep 25, 2023 · a076bea · a076bea
1 parent e7df629
commit a076bea
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 2 deletions.
diff --git a/deepfence_worker/controls/controls.go b/deepfence_worker/controls/controls.go
@@ -43,8 +43,13 @@ func RegisterControl[T ctl.StartVulnerabilityScanRequest | ctl.StartSecretScanRe
 func ApplyControl(ctx context.Context, req ctl.Action) error {
 	controls_guard.RLock()
 	defer controls_guard.RUnlock()
-	log.Info().Msgf("apply control req: %+v", req)
-	return controls[ctl.ActionID(req.ID)](ctx, []byte(req.RequestPayload))
+	f, has := controls[ctl.ActionID(req.ID)]
+	if has {
+		log.Info().Msgf("apply control req: %+v", req)
+		return f(ctx, []byte(req.RequestPayload))
+	}
+	log.Warn().Msgf("apply control req: %+v not implemented", req)
+	return nil
 }
 
 func init() {

diff --git a/deepfence_worker/tasks/malwarescan/malwarescan.go b/deepfence_worker/tasks/malwarescan/malwarescan.go
@@ -50,6 +50,8 @@ func NewMalwareScanner(ingest chan *kgo.Record) MalwareScan {
 }
 
 func (s MalwareScan) StopMalwareScan(ctx context.Context, task *asynq.Task) error {
+	defer cronjobs.ScanWorkloadAllocator.Free()
+
 	var params utils.MalwareScanParameters
 
 	log.Info().Msgf("StopMalwareScan, payload: %s ", string(task.Payload()))

diff --git a/deepfence_worker/tasks/sbom/generate_sbom.go b/deepfence_worker/tasks/sbom/generate_sbom.go
@@ -36,6 +36,8 @@ func NewSbomGenerator(ingest chan *kgo.Record) SbomGenerator {
 }
 
 func StopVulnerabilityScan(ctx context.Context, task *asynq.Task) error {
+	defer cronjobs.ScanWorkloadAllocator.Free()
+
 	log.Info().Msgf("StopVulnerabilityScan, payload: %s ", string(task.Payload()))
 	var params utils.SbomParameters
 	if err := json.Unmarshal(task.Payload(), &params); err != nil {

diff --git a/deepfence_worker/tasks/secretscan/secretscan.go b/deepfence_worker/tasks/secretscan/secretscan.go
@@ -38,6 +38,8 @@ func NewSecretScanner(ingest chan *kgo.Record) SecretScan {
 }
 
 func (s SecretScan) StopSecretScan(ctx context.Context, task *asynq.Task) error {
+	defer cronjobs.ScanWorkloadAllocator.Free()
+
 	var params utils.SecretScanParameters
 
 	log.Info().Msgf("StopSecretScan, payload: %s ", string(task.Payload()))