yoloyolo/classify/classifier.go

package classify

import (
	"github.com/gopacket/gopacket"
	"github.com/gopacket/gopacket/layers"
)

// Classifier classifies packets against a rule table.
type Classifier struct {
	rules []Rule
}

// NewClassifier creates a classifier with the given rules.
func NewClassifier(rules []Rule) *Classifier {
	return &Classifier{rules: rules}
}

// Classify inspects a gopacket.Packet and returns a ClassifiedPacket.
// It checks ICMP first, then extracts TCP/UDP port info and matches against rules.
// Per D-03: returns ClassUnknown1-4 via hash-bucketing if no rule matches.
func (c *Classifier) Classify(pkt gopacket.Packet) ClassifiedPacket {
	result := ClassifiedPacket{
		Length: len(pkt.Data()),
	}

	// Check ICMP
	if pkt.Layer(layers.LayerTypeICMPv4) != nil || pkt.Layer(layers.LayerTypeICMPv6) != nil {
		result.Protocol = "icmp"
		// Match against rules
		for _, rule := range c.rules {
			if rule.Protocol == "icmp" {
				result.Class = rule.Class
				return result
			}
		}
		result.Class = hashBucket(result.DstPort, result.Protocol)
		return result
	}

	// Check TCP
	if tcpLayer := pkt.Layer(layers.LayerTypeTCP); tcpLayer != nil {
		tcp := tcpLayer.(*layers.TCP)
		result.Protocol = "tcp"
		result.SrcPort = uint16(tcp.SrcPort)
		result.DstPort = uint16(tcp.DstPort)
		for _, rule := range c.rules {
			if rule.Protocol == "tcp" && (rule.DstPort == 0 || rule.DstPort == result.DstPort) {
				result.Class = rule.Class
				return result
			}
		}
		result.Class = hashBucket(result.DstPort, result.Protocol)
		return result
	}

	// Check UDP
	if udpLayer := pkt.Layer(layers.LayerTypeUDP); udpLayer != nil {
		udp := udpLayer.(*layers.UDP)
		result.Protocol = "udp"
		result.SrcPort = uint16(udp.SrcPort)
		result.DstPort = uint16(udp.DstPort)
		for _, rule := range c.rules {
			if rule.Protocol == "udp" && (rule.DstPort == 0 || rule.DstPort == result.DstPort) {
				result.Class = rule.Class
				return result
			}
		}
		result.Class = hashBucket(result.DstPort, result.Protocol)
		return result
	}

	// No recognized transport layer -> hash-bucket per D-03
	result.Class = hashBucket(result.DstPort, result.Protocol)
	return result
}

// hashBucket maps an unrecognized packet to one of 4 unknown traffic classes.
// Deterministic: same (dstPort, protocol) always maps to the same bucket.
// ARP and other non-transport packets have dstPort=0, protocol="" -> bucket 0 (ClassUnknown1).
func hashBucket(dstPort uint16, protocol string) TrafficClass {
	var protoNum uint16
	switch protocol {
	case "tcp":
		protoNum = 6
	case "udp":
		protoNum = 17
	case "icmp":
		protoNum = 1
	}
	h := uint32(dstPort)*31 + uint32(protoNum)*7
	switch h % 4 {
	case 0:
		return ClassUnknown1
	case 1:
		return ClassUnknown2
	case 2:
		return ClassUnknown3
	default:
		return ClassUnknown4
	}
}