mac resolution

engine/engine.go

@@ -2,6 +2,7 @@ package engine
 
 import (
 	"context"
+	"encoding/binary"
 	"runtime"
 
 	"git.difuse.io/Difuse/Mellaris/io"
@@ -24,6 +25,7 @@ func NewEngine(config Config) (Engine, error) {
 	if workerCount <= 0 {
 		workerCount = runtime.NumCPU()
 	}
+	macResolver := newSourceMACResolver()
 	var err error
 	workers := make([]*worker, workerCount)
 	for i := range workers {
@@ -32,6 +34,7 @@ func NewEngine(config Config) (Engine, error) {
 			ChanSize:                   config.WorkerQueueSize,
 			Logger:                     config.Logger,
 			Ruleset:                    config.Ruleset,
+			MACResolver:                macResolver,
 			TCPMaxBufferedPagesTotal:   config.WorkerTCPMaxBufferedPagesTotal,
 			TCPMaxBufferedPagesPerConn: config.WorkerTCPMaxBufferedPagesPerConn,
 			UDPMaxStreams:              config.WorkerUDPMaxStreams,
@@ -90,13 +93,8 @@ func (e *engine) Run(ctx context.Context) error {
 // dispatch dispatches a packet to a worker.
 func (e *engine) dispatch(p io.Packet) bool {
 	data := p.Data()
-	ipVersion := data[0] >> 4
-	var layerType gopacket.LayerType
-	if ipVersion == 4 {
-		layerType = layers.LayerTypeIPv4
-	} else if ipVersion == 6 {
-		layerType = layers.LayerTypeIPv6
-	} else {
+	layerType, srcMAC, dstMAC, ok := classifyPacket(data)
+	if !ok {
 		// Unsupported network layer
 		_ = e.io.SetVerdict(p, io.VerdictAcceptStream, nil)
 		return true
@@ -107,9 +105,41 @@ func (e *engine) dispatch(p io.Packet) bool {
 	e.workers[index].Feed(&workerPacket{
 		StreamID: p.StreamID(),
 		Packet:   packet,
+		SrcMAC:   srcMAC,
+		DstMAC:   dstMAC,
 		SetVerdict: func(v io.Verdict, b []byte) error {
 			return e.io.SetVerdict(p, v, b)
 		},
 	})
 	return true
 }
+
+// classifyPacket detects packet framing and returns a gopacket decode layer
+// plus best-effort source/destination MAC addresses when available.
+func classifyPacket(data []byte) (gopacket.LayerType, []byte, []byte, bool) {
+	if len(data) == 0 {
+		return 0, nil, nil, false
+	}
+
+	// Fast path for IP packets (NFQUEUE payloads are typically IP-only).
+	ipVersion := data[0] >> 4
+	if ipVersion == 4 {
+		return layers.LayerTypeIPv4, nil, nil, true
+	}
+	if ipVersion == 6 {
+		return layers.LayerTypeIPv6, nil, nil, true
+	}
+
+	// Ethernet frame path (for custom PacketIO implementations).
+	if len(data) >= 14 {
+		etherType := binary.BigEndian.Uint16(data[12:14])
+		if etherType == uint16(layers.EthernetTypeIPv4) || etherType == uint16(layers.EthernetTypeIPv6) {
+			return layers.LayerTypeEthernet,
+				append([]byte(nil), data[6:12]...),
+				append([]byte(nil), data[:6]...),
+				true
+		}
+	}
+
+	return 0, nil, nil, false
+}