Fix server controls workload logic

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

Hide file tree

Showing 4 changed files with 23 additions and 5 deletions.
diff --git a/deepfence_worker/controls/controls.go b/deepfence_worker/controls/controls.go
@@ -12,6 +12,10 @@ import (
 	sdkUtils "github.com/deepfence/ThreatMapper/deepfence_utils/utils"
 )
 
+const (
+	ContextKeyIsConsole = "server-controls-is-console"
+)
+
 var controls map[ctl.ActionID]func(ctx context.Context, req []byte) error
 var controls_guard sync.RWMutex
 
@@ -43,8 +47,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(context.WithValue(ctx, ContextKeyIsConsole, struct{}{}), []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
@@ -11,6 +11,7 @@ import (
 	"github.com/deepfence/YaraHunter/pkg/output"
 	"github.com/hibiken/asynq"
 
+	"github.com/deepfence/ThreatMapper/deepfence_worker/controls"
 	"github.com/deepfence/ThreatMapper/deepfence_worker/cronjobs"
 	workerUtils "github.com/deepfence/ThreatMapper/deepfence_worker/utils"
 	malwareScanConstants "github.com/deepfence/YaraHunter/constants"
@@ -76,7 +77,9 @@ func (s MalwareScan) StopMalwareScan(ctx context.Context, task *asynq.Task) erro
 }
 
 func (s MalwareScan) StartMalwareScan(ctx context.Context, task *asynq.Task) error {
-	defer cronjobs.ScanWorkloadAllocator.Free()
+	if ctx.Value(controls.ContextKeyIsConsole) != nil {
+		defer cronjobs.ScanWorkloadAllocator.Free()
+	}
 
 	tenantID, err := directory.ExtractNamespace(ctx)
 	if err != nil {

diff --git a/deepfence_worker/tasks/sbom/generate_sbom.go b/deepfence_worker/tasks/sbom/generate_sbom.go
@@ -13,6 +13,7 @@ import (
 	"github.com/deepfence/ThreatMapper/deepfence_utils/directory"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/utils"
+	"github.com/deepfence/ThreatMapper/deepfence_worker/controls"
 	"github.com/deepfence/ThreatMapper/deepfence_worker/cronjobs"
 	workerUtils "github.com/deepfence/ThreatMapper/deepfence_worker/utils"
 	"github.com/deepfence/package-scanner/sbom/syft"
@@ -59,7 +60,9 @@ func StopVulnerabilityScan(ctx context.Context, task *asynq.Task) error {
 }
 
 func (s SbomGenerator) GenerateSbom(ctx context.Context, task *asynq.Task) error {
-	defer cronjobs.ScanWorkloadAllocator.Free()
+	if ctx.Value(controls.ContextKeyIsConsole) != nil {
+		defer cronjobs.ScanWorkloadAllocator.Free()
+	}
 
 	var params utils.SbomParameters
 

diff --git a/deepfence_worker/tasks/secretscan/secretscan.go b/deepfence_worker/tasks/secretscan/secretscan.go
@@ -16,6 +16,7 @@ import (
 	"github.com/deepfence/ThreatMapper/deepfence_utils/directory"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/log"
 	"github.com/deepfence/ThreatMapper/deepfence_utils/utils"
+	"github.com/deepfence/ThreatMapper/deepfence_worker/controls"
 	"github.com/deepfence/ThreatMapper/deepfence_worker/cronjobs"
 	workerUtils "github.com/deepfence/ThreatMapper/deepfence_worker/utils"
 	pb "github.com/deepfence/agent-plugins-grpc/srcgo"
@@ -63,7 +64,9 @@ func (s SecretScan) StopSecretScan(ctx context.Context, task *asynq.Task) error
 }
 
 func (s SecretScan) StartSecretScan(ctx context.Context, task *asynq.Task) error {
-	defer cronjobs.ScanWorkloadAllocator.Free()
+	if ctx.Value(controls.ContextKeyIsConsole) != nil {
+		defer cronjobs.ScanWorkloadAllocator.Free()
+	}
 
 	tenantID, err := directory.ExtractNamespace(ctx)
 	if err != nil {