yoloyolo/classify/classifier_test.go

package classify_test

import (
	"testing"

	"github.com/gopacket/gopacket"
	"github.com/gopacket/gopacket/layers"
	"github.com/netsynth/netsynth/classify"
)

// buildTCPPacket builds a synthetic IPv4/TCP packet targeting the given dst port.
func buildTCPPacket(t *testing.T, dstPort uint16) gopacket.Packet {
	t.Helper()
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: false}

	eth := &layers.Ethernet{
		SrcMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x01},
		DstMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x02},
		EthernetType: layers.EthernetTypeIPv4,
	}
	ip := &layers.IPv4{
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolTCP,
		SrcIP:    []byte{192, 168, 1, 1},
		DstIP:    []byte{192, 168, 1, 2},
	}
	tcp := &layers.TCP{
		SrcPort: layers.TCPPort(54321),
		DstPort: layers.TCPPort(dstPort),
		SYN:     true,
	}
	if err := tcp.SetNetworkLayerForChecksum(ip); err != nil {
		t.Fatalf("SetNetworkLayerForChecksum: %v", err)
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, ip, tcp, gopacket.Payload{}); err != nil {
		t.Fatalf("SerializeLayers tcp: %v", err)
	}
	return gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
}

// buildUDPPacket builds a synthetic IPv4/UDP packet targeting the given dst port.
func buildUDPPacket(t *testing.T, dstPort uint16) gopacket.Packet {
	t.Helper()
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: false}

	eth := &layers.Ethernet{
		SrcMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x01},
		DstMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x02},
		EthernetType: layers.EthernetTypeIPv4,
	}
	ip := &layers.IPv4{
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolUDP,
		SrcIP:    []byte{192, 168, 1, 1},
		DstIP:    []byte{192, 168, 1, 2},
	}
	udp := &layers.UDP{
		SrcPort: layers.UDPPort(54321),
		DstPort: layers.UDPPort(dstPort),
	}
	if err := udp.SetNetworkLayerForChecksum(ip); err != nil {
		t.Fatalf("SetNetworkLayerForChecksum: %v", err)
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, ip, udp, gopacket.Payload{}); err != nil {
		t.Fatalf("SerializeLayers udp: %v", err)
	}
	return gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
}

// buildICMPPacket builds a synthetic IPv4/ICMPv4 packet.
func buildICMPPacket(t *testing.T) gopacket.Packet {
	t.Helper()
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: false}

	eth := &layers.Ethernet{
		SrcMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x01},
		DstMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x02},
		EthernetType: layers.EthernetTypeIPv4,
	}
	ip := &layers.IPv4{
		Version:  4,
		TTL:      64,
		Protocol: layers.IPProtocolICMPv4,
		SrcIP:    []byte{192, 168, 1, 1},
		DstIP:    []byte{192, 168, 1, 2},
	}
	icmp := &layers.ICMPv4{
		TypeCode: layers.CreateICMPv4TypeCode(layers.ICMPv4TypeEchoRequest, 0),
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, ip, icmp, gopacket.Payload{}); err != nil {
		t.Fatalf("SerializeLayers icmp: %v", err)
	}
	return gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
}

// buildUnknownPacket builds a synthetic Ethernet packet with no TCP/UDP/ICMP payload.
func buildUnknownPacket(t *testing.T) gopacket.Packet {
	t.Helper()
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{FixLengths: true}

	eth := &layers.Ethernet{
		SrcMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x01},
		DstMAC:       []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x02},
		EthernetType: layers.EthernetTypeARP,
	}
	if err := gopacket.SerializeLayers(buf, opts, eth, gopacket.Payload{0x01, 0x02, 0x03, 0x04}); err != nil {
		t.Fatalf("SerializeLayers unknown: %v", err)
	}
	return gopacket.NewPacket(buf.Bytes(), layers.LayerTypeEthernet, gopacket.Default)
}

func TestClassify(t *testing.T) {
	c := classify.NewClassifier(classify.DefaultRules)

	t.Run("TestClassifyICMP", func(t *testing.T) {
		pkt := buildICMPPacket(t)
		got := c.Classify(pkt)
		if got.Class != classify.ClassICMP {
			t.Errorf("ICMP packet: got class %q, want %q", got.Class, classify.ClassICMP)
		}
		if got.Protocol != "icmp" {
			t.Errorf("ICMP packet: got protocol %q, want %q", got.Protocol, "icmp")
		}
	})

	t.Run("TestClassifyDNS_UDP", func(t *testing.T) {
		pkt := buildUDPPacket(t, 53)
		got := c.Classify(pkt)
		if got.Class != classify.ClassDNS {
			t.Errorf("DNS/UDP packet: got class %q, want %q", got.Class, classify.ClassDNS)
		}
	})

	t.Run("TestClassifyDNS_TCP", func(t *testing.T) {
		pkt := buildTCPPacket(t, 53)
		got := c.Classify(pkt)
		if got.Class != classify.ClassDNS {
			t.Errorf("DNS/TCP packet: got class %q, want %q", got.Class, classify.ClassDNS)
		}
	})

	t.Run("TestClassifyHTTPS", func(t *testing.T) {
		pkt := buildTCPPacket(t, 443)
		got := c.Classify(pkt)
		if got.Class != classify.ClassHTTPS {
			t.Errorf("HTTPS packet: got class %q, want %q", got.Class, classify.ClassHTTPS)
		}
	})

	t.Run("TestClassifyHTTP", func(t *testing.T) {
		pkt := buildTCPPacket(t, 80)
		got := c.Classify(pkt)
		if got.Class != classify.ClassHTTP {
			t.Errorf("HTTP packet: got class %q, want %q", got.Class, classify.ClassHTTP)
		}
	})

	t.Run("TestClassifySSH", func(t *testing.T) {
		pkt := buildTCPPacket(t, 22)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSSH {
			t.Errorf("SSH packet: got class %q, want %q", got.Class, classify.ClassSSH)
		}
	})

	t.Run("TestClassifySMTP", func(t *testing.T) {
		pkt := buildTCPPacket(t, 25)
		got := c.Classify(pkt)
		if got.Class != classify.ClassSMTP {
			t.Errorf("SMTP packet: got class %q, want %q", got.Class, classify.ClassSMTP)
		}
	})

	t.Run("TestClassifyNTP", func(t *testing.T) {
		pkt := buildUDPPacket(t, 123)
		got := c.Classify(pkt)
		if got.Class != classify.ClassNTP {
			t.Errorf("NTP packet: got class %q, want %q", got.Class, classify.ClassNTP)
		}
	})

	t.Run("TestClassifyDHCP_port67", func(t *testing.T) {
		pkt := buildUDPPacket(t, 67)
		got := c.Classify(pkt)
		if got.Class != classify.ClassDHCP {
			t.Errorf("DHCP port 67 packet: got class %q, want %q", got.Class, classify.ClassDHCP)
		}
	})

	t.Run("TestClassifyDHCP_port68", func(t *testing.T) {
		pkt := buildUDPPacket(t, 68)
		got := c.Classify(pkt)
		if got.Class != classify.ClassDHCP {
			t.Errorf("DHCP port 68 packet: got class %q, want %q", got.Class, classify.ClassDHCP)
		}
	})

	t.Run("TestClassifyOtherTCP", func(t *testing.T) {
		pkt := buildTCPPacket(t, 8080)
		got := c.Classify(pkt)
		if got.Class != classify.ClassOtherTCP {
			t.Errorf("OtherTCP packet (port 8080): got class %q, want %q", got.Class, classify.ClassOtherTCP)
		}
	})

	t.Run("TestClassifyOtherUDP", func(t *testing.T) {
		pkt := buildUDPPacket(t, 9999)
		got := c.Classify(pkt)
		if got.Class != classify.ClassOtherUDP {
			t.Errorf("OtherUDP packet (port 9999): got class %q, want %q", got.Class, classify.ClassOtherUDP)
		}
	})

	t.Run("TestClassifyUnknown", func(t *testing.T) {
		pkt := buildUnknownPacket(t)
		got := c.Classify(pkt)
		if got.Class != classify.ClassUnknown {
			t.Errorf("Unknown packet: got class %q, want %q", got.Class, classify.ClassUnknown)
		}
	})

	t.Run("TestRulesAreOrderDependent", func(t *testing.T) {
		// A rule list where TCP port 443 maps to ClassHTTP (wrong) placed first,
		// then ClassHTTPS. The first matching rule should win.
		customRules := []classify.Rule{
			{Protocol: "tcp", DstPort: 443, Class: classify.ClassHTTP}, // first match
			{Protocol: "tcp", DstPort: 443, Class: classify.ClassHTTPS},
			{Protocol: "tcp", DstPort: 0, Class: classify.ClassOtherTCP},
		}
		custom := classify.NewClassifier(customRules)
		pkt := buildTCPPacket(t, 443)
		got := custom.Classify(pkt)
		if got.Class != classify.ClassHTTP {
			t.Errorf("Order-dependent rules: got class %q, want %q (first match should win)", got.Class, classify.ClassHTTP)
		}
	})
}