ruleset: try to fix reloader

2026-02-12 13:34:37 +05:30
parent beaaddad2b
commit e1c68ec7d0
3 changed files with 195 additions and 46 deletions
--- a/engine/reload_rules_test.go
+++ b/engine/reload_rules_test.go
@@ -84,6 +84,59 @@ func TestUDPStreamUsesUpdatedRuleset(t *testing.T) {
 	}
 }

+func TestUDPStreamReevaluatesAfterRulesetVersionChange(t *testing.T) {
+	node, err := snowflake.NewNode(0)
+	if err != nil {
+		t.Fatalf("create node: %v", err)
+	}
+	f := &udpStreamFactory{
+		WorkerID: 0,
+		Logger:   noopTestLogger{},
+		Node:     node,
+		Ruleset:  fixedRuleset{action: ruleset.ActionAllow},
+	}
+
+	ipFlow := gopacket.NewFlow(layers.EndpointIPv4, net.IPv4(10, 0, 0, 1).To4(), net.IPv4(10, 0, 0, 2).To4())
+	udp := &layers.UDP{
+		SrcPort: 12345,
+		DstPort: 53,
+		BaseLayer: layers.BaseLayer{
+			Payload: []byte("query"),
+		},
+	}
+
+	ctx1 := &udpContext{Verdict: udpVerdictAccept}
+	s := f.New(ipFlow, udp.TransportFlow(), udp, ctx1)
+	if !s.Accept(udp, false, ctx1) {
+		t.Fatalf("unexpected Accept=false before first feed")
+	}
+	s.Feed(udp, false, ctx1)
+	if ctx1.Verdict != udpVerdictAcceptStream {
+		t.Fatalf("verdict=%v want=%v", ctx1.Verdict, udpVerdictAcceptStream)
+	}
+
+	if err := f.UpdateRuleset(fixedRuleset{action: ruleset.ActionBlock}); err != nil {
+		t.Fatalf("update ruleset: %v", err)
+	}
+
+	ctx2 := &udpContext{Verdict: udpVerdictAccept}
+	if !s.Accept(udp, false, ctx2) {
+		t.Fatalf("expected Accept=true after ruleset update")
+	}
+	s.Feed(udp, false, ctx2)
+	if ctx2.Verdict != udpVerdictDropStream {
+		t.Fatalf("verdict=%v want=%v", ctx2.Verdict, udpVerdictDropStream)
+	}
+
+	ctx3 := &udpContext{Verdict: udpVerdictAccept}
+	if s.Accept(udp, false, ctx3) {
+		t.Fatalf("expected Accept=false with unchanged ruleset and no active entries")
+	}
+	if ctx3.Verdict != udpVerdictDropStream {
+		t.Fatalf("verdict=%v want=%v", ctx3.Verdict, udpVerdictDropStream)
+	}
+}
+
 func TestTCPStreamUsesUpdatedRuleset(t *testing.T) {
 	node, err := snowflake.NewNode(0)
 	if err != nil {
@@ -121,6 +174,72 @@ func TestTCPStreamUsesUpdatedRuleset(t *testing.T) {
 	}
 }

+func TestTCPStreamReevaluatesAfterRulesetVersionChange(t *testing.T) {
+	node, err := snowflake.NewNode(0)
+	if err != nil {
+		t.Fatalf("create node: %v", err)
+	}
+	f := &tcpStreamFactory{
+		WorkerID: 0,
+		Logger:   noopTestLogger{},
+		Node:     node,
+		Ruleset:  fixedRuleset{action: ruleset.ActionAllow},
+	}
+
+	ipFlow := gopacket.NewFlow(layers.EndpointIPv4, net.IPv4(10, 0, 0, 1).To4(), net.IPv4(10, 0, 0, 2).To4())
+	tcp := &layers.TCP{
+		SrcPort: 12345,
+		DstPort: 443,
+	}
+	ctx1 := &tcpContext{
+		PacketMetadata: &gopacket.PacketMetadata{},
+		Verdict:        tcpVerdictAccept,
+	}
+	rs := f.New(ipFlow, tcp.TransportFlow(), tcp, ctx1)
+	s, ok := rs.(*tcpStream)
+	if !ok {
+		t.Fatalf("unexpected stream type %T", rs)
+	}
+
+	start1 := false
+	if !s.Accept(tcp, gopacket.CaptureInfo{}, reassembly.TCPDirClientToServer, 0, &start1, ctx1) {
+		t.Fatalf("unexpected Accept=false before first feed")
+	}
+	s.ReassembledSG(fakeScatterGather{data: []byte("first")}, ctx1)
+	if ctx1.Verdict != tcpVerdictAcceptStream {
+		t.Fatalf("verdict=%v want=%v", ctx1.Verdict, tcpVerdictAcceptStream)
+	}
+
+	if err := f.UpdateRuleset(fixedRuleset{action: ruleset.ActionBlock}); err != nil {
+		t.Fatalf("update ruleset: %v", err)
+	}
+
+	ctx2 := &tcpContext{
+		PacketMetadata: &gopacket.PacketMetadata{},
+		Verdict:        tcpVerdictAccept,
+	}
+	start2 := false
+	if !s.Accept(tcp, gopacket.CaptureInfo{}, reassembly.TCPDirClientToServer, 0, &start2, ctx2) {
+		t.Fatalf("expected Accept=true after ruleset update")
+	}
+	s.ReassembledSG(fakeScatterGather{data: []byte("second")}, ctx2)
+	if ctx2.Verdict != tcpVerdictDropStream {
+		t.Fatalf("verdict=%v want=%v", ctx2.Verdict, tcpVerdictDropStream)
+	}
+
+	ctx3 := &tcpContext{
+		PacketMetadata: &gopacket.PacketMetadata{},
+		Verdict:        tcpVerdictAccept,
+	}
+	start3 := false
+	if s.Accept(tcp, gopacket.CaptureInfo{}, reassembly.TCPDirClientToServer, 0, &start3, ctx3) {
+		t.Fatalf("expected Accept=false with unchanged ruleset and no active entries")
+	}
+	if ctx3.Verdict != tcpVerdictDropStream {
+		t.Fatalf("verdict=%v want=%v", ctx3.Verdict, tcpVerdictDropStream)
+	}
+}
+
 type fakeScatterGather struct {
 	data []byte
 }