ISSUE-1830: Adding support for stopping the scan (#61)

pkg/jobs/status.go

@@ -21,40 +21,47 @@ func StartStatusReporter(ctx context.Context, scan_id string, scanner *scan.Scan
 
 	go func() {
 		defer stopScanJob()
-		ticker := time.NewTicker(30 * time.Second)
-		var err, abort error
+		ticker := time.NewTicker(1 * time.Second)
+		var err error
 		ts := time.Now()
+		log.Infof("StatusReporter started, scan_id: %s", scan_id)
 	loop:
 		for {
 			select {
 			case err = <-res:
 				break loop
 			case <-ctx.Done():
-				abort = ctx.Err()
+				err = ctx.Err()
 				break loop
 			case <-scanner.ScanStatusChan:
 				ts = time.Now()
 			case <-ticker.C:
+				if scanner.Stopped.Load() == true {
+					log.Errorf("Scanner job stopped, scan_id: %s", scan_id)
+					break loop
+				}
+
 				elapsed := int(time.Since(ts).Seconds())
 				if elapsed > threshold {
 					err = fmt.Errorf("Scan job aborted due to inactivity")
-					log.Error("Scanner job aborted as no update within threshold, Scan id:" + scan_id)
+					log.Errorf("Scanner job aborted due to inactivity, scan_id: %s" + scan_id)
 					scanner.Aborted.Store(true)
 					break loop
 				} else {
 					output.WriteScanStatus("IN_PROGRESS", scan_id, "")
 				}
 			}
 		}
-		if abort != nil {
-			output.WriteScanStatus("CANCELLED", scan_id, abort.Error())
-			return
-		}
-		if err != nil {
+
+		if scanner.Stopped.Load() == true {
+			output.WriteScanStatus("ERROR", scan_id, "Scan stopped by user")
+		} else if err != nil {
 			output.WriteScanStatus("ERROR", scan_id, err.Error())
-			return
+		} else {
+			output.WriteScanStatus("COMPLETE", scan_id, "")
 		}
-		output.WriteScanStatus("COMPLETE", scan_id, "")
+
+		log.Infof("StatusReporter finished, scan_id: %s", scan_id)
 	}()
 	return res
 }

pkg/scan/process_image.go

@@ -890,6 +890,10 @@ func CheckScanStatus(s *Scanner) error {
 			close(s.ScanStatusChan)
 			log.Error("Scan aborted due to inactivity, scanid:", s.ScanID)
 			return fmt.Errorf("Scan aborted due to inactivity")
+		} else if s.Stopped.Load() == true {
+			close(s.ScanStatusChan)
+			log.Error("Scan stopped by user request, scanid:", s.ScanID)
+			return fmt.Errorf("Scan stopped by user request")
 		} else {
 			s.ScanStatusChan <- true
 		}

pkg/scan/scanner.go

@@ -12,8 +12,9 @@ func New(opts *config.Options, yaraconfig *config.Config,
 
 	statusChan := make(chan bool)
 	obj := Scanner{opts, yaraconfig, yaraScannerIn, scanID, statusChan,
-		atomic.Bool{}, atomic.Bool{}}
+		atomic.Bool{}, atomic.Bool{}, atomic.Bool{}}
 	obj.Aborted.Store(false)
+	obj.Stopped.Store(false)
 	obj.ReportStatus.Store(true)
 	return &obj
 }
@@ -26,6 +27,7 @@ type Scanner struct {
 	ScanID         string
 	ScanStatusChan chan bool
 	Aborted        atomic.Bool
+	Stopped        atomic.Bool
 	ReportStatus   atomic.Bool
 }
 

pkg/server/grpc.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"os/signal"
 	"strings"
+	"sync"
 	"syscall"
 
 	"github.com/deepfence/YaraHunter/constants"
@@ -41,6 +42,8 @@ type gRPCServer struct {
 	pb.UnimplementedMalwareScannerServer
 	pb.UnimplementedAgentPluginServer
 	pb.UnimplementedScannersServer
+
+	scanMap sync.Map
 }
 
 func (s *gRPCServer) ReportJobsStatus(context.Context, *pb.Empty) (*pb.JobReports, error) {
@@ -49,6 +52,33 @@ func (s *gRPCServer) ReportJobsStatus(context.Context, *pb.Empty) (*pb.JobReport
 	}, nil
 }
 
+func (s *gRPCServer) StopScan(c context.Context, req *pb.StopScanRequest) (*pb.StopScanResult, error) {
+	scanID := req.ScanId
+	result := &pb.StopScanResult{
+		Success:     true,
+		Description: "",
+	}
+
+	obj, found := s.scanMap.Load(scanID)
+	if !found {
+		msg := "Failed to Stop scan"
+		log.Info("%s, may have already completed, scan_id: %s", msg, scanID)
+		result.Success = false
+		result.Description = "Failed to Stop scan"
+		return result, nil
+	} else {
+		msg := "Stop request submitted"
+		log.Infof("%s, scan_id: %s", msg, scanID)
+		result.Success = true
+		result.Description = msg
+	}
+
+	scanner := obj.(*scan.Scanner)
+	scanner.Stopped.Store(true)
+
+	return result, nil
+}
+
 func (s *gRPCServer) GetName(context.Context, *pb.Empty) (*pb.Name, error) {
 	return &pb.Name{Str: s.plugin_name}, nil
 }
@@ -63,8 +93,10 @@ func (s *gRPCServer) FindMalwareInfo(c context.Context, r *pb.MalwareRequest) (*
 
 		yaraScanner, err := s.yaraRules.NewScanner()
 		scanner := scan.New(s.options, s.yaraConfig, yaraScanner, r.ScanId)
+		s.scanMap.Store(scanner.ScanID, scanner)
 		res := jobs.StartStatusReporter(context.Background(), r.ScanId, scanner)
 		defer func() {
+			s.scanMap.Delete(scanner.ScanID)
 			res <- err
 			close(res)
 		}()
@@ -137,6 +169,8 @@ func RunGrpcServer(opts *config.Options, config *config.Config, plugin_name stri
 		return err
 	}
 
+	impl.scanMap = sync.Map{}
+
 	// compile yara rules
 	impl.yaraRules = yararules.New(*opts.RulesPath)
 	err = impl.yaraRules.Compile(constants.Filescan, *opts.FailOnCompileWarning)