feat(tcp, udp): implement stream state release after verdict and adjust flush interval logic

2026-09-12 17:20:24 +08:00
parent cc6ad72706
commit 8f8cc1d4d9
4 changed files with 254 additions and 12 deletions
@@ -0,0 +1,199 @@
+package engine
+
+import (
+	"net"
+	"testing"
+	"time"
+
+	"github.com/apernet/OpenGFW/analyzer"
+	"github.com/apernet/OpenGFW/analyzer/tcp"
+	"github.com/apernet/OpenGFW/analyzer/udp"
+	"github.com/apernet/OpenGFW/io"
+	"github.com/apernet/OpenGFW/ruleset"
+	"github.com/apernet/OpenGFW/ruleset/builtins/geo"
+
+	"github.com/bwmarrin/snowflake"
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+	"github.com/google/gopacket/reassembly"
+)
+
+type nullEngineLogger struct{}
+
+func (nullEngineLogger) WorkerStart(int)                                          {}
+func (nullEngineLogger) WorkerStop(int)                                           {}
+func (nullEngineLogger) TCPStreamNew(int, ruleset.StreamInfo)                     {}
+func (nullEngineLogger) TCPStreamPropUpdate(ruleset.StreamInfo, bool)             {}
+func (nullEngineLogger) TCPStreamAction(ruleset.StreamInfo, ruleset.Action, bool) {}
+func (nullEngineLogger) TCPFlush(int, int, int)                                   {}
+func (nullEngineLogger) UDPStreamNew(int, ruleset.StreamInfo)                     {}
+func (nullEngineLogger) UDPStreamPropUpdate(ruleset.StreamInfo, bool)             {}
+func (nullEngineLogger) UDPStreamAction(ruleset.StreamInfo, ruleset.Action, bool) {}
+func (nullEngineLogger) ModifyError(ruleset.StreamInfo, error)                    {}
+func (nullEngineLogger) AnalyzerDebugf(int64, string, string, ...interface{})     {}
+func (nullEngineLogger) AnalyzerInfof(int64, string, string, ...interface{})      {}
+func (nullEngineLogger) AnalyzerErrorf(int64, string, string, ...interface{})     {}
+
+type nullRulesetLogger struct{}
+
+func (nullRulesetLogger) Log(ruleset.StreamInfo, string)               {}
+func (nullRulesetLogger) MatchError(ruleset.StreamInfo, string, error) {}
+
+func testRuleset(t *testing.T, rules ...ruleset.ExprRule) ruleset.Ruleset {
+	t.Helper()
+	ans := []analyzer.Analyzer{
+		&tcp.FETAnalyzer{}, &tcp.HTTPAnalyzer{}, &tcp.TLSAnalyzer{},
+		&udp.QUICAnalyzer{}, &udp.WireGuardAnalyzer{},
+	}
+	rs, err := ruleset.CompileExprRules(rules, ans, nil, &ruleset.BuiltinConfig{Logger: nullRulesetLogger{}, GeoMatcher: geo.NewGeoMatcher("", "")})
+	if err != nil {
+		t.Fatal(err)
+	}
+	return rs
+}
+
+// capturingFactory 记下工厂创建的每条 tcpStream，让测试能检查判定之后它的状态
+type capturingFactory struct {
+	inner   *tcpStreamFactory
+	streams []*tcpStream
+}
+
+func (f *capturingFactory) New(ipFlow, tcpFlow gopacket.Flow, tcp *layers.TCP, ac reassembly.AssemblerContext) reassembly.Stream {
+	s := f.inner.New(ipFlow, tcpFlow, tcp, ac).(*tcpStream)
+	f.streams = append(f.streams, s)
+	return s
+}
+
+func newTestAssembler(t *testing.T, rs ruleset.Ruleset) (*reassembly.Assembler, *capturingFactory) {
+	t.Helper()
+	node, err := snowflake.NewNode(1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	cf := &capturingFactory{inner: &tcpStreamFactory{WorkerID: 1, Logger: nullEngineLogger{}, Node: node, Ruleset: rs}}
+	asm := reassembly.NewAssembler(reassembly.NewStreamPool(cf))
+	asm.MaxBufferedPagesTotal = defaultTCPMaxBufferedPagesTotal
+	asm.MaxBufferedPagesPerConnection = defaultTCPMaxBufferedPagesPerConnection
+	return asm, cf
+}
+
+var testIPFlow = gopacket.NewFlow(layers.EndpointIPv4, net.IPv4(10, 0, 0, 1).To4(), net.IPv4(203, 0, 113, 1).To4())
+
+func tcpSegment(seq uint32, syn bool, payload []byte) *layers.TCP {
+	seg := &layers.TCP{SrcPort: 40000, DstPort: 80, Seq: seq, SYN: syn, ACK: !syn, PSH: len(payload) > 0, Window: 65535}
+	seg.Payload = payload
+	return seg
+}
+
+func feedTCP(asm *reassembly.Assembler, seg *layers.TCP, ts time.Time) tcpVerdict {
+	ctx := &tcpContext{PacketMetadata: &gopacket.PacketMetadata{CaptureInfo: gopacket.CaptureInfo{Timestamp: ts}}, Verdict: tcpVerdictAccept}
+	asm.AssembleWithContext(testIPFlow, seg, ctx)
+	return ctx.Verdict
+}
+
+var httpRequest = []byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")
+
+func TestTCPStreamReleasesStateAfterVerdict(t *testing.T) {
+	rs := testRuleset(t, ruleset.ExprRule{Name: "block-http", Action: "block", Expr: "http?.req != nil"})
+	asm, cf := newTestAssembler(t, rs)
+	ts := time.Now()
+
+	feedTCP(asm, tcpSegment(1000, true, nil), ts)
+	if v := feedTCP(asm, tcpSegment(1001, false, httpRequest), ts); v != tcpVerdictDropStream {
+		t.Fatalf("verdict = %v, want drop stream", v)
+	}
+	if len(cf.streams) != 1 {
+		t.Fatalf("streams created = %d, want 1", len(cf.streams))
+	}
+	s := cf.streams[0]
+	if !s.released || s.activeEntries != nil || s.info.Props != nil || s.ruleset != nil {
+		t.Fatalf("stream still holds state after verdict: released=%v entries=%d props=%v ruleset=%v", s.released, len(s.activeEntries), s.info.Props, s.ruleset != nil)
+	}
+
+	// 后续报文沿用判定，且不能碰已经释放的状态
+	if v := feedTCP(asm, tcpSegment(1001+uint32(len(httpRequest)), false, []byte("more")), ts); v != tcpVerdictDropStream {
+		t.Fatalf("verdict on a later segment = %v, want drop stream", v)
+	}
+}
+
+func TestTCPStreamReleasedWhenAllAnalyzersFinish(t *testing.T) {
+	// 规则只引用 tls；一个 HTTP 请求让 tls 分析器直接放弃，没有裁决，引擎应当发 AcceptStream 并释放
+	rs := testRuleset(t, ruleset.ExprRule{Name: "block-tls", Action: "block", Expr: "tls?.req != nil"})
+	asm, cf := newTestAssembler(t, rs)
+	ts := time.Now()
+
+	feedTCP(asm, tcpSegment(1000, true, nil), ts)
+	if v := feedTCP(asm, tcpSegment(1001, false, httpRequest), ts); v != tcpVerdictAcceptStream {
+		t.Fatalf("verdict = %v, want accept stream", v)
+	}
+	if s := cf.streams[0]; !s.released || s.info.Props != nil {
+		t.Fatalf("stream not released after all analyzers finished: released=%v props=%v", s.released, s.info.Props)
+	}
+}
+
+// 判定之前缓冲的乱序页会在超时 flush 时被补送到 ReassembledSG，那时流已经释放，不能崩
+func TestTCPStreamSurvivesFlushDeliveringBufferedPageAfterRelease(t *testing.T) {
+	rs := testRuleset(t, ruleset.ExprRule{Name: "block-http", Action: "block", Expr: "http?.req != nil"})
+	asm, cf := newTestAssembler(t, rs)
+	ts := time.Now()
+
+	feedTCP(asm, tcpSegment(1000, true, nil), ts)
+	// 乱序段：序号跳过一大截，被缓冲起来等空洞补齐
+	feedTCP(asm, tcpSegment(9000, false, []byte("out of order")), ts)
+	if v := feedTCP(asm, tcpSegment(1001, false, httpRequest), ts); v != tcpVerdictDropStream {
+		t.Fatalf("verdict = %v, want drop stream", v)
+	}
+	s := cf.streams[0]
+	if !s.released {
+		t.Fatal("stream not released after verdict")
+	}
+
+	// 超时 flush 会带着空洞把那一页送进来
+	flushed, closed := asm.FlushCloseOlderThan(ts.Add(time.Hour))
+	if flushed == 0 || closed == 0 {
+		t.Fatalf("flush did nothing: flushed=%d closed=%d", flushed, closed)
+	}
+	if s.info.Props != nil || s.activeEntries != nil {
+		t.Fatal("flush resurrected released state")
+	}
+}
+
+func TestUDPStreamReleasesStateAfterVerdict(t *testing.T) {
+	rs := testRuleset(t, ruleset.ExprRule{Name: "allow-443", Action: "allow", Expr: "port.dst == 443"})
+	w, err := newWorker(workerConfig{ID: 1, Logger: nullEngineLogger{}, Ruleset: rs})
+	if err != nil {
+		t.Fatal(err)
+	}
+	pkt := &layers.UDP{SrcPort: 40000, DstPort: 443}
+	pkt.Payload = []byte("hello")
+
+	if v, _ := w.handleUDP(7, testIPFlow, pkt); v != io.VerdictAcceptStream {
+		t.Fatalf("verdict = %v, want accept stream", v)
+	}
+	value, ok := w.udpStreamManager.streams.Get(7)
+	if !ok {
+		t.Fatal("stream missing from the LRU")
+	}
+	if s := value.Stream; s.activeEntries != nil || s.info.Props != nil || s.ruleset != nil {
+		t.Fatalf("udp stream still holds state after verdict: entries=%d props=%v", len(s.activeEntries), s.info.Props)
+	}
+	if v, _ := w.handleUDP(7, testIPFlow, pkt); v != io.VerdictAcceptStream {
+		t.Fatalf("verdict on a later packet = %v, want accept stream", v)
+	}
+}
+
+func TestTCPFlushIntervalFollowsTimeout(t *testing.T) {
+	cases := []struct {
+		timeout, want time.Duration
+	}{
+		{30 * time.Second, 15 * time.Second},
+		{10 * time.Minute, time.Minute},
+		{time.Second, time.Second},
+		{0, time.Second},
+	}
+	for _, c := range cases {
+		if got := tcpFlushIntervalFor(c.timeout); got != c.want {
+			t.Errorf("timeout %v: interval = %v, want %v", c.timeout, got, c.want)
+		}
+	}
+}
@@ -124,8 +124,8 @@ type tcpStream struct {
 	logger        Logger
 	ruleset       ruleset.Ruleset
 	activeEntries []*tcpStreamEntry
-	doneEntries   []*tcpStreamEntry
 	lastVerdict   tcpVerdict
+	released      bool // 最终判定已发出，分析器和属性表已经释放
 	firstSegSeen  bool // true once the client's first data segment was measured
 	synAckWindow  int  // window of the observed SYN-ACK, -1 if none seen yet
 	synAckWscale  int  // its window scale factor, -1 if the option was absent
@@ -139,6 +139,11 @@ type tcpStreamEntry struct {
 }

 func (s *tcpStream) Accept(tcp *layers.TCP, ci gopacket.CaptureInfo, dir reassembly.TCPFlowDirection, nextSeq reassembly.Sequence, start *bool, ac reassembly.AssemblerContext) bool {
+	if s.released {
+		// 最终判定已发出，后续报文一律沿用；Props 已经释放，不能再碰
+		ac.(*tcpContext).Verdict = s.lastVerdict
+		return false
+	}
 	if tcp.SYN && tcp.ACK && s.synAckWindow < 0 {
 		s.synAckWindow, s.synAckWscale = int(tcp.Window), tcpWindowScale(tcp)
 		if s.synAckWindow <= tcpSynAckWindowMax {
@@ -210,6 +215,10 @@ func (s *tcpStream) recordFirstSegLen(n int) bool {
 }

 func (s *tcpStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.AssemblerContext) {
+	if s.released {
+		// 判定已发出。还会走到这里只可能是超时 flush 补送的乱序页，没有东西可处理
+		return
+	}
 	dir, start, end, skip := sg.Info()
 	rev := dir == reassembly.TCPDirServerToClient
 	avail, _ := sg.Lengths()
@@ -229,7 +238,6 @@ func (s *tcpStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.Ass
 		updated = updated || up1 || up2
 		if done {
 			s.activeEntries = append(s.activeEntries[:i], s.activeEntries[i+1:]...)
-			s.doneEntries = append(s.doneEntries, entry)
 		}
 	}
 	ctx := ac.(*tcpContext)
@@ -254,15 +262,20 @@ func (s *tcpStream) ReassembledSG(sg reassembly.ScatterGather, ac reassembly.Ass
 		ctx.Verdict = tcpVerdictAcceptStream
 		s.logger.TCPStreamAction(s.info, ruleset.ActionAllow, true)
 	}
+	if len(s.activeEntries) == 0 {
+		// 不会再有分析器读数据，也不会再匹配规则：从这一刻起 Accept 只看 lastVerdict
+		s.release()
+	}
 }

 func (s *tcpStream) ReassemblyComplete(ac reassembly.AssemblerContext) bool {
 	s.closeActiveEntries()
+	s.release()
 	return true
 }

 func (s *tcpStream) closeActiveEntries() {
-	// Signal close to all active entries & move them to doneEntries
+	// Signal close to all active entries
 	updated := false
 	for _, entry := range s.activeEntries {
 		update := entry.Stream.Close(false)
@@ -272,10 +285,22 @@ func (s *tcpStream) closeActiveEntries() {
 	if updated {
 		s.logger.TCPStreamPropUpdate(s.info, true)
 	}
-	s.doneEntries = append(s.doneEntries, s.activeEntries...)
 	s.activeEntries = nil
 }

+// release 在最终判定发出后丢掉不会再用到的状态。
+//
+// 判定之后 Accept 只看 lastVerdict，分析器和属性表都不会再被读。但这条流本身还要在
+// 连接池里挂到超时被 flush，而 gopacket 回收连接槽位时并不清掉 stream 指针，
+// 要等下一条连接复用那个槽位才会覆盖。不主动放手的话，每条判定完的流会带着
+// 全部分析器对象一直活到那时候；节点上的 heap profile 里这部分占到了八成。
+func (s *tcpStream) release() {
+	s.released = true
+	s.activeEntries = nil
+	s.info.Props = nil
+	s.ruleset = nil
+}
+
 func (s *tcpStream) feedEntry(entry *tcpStreamEntry, rev, start, end bool, skip int, data []byte) (update *analyzer.PropUpdate, closeUpdate *analyzer.PropUpdate, done bool) {
 	if !entry.HasLimit {
 		update, done = entry.Stream.Feed(rev, start, end, skip, data)
@@ -287,7 +312,7 @@ func (s *tcpStream) feedEntry(entry *tcpStreamEntry, rev, start, end bool, skip
 		update, done = entry.Stream.Feed(rev, start, end, skip, qData)
 		entry.Quota -= len(qData)
 		if entry.Quota <= 0 {
-			// Quota exhausted, signal close & move to doneEntries
+			// Quota exhausted, signal close & drop the entry
 			closeUpdate = entry.Stream.Close(true)
 			done = true
 		}
@@ -160,7 +160,6 @@ type udpStream struct {
 	logger        Logger
 	ruleset       ruleset.Ruleset
 	activeEntries []*udpStreamEntry
-	doneEntries   []*udpStreamEntry
 	lastVerdict   udpVerdict
 }

@@ -194,7 +193,6 @@ func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) {
 		updated = updated || up1 || up2
 		if done {
 			s.activeEntries = append(s.activeEntries[:i], s.activeEntries[i+1:]...)
-			s.doneEntries = append(s.doneEntries, entry)
 		}
 	}
 	if updated || s.virgin {
@@ -236,14 +234,19 @@ func (s *udpStream) Feed(udp *layers.UDP, rev bool, uc *udpContext) {
 		uc.Verdict = udpVerdictAcceptStream
 		s.logger.UDPStreamAction(s.info, ruleset.ActionAllow, true)
 	}
+	if len(s.activeEntries) == 0 {
+		// 同 tcpStream.release：这条流会在 LRU 里一直待到被挤出去，别让它带着分析器一起等
+		s.release()
+	}
 }

 func (s *udpStream) Close() {
 	s.closeActiveEntries()
+	s.release()
 }

 func (s *udpStream) closeActiveEntries() {
-	// Signal close to all active entries & move them to doneEntries
+	// Signal close to all active entries
 	updated := false
 	for _, entry := range s.activeEntries {
 		update := entry.Stream.Close(false)
@@ -253,16 +256,22 @@ func (s *udpStream) closeActiveEntries() {
 	if updated {
 		s.logger.UDPStreamPropUpdate(s.info, true)
 	}
-	s.doneEntries = append(s.doneEntries, s.activeEntries...)
 	s.activeEntries = nil
 }

+// release 在最终判定发出后丢掉不会再用到的状态，Accept 此后只看 lastVerdict
+func (s *udpStream) release() {
+	s.activeEntries = nil
+	s.info.Props = nil
+	s.ruleset = nil
+}
+
 func (s *udpStream) feedEntry(entry *udpStreamEntry, rev bool, data []byte) (update *analyzer.PropUpdate, closeUpdate *analyzer.PropUpdate, done bool) {
 	update, done = entry.Stream.Feed(rev, data)
 	if entry.HasLimit {
 		entry.Quota -= len(data)
 		if entry.Quota <= 0 {
-			// Quota exhausted, signal close & move to doneEntries
+			// Quota exhausted, signal close & drop the entry
 			closeUpdate = entry.Stream.Close(true)
 			done = true
 		}
@@ -20,9 +20,18 @@ const (
 	defaultTCPTimeout                       = 10 * time.Minute
 	defaultUDPMaxStreams                    = 4096

-	tcpFlushInterval = 1 * time.Minute
+	// 超时 flush 的间隔范围。判定后的流走 connmark 旁路，FIN 到不了用户态，
+	// 只能靠定时 flush 把它们从连接池里清掉。间隔取超时的一半：
+	// 之前固定 1 分钟，30 秒的超时实际会让流在池里活到 90 秒。
+	tcpMaxFlushInterval = 1 * time.Minute
+	tcpMinFlushInterval = 1 * time.Second
 )

+// tcpFlushIntervalFor 返回给定 TCP 超时下的 flush 间隔
+func tcpFlushIntervalFor(timeout time.Duration) time.Duration {
+	return min(tcpMaxFlushInterval, max(timeout/2, tcpMinFlushInterval))
+}
+
 type workerPacket struct {
 	StreamID   uint32
 	Packet     gopacket.Packet
@@ -122,7 +131,7 @@ func (w *worker) Run(ctx context.Context) {
 	w.logger.WorkerStart(w.id)
 	defer w.logger.WorkerStop(w.id)

-	tcpFlushTicker := time.NewTicker(tcpFlushInterval)
+	tcpFlushTicker := time.NewTicker(tcpFlushIntervalFor(w.tcpTimeout))
 	defer tcpFlushTicker.Stop()

 	for {