mac resolution

engine/worker.go

@@ -2,6 +2,7 @@ package engine
 
 import (
 	"context"
+	"net"
 
 	"git.difuse.io/Difuse/Mellaris/io"
 	"git.difuse.io/Difuse/Mellaris/ruleset"
@@ -22,13 +23,16 @@ const (
 type workerPacket struct {
 	StreamID   uint32
 	Packet     gopacket.Packet
+	SrcMAC     net.HardwareAddr
+	DstMAC     net.HardwareAddr
 	SetVerdict func(io.Verdict, []byte) error
 }
 
 type worker struct {
-	id         int
-	packetChan chan *workerPacket
-	logger     Logger
+	id          int
+	packetChan  chan *workerPacket
+	logger      Logger
+	macResolver *sourceMACResolver
 
 	tcpStreamFactory *tcpStreamFactory
 	tcpStreamPool    *reassembly.StreamPool
@@ -45,6 +49,7 @@ type workerConfig struct {
 	ChanSize                   int
 	Logger                     Logger
 	Ruleset                    ruleset.Ruleset
+	MACResolver                *sourceMACResolver
 	TCPMaxBufferedPagesTotal   int
 	TCPMaxBufferedPagesPerConn int
 	UDPMaxStreams              int
@@ -95,6 +100,7 @@ func newWorker(config workerConfig) (*worker, error) {
 		id:                 config.ID,
 		packetChan:         make(chan *workerPacket, config.ChanSize),
 		logger:             config.Logger,
+		macResolver:        config.MACResolver,
 		tcpStreamFactory:   tcpSF,
 		tcpStreamPool:      tcpStreamPool,
 		tcpAssembler:       tcpAssembler,
@@ -120,7 +126,7 @@ func (w *worker) Run(ctx context.Context) {
 				// Closed
 				return
 			}
-			v, b := w.handle(wPkt.StreamID, wPkt.Packet)
+			v, b := w.handle(wPkt.StreamID, wPkt.Packet, wPkt.SrcMAC, wPkt.DstMAC)
 			_ = wPkt.SetVerdict(v, b)
 		}
 	}
@@ -133,18 +139,21 @@ func (w *worker) UpdateRuleset(r ruleset.Ruleset) error {
 	return w.udpStreamFactory.UpdateRuleset(r)
 }
 
-func (w *worker) handle(streamID uint32, p gopacket.Packet) (io.Verdict, []byte) {
+func (w *worker) handle(streamID uint32, p gopacket.Packet, srcMAC, dstMAC net.HardwareAddr) (io.Verdict, []byte) {
 	netLayer, trLayer := p.NetworkLayer(), p.TransportLayer()
 	if netLayer == nil || trLayer == nil {
 		// Invalid packet
 		return io.VerdictAccept, nil
 	}
 	ipFlow := netLayer.NetworkFlow()
+	if len(srcMAC) == 0 && w.macResolver != nil {
+		srcMAC = w.macResolver.Resolve(net.IP(ipFlow.Src().Raw()))
+	}
 	switch tr := trLayer.(type) {
 	case *layers.TCP:
-		return w.handleTCP(ipFlow, p.Metadata(), tr), nil
+		return w.handleTCP(ipFlow, srcMAC, dstMAC, p.Metadata(), tr), nil
 	case *layers.UDP:
-		v, modPayload := w.handleUDP(streamID, ipFlow, tr)
+		v, modPayload := w.handleUDP(streamID, ipFlow, srcMAC, dstMAC, tr)
 		if v == io.VerdictAcceptModify && modPayload != nil {
 			tr.Payload = modPayload
 			_ = tr.SetNetworkLayerForChecksum(netLayer)
@@ -167,18 +176,22 @@ func (w *worker) handle(streamID uint32, p gopacket.Packet) (io.Verdict, []byte)
 	}
 }
 
-func (w *worker) handleTCP(ipFlow gopacket.Flow, pMeta *gopacket.PacketMetadata, tcp *layers.TCP) io.Verdict {
+func (w *worker) handleTCP(ipFlow gopacket.Flow, srcMAC, dstMAC net.HardwareAddr, pMeta *gopacket.PacketMetadata, tcp *layers.TCP) io.Verdict {
 	ctx := &tcpContext{
 		PacketMetadata: pMeta,
 		Verdict:        tcpVerdictAccept,
+		SrcMAC:         srcMAC,
+		DstMAC:         dstMAC,
 	}
 	w.tcpAssembler.AssembleWithContext(ipFlow, tcp, ctx)
 	return io.Verdict(ctx.Verdict)
 }
 
-func (w *worker) handleUDP(streamID uint32, ipFlow gopacket.Flow, udp *layers.UDP) (io.Verdict, []byte) {
+func (w *worker) handleUDP(streamID uint32, ipFlow gopacket.Flow, srcMAC, dstMAC net.HardwareAddr, udp *layers.UDP) (io.Verdict, []byte) {
 	ctx := &udpContext{
 		Verdict: udpVerdictAccept,
+		SrcMAC:  srcMAC,
+		DstMAC:  dstMAC,
 	}
 	w.udpStreamManager.MatchWithContext(streamID, ipFlow, udp, ctx)
 	return io.Verdict(ctx.Verdict), ctx.Packet