chore: Reorganize the codes

.gitignore

@@ -15,3 +15,4 @@
 # vendor/
 test.pcapng
 ptcpdump
+ptcpdump.pcapng

bpf/bpf.go

@@ -49,17 +49,22 @@ func (b *BPF) Load() error {
 
 func (b *BPF) Close() {
 	for _, lk := range b.links {
-		lk.Close()
+		if err := lk.Close(); err != nil {
+			log.Printf("[bpf] close link %v failed: %+v", lk, err)
+		}
 	}
 	for i := len(b.closeFuncs) - 1; i > 0; i-- {
 		f := b.closeFuncs[i]
 		f()
 	}
-	b.objs.Close()
+	if err := b.objs.Close(); err != nil {
+		log.Printf("[bpf] close objects failed: %+v", err)
+	}
 }
 
 func (b *BPF) AttachKprobes() error {
-	lk, err := link.Kprobe("security_sk_classify_flow", b.objs.KprobeSecuritySkClassifyFlow, &link.KprobeOptions{})
+	lk, err := link.Kprobe("security_sk_classify_flow",
+		b.objs.KprobeSecuritySkClassifyFlow, &link.KprobeOptions{})
 	if err != nil {
 		return xerrors.Errorf("attach kprobe/security_sk_classify_flow: %w", err)
 	}

bpf/ptcpdump.c

@@ -1,3 +1,4 @@
+//go:build ignore
 // +build ignore
 
 #include "vmlinux.h"

bpf/tools.go

@@ -0,0 +1,6 @@
+//go:build ignore
+// +build ignore
+
+package bpf
+
+import _ "github.com/cilium/ebpf/cmd/bpf2go"

internal/event/net.go

@@ -0,0 +1,59 @@
+package event
+
+import (
+	"bytes"
+	"encoding/binary"
+	"github.com/mozillazg/ptcpdump/bpf"
+	"golang.org/x/xerrors"
+	"unsafe"
+)
+
+type packetType int
+
+const (
+	packetTypeIngress packetType = 0
+	packetTypeEgress  packetType = 1
+)
+
+type Packet struct {
+	Type packetType
+	Pid  int
+	Comm string
+
+	Data []byte
+}
+
+func ParsePacketEvent(rawSample []byte) (*Packet, error) {
+	var p Packet
+	event := bpf.BpfPacketEventT{}
+	if err := binary.Read(bytes.NewBuffer(rawSample), binary.LittleEndian, &event.Meta); err != nil {
+		return nil, xerrors.Errorf("parse meta: %w", err)
+	}
+	copy(event.Payload[:], rawSample[unsafe.Offsetof(event.Payload):])
+
+	p.Pid = int(event.Meta.Pid)
+	p.Comm = strComm(event.Meta.Comm)
+	if event.Meta.PacketType == 1 {
+		p.Type = packetTypeEgress
+	}
+	p.Data = make([]byte, event.Meta.PayloadLen)
+	copy(p.Data[:], event.Payload[:event.Meta.PayloadLen])
+
+	return &p, nil
+}
+
+func (p Packet) Ingress() bool {
+	return p.Type == packetTypeIngress
+}
+
+func (p Packet) Egress() bool {
+	return p.Type == packetTypeEgress
+}
+
+func strComm(comm [16]int8) string {
+	b := make([]byte, len(comm))
+	for i, c := range comm {
+		b[i] = byte(c)
+	}
+	return string(b)
+}

internal/writer/pcapng.go

@@ -0,0 +1,41 @@
+package writer
+
+import (
+	"fmt"
+	"github.com/gopacket/gopacket"
+	"github.com/gopacket/gopacket/pcapgo"
+	"github.com/mozillazg/ptcpdump/internal/event"
+	"golang.org/x/xerrors"
+	"time"
+)
+
+type PcapNGWriter struct {
+	pw *pcapgo.NgWriter
+}
+
+func NewPcapNGWriter(pw *pcapgo.NgWriter) *PcapNGWriter {
+	return &PcapNGWriter{pw: pw}
+}
+
+func (w *PcapNGWriter) Write(p *event.Packet) error {
+	payloadLen := len(p.Data)
+	info := gopacket.CaptureInfo{
+		Timestamp:      time.Now(),
+		CaptureLength:  payloadLen,
+		Length:         payloadLen,
+		InterfaceIndex: 0,
+	}
+	opts := pcapgo.NgPacketOptions{
+		Comment: fmt.Sprintf("PID: %d\nCOMMAND: %s", p.Pid, p.Comm),
+	}
+
+	if err := w.pw.WritePacketWithOptions(info, p.Data, opts); err != nil {
+		return xerrors.Errorf("writing packet: %w", err)
+	}
+
+	return nil
+}
+
+func (w *PcapNGWriter) Flush() error {
+	return w.pw.Flush()
+}

internal/writer/stdout.go

@@ -0,0 +1,48 @@
+package writer
+
+import (
+	"github.com/gopacket/gopacket"
+	"github.com/gopacket/gopacket/layers"
+	"github.com/mozillazg/ptcpdump/internal/event"
+	"log"
+)
+
+type StdoutWriter struct {
+}
+
+func NewStdoutWriter() *StdoutWriter {
+	return &StdoutWriter{}
+}
+
+func (w *StdoutWriter) Write(p *event.Packet) error {
+	packetType := "=>·  "
+	if p.Egress() {
+		packetType = "  ·=>"
+	}
+
+	// Decode a packet
+	packet := gopacket.NewPacket(p.Data, layers.LayerTypeEthernet, gopacket.Default)
+	var ipv4 *layers.IPv4
+	if ipv4Layer := packet.Layer(layers.LayerTypeIPv4); ipv4Layer != nil {
+		ipv4, _ = ipv4Layer.(*layers.IPv4)
+	}
+	if ipv4 == nil {
+		return nil
+	}
+	if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
+		tcp, _ := tcpLayer.(*layers.TCP)
+		log.Printf("%s %s-%d %s:%d => %s:%d",
+			packetType, p.Comm, p.Pid,
+			ipv4.SrcIP.String(), tcp.SrcPort,
+			ipv4.DstIP.String(), tcp.DstPort)
+	}
+	if udpLayer := packet.Layer(layers.LayerTypeUDP); udpLayer != nil {
+		udp, _ := udpLayer.(*layers.UDP)
+		log.Printf("%s %s-%d %s:%d => %s:%d",
+			packetType, p.Comm, p.Pid,
+			ipv4.SrcIP.String(), udp.SrcPort,
+			ipv4.DstIP.String(), udp.DstPort)
+	}
+
+	return nil
+}

internal/writer/writer.go

@@ -0,0 +1,4 @@
+package writer
+
+type PacketWriter struct {
+}

main.go

@@ -1,29 +1,25 @@
 package main
 
 import (
-	"bytes"
 	"context"
-	"encoding/binary"
 	"errors"
 	"fmt"
-	"io"
-	"log"
-	"math"
-	"net"
-	"os"
-	"os/signal"
-	"syscall"
-	"time"
-	"unsafe"
-
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/perf"
 	"github.com/cilium/ebpf/rlimit"
-	"github.com/gopacket/gopacket"
 	"github.com/gopacket/gopacket/layers"
 	"github.com/gopacket/gopacket/pcapgo"
 	"github.com/mozillazg/ptcpdump/bpf"
+	"github.com/mozillazg/ptcpdump/internal/event"
+	"github.com/mozillazg/ptcpdump/internal/writer"
 	"golang.org/x/xerrors"
+	"io"
+	"log"
+	"math"
+	"net"
+	"os"
+	"os/signal"
+	"syscall"
 )
 
 func logErr(err error) {
@@ -36,71 +32,23 @@ func logErr(err error) {
 	log.Printf("%+v", err)
 }
 
-func parseEvent(pcapWriter *pcapgo.NgWriter, rawSample []byte) {
-	event := bpf.BpfPacketEventT{}
-	if err := binary.Read(bytes.NewBuffer(rawSample), binary.LittleEndian, &event.Meta); err != nil {
-		log.Printf("parse event failed: %+v", err)
+func parseEvent(pcapWriter *writer.PcapNGWriter, rawSample []byte) {
+	pevent, err := event.ParsePacketEvent(rawSample)
+	if err != nil {
+		logErr(err)
 		return
 	}
-	copy(event.Payload[:], rawSample[unsafe.Offsetof(event.Payload):])
-
-	packetType := "=>·  "
-	if event.Meta.PacketType == 1 {
-		packetType = "  ·=>"
-	}
-
-	// Decode a packet
-	data := event.Payload[:]
-	packet := gopacket.NewPacket(data, layers.LayerTypeEthernet, gopacket.Default)
-	var ipv4 *layers.IPv4
-	if ipv4Layer := packet.Layer(layers.LayerTypeIPv4); ipv4Layer != nil {
-		ipv4, _ = ipv4Layer.(*layers.IPv4)
-	}
-	if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
-		tcp, _ := tcpLayer.(*layers.TCP)
-		log.Printf("%s %s-%d %s:%d => %s:%d",
-			packetType, strComm(event.Meta.Comm), event.Meta.Pid,
-			ipv4.SrcIP.String(), tcp.SrcPort,
-			ipv4.DstIP.String(), tcp.DstPort)
-	}
-	if udpLayer := packet.Layer(layers.LayerTypeUDP); udpLayer != nil {
-		udp, _ := udpLayer.(*layers.UDP)
-		log.Printf("%s %s-%d %s:%d => %s:%d",
-			packetType, strComm(event.Meta.Comm), event.Meta.Pid,
-			ipv4.SrcIP.String(), udp.SrcPort,
-			ipv4.DstIP.String(), udp.DstPort)
-	}
-
-	payloadLen := int(event.Meta.PayloadLen)
-	info := gopacket.CaptureInfo{
-		Timestamp:      time.Now(),
-		CaptureLength:  payloadLen,
-		Length:         payloadLen,
-		InterfaceIndex: 0,
-	}
-	packetData := make([]byte, payloadLen)
-	copy(packetData, data[:payloadLen])
-	log.Printf("len1: %d, len2: %d", len(packetData), len(data[:payloadLen]))
-	opts := pcapgo.NgPacketOptions{
-		Comment: fmt.Sprintf("PID: %d\nCOMMAND: %s", event.Meta.Pid, strComm(event.Meta.Comm)),
-	}
 
-	if err := pcapWriter.WritePacketWithOptions(info, packetData, opts); err != nil {
-		// if err := pcapWriter.WritePacket(info, packetData); err != nil {
-		log.Printf("Error writing packet: %+v", err)
+	if err := writer.NewStdoutWriter().Write(pevent); err != nil {
+		logErr(err)
 	}
-	pcapWriter.Flush()
-}
 
-func strComm(comm [16]int8) string {
-	b := []byte{}
-	for _, c := range comm {
-		b = append(b, byte(c))
+	if err := pcapWriter.Write(pevent); err != nil {
+		logErr(err)
 	}
-	return string(b)
 }
 
-func newPcapWriter(w io.Writer, ifaceNames []string) (*pcapgo.NgWriter, map[string]int, error) {
+func newPcapWriter(w io.Writer, ifaceNames []string) (*writer.PcapNGWriter, map[string]int, error) {
 	if len(ifaceNames) == 0 {
 		return nil, nil, xerrors.New("can't create pcap with no ifaceNames")
 	}
@@ -135,20 +83,20 @@ func newPcapWriter(w io.Writer, ifaceNames []string) (*pcapgo.NgWriter, map[stri
 		return nil, nil, xerrors.Errorf("writing pcap header: %w", err)
 	}
 
-	return pcapWriter, nameIfcs, nil
+	return writer.NewPcapNGWriter(pcapWriter), nameIfcs, nil
 }
 
 func main() {
 	if err := rlimit.RemoveMemlock(); err != nil {
 		logErr(err)
 		return
 	}
-	pcapFile, err := os.Create("test.pcapng")
+	pcapFile, err := os.Create("ptcpdump.pcapng")
 	if err != nil {
 		logErr(err)
 		return
 	}
-	pcapWriter, _, err := newPcapWriter(pcapFile, []string{"lo", "enp0s3", "docker0"})
+	pcapWriter, _, err := newPcapWriter(pcapFile, []string{"lo", "enp0s3", "docker0", "wlp4s0", "enp5s0"})
 	if err != nil {
 		logErr(err)
 		return
@@ -171,7 +119,7 @@ func main() {
 		return
 	}
 
-	for _, ifaceName := range []string{"lo", "enp0s3", "docker0"} {
+	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)