flows: implement ipv

engine/worker.go

@@ -150,56 +150,67 @@ func (w *worker) handle(wp *workerPacket) (io.Verdict, []byte) {
 		return io.VerdictAccept, nil
 	}
 
-	ipVersion := data[0] >> 4
-	if ipVersion == 4 {
-		l3, transport, ok := ParseL3(data)
-		if !ok {
-			return io.VerdictAccept, nil
-		}
-		switch l3.Protocol {
-		case 6: // TCP
-			tcp, payload, ok := ParseTCP(transport)
-			if !ok {
-				return io.VerdictAccept, nil
-			}
-			verdict := w.tcpFlowMgr.handle(
-				wp.StreamID, l3, tcp, payload,
-				wp.SrcMAC, wp.DstMAC,
-			)
-			return verdict, nil
-
-		case 17: // UDP
-			udp, payload, ok := ParseUDP(transport)
-			if !ok {
-				return io.VerdictAccept, nil
-			}
-			v, modPayload := w.handleUDP(
-				wp.StreamID, l3, udp, payload,
-				wp.SrcMAC, wp.DstMAC,
-			)
-			if v == io.VerdictAcceptModify && modPayload != nil {
-				return w.serializeModifiedUDP(data, l3, udp, transport, modPayload)
-			}
-			return v, nil
-
-		default:
-			return io.VerdictAccept, nil
-		}
+	if v, b, ok := w.handleIPPacket(wp, data); ok {
+		return v, b
 	}
 
-	// Ethernet frame path (for custom PacketIO)
-	if ipVersion == 6 {
-		// TODO: IPv6 support with raw parsing
-		return io.VerdictAccept, nil
+	// Ethernet frame fallback path (for custom PacketIO implementations).
+	if l3Payload, ok := extractL3PayloadFromEthernet(data); ok {
+		if v, b, ok := w.handleIPPacket(wp, l3Payload); ok {
+			return v, b
+		}
 	}
 
 	return io.VerdictAccept, nil
 }
 
+func (w *worker) handleIPPacket(wp *workerPacket, data []byte) (io.Verdict, []byte, bool) {
+	l3, transport, ok := ParseL3(data)
+	if !ok {
+		return io.VerdictAccept, nil, false
+	}
+	switch l3.Protocol {
+	case 6: // TCP
+		tcp, payload, ok := ParseTCP(transport)
+		if !ok {
+			return io.VerdictAccept, nil, true
+		}
+		verdict := w.tcpFlowMgr.handle(
+			wp.StreamID, l3, tcp, payload,
+			wp.SrcMAC, wp.DstMAC,
+		)
+		return verdict, nil, true
+	case 17: // UDP
+		udp, payload, ok := ParseUDP(transport)
+		if !ok {
+			return io.VerdictAccept, nil, true
+		}
+		v, modPayload := w.handleUDP(
+			wp.StreamID, l3, udp, payload,
+			wp.SrcMAC, wp.DstMAC,
+		)
+		if v == io.VerdictAcceptModify && modPayload != nil {
+			mv, mb := w.serializeModifiedUDP(data, l3, modPayload)
+			return mv, mb, true
+		}
+		return v, nil, true
+	default:
+		return io.VerdictAccept, nil, true
+	}
+}
+
 func (w *worker) handleUDP(streamID uint32, l3 L3Info, udp UDPInfo, payload []byte, srcMAC, dstMAC net.HardwareAddr) (io.Verdict, []byte) {
-	ipSrc := net.IP(l3.SrcIP[:])
-	ipDst := net.IP(l3.DstIP[:])
-	ipFlow := gopacket.NewFlow(layers.EndpointIPv4, ipSrc.To4(), ipDst.To4())
+	ipSrc := l3.SrcIPAddr()
+	ipDst := l3.DstIPAddr()
+	endpointType := layers.EndpointIPv4
+	flowSrc := ipSrc.To4()
+	flowDst := ipDst.To4()
+	if l3.Version == 6 {
+		endpointType = layers.EndpointIPv6
+		flowSrc = ipSrc.To16()
+		flowDst = ipDst.To16()
+	}
+	ipFlow := gopacket.NewFlow(endpointType, flowSrc, flowDst)
 
 	if len(srcMAC) == 0 && w.macResolver != nil {
 		srcMAC = w.macResolver.Resolve(ipSrc)
@@ -219,8 +230,12 @@ func (w *worker) handleUDP(streamID uint32, l3 L3Info, udp UDPInfo, payload []by
 	return io.Verdict(uc.Verdict), uc.Packet
 }
 
-func (w *worker) serializeModifiedUDP(fullData []byte, l3 L3Info, udp UDPInfo, transport []byte, modPayload []byte) (io.Verdict, []byte) {
-	ipPkt := gopacket.NewPacket(fullData, layers.LayerTypeIPv4, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
+func (w *worker) serializeModifiedUDP(fullData []byte, l3 L3Info, modPayload []byte) (io.Verdict, []byte) {
+	layerType := layers.LayerTypeIPv4
+	if l3.Version == 6 {
+		layerType = layers.LayerTypeIPv6
+	}
+	ipPkt := gopacket.NewPacket(fullData, layerType, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
 	netLayer := ipPkt.NetworkLayer()
 	trLayer := ipPkt.TransportLayer()
 	if netLayer == nil || trLayer == nil {
@@ -240,3 +255,28 @@ func (w *worker) serializeModifiedUDP(fullData []byte, l3 L3Info, udp UDPInfo, t
 	}
 	return io.VerdictAcceptModify, w.modSerializeBuffer.Bytes()
 }
+
+func extractL3PayloadFromEthernet(data []byte) ([]byte, bool) {
+	if len(data) < 14 {
+		return nil, false
+	}
+	offset := 12
+	etherType := uint16(data[offset])<<8 | uint16(data[offset+1])
+	offset += 2
+
+	for etherType == 0x8100 || etherType == 0x88A8 {
+		if len(data) < offset+4 {
+			return nil, false
+		}
+		etherType = uint16(data[offset+2])<<8 | uint16(data[offset+3])
+		offset += 4
+	}
+
+	if etherType != 0x0800 && etherType != 0x86DD {
+		return nil, false
+	}
+	if len(data) <= offset {
+		return nil, false
+	}
+	return data[offset:], true
+}