From d58df54cd60264f76a6354545cd59b8fae3ab687 Mon Sep 17 00:00:00 2001
From: starainrt <i@b612.me>
Date: Thu, 28 Nov 2024 18:36:19 +0800
Subject: [PATCH] =?UTF-8?q?bug=20fix:=E5=B9=B6=E5=8F=91=E8=AF=BB=E5=86=99m?=
 =?UTF-8?q?ap=E9=97=AE=E9=A2=98;add:=E6=96=B0=E5=A2=9EStop()=E5=87=BD?=
 =?UTF-8?q?=E6=95=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 bcap.go            |  6 ++++++
 libpcap/libpcap.go | 20 ++++++++++++++++----
 nfq/nfqueue.go     |  6 ++++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/bcap.go b/bcap.go
index 264908b..1fd2e8b 100644
--- a/bcap.go
+++ b/bcap.go
@@ -40,6 +40,8 @@ type PacketInfo struct {
 }
 
 func (p *Packets) Key(key string) PacketInfo {
+	p.RLock()
+	defer p.RUnlock()
 	return p.cu[key]
 }
 
@@ -155,7 +157,9 @@ func (p *Packets) parseTcp(info PacketInfo, packet gopacket.Packet, layer gopack
 	info.finState = tcp.FIN
 	info.synState = tcp.SYN
 	info.tcpWindow = tcp.Window
+	p.RLock()
 	lastPacket := p.cu[info.Key]
+	p.RUnlock()
 	if lastPacket.Key != info.Key {
 		lastPacket = PacketInfo{
 			Key:           info.Key,
@@ -180,7 +184,9 @@ func (p *Packets) parseTcp(info PacketInfo, packet gopacket.Packet, layer gopack
 		p.cu[info.Key] = lastPacket
 		p.Unlock()
 	}
+	p.RLock()
 	lastReverse := p.cu[info.ReverseKey]
+	p.RUnlock()
 	if !lastPacket.isFirst {
 		info.comment = lastPacket.comment
 		if lastPacket.SrcMac != nil && len(info.SrcMac) == 0 {
diff --git a/libpcap/libpcap.go b/libpcap/libpcap.go
index 316f0c9..5f631d1 100644
--- a/libpcap/libpcap.go
+++ b/libpcap/libpcap.go
@@ -1,6 +1,7 @@
 package libpcap
 
 import (
+	"context"
 	"fmt"
 	"github.com/google/gopacket"
 	"github.com/google/gopacket/pcap"
@@ -12,12 +13,18 @@ type NetCatch struct {
 	sentence string
 	fn       func(gopacket.Packet)
 	*pcap.Handle
+	ctx    context.Context
+	stopFn context.CancelFunc
 }
 
 func (n *NetCatch) SetRecall(fn func(p gopacket.Packet)) {
 	n.fn = fn
 }
 
+func (n *NetCatch) Stop() {
+	n.stopFn()
+}
+
 func FindAllDevs() ([]pcap.Interface, error) {
 	return pcap.FindAllDevs()
 }
@@ -45,6 +52,7 @@ func NewCatch(host string, sentence string) (*NetCatch, error) {
 	nc.host = host
 	nc.eth = eth
 	nc.sentence = sentence
+	nc.ctx, nc.stopFn = context.WithCancel(context.Background())
 	return nc, nil
 }
 
@@ -68,10 +76,14 @@ func (n *NetCatch) Run() error {
 		return err
 	}
 	pks := gopacket.NewPacketSource(handle, handle.LinkType())
-	for packet := range pks.Packets() {
-		if n.fn != nil {
-			n.fn(packet)
+	for {
+		select {
+		case packet := <-pks.Packets():
+			if n.fn != nil {
+				n.fn(packet)
+			}
+		case <-n.ctx.Done():
+			return nil
 		}
 	}
-	return nil
 }
diff --git a/nfq/nfqueue.go b/nfq/nfqueue.go
index 1d71f3a..c8468d5 100644
--- a/nfq/nfqueue.go
+++ b/nfq/nfqueue.go
@@ -38,6 +38,12 @@ 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,