Merge pull request #136 from apernet/wip-pcapgo

Add pcap back

cmd/root.go

@@ -17,6 +17,7 @@ import (
 	"github.com/apernet/OpenGFW/modifier"
 	modUDP "github.com/apernet/OpenGFW/modifier/udp"
 	"github.com/apernet/OpenGFW/ruleset"
+	"github.com/apernet/OpenGFW/ruleset/builtins/geo"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -42,6 +43,7 @@ var logger *zap.Logger
 // Flags
 var (
 	cfgFile   string
+	pcapFile  string
 	logLevel  string
 	logFormat string
 )
@@ -117,6 +119,7 @@ func init() {
 
 func initFlags() {
 	rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "", "config file")
+	rootCmd.PersistentFlags().StringVarP(&pcapFile, "pcap", "p", "", "pcap file (optional)")
 	rootCmd.PersistentFlags().StringVarP(&logLevel, "log-level", "l", envOrDefaultString(appLogLevelEnv, "info"), "log level")
 	rootCmd.PersistentFlags().StringVarP(&logFormat, "log-format", "f", envOrDefaultString(appLogFormatEnv, "console"), "log format")
 }
@@ -166,6 +169,7 @@ type cliConfig struct {
 	IO      cliConfigIO      `mapstructure:"io"`
 	Workers cliConfigWorkers `mapstructure:"workers"`
 	Ruleset cliConfigRuleset `mapstructure:"ruleset"`
+	Replay  cliConfigReplay  `mapstructure:"replay"`
 }
 
 type cliConfigIO struct {
@@ -176,6 +180,10 @@ type cliConfigIO struct {
 	RST         bool   `mapstructure:"rst"`
 }
 
+type cliConfigReplay struct {
+	Realtime bool `mapstructure:"realtime"`
+}
+
 type cliConfigWorkers struct {
 	Count                      int           `mapstructure:"count"`
 	QueueSize                  int           `mapstructure:"queueSize"`
@@ -196,17 +204,30 @@ func (c *cliConfig) fillLogger(config *engine.Config) error {
 }
 
 func (c *cliConfig) fillIO(config *engine.Config) error {
-	nfio, err := io.NewNFQueuePacketIO(io.NFQueuePacketIOConfig{
-		QueueSize:   c.IO.QueueSize,
-		ReadBuffer:  c.IO.ReadBuffer,
-		WriteBuffer: c.IO.WriteBuffer,
-		Local:       c.IO.Local,
-		RST:         c.IO.RST,
-	})
+	var ioImpl io.PacketIO
+	var err error
+	if pcapFile != "" {
+		// Setup IO for pcap file replay
+		logger.Info("replaying from pcap file", zap.String("pcap file", pcapFile))
+		ioImpl, err = io.NewPcapPacketIO(io.PcapPacketIOConfig{
+			PcapFile: pcapFile,
+			Realtime: c.Replay.Realtime,
+		})
+	} else {
+		// Setup IO for nfqueue
+		ioImpl, err = io.NewNFQueuePacketIO(io.NFQueuePacketIOConfig{
+			QueueSize:   c.IO.QueueSize,
+			ReadBuffer:  c.IO.ReadBuffer,
+			WriteBuffer: c.IO.WriteBuffer,
+			Local:       c.IO.Local,
+			RST:         c.IO.RST,
+		})
+	}
+
 	if err != nil {
 		return configError{Field: "io", Err: err}
 	}
-	config.IO = nfio
+	config.IO = ioImpl
 	return nil
 }
 
@@ -259,8 +280,7 @@ func runMain(cmd *cobra.Command, args []string) {
 	}
 	rsConfig := &ruleset.BuiltinConfig{
 		Logger:               &rulesetLogger{},
-		GeoSiteFilename:      config.Ruleset.GeoSite,
-		GeoIpFilename:        config.Ruleset.GeoIp,
+		GeoMatcher:           geo.NewGeoMatcher(config.Ruleset.GeoSite, config.Ruleset.GeoIp),
 		ProtectedDialContext: engineConfig.IO.ProtectedDialContext,
 	}
 	rs, err := ruleset.CompileExprRules(rawRs, analyzers, modifiers, rsConfig)

engine/engine.go

@@ -58,12 +58,17 @@ func (e *engine) UpdateRuleset(r ruleset.Ruleset) error {
 }
 
 func (e *engine) Run(ctx context.Context) error {
+	workerCtx, workerCancel := context.WithCancel(ctx)
+	defer workerCancel() // Stop workers
+
+	// Register IO shutdown
 	ioCtx, ioCancel := context.WithCancel(ctx)
-	defer ioCancel() // Stop workers & IO
+	e.io.SetCancelFunc(ioCancel)
+	defer ioCancel() // Stop IO
 
 	// Start workers
 	for _, w := range e.workers {
-		go w.Run(ioCtx)
+		go w.Run(workerCtx)
 	}
 
 	// Register IO callback
@@ -85,6 +90,8 @@ func (e *engine) Run(ctx context.Context) error {
 		return err
 	case <-ctx.Done():
 		return nil
+	case <-ioCtx.Done():
+		return nil
 	}
 }
 
@@ -102,9 +109,11 @@ func (e *engine) dispatch(p io.Packet) bool {
 		_ = e.io.SetVerdict(p, io.VerdictAcceptStream, nil)
 		return true
 	}
+	// Convert to gopacket.Packet
+	packet := gopacket.NewPacket(data, layerType, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
+	packet.Metadata().Timestamp = p.Timestamp()
 	// Load balance by stream ID
 	index := p.StreamID() % uint32(len(e.workers))
-	packet := gopacket.NewPacket(data, layerType, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
 	e.workers[index].Feed(&workerPacket{
 		StreamID: p.StreamID(),
 		Packet:   packet,

io/interface.go

@@ -3,6 +3,7 @@ package io
 import (
 	"context"
 	"net"
+	"time"
 )
 
 type Verdict int
@@ -24,6 +25,8 @@ const (
 type Packet interface {
 	// StreamID is the ID of the stream the packet belongs to.
 	StreamID() uint32
+	// Timestamp is the time the packet was received.
+	Timestamp() time.Time
 	// Data is the raw packet data, starting with the IP header.
 	Data() []byte
 }
@@ -45,6 +48,9 @@ type PacketIO interface {
 	ProtectedDialContext(ctx context.Context, network, address string) (net.Conn, error)
 	// Close closes the packet IO.
 	Close() error
+	// SetCancelFunc gives packet IO access to context cancel function, enabling it to
+	// trigger a shutdown
+	SetCancelFunc(cancelFunc context.CancelFunc) error
 }
 
 type ErrInvalidPacket struct {

io/nfqueue.go

@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"strings"
 	"syscall"
+	"time"
 
 	"github.com/coreos/go-iptables/iptables"
 	"github.com/florianl/go-nfqueue"
@@ -189,6 +190,12 @@ func (n *nfqueuePacketIO) Register(ctx context.Context, cb PacketCallback) error
 				streamID: ctIDFromCtBytes(*a.Ct),
 				data:     *a.Payload,
 			}
+			// Use timestamp from attribute if available, otherwise use current time as fallback
+			if a.Timestamp != nil {
+				p.timestamp = *a.Timestamp
+			} else {
+				p.timestamp = time.Now()
+			}
 			return okBoolToInt(cb(p, nil))
 		},
 		func(e error) int {
@@ -274,6 +281,11 @@ func (n *nfqueuePacketIO) Close() error {
 	return n.n.Close()
 }
 
+// nfqueue IO does not issue shutdown
+func (n *nfqueuePacketIO) SetCancelFunc(cancelFunc context.CancelFunc) error {
+	return nil
+}
+
 func (n *nfqueuePacketIO) setupNft(local, rst, remove bool) error {
 	rules, err := generateNftRules(local, rst)
 	if err != nil {
@@ -312,15 +324,20 @@ func (n *nfqueuePacketIO) setupIpt(local, rst, remove bool) error {
 var _ Packet = (*nfqueuePacket)(nil)
 
 type nfqueuePacket struct {
-	id       uint32
-	streamID uint32
-	data     []byte
+	id        uint32
+	streamID  uint32
+	timestamp time.Time
+	data      []byte
 }
 
 func (p *nfqueuePacket) StreamID() uint32 {
 	return p.streamID
 }
 
+func (p *nfqueuePacket) Timestamp() time.Time {
+	return p.timestamp
+}
+
 func (p *nfqueuePacket) Data() []byte {
 	return p.data
 }

io/pcap.go

@@ -0,0 +1,136 @@
+package io
+
+import (
+	"context"
+	"hash/crc32"
+	"io"
+	"net"
+	"os"
+	"sort"
+	"strings"
+	"time"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/pcapgo"
+)
+
+var _ PacketIO = (*pcapPacketIO)(nil)
+
+type pcapPacketIO struct {
+	pcapFile   io.ReadCloser
+	pcap       *pcapgo.Reader
+	timeOffset *time.Duration
+	ioCancel   context.CancelFunc
+	config     PcapPacketIOConfig
+
+	dialer *net.Dialer
+}
+
+type PcapPacketIOConfig struct {
+	PcapFile string
+	Realtime bool
+}
+
+func NewPcapPacketIO(config PcapPacketIOConfig) (PacketIO, error) {
+	pcapFile, err := os.Open(config.PcapFile)
+	if err != nil {
+		return nil, err
+	}
+
+	handle, err := pcapgo.NewReader(pcapFile)
+	if err != nil {
+		return nil, err
+	}
+
+	return &pcapPacketIO{
+		pcapFile:   pcapFile,
+		pcap:       handle,
+		timeOffset: nil,
+		ioCancel:   nil,
+		config:     config,
+		dialer:     &net.Dialer{},
+	}, nil
+}
+
+func (p *pcapPacketIO) Register(ctx context.Context, cb PacketCallback) error {
+	go func() {
+		packetSource := gopacket.NewPacketSource(p.pcap, p.pcap.LinkType())
+		for packet := range packetSource.Packets() {
+			p.wait(packet)
+
+			networkLayer := packet.NetworkLayer()
+			if networkLayer != nil {
+				src, dst := networkLayer.NetworkFlow().Endpoints()
+				endpoints := []string{src.String(), dst.String()}
+				sort.Strings(endpoints)
+				id := crc32.Checksum([]byte(strings.Join(endpoints, ",")), crc32.IEEETable)
+
+				cb(&pcapPacket{
+					streamID:  id,
+					timestamp: packet.Metadata().Timestamp,
+					data:      packet.LinkLayer().LayerPayload(),
+				}, nil)
+			}
+		}
+		// Give the workers a chance to finish everything
+		time.Sleep(time.Second)
+		// Stop the engine when all packets are finished
+		p.ioCancel()
+	}()
+
+	return nil
+}
+
+// A normal dialer is sufficient as pcap IO does not mess up with the networking
+func (p *pcapPacketIO) ProtectedDialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	return p.dialer.DialContext(ctx, network, address)
+}
+
+func (p *pcapPacketIO) SetVerdict(pkt Packet, v Verdict, newPacket []byte) error {
+	return nil
+}
+
+func (p *pcapPacketIO) SetCancelFunc(cancelFunc context.CancelFunc) error {
+	p.ioCancel = cancelFunc
+	return nil
+}
+
+func (p *pcapPacketIO) Close() error {
+	return p.pcapFile.Close()
+}
+
+// Intentionally slow down the replay
+// In realtime mode, this is to match the timestamps in the capture
+func (p *pcapPacketIO) wait(packet gopacket.Packet) {
+	if !p.config.Realtime {
+		return
+	}
+
+	if p.timeOffset == nil {
+		offset := time.Since(packet.Metadata().Timestamp)
+		p.timeOffset = &offset
+	} else {
+		t := time.Until(packet.Metadata().Timestamp.Add(*p.timeOffset))
+		time.Sleep(t)
+	}
+}
+
+var _ Packet = (*pcapPacket)(nil)
+
+type pcapPacket struct {
+	streamID  uint32
+	timestamp time.Time
+	data      []byte
+}
+
+func (p *pcapPacket) StreamID() uint32 {
+	return p.streamID
+}
+
+func (p *pcapPacket) Timestamp() time.Time {
+	return p.timestamp
+}
+
+func (p *pcapPacket) Data() []byte {
+	return p.data
+}

ruleset/expr.go

@@ -20,7 +20,6 @@ import (
 	"github.com/apernet/OpenGFW/analyzer"
 	"github.com/apernet/OpenGFW/modifier"
 	"github.com/apernet/OpenGFW/ruleset/builtins"
-	"github.com/apernet/OpenGFW/ruleset/builtins/geo"
 )
 
 // ExprRule is the external representation of an expression rule.
@@ -302,23 +301,22 @@ type Function struct {
 }
 
 func buildFunctionMap(config *BuiltinConfig) map[string]*Function {
-	geoMatcher := geo.NewGeoMatcher(config.GeoSiteFilename, config.GeoIpFilename)
 	return map[string]*Function{
 		"geoip": {
-			InitFunc:  geoMatcher.LoadGeoIP,
+			InitFunc:  config.GeoMatcher.LoadGeoIP,
 			PatchFunc: nil,
 			Func: func(params ...any) (any, error) {
-				return geoMatcher.MatchGeoIp(params[0].(string), params[1].(string)), nil
+				return config.GeoMatcher.MatchGeoIp(params[0].(string), params[1].(string)), nil
 			},
-			Types: []reflect.Type{reflect.TypeOf(geoMatcher.MatchGeoIp)},
+			Types: []reflect.Type{reflect.TypeOf(config.GeoMatcher.MatchGeoIp)},
 		},
 		"geosite": {
-			InitFunc:  geoMatcher.LoadGeoSite,
+			InitFunc:  config.GeoMatcher.LoadGeoSite,
 			PatchFunc: nil,
 			Func: func(params ...any) (any, error) {
-				return geoMatcher.MatchGeoSite(params[0].(string), params[1].(string)), nil
+				return config.GeoMatcher.MatchGeoSite(params[0].(string), params[1].(string)), nil
 			},
-			Types: []reflect.Type{reflect.TypeOf(geoMatcher.MatchGeoSite)},
+			Types: []reflect.Type{reflect.TypeOf(config.GeoMatcher.MatchGeoSite)},
 		},
 		"cidr": {
 			InitFunc: nil,

ruleset/interface.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/apernet/OpenGFW/analyzer"
 	"github.com/apernet/OpenGFW/modifier"
+	"github.com/apernet/OpenGFW/ruleset/builtins/geo"
 )
 
 type Action int
@@ -102,7 +103,6 @@ type Logger interface {
 
 type BuiltinConfig struct {
 	Logger               Logger
-	GeoSiteFilename      string
-	GeoIpFilename        string
+	GeoMatcher           *geo.GeoMatcher
 	ProtectedDialContext func(ctx context.Context, network, address string) (net.Conn, error)
 }