first commit

nfq/nfqueue.go

@@ -0,0 +1,85 @@
+package nfq
+
+import (
+	"context"
+	"fmt"
+	"github.com/florianl/go-nfqueue/v2"
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+	"time"
+)
+
+type NfQueue struct {
+	queid    uint16
+	maxqueue uint32
+	ctx      context.Context
+	stopFn   context.CancelFunc
+	recallFn func(id uint32, q *nfqueue.Nfqueue, p Packet)
+}
+
+type Packet struct {
+	Packet gopacket.Packet
+	Attr   nfqueue.Attribute
+}
+
+func NewNfQueue(ctx context.Context, queid uint16, maxqueue uint32) *NfQueue {
+	var q = new(NfQueue)
+	if ctx == nil {
+		q.ctx, q.stopFn = context.WithCancel(context.Background())
+	} else {
+		q.ctx, q.stopFn = context.WithCancel(ctx)
+	}
+	q.queid = queid
+	q.maxqueue = maxqueue
+	return q
+}
+
+func (n *NfQueue) SetRecall(fn func(id uint32, q *nfqueue.Nfqueue, p Packet)) {
+	n.recallFn = fn
+}
+
+func (n *NfQueue) Run() error {
+	cfg := nfqueue.Config{
+		NfQueue:      n.queid,
+		MaxQueueLen:  n.maxqueue,
+		Copymode:     nfqueue.NfQnlCopyPacket,
+		WriteTimeout: time.Second * 10,
+	}
+	nfq, err := nfqueue.Open(&cfg)
+	if err != nil {
+		return fmt.Errorf("failed to open nfqueue, err:", err)
+	}
+
+	if err := nfq.RegisterWithErrorFunc(n.ctx, func(a nfqueue.Attribute) int {
+		return n.handlePacket(nfq, a)
+	}, func(e error) int {
+		return 0
+	}); err != nil {
+		return fmt.Errorf("failed to register handlers, err:", err)
+	}
+	<-n.ctx.Done()
+	return nil
+}
+
+func (n *NfQueue) handlePacket(q *nfqueue.Nfqueue, a nfqueue.Attribute) int {
+	if a.Payload != nil && len(*a.Payload) != 0 {
+		var packet gopacket.Packet
+		data := *a.Payload
+		if data[0]&0xf0 == 0x40 {
+			packet = gopacket.NewPacket(data, layers.LayerTypeIPv4, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
+		} else {
+			packet = gopacket.NewPacket(data, layers.LayerTypeIPv6, gopacket.DecodeOptions{Lazy: true, NoCopy: true})
+		}
+		if n.recallFn != nil {
+			n.recallFn(*a.PacketID, q, Packet{
+				Packet: packet,
+				Attr:   a,
+			})
+		} else {
+			q.SetVerdict(*a.PacketID, nfqueue.NfAccept)
+		}
+		return 0
+	}
+	q.SetVerdict(*a.PacketID, nfqueue.NfAccept)
+	return 0
+}