mirror of
https://github.com/bettercap/bettercap
synced 2025-07-11 07:37:00 -07:00
fix: using sync.Map to avoid race conditions on the packets.Queue
This commit is contained in:
parent
b676d68b4c
commit
64a5ce2b58
4 changed files with 34 additions and 71 deletions
|
@ -61,8 +61,11 @@ func (mod *Discovery) getRow(e *network.Endpoint, withMeta bool) [][]string {
|
||||||
|
|
||||||
var traffic *packets.Traffic
|
var traffic *packets.Traffic
|
||||||
var found bool
|
var found bool
|
||||||
if traffic, found = mod.Session.Queue.Traffic[e.IpAddress]; !found {
|
var v interface{}
|
||||||
|
if v, found = mod.Session.Queue.Traffic.Load(e.IpAddress); !found {
|
||||||
traffic = &packets.Traffic{}
|
traffic = &packets.Traffic{}
|
||||||
|
} else {
|
||||||
|
traffic = v.(*packets.Traffic)
|
||||||
}
|
}
|
||||||
|
|
||||||
seen := e.LastSeen.Format("15:04:05")
|
seen := e.LastSeen.Format("15:04:05")
|
||||||
|
@ -203,9 +206,6 @@ func (mod *Discovery) colNames(hasMeta bool) []string {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (mod *Discovery) showStatusBar() {
|
func (mod *Discovery) showStatusBar() {
|
||||||
mod.Session.Queue.Stats.RLock()
|
|
||||||
defer mod.Session.Queue.Stats.RUnlock()
|
|
||||||
|
|
||||||
parts := []string{
|
parts := []string{
|
||||||
fmt.Sprintf("%s %s", tui.Red("↑"), humanize.Bytes(mod.Session.Queue.Stats.Sent)),
|
fmt.Sprintf("%s %s", tui.Red("↑"), humanize.Bytes(mod.Session.Queue.Stats.Sent)),
|
||||||
fmt.Sprintf("%s %s", tui.Green("↓"), humanize.Bytes(mod.Session.Queue.Stats.Received)),
|
fmt.Sprintf("%s %s", tui.Green("↓"), humanize.Bytes(mod.Session.Queue.Stats.Received)),
|
||||||
|
|
|
@ -41,24 +41,19 @@ func (a BySeenSorter) Less(i, j int) bool { return a[i].LastSeen.Before(a[j].Las
|
||||||
|
|
||||||
type BySentSorter []*network.Endpoint
|
type BySentSorter []*network.Endpoint
|
||||||
|
|
||||||
|
func trafficOf(ip string) *packets.Traffic {
|
||||||
|
if v, found := session.I.Queue.Traffic.Load(ip); !found {
|
||||||
|
return &packets.Traffic{}
|
||||||
|
} else {
|
||||||
|
return v.(*packets.Traffic)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (a BySentSorter) Len() int { return len(a) }
|
func (a BySentSorter) Len() int { return len(a) }
|
||||||
func (a BySentSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
|
func (a BySentSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
|
||||||
func (a BySentSorter) Less(i, j int) bool {
|
func (a BySentSorter) Less(i, j int) bool {
|
||||||
session.I.Queue.Lock()
|
aTraffic := trafficOf(a[i].IpAddress)
|
||||||
defer session.I.Queue.Unlock()
|
bTraffic := trafficOf(a[j].IpAddress)
|
||||||
|
|
||||||
var found bool = false
|
|
||||||
var aTraffic *packets.Traffic = nil
|
|
||||||
var bTraffic *packets.Traffic = nil
|
|
||||||
|
|
||||||
if aTraffic, found = session.I.Queue.Traffic[a[i].IpAddress]; !found {
|
|
||||||
aTraffic = &packets.Traffic{}
|
|
||||||
}
|
|
||||||
|
|
||||||
if bTraffic, found = session.I.Queue.Traffic[a[j].IpAddress]; !found {
|
|
||||||
bTraffic = &packets.Traffic{}
|
|
||||||
}
|
|
||||||
|
|
||||||
return bTraffic.Sent > aTraffic.Sent
|
return bTraffic.Sent > aTraffic.Sent
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -67,20 +62,7 @@ type ByRcvdSorter []*network.Endpoint
|
||||||
func (a ByRcvdSorter) Len() int { return len(a) }
|
func (a ByRcvdSorter) Len() int { return len(a) }
|
||||||
func (a ByRcvdSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
|
func (a ByRcvdSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
|
||||||
func (a ByRcvdSorter) Less(i, j int) bool {
|
func (a ByRcvdSorter) Less(i, j int) bool {
|
||||||
session.I.Queue.Lock()
|
aTraffic := trafficOf(a[i].IpAddress)
|
||||||
defer session.I.Queue.Unlock()
|
bTraffic := trafficOf(a[j].IpAddress)
|
||||||
|
|
||||||
var found bool = false
|
|
||||||
var aTraffic *packets.Traffic = nil
|
|
||||||
var bTraffic *packets.Traffic = nil
|
|
||||||
|
|
||||||
if aTraffic, found = session.I.Queue.Traffic[a[i].IpAddress]; !found {
|
|
||||||
aTraffic = &packets.Traffic{}
|
|
||||||
}
|
|
||||||
|
|
||||||
if bTraffic, found = session.I.Queue.Traffic[a[j].IpAddress]; !found {
|
|
||||||
bTraffic = &packets.Traffic{}
|
|
||||||
}
|
|
||||||
|
|
||||||
return bTraffic.Received > aTraffic.Received
|
return bTraffic.Received > aTraffic.Received
|
||||||
}
|
}
|
||||||
|
|
|
@ -296,9 +296,6 @@ func (mod *WiFiModule) colNames(nrows int) []string {
|
||||||
}
|
}
|
||||||
|
|
||||||
func (mod *WiFiModule) showStatusBar() {
|
func (mod *WiFiModule) showStatusBar() {
|
||||||
mod.Session.Queue.Stats.RLock()
|
|
||||||
defer mod.Session.Queue.Stats.RUnlock()
|
|
||||||
|
|
||||||
parts := []string{
|
parts := []string{
|
||||||
fmt.Sprintf("%s (ch. %d)", mod.iface.Name(), network.GetInterfaceChannel(mod.iface.Name())),
|
fmt.Sprintf("%s (ch. %d)", mod.iface.Name(), network.GetInterfaceChannel(mod.iface.Name())),
|
||||||
fmt.Sprintf("%s %s", tui.Red("↑"), humanize.Bytes(mod.Session.Queue.Stats.Sent)),
|
fmt.Sprintf("%s %s", tui.Red("↑"), humanize.Bytes(mod.Session.Queue.Stats.Sent)),
|
||||||
|
|
|
@ -4,6 +4,7 @@ import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"net"
|
"net"
|
||||||
"sync"
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
|
||||||
"github.com/bettercap/bettercap/network"
|
"github.com/bettercap/bettercap/network"
|
||||||
|
|
||||||
|
@ -25,8 +26,6 @@ type Traffic struct {
|
||||||
}
|
}
|
||||||
|
|
||||||
type Stats struct {
|
type Stats struct {
|
||||||
sync.RWMutex
|
|
||||||
|
|
||||||
Sent uint64
|
Sent uint64
|
||||||
Received uint64
|
Received uint64
|
||||||
PktReceived uint64
|
PktReceived uint64
|
||||||
|
@ -39,10 +38,9 @@ type Queue struct {
|
||||||
sync.RWMutex
|
sync.RWMutex
|
||||||
|
|
||||||
Activities chan Activity `json:"-"`
|
Activities chan Activity `json:"-"`
|
||||||
|
Stats Stats
|
||||||
Stats Stats
|
Protos sync.Map
|
||||||
Protos map[string]uint64
|
Traffic sync.Map
|
||||||
Traffic map[string]*Traffic
|
|
||||||
|
|
||||||
iface *network.Endpoint
|
iface *network.Endpoint
|
||||||
handle *pcap.Handle
|
handle *pcap.Handle
|
||||||
|
@ -55,8 +53,8 @@ type Queue struct {
|
||||||
|
|
||||||
func NewQueue(iface *network.Endpoint) (q *Queue, err error) {
|
func NewQueue(iface *network.Endpoint) (q *Queue, err error) {
|
||||||
q = &Queue{
|
q = &Queue{
|
||||||
Protos: make(map[string]uint64),
|
Protos: sync.Map{},
|
||||||
Traffic: make(map[string]*Traffic),
|
Traffic: sync.Map{},
|
||||||
Activities: make(chan Activity),
|
Activities: make(chan Activity),
|
||||||
|
|
||||||
writes: &sync.WaitGroup{},
|
writes: &sync.WaitGroup{},
|
||||||
|
@ -102,14 +100,12 @@ func (q *Queue) trackProtocols(pkt gopacket.Packet) {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
|
|
||||||
q.Lock()
|
|
||||||
name := proto.String()
|
name := proto.String()
|
||||||
if _, found := q.Protos[name]; !found {
|
if v, found := q.Protos.Load(name); !found {
|
||||||
q.Protos[name] = 1
|
q.Protos.Store(name, 1)
|
||||||
} else {
|
} else {
|
||||||
q.Protos[name]++
|
q.Protos.Store(name, v.(int)+1)
|
||||||
}
|
}
|
||||||
q.Unlock()
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -122,46 +118,34 @@ func (q *Queue) trackActivity(eth *layers.Ethernet, ip4 *layers.IPv4, address ne
|
||||||
Source: isSent,
|
Source: isSent,
|
||||||
}
|
}
|
||||||
|
|
||||||
q.Lock()
|
|
||||||
defer q.Unlock()
|
|
||||||
|
|
||||||
// initialize or update stats
|
// initialize or update stats
|
||||||
addr := address.String()
|
addr := address.String()
|
||||||
if _, found := q.Traffic[addr]; !found {
|
if v, found := q.Traffic.Load(addr); !found {
|
||||||
if isSent {
|
if isSent {
|
||||||
q.Traffic[addr] = &Traffic{Sent: pktSize}
|
q.Traffic.Store(addr, &Traffic{Sent: pktSize})
|
||||||
} else {
|
} else {
|
||||||
q.Traffic[addr] = &Traffic{Received: pktSize}
|
q.Traffic.Store(addr, &Traffic{Received: pktSize})
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
if isSent {
|
if isSent {
|
||||||
q.Traffic[addr].Sent += pktSize
|
v.(*Traffic).Sent += pktSize
|
||||||
} else {
|
} else {
|
||||||
q.Traffic[addr].Received += pktSize
|
v.(*Traffic).Received += pktSize
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (q *Queue) TrackPacket(size uint64) {
|
func (q *Queue) TrackPacket(size uint64) {
|
||||||
q.Stats.Lock()
|
atomic.AddUint64(&q.Stats.PktReceived, 1)
|
||||||
defer q.Stats.Unlock()
|
atomic.AddUint64(&q.Stats.Received, size)
|
||||||
|
|
||||||
q.Stats.PktReceived++
|
|
||||||
q.Stats.Received += size
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (q *Queue) TrackSent(size uint64) {
|
func (q *Queue) TrackSent(size uint64) {
|
||||||
q.Stats.Lock()
|
atomic.AddUint64(&q.Stats.Sent, size)
|
||||||
defer q.Stats.Unlock()
|
|
||||||
|
|
||||||
q.Stats.Sent += size
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (q *Queue) TrackError() {
|
func (q *Queue) TrackError() {
|
||||||
q.Stats.Lock()
|
atomic.AddUint64(&q.Stats.Errors, 1)
|
||||||
defer q.Stats.Unlock()
|
|
||||||
|
|
||||||
q.Stats.Errors++
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (q *Queue) getPacketMeta(pkt gopacket.Packet) map[string]string {
|
func (q *Queue) getPacketMeta(pkt gopacket.Packet) map[string]string {
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue