Improve NgWriter interface handling

Use CaptureInfoToID callback in NgWriterOptions to map captures to
specific interfaces when multiple hosts and network namespaces are
involved, to avoid conflicts caused by interface indexes often being the
same across multiple hosts.

Signed-off-by: Chance Zibolski <chance.zibolski@gmail.com>

cmd/output_handler.go

@@ -56,7 +56,7 @@ type outputFileHandler struct {
 	// Unfortunately, since pcapng didn't consider multi-host captures, this
 	// means the same interface index across different hosts may clash.
 	// TODO: Update NgWriter to key by more than the interface index.
-	interfaceConfigured map[int]struct{}
+	interfaceConfigured map[capture.CaptureInterface]struct{}
 }
 
 func newOutputFileHandler(outputPath string, isDir bool, linkType layers.LinkType, snaplen uint32, outputFormat capperpb.PcapOutputFormat) *outputFileHandler {
@@ -67,12 +67,12 @@ func newOutputFileHandler(outputPath string, isDir bool, linkType layers.LinkTyp
 		linkType:            linkType,
 		snaplen:             snaplen,
 		outputFormat:        outputFormat,
-		interfaceConfigured: make(map[int]struct{}),
+		interfaceConfigured: make(map[capture.CaptureInterface]struct{}),
 	}
 }
 
 func (h *outputFileHandler) HandlePacket(p gopacket.Packet) error {
-	ad, err := getCapperAncillaryData(p)
+	ad, err := capture.GetCapperAncillaryData(p.Metadata().CaptureInfo)
 	if err != nil {
 		return fmt.Errorf("error getting packet ancillary data: %w", err)
 	}
@@ -91,16 +91,25 @@ func (h *outputFileHandler) HandlePacket(p gopacket.Packet) error {
 		h.writers[identifier] = packetWriter
 	} else {
 		// We already have a writer, check if we need to update it.
-		// If we're outputting to a directory, then we have a writer for each interface, so there's no need to update it,
-		// And we only need to add interfaces for pcapng format.
-		if !h.isDir && h.outputFormat == capperpb.PcapOutputFormat_OUTPUT_FORMAT_PCAPNG {
-			if _, configured := h.interfaceConfigured[p.Metadata().InterfaceIndex]; !configured {
+		// We only need to add interfaces for pcapng format.
+		// We don't do this when the writer already exists because this is handled
+		// for the first interface automatically as part of capture.NewPcapNgWriter.
+		if h.outputFormat == capperpb.PcapOutputFormat_OUTPUT_FORMAT_PCAPNG {
+			captureIface := capture.CaptureInterface{
+				Name:       ad.IfaceName,
+				Index:      uint64(p.Metadata().InterfaceIndex),
+				Hostname:   ad.NodeName,
+				Netns:      ad.Netns,
+				NetnsInode: ad.NetnsInode,
+				LinkType:   layers.LinkType(ad.LinkType),
+			}
+			if _, configured := h.interfaceConfigured[captureIface]; !configured {
 				ngWriter := packetWriter.(*capture.PcapNgWriter)
-				_, err := ngWriter.AddInterface(ad.IfaceName, p.Metadata().InterfaceIndex, layers.LinkType(ad.LinkType))
+				_, err := ngWriter.AddInterface(captureIface)
 				if err != nil {
 					return err
 				}
-				h.interfaceConfigured[p.Metadata().InterfaceIndex] = struct{}{}
+				h.interfaceConfigured[captureIface] = struct{}{}
 			}
 		}
 	}
@@ -131,16 +140,19 @@ func (h *outputFileHandler) newPacketWriter(identifier string, interfaceIndex in
 	switch h.outputFormat {
 	case capperpb.PcapOutputFormat_OUTPUT_FORMAT_PCAPNG:
 		var err error
-		packetWriter, err = capture.NewPcapNgWriter(
-			w, h.linkType, h.snaplen,
-			ad.GetIfaceName(), interfaceIndex,
-			ad.GetHardware(), ad.GetOperatingSystem(),
-			ad.GetNodeName(),
-		)
+		captureIface := capture.CaptureInterface{
+			Name:       ad.IfaceName,
+			Index:      uint64(interfaceIndex),
+			Hostname:   ad.NodeName,
+			Netns:      ad.Netns,
+			NetnsInode: ad.NetnsInode,
+			LinkType:   layers.LinkType(ad.LinkType),
+		}
+		packetWriter, err = capture.NewPcapNgWriter(w, captureIface, h.snaplen, ad.GetHardware(), ad.GetOperatingSystem())
 		if err != nil {
 			return nil, err
 		}
-		h.interfaceConfigured[interfaceIndex] = struct{}{}
+		h.interfaceConfigured[captureIface] = struct{}{}
 	case capperpb.PcapOutputFormat_OUTPUT_FORMAT_PCAP:
 		fallthrough
 	default:
@@ -159,22 +171,3 @@ func (h *outputFileHandler) Flush() error {
 	}
 	return err
 }
-
-func getCapperAncillaryData(p gopacket.Packet) (*capperpb.AncillaryPacketData, error) {
-	pancillaryData := p.Metadata().AncillaryData
-	if len(pancillaryData) == 0 {
-		return nil, fmt.Errorf("no gopacket AncillaryData")
-	}
-	var ancillaryData *capperpb.AncillaryPacketData
-	for _, ad := range pancillaryData {
-		var ok bool
-		ancillaryData, ok = ad.(*capperpb.AncillaryPacketData)
-		if ok {
-			break
-		}
-	}
-	if ancillaryData == nil {
-		return nil, fmt.Errorf("no capper AncillaryPacketData found in gopacket AncillaryData")
-	}
-	return ancillaryData, nil
-}

cmd/server.go

@@ -376,7 +376,7 @@ func (s *server) capture(ctx context.Context, ifaces []string, netns string, con
 // bytes to the given Capper_CaptureServer stream.
 func newStreamPacketHandler(linkType layers.LinkType, snaplen uint32, stream capperpb.Capper_CaptureServer) (capture.PacketHandler, error) {
 	streamHandler := capture.PacketHandlerFunc(func(p gopacket.Packet) error {
-		ad, err := getCapperAncillaryData(p)
+		ad, err := capture.GetCapperAncillaryData(p.Metadata().CaptureInfo)
 		if err != nil {
 			return fmt.Errorf("error getting packet AncillaryData: %w", err)
 		}

go.mod

@@ -73,4 +73,4 @@ require (
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240227224415-6ceb2ff114de // indirect
 )
 
-replace github.com/gopacket/gopacket => github.com/chancez/gopacket v1.2.1-0.20240719041229-5f4144d9b475
+replace github.com/gopacket/gopacket => github.com/chancez/gopacket v1.2.1-0.20240722032546-3e0a9b78ef2d

pkg/capture/ancillary_data.go

@@ -0,0 +1,27 @@
+package capture
+
+import (
+	"fmt"
+
+	capperpb "github.com/chancez/capper/proto/capper"
+	"github.com/gopacket/gopacket"
+)
+
+func GetCapperAncillaryData(ci gopacket.CaptureInfo) (*capperpb.AncillaryPacketData, error) {
+	pancillaryData := ci.AncillaryData
+	if len(pancillaryData) == 0 {
+		return nil, fmt.Errorf("no gopacket AncillaryData")
+	}
+	var ancillaryData *capperpb.AncillaryPacketData
+	for _, ad := range pancillaryData {
+		var ok bool
+		ancillaryData, ok = ad.(*capperpb.AncillaryPacketData)
+		if ok {
+			break
+		}
+	}
+	if ancillaryData == nil {
+		return nil, fmt.Errorf("no capper AncillaryPacketData found in gopacket AncillaryData")
+	}
+	return ancillaryData, nil
+}

pkg/capture/capture.go

@@ -98,6 +98,7 @@ type CaptureInterface struct {
 	Hostname   string
 	NetnsInode uint64
 	Netns      string
+	LinkType   layers.LinkType
 }
 
 func getInterface(ifaceName string, netns string) (CaptureInterface, error) {
@@ -185,6 +186,8 @@ func NewBasic(ctx context.Context, log *slog.Logger, ifaceName, netns string, co
 		return nil, fmt.Errorf("error creating handle: %w", err)
 	}
 
+	iface.LinkType = handle.LinkType()
+
 	return &BasicCapture{
 		log:    log,
 		clock:  clock,
@@ -232,7 +235,7 @@ func (c *BasicCapture) Start(ctx context.Context, handler PacketHandler) error {
 	packetSource := gopacket.NewPacketSource(c.handle, c.handle.LinkType())
 	for packet := range packetSource.PacketsCtx(packetsCtx) {
 		packet.Metadata().AncillaryData = append(packet.Metadata().AncillaryData, &capperpb.AncillaryPacketData{
-			LinkType:        int64(c.handle.LinkType()),
+			LinkType:        int64(c.iface.LinkType),
 			NodeName:        c.iface.Hostname,
 			Netns:           c.iface.Netns,
 			NetnsInode:      c.iface.NetnsInode,

pkg/capture/writer.go

@@ -44,38 +44,56 @@ func (w *PcapWriter) Flush() error {
 
 type PcapNgWriter struct {
 	ngWriter *pcapgo.NgWriter
-	linkType layers.LinkType
 	snaplen  uint32
 	os       string
-	hostname string
+
+	interfaceToID map[CaptureInterface]int
 }
 
-func NewPcapNgWriter(w io.Writer, linkType layers.LinkType, snaplen uint32, ifaceName string, ifaceIndex int, hardware, os, hostname string) (*PcapNgWriter, error) {
+func NewPcapNgWriter(w io.Writer, iface CaptureInterface, snaplen uint32, hardware, os string) (*PcapNgWriter, error) {
+	interfaceToID := make(map[CaptureInterface]int)
 	intf := pcapgo.NgInterface{
-		Name:        ifaceName,
-		Index:       ifaceIndex,
-		LinkType:    linkType,
+		Name:        iface.Name,
+		Index:       int(iface.Index),
+		LinkType:    iface.LinkType,
 		SnapLength:  snaplen,
 		OS:          os,
-		Description: fmt.Sprintf("iface: %s hostname: %q", ifaceName, hostname),
+		Description: fmt.Sprintf("iface: %s hostname: %q", iface.Name, iface.Hostname),
 	}
+	interfaceToID[iface] = 0
+
 	ngOpts := pcapgo.NgWriterOptions{
 		SectionInfo: pcapgo.NgSectionInfo{
 			Hardware:    hardware,
 			OS:          os,
 			Application: "capper",
 		},
+		CaptureInfoToID: func(ci gopacket.CaptureInfo, data []byte) (id int, ok bool) {
+			ad, err := GetCapperAncillaryData(ci)
+			if err != nil {
+				return -1, false
+			}
+			iface := CaptureInterface{
+				Name:       ad.IfaceName,
+				Index:      uint64(ci.InterfaceIndex),
+				Hostname:   ad.NodeName,
+				Netns:      ad.Netns,
+				NetnsInode: ad.NetnsInode,
+				LinkType:   layers.LinkType(ad.LinkType),
+			}
+			id, ok = interfaceToID[iface]
+			return id, ok
+		},
 	}
 	ngWriter, err := pcapgo.NewNgWriterInterface(w, intf, ngOpts)
 	if err != nil {
 		return nil, err
 	}
 	return &PcapNgWriter{
-		ngWriter: ngWriter,
-		linkType: linkType,
-		snaplen:  snaplen,
-		os:       os,
-		hostname: hostname,
+		ngWriter:      ngWriter,
+		snaplen:       snaplen,
+		os:            os,
+		interfaceToID: interfaceToID,
 	}, nil
 }
 
@@ -87,13 +105,17 @@ func (w *PcapNgWriter) Flush() error {
 	return w.ngWriter.Flush()
 }
 
-func (w *PcapNgWriter) AddInterface(name string, index int, linkType layers.LinkType) (int, error) {
-	return w.ngWriter.AddInterface(pcapgo.NgInterface{
-		Name:        name,
-		Index:       index,
-		LinkType:    linkType,
+func (w *PcapNgWriter) AddInterface(iface CaptureInterface) (int, error) {
+	id, err := w.ngWriter.AddInterface(pcapgo.NgInterface{
+		Name:        iface.Name,
+		Index:       int(iface.Index),
+		LinkType:    iface.LinkType,
 		SnapLength:  w.snaplen,
 		OS:          w.os,
-		Description: fmt.Sprintf("iface: %s hostname: %q", name, w.hostname),
+		Description: fmt.Sprintf("iface: %s hostname: %q", iface.Name, iface.Hostname),
 	})
+	if err == nil {
+		w.interfaceToID[iface] = id
+	}
+	return id, err
 }

vendor/modules.txt

@@ -182,7 +182,7 @@ github.com/google/go-cmp/cmp/internal/value
 # github.com/google/uuid v1.6.0
 ## explicit
 github.com/google/uuid
-# github.com/gopacket/gopacket v1.2.0 => github.com/chancez/gopacket v1.2.1-0.20240719041229-5f4144d9b475
+# github.com/gopacket/gopacket v1.2.0 => github.com/chancez/gopacket v1.2.1-0.20240722032546-3e0a9b78ef2d
 ## explicit; go 1.20
 github.com/gopacket/gopacket
 github.com/gopacket/gopacket/layers
@@ -469,4 +469,4 @@ google.golang.org/protobuf/types/known/durationpb
 google.golang.org/protobuf/types/known/emptypb
 google.golang.org/protobuf/types/known/fieldmaskpb
 google.golang.org/protobuf/types/known/timestamppb
-# github.com/gopacket/gopacket => github.com/chancez/gopacket v1.2.1-0.20240719041229-5f4144d9b475
+# github.com/gopacket/gopacket => github.com/chancez/gopacket v1.2.1-0.20240722032546-3e0a9b78ef2d