Scan engine refactoring

command/arp.go

@@ -60,11 +60,14 @@ var arpCmd = &cobra.Command{
 
 		m := newARPScanMethod(ctx)
 
-		return startEngine(ctx, newEngineConfig(
-			withLogger(logger),
-			withScanRange(r),
-			withScanMethod(m),
-			withBPFFilter(arp.BPFFilter)))
+		return startPacketScanEngine(ctx, newPacketScanConfig(
+			withPacketScanMethod(m),
+			withPacketBPFFilter(arp.BPFFilter),
+			withPacketEngineConfig(newEngineConfig(
+				withLogger(logger),
+				withScanRange(r),
+			)),
+		))
 	},
 }
 

command/root.go

@@ -306,8 +306,6 @@ type bpfFilterFunc func(r *scan.Range) (filter string, maxPacketLength int)
 type engineConfig struct {
 	logger     log.Logger
 	scanRange  *scan.Range
-	scanMethod resultScanMethod
-	bpfFilter  bpfFilterFunc
 	rateCount  int
 	rateWindow time.Duration
 	exitDelay  time.Duration
@@ -327,18 +325,6 @@ func withScanRange(r *scan.Range) engineConfigOption {
 	}
 }
 
-func withScanMethod(sm resultScanMethod) engineConfigOption {
-	return func(c *engineConfig) {
-		c.scanMethod = sm
-	}
-}
-
-func withBPFFilter(bpfFilter bpfFilterFunc) engineConfigOption {
-	return func(c *engineConfig) {
-		c.bpfFilter = bpfFilter
-	}
-}
-
 func newEngineConfig(opts ...engineConfigOption) *engineConfig {
 	c := &engineConfig{
 		rateCount:  cliRateCount,
@@ -351,18 +337,42 @@ func newEngineConfig(opts ...engineConfigOption) *engineConfig {
 	return c
 }
 
-type resultScanMethod interface {
-	scan.Method
-	Results() <-chan scan.Result
+type packetScanConfig struct {
+	*engineConfig
+	scanMethod scan.PacketMethod
+	bpfFilter  bpfFilterFunc
 }
 
-func startEngine(ctx context.Context, conf *engineConfig) error {
-	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
+type packetScanConfigOption func(c *packetScanConfig)
+
+func withPacketEngineConfig(conf *engineConfig) packetScanConfigOption {
+	return func(c *packetScanConfig) {
+		c.engineConfig = conf
+	}
+}
+
+func withPacketScanMethod(sm scan.PacketMethod) packetScanConfigOption {
+	return func(c *packetScanConfig) {
+		c.scanMethod = sm
+	}
+}
+
+func withPacketBPFFilter(bpfFilter bpfFilterFunc) packetScanConfigOption {
+	return func(c *packetScanConfig) {
+		c.bpfFilter = bpfFilter
+	}
+}
 
+func newPacketScanConfig(opts ...packetScanConfigOption) *packetScanConfig {
+	c := &packetScanConfig{}
+	for _, o := range opts {
+		o(c)
+	}
+	return c
+}
+
+func startPacketScanEngine(ctx context.Context, conf *packetScanConfig) error {
 	r := conf.scanRange
-	m := conf.scanMethod
-	logger := conf.logger
 
 	// setup network interface to read/write packets
 	afps, err := afpacket.NewPacketSource(r.Interface.Name)
@@ -380,18 +390,26 @@ func startEngine(ctx context.Context, conf *engineConfig) error {
 		rw = packet.NewRateLimitReadWriter(afps,
 			ratelimit.New(conf.rateCount, ratelimit.Per(conf.rateWindow)))
 	}
+	engine := scan.SetupPacketEngine(rw, conf.scanMethod)
+	return startScanEngine(ctx, engine, conf.engineConfig)
+}
+
+func startScanEngine(ctx context.Context, engine scan.EngineResulter, conf *engineConfig) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	logger := conf.logger
 
 	// setup result logging
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		logger.LogResults(ctx, m.Results())
+		logger.LogResults(ctx, engine.Results())
 	}()
 
 	// start scan
-	engine := scan.SetupEngine(rw, m)
-	done, errc := engine.Start(ctx, r)
+	done, errc := engine.Start(ctx, conf.scanRange)
 	go func() {
 		defer cancel()
 		<-done

command/tcp.go

@@ -88,11 +88,14 @@ var tcpCmd = &cobra.Command{
 			withTCPPacketFlags(tcp.AllFlags),
 		)
 
-		return startEngine(ctx, newEngineConfig(
-			withLogger(conf.logger),
-			withScanRange(conf.scanRange),
-			withScanMethod(m),
-			withBPFFilter(tcp.BPFFilter)))
+		return startPacketScanEngine(ctx, newPacketScanConfig(
+			withPacketScanMethod(m),
+			withPacketBPFFilter(tcp.BPFFilter),
+			withPacketEngineConfig(newEngineConfig(
+				withLogger(conf.logger),
+				withScanRange(conf.scanRange),
+			)),
+		))
 	},
 }
 

command/tcp_fin.go

@@ -43,10 +43,13 @@ var tcpfinCmd = &cobra.Command{
 			withTCPPacketFlags(tcp.AllFlags),
 		)
 
-		return startEngine(ctx, newEngineConfig(
-			withLogger(conf.logger),
-			withScanRange(conf.scanRange),
-			withScanMethod(m),
-			withBPFFilter(tcp.BPFFilter)))
+		return startPacketScanEngine(ctx, newPacketScanConfig(
+			withPacketScanMethod(m),
+			withPacketBPFFilter(tcp.BPFFilter),
+			withPacketEngineConfig(newEngineConfig(
+				withLogger(conf.logger),
+				withScanRange(conf.scanRange),
+			)),
+		))
 	},
 }

command/tcp_null.go

@@ -43,10 +43,13 @@ var tcpnullCmd = &cobra.Command{
 			withTCPPacketFlags(tcp.AllFlags),
 		)
 
-		return startEngine(ctx, newEngineConfig(
-			withLogger(conf.logger),
-			withScanRange(conf.scanRange),
-			withScanMethod(m),
-			withBPFFilter(tcp.BPFFilter)))
+		return startPacketScanEngine(ctx, newPacketScanConfig(
+			withPacketScanMethod(m),
+			withPacketBPFFilter(tcp.BPFFilter),
+			withPacketEngineConfig(newEngineConfig(
+				withLogger(conf.logger),
+				withScanRange(conf.scanRange),
+			)),
+		))
 	},
 }

command/tcp_syn.go

@@ -51,10 +51,13 @@ func startTCPSYNScan(ctx context.Context, subnet string) (err error) {
 		withTCPPacketFlags(tcp.EmptyFlags),
 	)
 
-	return startEngine(ctx, newEngineConfig(
-		withLogger(conf.logger),
-		withScanRange(conf.scanRange),
-		withScanMethod(m),
+	return startPacketScanEngine(ctx, newPacketScanConfig(
+		withPacketScanMethod(m),
 		// TODO SYN,ACK filter
-		withBPFFilter(tcp.BPFFilter)))
+		withPacketBPFFilter(tcp.BPFFilter),
+		withPacketEngineConfig(newEngineConfig(
+			withLogger(conf.logger),
+			withScanRange(conf.scanRange),
+		)),
+	))
 }

command/tcp_xmas.go

@@ -43,10 +43,13 @@ var tcpxmasCmd = &cobra.Command{
 			withTCPPacketFlags(tcp.AllFlags),
 		)
 
-		return startEngine(ctx, newEngineConfig(
-			withLogger(conf.logger),
-			withScanRange(conf.scanRange),
-			withScanMethod(m),
-			withBPFFilter(tcp.BPFFilter)))
+		return startPacketScanEngine(ctx, newPacketScanConfig(
+			withPacketScanMethod(m),
+			withPacketBPFFilter(tcp.BPFFilter),
+			withPacketEngineConfig(newEngineConfig(
+				withLogger(conf.logger),
+				withScanRange(conf.scanRange),
+			)),
+		))
 	},
 }

command/udp.go

@@ -43,11 +43,14 @@ var udpCmd = &cobra.Command{
 
 		m := newUDPScanMethod(ctx, conf)
 
-		return startEngine(ctx, newEngineConfig(
-			withLogger(conf.logger),
-			withScanRange(conf.scanRange),
-			withScanMethod(m),
-			withBPFFilter(icmp.BPFFilter)))
+		return startPacketScanEngine(ctx, newPacketScanConfig(
+			withPacketScanMethod(m),
+			withPacketBPFFilter(icmp.BPFFilter),
+			withPacketEngineConfig(newEngineConfig(
+				withLogger(conf.logger),
+				withScanRange(conf.scanRange),
+			)),
+		))
 	},
 }
 

pkg/scan/arp/arp.go

@@ -24,7 +24,7 @@ type ScanMethod struct {
 }
 
 // Assert that arp.ScanMethod conforms to the scan.Method interface
-var _ scan.Method = (*ScanMethod)(nil)
+var _ scan.PacketMethod = (*ScanMethod)(nil)
 
 //easyjson:json
 type ScanResult struct {

pkg/scan/engine.go

@@ -1,4 +1,4 @@
-//go:generate mockgen -package scan -destination=mock_engine_test.go -source=engine.go
+//go:generate mockgen -package scan -destination=mock_engine_test.go . PacketSource
 
 package scan
 
@@ -21,11 +21,31 @@ type Range struct {
 	SrcSubnet *net.IPNet
 	SrcIP     net.IP
 	SrcMAC    net.HardwareAddr
-	StartPort uint16
-	EndPort   uint16
 	Ports     []*PortRange
 }
 
+type Engine interface {
+	Start(ctx context.Context, r *Range) (done <-chan interface{}, errc <-chan error)
+}
+
+type Resulter interface {
+	Results() <-chan Result
+}
+
+type EngineResulter interface {
+	Engine
+	Resulter
+}
+
+type engineResulter struct {
+	Engine
+	Resulter
+}
+
+func NewEngineResulter(e Engine, r Resulter) EngineResulter {
+	return &engineResulter{Engine: e, Resulter: r}
+}
+
 type PacketSource interface {
 	Packets(ctx context.Context, r *Range) <-chan *packet.BufferData
 }
@@ -50,17 +70,17 @@ func (s *packetSource) Packets(ctx context.Context, r *Range) <-chan *packet.Buf
 	return s.pktgen.Packets(ctx, requests)
 }
 
-type Engine struct {
+type PacketEngine struct {
 	src PacketSource
-	rcv packet.Receiver
 	snd packet.Sender
+	rcv packet.Receiver
 }
 
-func NewEngine(ps PacketSource, s packet.Sender, r packet.Receiver) *Engine {
-	return &Engine{src: ps, snd: s, rcv: r}
+func NewPacketEngine(ps PacketSource, s packet.Sender, r packet.Receiver) *PacketEngine {
+	return &PacketEngine{src: ps, snd: s, rcv: r}
 }
 
-func (e *Engine) Start(ctx context.Context, r *Range) (<-chan interface{}, <-chan error) {
+func (e *PacketEngine) Start(ctx context.Context, r *Range) (<-chan interface{}, <-chan error) {
 	packets := e.src.Packets(ctx, r)
 	done, errc1 := e.snd.SendPackets(ctx, packets)
 	errc2 := e.rcv.ReceivePackets(ctx)
@@ -101,14 +121,15 @@ func mergeErrChan(ctx context.Context, channels ...<-chan error) <-chan error {
 	return out
 }
 
-type Method interface {
+type PacketMethod interface {
 	PacketSource
 	packet.Processor
+	Resulter
 }
 
-func SetupEngine(rw packet.ReadWriter, m Method) *Engine {
+func SetupPacketEngine(rw packet.ReadWriter, m PacketMethod) EngineResulter {
 	sender := packet.NewSender(rw)
 	receiver := packet.NewReceiver(rw, m)
-	engine := NewEngine(m, sender, receiver)
-	return engine
+	engine := NewPacketEngine(m, sender, receiver)
+	return NewEngineResulter(engine, m)
 }

pkg/scan/engine_test.go

@@ -78,7 +78,7 @@ func TestMergeErrChanContextExit(t *testing.T) {
 	assert.Equal(t, 0, len(result), "error slice is not empty")
 }
 
-func TestEngineStartCollectsAllErrors(t *testing.T) {
+func TestPacketEngineStartCollectsAllErrors(t *testing.T) {
 	t.Parallel()
 
 	ctrl := gomock.NewController(t)
@@ -99,7 +99,7 @@ func TestEngineStartCollectsAllErrors(t *testing.T) {
 	close(errc2)
 	r := NewMockReceiver(ctrl)
 	r.EXPECT().ReceivePackets(notNil).Return(errc2)
-	e := NewEngine(ps, s, r)
+	e := NewPacketEngine(ps, s, r)
 
 	_, out := e.Start(context.Background(), &Range{
 		DstSubnet: &net.IPNet{