fix: using sync.Map to avoid race conditions on the packets.Queue

This commit is contained in:
evilsocket 2019-03-17 13:12:31 +01:00
parent b676d68b4c
commit 64a5ce2b58
No known key found for this signature in database
GPG key ID: 1564D7F30393A456
4 changed files with 34 additions and 71 deletions

View file

@ -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)),

View file

@ -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
} }

View file

@ -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)),

View file

@ -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 {