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,