flows: implement ipv

2026-05-14 09:41:07 +05:30
parent 7a3f6e945d
commit 1e5c1dea75
5 changed files with 435 additions and 69 deletions
@@ -0,0 +1,114 @@
+package engine
+
+import (
+	"net"
+	"testing"
+
+	"git.difuse.io/Difuse/Mellaris/io"
+	"git.difuse.io/Difuse/Mellaris/ruleset"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+)
+
+func TestWorkerHandleIPv6TCP(t *testing.T) {
+	w, err := newWorker(workerConfig{
+		ID:         0,
+		Logger:     noopTestLogger{},
+		Ruleset:    fixedRuleset{action: ruleset.ActionBlock},
+		ResultChan: make(chan workerResult, 1),
+	})
+	if err != nil {
+		t.Fatalf("new worker: %v", err)
+	}
+
+	src := net.ParseIP("2001:db8::11").To16()
+	dst := net.ParseIP("2001:db8::22").To16()
+	data := serializeIPv6TCP(t, src, dst, 42310, 443, 1000)
+
+	v, _ := w.handle(&workerPacket{
+		StreamID: 11,
+		Data:     data,
+	})
+	if v != io.VerdictDropStream {
+		t.Fatalf("verdict=%v want=%v", v, io.VerdictDropStream)
+	}
+}
+
+func TestWorkerHandleIPv6UDP(t *testing.T) {
+	w, err := newWorker(workerConfig{
+		ID:         0,
+		Logger:     noopTestLogger{},
+		Ruleset:    fixedRuleset{action: ruleset.ActionBlock},
+		ResultChan: make(chan workerResult, 1),
+	})
+	if err != nil {
+		t.Fatalf("new worker: %v", err)
+	}
+
+	src := net.ParseIP("2001:db8::33").To16()
+	dst := net.ParseIP("2001:db8::44").To16()
+	data := serializeIPv6UDP(t, src, dst, 50000, 53, []byte("dns"))
+
+	v, _ := w.handle(&workerPacket{
+		StreamID: 12,
+		Data:     data,
+	})
+	if v != io.VerdictDropStream {
+		t.Fatalf("verdict=%v want=%v", v, io.VerdictDropStream)
+	}
+}
+
+func serializeIPv6TCP(t *testing.T, src, dst net.IP, srcPort, dstPort uint16, seq uint32) []byte {
+	t.Helper()
+	ip6 := &layers.IPv6{
+		Version:    6,
+		HopLimit:   64,
+		NextHeader: layers.IPProtocolTCP,
+		SrcIP:      src,
+		DstIP:      dst,
+	}
+	tcp := &layers.TCP{
+		SrcPort: layers.TCPPort(srcPort),
+		DstPort: layers.TCPPort(dstPort),
+		Seq:     seq,
+		SYN:     true,
+	}
+	if err := tcp.SetNetworkLayerForChecksum(ip6); err != nil {
+		t.Fatalf("set tcp checksum network layer: %v", err)
+	}
+	buf := gopacket.NewSerializeBuffer()
+	if err := gopacket.SerializeLayers(buf, gopacket.SerializeOptions{
+		FixLengths:       true,
+		ComputeChecksums: true,
+	}, ip6, tcp); err != nil {
+		t.Fatalf("serialize ipv6 tcp: %v", err)
+	}
+	return append([]byte(nil), buf.Bytes()...)
+}
+
+func serializeIPv6UDP(t *testing.T, src, dst net.IP, srcPort, dstPort uint16, payload []byte) []byte {
+	t.Helper()
+	ip6 := &layers.IPv6{
+		Version:    6,
+		HopLimit:   64,
+		NextHeader: layers.IPProtocolUDP,
+		SrcIP:      src,
+		DstIP:      dst,
+	}
+	udp := &layers.UDP{
+		SrcPort: layers.UDPPort(srcPort),
+		DstPort: layers.UDPPort(dstPort),
+	}
+	if err := udp.SetNetworkLayerForChecksum(ip6); err != nil {
+		t.Fatalf("set udp checksum network layer: %v", err)
+	}
+	buf := gopacket.NewSerializeBuffer()
+	if err := gopacket.SerializeLayers(buf, gopacket.SerializeOptions{
+		FixLengths:       true,
+		ComputeChecksums: true,
+	}, ip6, udp, gopacket.Payload(payload)); err != nil {
+		t.Fatalf("serialize ipv6 udp: %v", err)
+	}
+	return append([]byte(nil), buf.Bytes()...)
+}