Merge pull request #135 from apernet/revert-pcap

Revert pcap

cmd/root.go

@@ -7,6 +7,7 @@ import (
 	"os/signal"
 	"strings"
 	"syscall"
+	"time"
 
 	"github.com/apernet/OpenGFW/analyzer"
 	"github.com/apernet/OpenGFW/analyzer/tcp"
@@ -16,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"
@@ -176,11 +178,12 @@ type cliConfigIO struct {
 }
 
 type cliConfigWorkers struct {
-	Count                      int `mapstructure:"count"`
-	QueueSize                  int `mapstructure:"queueSize"`
-	TCPMaxBufferedPagesTotal   int `mapstructure:"tcpMaxBufferedPagesTotal"`
-	TCPMaxBufferedPagesPerConn int `mapstructure:"tcpMaxBufferedPagesPerConn"`
-	UDPMaxStreams              int `mapstructure:"udpMaxStreams"`
+	Count                      int           `mapstructure:"count"`
+	QueueSize                  int           `mapstructure:"queueSize"`
+	TCPMaxBufferedPagesTotal   int           `mapstructure:"tcpMaxBufferedPagesTotal"`
+	TCPMaxBufferedPagesPerConn int           `mapstructure:"tcpMaxBufferedPagesPerConn"`
+	TCPTimeout                 time.Duration `mapstructure:"tcpTimeout"`
+	UDPMaxStreams              int           `mapstructure:"udpMaxStreams"`
 }
 
 type cliConfigRuleset struct {
@@ -213,6 +216,7 @@ func (c *cliConfig) fillWorkers(config *engine.Config) error {
 	config.WorkerQueueSize = c.Workers.QueueSize
 	config.WorkerTCPMaxBufferedPagesTotal = c.Workers.TCPMaxBufferedPagesTotal
 	config.WorkerTCPMaxBufferedPagesPerConn = c.Workers.TCPMaxBufferedPagesPerConn
+	config.WorkerTCPTimeout = c.Workers.TCPTimeout
 	config.WorkerUDPMaxStreams = c.Workers.UDPMaxStreams
 	return nil
 }
@@ -256,8 +260,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)
@@ -340,12 +343,26 @@ func (l *engineLogger) TCPStreamPropUpdate(info ruleset.StreamInfo, close bool)
 }
 
 func (l *engineLogger) TCPStreamAction(info ruleset.StreamInfo, action ruleset.Action, noMatch bool) {
-	logger.Info("TCP stream action",
-		zap.Int64("id", info.ID),
-		zap.String("src", info.SrcString()),
-		zap.String("dst", info.DstString()),
-		zap.String("action", action.String()),
-		zap.Bool("noMatch", noMatch))
+	if noMatch {
+		logger.Debug("TCP stream no match",
+			zap.Int64("id", info.ID),
+			zap.String("src", info.SrcString()),
+			zap.String("dst", info.DstString()),
+			zap.String("action", action.String()))
+	} else {
+		logger.Info("TCP stream action",
+			zap.Int64("id", info.ID),
+			zap.String("src", info.SrcString()),
+			zap.String("dst", info.DstString()),
+			zap.String("action", action.String()))
+	}
+}
+
+func (l *engineLogger) TCPFlush(workerID, flushed, closed int) {
+	logger.Debug("TCP flush",
+		zap.Int("workerID", workerID),
+		zap.Int("flushed", flushed),
+		zap.Int("closed", closed))
 }
 
 func (l *engineLogger) UDPStreamNew(workerID int, info ruleset.StreamInfo) {
@@ -366,12 +383,19 @@ func (l *engineLogger) UDPStreamPropUpdate(info ruleset.StreamInfo, close bool)
 }
 
 func (l *engineLogger) UDPStreamAction(info ruleset.StreamInfo, action ruleset.Action, noMatch bool) {
-	logger.Info("UDP stream action",
-		zap.Int64("id", info.ID),
-		zap.String("src", info.SrcString()),
-		zap.String("dst", info.DstString()),
-		zap.String("action", action.String()),
-		zap.Bool("noMatch", noMatch))
+	if noMatch {
+		logger.Debug("UDP stream no match",
+			zap.Int64("id", info.ID),
+			zap.String("src", info.SrcString()),
+			zap.String("dst", info.DstString()),
+			zap.String("action", action.String()))
+	} else {
+		logger.Info("UDP stream action",
+			zap.Int64("id", info.ID),
+			zap.String("src", info.SrcString()),
+			zap.String("dst", info.DstString()),
+			zap.String("action", action.String()))
+	}
 }
 
 func (l *engineLogger) ModifyError(info ruleset.StreamInfo, err error) {

engine/engine.go

@@ -34,6 +34,7 @@ func NewEngine(config Config) (Engine, error) {
 			Ruleset:                    config.Ruleset,
 			TCPMaxBufferedPagesTotal:   config.WorkerTCPMaxBufferedPagesTotal,
 			TCPMaxBufferedPagesPerConn: config.WorkerTCPMaxBufferedPagesPerConn,
+			TCPTimeout:                 config.WorkerTCPTimeout,
 			UDPMaxStreams:              config.WorkerUDPMaxStreams,
 		})
 		if err != nil {
@@ -101,9 +102,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,

engine/interface.go

@@ -2,6 +2,7 @@ package engine
 
 import (
 	"context"
+	"time"
 
 	"github.com/apernet/OpenGFW/io"
 	"github.com/apernet/OpenGFW/ruleset"
@@ -25,6 +26,7 @@ type Config struct {
 	WorkerQueueSize                  int
 	WorkerTCPMaxBufferedPagesTotal   int
 	WorkerTCPMaxBufferedPagesPerConn int
+	WorkerTCPTimeout                 time.Duration
 	WorkerUDPMaxStreams              int
 }
 
@@ -36,6 +38,7 @@ type Logger interface {
 	TCPStreamNew(workerID int, info ruleset.StreamInfo)
 	TCPStreamPropUpdate(info ruleset.StreamInfo, close bool)
 	TCPStreamAction(info ruleset.StreamInfo, action ruleset.Action, noMatch bool)
+	TCPFlush(workerID, flushed, closed int)
 
 	UDPStreamNew(workerID int, info ruleset.StreamInfo)
 	UDPStreamPropUpdate(info ruleset.StreamInfo, close bool)

engine/worker.go

@@ -2,6 +2,7 @@ package engine
 
 import (
 	"context"
+	"time"
 
 	"github.com/apernet/OpenGFW/io"
 	"github.com/apernet/OpenGFW/ruleset"
@@ -14,9 +15,12 @@ import (
 
 const (
 	defaultChanSize                         = 64
-	defaultTCPMaxBufferedPagesTotal         = 4096
-	defaultTCPMaxBufferedPagesPerConnection = 64
+	defaultTCPMaxBufferedPagesTotal         = 65536
+	defaultTCPMaxBufferedPagesPerConnection = 16
+	defaultTCPTimeout                       = 10 * time.Minute
 	defaultUDPMaxStreams                    = 4096
+
+	tcpFlushInterval = 1 * time.Minute
 )
 
 type workerPacket struct {
@@ -33,6 +37,7 @@ type worker struct {
 	tcpStreamFactory *tcpStreamFactory
 	tcpStreamPool    *reassembly.StreamPool
 	tcpAssembler     *reassembly.Assembler
+	tcpTimeout       time.Duration
 
 	udpStreamFactory *udpStreamFactory
 	udpStreamManager *udpStreamManager
@@ -47,6 +52,7 @@ type workerConfig struct {
 	Ruleset                    ruleset.Ruleset
 	TCPMaxBufferedPagesTotal   int
 	TCPMaxBufferedPagesPerConn int
+	TCPTimeout                 time.Duration
 	UDPMaxStreams              int
 }
 
@@ -60,6 +66,9 @@ func (c *workerConfig) fillDefaults() {
 	if c.TCPMaxBufferedPagesPerConn <= 0 {
 		c.TCPMaxBufferedPagesPerConn = defaultTCPMaxBufferedPagesPerConnection
 	}
+	if c.TCPTimeout <= 0 {
+		c.TCPTimeout = defaultTCPTimeout
+	}
 	if c.UDPMaxStreams <= 0 {
 		c.UDPMaxStreams = defaultUDPMaxStreams
 	}
@@ -98,6 +107,7 @@ func newWorker(config workerConfig) (*worker, error) {
 		tcpStreamFactory:   tcpSF,
 		tcpStreamPool:      tcpStreamPool,
 		tcpAssembler:       tcpAssembler,
+		tcpTimeout:         config.TCPTimeout,
 		udpStreamFactory:   udpSF,
 		udpStreamManager:   udpSM,
 		modSerializeBuffer: gopacket.NewSerializeBuffer(),
@@ -111,6 +121,10 @@ func (w *worker) Feed(p *workerPacket) {
 func (w *worker) Run(ctx context.Context) {
 	w.logger.WorkerStart(w.id)
 	defer w.logger.WorkerStop(w.id)
+
+	tcpFlushTicker := time.NewTicker(tcpFlushInterval)
+	defer tcpFlushTicker.Stop()
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -122,6 +136,8 @@ func (w *worker) Run(ctx context.Context) {
 			}
 			v, b := w.handle(wPkt.StreamID, wPkt.Packet)
 			_ = wPkt.SetVerdict(v, b)
+		case <-tcpFlushTicker.C:
+			w.flushTCP(w.tcpTimeout)
 		}
 	}
 }
@@ -176,6 +192,11 @@ func (w *worker) handleTCP(ipFlow gopacket.Flow, pMeta *gopacket.PacketMetadata,
 	return io.Verdict(ctx.Verdict)
 }
 
+func (w *worker) flushTCP(timeout time.Duration) {
+	flushed, closed := w.tcpAssembler.FlushCloseOlderThan(time.Now().Add(-timeout))
+	w.logger.TCPFlush(w.id, flushed, closed)
+}
+
 func (w *worker) handleUDP(streamID uint32, ipFlow gopacket.Flow, udp *layers.UDP) (io.Verdict, []byte) {
 	ctx := &udpContext{
 		Verdict: udpVerdictAccept,

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
 }

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 {
@@ -312,15 +319,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
 }

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)
 }