Mellaris/engine/reload_rules_test.go

package engine

import (
	"net"
	"testing"

	"git.difuse.io/Difuse/Mellaris/analyzer"
	"git.difuse.io/Difuse/Mellaris/ruleset"

	"github.com/bwmarrin/snowflake"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/reassembly"
)

type fixedRuleset struct {
	action ruleset.Action
}

func (r fixedRuleset) Analyzers(ruleset.StreamInfo) []analyzer.Analyzer {
	return nil
}

func (r fixedRuleset) Match(ruleset.StreamInfo) ruleset.MatchResult {
	return ruleset.MatchResult{Action: r.action}
}

type noopTestLogger struct{}

func (noopTestLogger) WorkerStart(int) {}
func (noopTestLogger) WorkerStop(int)  {}

func (noopTestLogger) TCPStreamNew(int, ruleset.StreamInfo) {}
func (noopTestLogger) TCPStreamPropUpdate(ruleset.StreamInfo, bool) {
}
func (noopTestLogger) TCPStreamAction(ruleset.StreamInfo, ruleset.Action, bool) {
}

func (noopTestLogger) UDPStreamNew(int, ruleset.StreamInfo) {}
func (noopTestLogger) UDPStreamPropUpdate(ruleset.StreamInfo, bool) {
}
func (noopTestLogger) UDPStreamAction(ruleset.StreamInfo, ruleset.Action, bool) {
}

func (noopTestLogger) ModifyError(ruleset.StreamInfo, error) {}

func (noopTestLogger) AnalyzerDebugf(int64, string, string, ...interface{}) {}
func (noopTestLogger) AnalyzerInfof(int64, string, string, ...interface{})  {}
func (noopTestLogger) AnalyzerErrorf(int64, string, string, ...interface{}) {}

func TestUDPStreamUsesUpdatedRuleset(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"),
		},
	}
	ctx := &udpContext{Verdict: udpVerdictAccept}
	s := f.New(ipFlow, udp.TransportFlow(), udp, ctx)

	if err := f.UpdateRuleset(fixedRuleset{action: ruleset.ActionBlock}); err != nil {
		t.Fatalf("update ruleset: %v", err)
	}

	if !s.Accept(udp, false, ctx) {
		t.Fatalf("unexpected Accept=false for virgin stream")
	}
	s.Feed(udp, false, ctx)
	if ctx.Verdict != udpVerdictDropStream {
		t.Fatalf("verdict=%v want=%v", ctx.Verdict, udpVerdictDropStream)
	}
}

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 {
		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,
	}
	ctx := &tcpContext{
		PacketMetadata: &gopacket.PacketMetadata{},
		Verdict:        tcpVerdictAccept,
	}
	rs := f.New(ipFlow, tcp.TransportFlow(), tcp, ctx)
	s, ok := rs.(*tcpStream)
	if !ok {
		t.Fatalf("unexpected stream type %T", rs)
	}

	if err := f.UpdateRuleset(fixedRuleset{action: ruleset.ActionBlock}); err != nil {
		t.Fatalf("update ruleset: %v", err)
	}

	s.ReassembledSG(fakeScatterGather{data: []byte("payload")}, ctx)
	if ctx.Verdict != tcpVerdictDropStream {
		t.Fatalf("verdict=%v want=%v", ctx.Verdict, tcpVerdictDropStream)
	}
}

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
}

func (s fakeScatterGather) Lengths() (int, int) {
	return len(s.data), 0
}

func (s fakeScatterGather) Fetch(length int) []byte {
	if length < 0 {
		return nil
	}
	if length > len(s.data) {
		length = len(s.data)
	}
	return s.data[:length]
}

func (fakeScatterGather) KeepFrom(int) {}

func (fakeScatterGather) CaptureInfo(int) gopacket.CaptureInfo {
	return gopacket.CaptureInfo{}
}

func (fakeScatterGather) Info() (reassembly.TCPFlowDirection, bool, bool, int) {
	return reassembly.TCPDirClientToServer, true, false, 0
}

func (fakeScatterGather) Stats() reassembly.TCPAssemblyStats {
	return reassembly.TCPAssemblyStats{}
}