feat: Clean up pcap handles and allow alternative sources (#103)

## Which problem is this PR solving?
The assembler can only currently use the gopacket.pcaphandle as it's
packet source. This PR updates the assembler's New function to break the
direct dependency on pcap and allow config to select the packet source.
This does not add additional sources.

## Short description of the changes
- Add source (defualt "pcap") and bpf filter (default "tcp") options to
config
- Move pcap handle creation to separate function and set bpf filter
- Update the packet source lazy option to use config instead of direct
value
- Move started log to Start func

## How to verify that this has the expected result
The agent should continue to work as it is.

Co-authored-by: Jamie Danielson <jamieedanielson@gmail.com>

assemblers/config.go

@@ -28,6 +28,8 @@ var fname = flag.String("r", "", "Filename to read from, overrides -i")
 var snaplen = flag.Int("s", 65536, "Snap length (number of bytes max to read per packet")
 var tstype = flag.String("timestamp_type", "", "Type of timestamps to use")
 var promisc = flag.Bool("promisc", true, "Set promiscuous mode")
+var packetSource = flag.String("source", "pcap", "Packet source (defaults to pcap)")
+var bpfFilter = flag.String("filter", "tcp", "BPF filter")
 
 type config struct {
 	Maxcount         int
@@ -48,6 +50,8 @@ type config struct {
 	Promiscuous      bool
 	CloseTimeout     time.Duration
 	Timeout          time.Duration
+	packetSource     string
+	bpfFilter        string
 }
 
 func NewConfig() *config {
@@ -69,6 +73,8 @@ func NewConfig() *config {
 		TsType:           *tstype,
 		Promiscuous:      *promisc,
 		Timeout:          timeout,
+		packetSource:     *packetSource,
+		bpfFilter:        *bpfFilter,
 	}
 
 	if c.Debug {

assemblers/tcp_assembler.go

@@ -1,8 +1,6 @@
 package assemblers
 
 import (
-	"flag"
-	"strings"
 	"time"
 
 	"github.com/google/gopacket"
@@ -42,7 +40,6 @@ func (c *Context) GetCaptureInfo() gopacket.CaptureInfo {
 
 type tcpAssembler struct {
 	config        *config
-	handle        *pcap.Handle
 	packetSource  *gopacket.PacketSource
 	streamFactory *tcpStreamFactory
 	streamPool    *reassembly.StreamPool
@@ -51,51 +48,29 @@ type tcpAssembler struct {
 }
 
 func NewTcpAssembler(config config, httpEvents chan HttpEvent) tcpAssembler {
-	var handle *pcap.Handle
+	var packetSource *gopacket.PacketSource
 	var err error
 
-	// Set up pcap packet capture
-	if *fname != "" {
-		log.Info().
-			Str("filename", *fname).
-			Msg("Reading from pcap dump")
-		handle, err = pcap.OpenOffline(*fname)
-	} else {
-		log.Info().
-			Str("interface", *iface).
-			Msg("Starting capture")
-		handle, err = pcap.OpenLive(*iface, int32(*snaplen), true, pcap.BlockForever)
-	}
-	if err != nil {
-		log.Fatal().
-			Err(err).
-			Msg("Failed to open a pcap handle")
-	}
-	if len(flag.Args()) > 0 {
-		bpffilter := strings.Join(flag.Args(), " ")
-		log.Info().
-			Str("bpf_filter", bpffilter).
-			Msg("Using BPF filter")
-		if err = handle.SetBPFFilter(bpffilter); err != nil {
-			log.Fatal().
-				Err(err).
-				Str("bpf_filter", bpffilter).
-				Msg("BPF filter error")
+	switch config.packetSource {
+	case "pcap":
+		packetSource, err = newPcapPacketSource(config)
+		if err != nil {
+			log.Fatal().Err(err).Msg("Failed to setup pcap handle")
 		}
+	// TODO: other data sources (eg afpacket, pfring, etc)
+	default:
+		log.Fatal().Str("packet_source", config.packetSource).Msg("Unknown packet source")
 	}
 
-	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
-	packetSource.Lazy = *lazy
+	packetSource.Lazy = config.Lazy
 	packetSource.NoCopy = true
-	log.Info().Msg("Starting to read packets")
 
 	streamFactory := NewTcpStreamFactory(httpEvents)
 	streamPool := reassembly.NewStreamPool(&streamFactory)
 	assembler := reassembly.NewAssembler(streamPool)
 
 	return tcpAssembler{
 		config:        &config,
-		handle:        handle,
 		packetSource:  packetSource,
 		streamFactory: &streamFactory,
 		streamPool:    streamPool,
@@ -180,17 +155,14 @@ func (h *tcpAssembler) Start() {
 
 func (h *tcpAssembler) Stop() {
 	closed := h.assembler.FlushAll()
-	// Debug("Final flush: %d closed", closed)
-	log.Debug().
-		Int("closed", closed).
-		Msg("Final flush")
 	if zerolog.GlobalLevel() >= zerolog.DebugLevel {
 		// this uses stdlib's log, but oh well
 		h.streamPool.Dump()
 	}
 
 	h.streamFactory.WaitGoRoutines()
 	log.Debug().
+		Int("closed", closed).
 		Str("assember_page_usage", h.assembler.Dump()).
 		Int("IPdefrag", stats.ipdefrag).
 		Int("missed_bytes", stats.missedBytes).
@@ -207,5 +179,34 @@ func (h *tcpAssembler) Stop() {
 		Int("biggest_chunk_bytes", stats.biggestChunkBytes).
 		Int("overlap_packets", stats.overlapPackets).
 		Int("overlap_bytes", stats.overlapBytes).
-		Msg("Stop")
+		Msg("Stopping TCP assembler")
+}
+
+func newPcapPacketSource(config config) (*gopacket.PacketSource, error) {
+	log.Info().
+		Str("interface", config.Interface).
+		Int("snaplen", config.Snaplen).
+		Bool("promiscuous", config.Promiscuous).
+		Str("bpf_filter", config.bpfFilter).
+		Msg("Configuring pcap packet source")
+	handle, err := pcap.OpenLive(config.Interface, int32(config.Snaplen), config.Promiscuous, time.Second)
+	if err != nil {
+		log.Fatal().
+			Err(err).
+			Msg("Failed to open a pcap handle")
+		return nil, err
+	}
+	if config.bpfFilter != "" {
+		if err = handle.SetBPFFilter(config.bpfFilter); err != nil {
+			log.Fatal().
+				Err(err).
+				Msg("Error setting BPF filter")
+			return nil, err
+		}
+	}
+
+	return gopacket.NewPacketSource(
+		handle,
+		handle.LinkType(),
+	), nil
 }