stackrox · vikin91 · Nov 29, 2022 · Nov 14, 2022 · Nov 14, 2022 · Nov 15, 2022
@@ -1,24 +1,25 @@
 package compliance
 
 import (
-	"github.com/stackrox/rox/generated/internalapi/central"
 	"github.com/stackrox/rox/generated/storage"
 	"github.com/stackrox/rox/pkg/concurrency"
+	"github.com/stackrox/rox/pkg/sync"
 	"github.com/stackrox/rox/sensor/common"
 )
 
 // NodeScanHandler is responsible for handling the arriving NodeScanV2 messages, processing then, and sending them to central
 type NodeScanHandler interface {
 	common.SensorComponent
+	Stopped() concurrency.ReadOnlyErrorSignal
 }
 
 // NewNodeScanHandler returns a new instance of a NodeScanHandler
 func NewNodeScanHandler(ch <-chan *storage.NodeScanV2) NodeScanHandler {
 	return &nodeScanHandlerImpl{
 		nodeScans: ch,
-		toCentral: make(chan *central.MsgFromSensor),
-
-		stopC:    concurrency.NewErrorSignal(),
-		stoppedC: concurrency.NewErrorSignal(),
+		toCentral: nil,
+		stopC:     concurrency.NewErrorSignal(),
+		lock:      &sync.Mutex{},
+		stoppedC:  concurrency.NewErrorSignal(),
 	}
 }
@@ -6,31 +6,52 @@ import (
 	"github.com/stackrox/rox/generated/storage"
 	"github.com/stackrox/rox/pkg/centralsensor"
 	"github.com/stackrox/rox/pkg/concurrency"
+	"github.com/stackrox/rox/pkg/sync"
+)
+
+var (
+	errInputChanClosed   = errors.New("channel receiving node scans v2 is closed")
+	errStartMoreThanOnce = errors.New("unable to start the component more than once")
 )
 
 type nodeScanHandlerImpl struct {
 	nodeScans <-chan *storage.NodeScanV2
-	toCentral chan *central.MsgFromSensor
+	toCentral <-chan *central.MsgFromSensor
 
-	stopC    concurrency.ErrorSignal
+	// lock prevents the race condition between Start() [writer] and ResponsesC() [reader]
+	lock *sync.Mutex
+	// stopC is a command that tells this component to stop
+	stopC concurrency.ErrorSignal
+	// stoppedC is signaled when the goroutine inside of run() finishes
 	stoppedC concurrency.ErrorSignal
 }
 
+func (c *nodeScanHandlerImpl) Stopped() concurrency.ReadOnlyErrorSignal {
+	return &c.stoppedC
+}
+
 func (c *nodeScanHandlerImpl) Capabilities() []centralsensor.SensorCapability {
 	return []centralsensor.SensorCapability{centralsensor.NodeScanningCap}
 }
 
+// ResponsesC returns a channel with messages to Central. It must be called after Start() for the channel to be not nil
 func (c *nodeScanHandlerImpl) ResponsesC() <-chan *central.MsgFromSensor {
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.toCentral == nil {
+		log.Panic("Start must be called before ResponsesC")
+	}
 	return c.toCentral
 }
 
 func (c *nodeScanHandlerImpl) Start() error {
-	// c.run closes chan toCentral on exit, so we do not want to close twice in case of a restart
-	if !c.stopC.IsDone() {
-		go c.run()
-		return nil
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	if c.toCentral != nil {
+		return errStartMoreThanOnce
 	}
-	return errors.New("stopped handlers cannot be restarted")
+	c.toCentral = c.run()
+	return nil
 }
 
 func (c *nodeScanHandlerImpl) Stop(err error) {
@@ -43,40 +64,46 @@ func (c *nodeScanHandlerImpl) ProcessMessage(_ *central.MsgToSensor) error {
 	return nil
 }
 
-func (c *nodeScanHandlerImpl) run() {
-	defer c.stoppedC.SignalWithError(c.stopC.Err())
-	defer close(c.toCentral)
-
-	for !c.stopC.IsDone() {
-		select {
-		case <-c.stopC.Done():
-			return
-		case scan, ok := <-c.nodeScans:
-			if !ok {
-				c.stopC.SignalWithError(errors.New("channel receiving node scans v2 is closed"))
+// run handles the messages from Compliance and forwards them to Central
+// This is the only goroutine that writes into the toCentral channel, thus it is responsible for creating and closing that chan
+func (c *nodeScanHandlerImpl) run() <-chan *central.MsgFromSensor {
+	toC := make(chan *central.MsgFromSensor)
+	go func() {
+		defer func() {
+			c.stoppedC.SignalWithError(c.stopC.Err())
+		}()
+		defer close(toC)
+		for !c.stopC.IsDone() {
+			select {
+			case <-c.stopC.Done():
 				return
+			case scan, ok := <-c.nodeScans:
+				if !ok {
+					c.stopC.SignalWithError(errInputChanClosed)
+					return
+				}
+				// TODO(ROX-12943): Do something with the scan, e.g., attach NodeID
+				c.sendScan(toC, scan)
 			}
-			// TODO(ROX-12943): Do something with the scan, e.g., attach NodeID
-			c.sendScan(scan)
 		}
-	}
+	}()
+	return toC
 }
 
-func (c *nodeScanHandlerImpl) sendScan(scan *storage.NodeScanV2) {
-	if scan != nil {
-		select {
-		case <-c.stopC.Done():
-			return
-		case c.toCentral <- &central.MsgFromSensor{
-			Msg: &central.MsgFromSensor_Event{
-				Event: &central.SensorEvent{
-					Resource: &central.SensorEvent_NodeScanV2{
-						NodeScanV2: scan,
-					},
+func (c *nodeScanHandlerImpl) sendScan(toC chan *central.MsgFromSensor, scan *storage.NodeScanV2) {
+	if scan == nil {
+		return
+	}
+	select {
+	case <-c.stopC.Done():
+	case toC <- &central.MsgFromSensor{
+		Msg: &central.MsgFromSensor_Event{
+			Event: &central.SensorEvent{
+				Resource: &central.SensorEvent_NodeScanV2{
+					NodeScanV2: scan,
 				},
 			},
-		}:
-			return
-		}
+		},
+	}:
 	}
 }