bcap/nfq/nfqueue.go

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) Stop() {
	if n.stopFn != nil {
		n.stopFn()
	}
}

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
}