add -i flag

bpf/bpf.go

@@ -3,7 +3,6 @@ package bpf
 import (
 	"encoding/binary"
 	"log"
-	"net"
 	"unsafe"
 
 	"github.com/cilium/ebpf"
@@ -83,22 +82,22 @@ func (b *BPF) AttachTracepoints() error {
 	return nil
 }
 
-func (b *BPF) AttachTcHooks(dev *net.Interface) error {
-	closeFunc, err := ensureTcQdisc(dev)
+func (b *BPF) AttachTcHooks(ifindex int) error {
+	closeFunc, err := ensureTcQdisc(ifindex)
 	if err != nil {
 		if closeFunc != nil {
 			closeFunc()
 		}
 		return xerrors.Errorf("attach tc hooks: %w", err)
 	}
 
-	c1, err := attachTcHook(dev, b.objs.TcEgress, false)
+	c1, err := attachTcHook(ifindex, b.objs.TcEgress, false)
 	if err != nil {
 		closeFunc()
 		return xerrors.Errorf("attach tc hooks: %w", err)
 	}
 
-	c2, err := attachTcHook(dev, b.objs.TcIngress, true)
+	c2, err := attachTcHook(ifindex, b.objs.TcIngress, true)
 	if err != nil {
 		c1()
 		closeFunc()
@@ -125,7 +124,7 @@ func (b *BPF) NewExecEventReader() (*ringbuf.Reader, error) {
 	return reader, nil
 }
 
-func attachTcHook(dev *net.Interface, prog *ebpf.Program, ingress bool) (func(), error) {
+func attachTcHook(ifindex int, prog *ebpf.Program, ingress bool) (func(), error) {
 	tcnl, err := tc.Open(&tc.Config{})
 	if err != nil {
 		return nil, err
@@ -146,7 +145,7 @@ func attachTcHook(dev *net.Interface, prog *ebpf.Program, ingress bool) (func(),
 	filter := tc.Object{
 		tc.Msg{
 			Family:  unix.AF_UNSPEC,
-			Ifindex: uint32(dev.Index),
+			Ifindex: uint32(ifindex),
 			Handle:  0,
 			Parent:  core.BuildHandle(tc.HandleRoot, parent),
 			Info:    1<<16 | uint32(htons(unix.ETH_P_ALL)),
@@ -172,7 +171,7 @@ func attachTcHook(dev *net.Interface, prog *ebpf.Program, ingress bool) (func(),
 	return newCloseFunc, nil
 }
 
-func ensureTcQdisc(dev *net.Interface) (func(), error) {
+func ensureTcQdisc(ifindex int) (func(), error) {
 	tcnl, err := tc.Open(&tc.Config{})
 	if err != nil {
 		return nil, err
@@ -186,7 +185,7 @@ func ensureTcQdisc(dev *net.Interface) (func(), error) {
 	qdisc := tc.Object{
 		Msg: tc.Msg{
 			Family:  unix.AF_UNSPEC,
-			Ifindex: uint32(dev.Index),
+			Ifindex: uint32(ifindex),
 			Handle:  core.BuildHandle(tc.HandleRoot, 0x0000),
 			Parent:  tc.HandleIngress,
 		},

go.mod

@@ -23,6 +23,8 @@ require (
 	github.com/shoenig/go-m1cpu v0.1.6 // indirect
 	github.com/tklauser/go-sysconf v0.3.12 // indirect
 	github.com/tklauser/numcpus v0.6.1 // indirect
+	github.com/vishvananda/netlink v1.1.0 // indirect
+	github.com/vishvananda/netns v0.0.0-20211101163701-50045581ed74 // indirect
 	github.com/yusufpapurcu/wmi v1.2.4 // indirect
 	golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2 // indirect
 	golang.org/x/net v0.21.0 // indirect

go.sum

@@ -100,6 +100,7 @@ github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+F
 github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY=
 github.com/vishvananda/netlink v1.1.0 h1:1iyaYNBLmP6L0220aDnYQpo1QEV4t4hJ+xEEhhJH8j0=
 github.com/vishvananda/netlink v1.1.0/go.mod h1:cTgwzPIzzgDAYoQrMm0EdrjRUBkTqKYppBueQtXaqoE=
+github.com/vishvananda/netns v0.0.0-20191106174202-0a2b9b5464df/go.mod h1:JP3t17pCcGlemwknint6hfoeCVQrEMVwxRLRjXpq+BU=
 github.com/vishvananda/netns v0.0.0-20211101163701-50045581ed74 h1:gga7acRE695APm9hlsSMoOoE65U4/TcqNj90mc69Rlg=
 github.com/vishvananda/netns v0.0.0-20211101163701-50045581ed74/go.mod h1:DD4vA1DwXk04H54A1oHXtwZmA0grkVMdPxx/VGLCah0=
 github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
@@ -130,10 +131,12 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5h
 golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190411185658-b44545bcd369/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20190606203320-7fc4e5ec1444/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190826190057-c7b8b68b1456/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191008105621-543471e840be/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200202164722-d101bd2416d5/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20200217220822-9197077df867/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201009025420-dfb3f7c4e634/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201118182958-a01c418693c7/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

internal/dev/dev.go

@@ -0,0 +1,40 @@
+package dev
+
+import (
+	"github.com/vishvananda/netlink"
+	"golang.org/x/xerrors"
+)
+
+type Device struct {
+	Name    string
+	Ifindex int
+}
+
+func GetDevices(name string) (map[int]Device, error) {
+	var links []netlink.Link
+	var err error
+	ifindexMap := make(map[int]Device)
+
+	if name == "any" {
+		if links, err = netlink.LinkList(); err != nil {
+			return nil, xerrors.Errorf(": %w", err)
+		}
+	} else {
+		link, err := netlink.LinkByName(name)
+		if err != nil {
+			return nil, xerrors.Errorf("get device by name (%s): %w", name, err)
+		}
+		links = append(links, link)
+	}
+
+	for _, link := range links {
+		linkAttrs := link.Attrs()
+		dev := Device{
+			Name:    linkAttrs.Name,
+			Ifindex: linkAttrs.Index,
+		}
+		ifindexMap[dev.Ifindex] = dev
+	}
+
+	return ifindexMap, nil
+}

internal/event/net.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"github.com/mozillazg/ptcpdump/bpf"
+	"github.com/mozillazg/ptcpdump/internal/dev"
 	"golang.org/x/xerrors"
 	"unsafe"
 )
@@ -16,14 +17,15 @@ const (
 )
 
 type Packet struct {
-	Type packetType
-	Pid  int
-	Comm string
+	Type   packetType
+	Device dev.Device
+	Pid    int
+	Comm   string
 
 	Data []byte
 }
 
-func ParsePacketEvent(rawSample []byte) (*Packet, error) {
+func ParsePacketEvent(devices map[int]dev.Device, rawSample []byte) (*Packet, error) {
 	var p Packet
 	event := bpf.BpfPacketEventT{}
 	if err := binary.Read(bytes.NewBuffer(rawSample), binary.LittleEndian, &event.Meta); err != nil {
@@ -36,6 +38,7 @@ func ParsePacketEvent(rawSample []byte) (*Packet, error) {
 	if event.Meta.PacketType == 1 {
 		p.Type = packetTypeEgress
 	}
+	p.Device = devices[int(event.Meta.Ifindex)]
 	p.Data = make([]byte, event.Meta.PayloadLen)
 	copy(p.Data[:], event.Payload[:event.Meta.PayloadLen])
 

internal/writer/pcapng.go

@@ -26,7 +26,7 @@ func (w *PcapNGWriter) Write(e *event.Packet) error {
 		Timestamp:      time.Now(),
 		CaptureLength:  payloadLen,
 		Length:         payloadLen,
-		InterfaceIndex: 0,
+		InterfaceIndex: e.Device.Ifindex,
 	}
 	p := w.pcache.Get(e.Pid)
 	if p.Pid == 0 {

main.go

@@ -4,27 +4,27 @@ import (
 	"context"
 	"errors"
 	"flag"
-	"fmt"
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/perf"
 	"github.com/cilium/ebpf/rlimit"
 	"github.com/gopacket/gopacket/layers"
 	"github.com/gopacket/gopacket/pcapgo"
 	"github.com/mozillazg/ptcpdump/bpf"
+	"github.com/mozillazg/ptcpdump/internal/dev"
 	"github.com/mozillazg/ptcpdump/internal/event"
 	"github.com/mozillazg/ptcpdump/internal/metadata"
 	"github.com/mozillazg/ptcpdump/internal/writer"
 	"golang.org/x/xerrors"
 	"io"
 	"log"
 	"math"
-	"net"
 	"os"
 	"os/signal"
 	"syscall"
 )
 
 type Options struct {
+	iface         string
 	writeFilePath string
 }
 
@@ -68,54 +68,57 @@ func parseExecEvent(pcache *metadata.ProcessCache, rawSample []byte) {
 	pcache.AddItem(*e)
 }
 
-func newPcapWriter(w io.Writer, ifaceNames []string, pcache *metadata.ProcessCache) (*writer.PcapNGWriter, map[string]int, error) {
-	if len(ifaceNames) == 0 {
-		return nil, nil, xerrors.New("can't create pcap with no ifaceNames")
+func newPcapWriter(w io.Writer, devices []dev.Device, pcache *metadata.ProcessCache) (*writer.PcapNGWriter, error) {
+	if len(devices) == 0 {
+		return nil, xerrors.New("can't create pcap with no interface")
 	}
 
 	var interfaces []pcapgo.NgInterface
-	nameIfcs := make(map[string]int)
-
-	for id, name := range ifaceNames {
+	for _, dev := range devices {
 		interfaces = append(interfaces, pcapgo.NgInterface{
-			Name:       fmt.Sprintf("ptcpdump-%s", name),
-			Comment:    "ptcpdump iface name",
+			Name:       dev.Name,
+			Comment:    "ptcpdump interface name",
 			LinkType:   layers.LinkTypeEthernet,
 			SnapLength: uint32(math.MaxUint16),
 		})
-		nameIfcs[name] = id
 	}
 
 	pcapWriter, err := pcapgo.NewNgWriterInterface(w, interfaces[0], pcapgo.NgWriterOptions{})
 	if err != nil {
-		return nil, nil, err
+		return nil, xerrors.Errorf(": %w", err)
 	}
-
 	for _, ifc := range interfaces[1:] {
 		_, err := pcapWriter.AddInterface(ifc)
 		if err != nil {
-			return nil, nil, err
+			return nil, xerrors.Errorf(": %w", err)
 		}
 	}
 
 	// Flush the header out in case we're writing to stdout, this lets tcpdump print a reassuring message
 	if err := pcapWriter.Flush(); err != nil {
-		return nil, nil, xerrors.Errorf("writing pcap header: %w", err)
+		return nil, xerrors.Errorf("writing pcap header: %w", err)
 	}
 
-	return writer.NewPcapNGWriter(pcapWriter, pcache), nameIfcs, nil
+	return writer.NewPcapNGWriter(pcapWriter, pcache), nil
 }
 
 func setupFlags() *Options {
 	opts := &Options{}
 	flag.StringVar(&opts.writeFilePath, "w", "",
 		"Write the raw packets to file rather than parsing and printing them out. e.g. ptcpdump.pcapng")
+	flag.StringVar(&opts.iface, "i", "eth0", "")
 	flag.Parse()
 	return opts
 }
 
 func main() {
 	opts := setupFlags()
+	devices, err := dev.GetDevices(opts.iface)
+	if err != nil {
+		logErr(err)
+		return
+	}
+
 	if err := rlimit.RemoveMemlock(); err != nil {
 		logErr(err)
 		return
@@ -130,7 +133,7 @@ func main() {
 			logErr(err)
 			return
 		}
-		pcapWriter, _, err := newPcapWriter(pcapFile, []string{"lo", "enp0s3", "docker0", "wlp4s0", "enp5s0"}, pcache)
+		pcapWriter, err := newPcapWriter(pcapFile, devices, pcache)
 		if err != nil {
 			logErr(err)
 			return
@@ -165,13 +168,8 @@ func main() {
 		return
 	}
 
-	for _, ifaceName := range []string{"lo", "enp0s3", "docker0", "wlp4s0", "enp5s0"} {
-		dev, err := net.InterfaceByName(ifaceName)
-		if err != nil {
-			log.Printf("get interface by name %s failed: %+v", ifaceName, err)
-			continue
-		}
-		if err := bf.AttachTcHooks(dev); err != nil {
+	for _, dev := range devices {
+		if err := bf.AttachTcHooks(dev.Ifindex); err != nil {
 			logErr(err)
 			return
 		}