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