From ac4b1f6e9e4faa54533dbc5b1cb68d3fbe3b3bc3 Mon Sep 17 00:00:00 2001 From: bonedaddy Date: Tue, 22 Dec 2020 17:10:33 -0800 Subject: [PATCH 1/3] network: optimize wifi locking and include memory allocation optimization --- network/wifi.go | 38 ++++++++++++++++++++++---------------- network/wifi_ap.go | 34 +++++++++++++++++----------------- network/wifi_handshake.go | 18 +++++++++--------- 3 files changed, 48 insertions(+), 42 deletions(-) diff --git a/network/wifi.go b/network/wifi.go index 0fc57d35..0a452d20 100644 --- a/network/wifi.go +++ b/network/wifi.go @@ -43,7 +43,7 @@ type APNewCallback func(ap *AccessPoint) type APLostCallback func(ap *AccessPoint) type WiFi struct { - sync.Mutex + sync.RWMutex aliases *data.UnsortedKV aps map[string]*AccessPoint @@ -67,8 +67,12 @@ func NewWiFi(iface *Endpoint, aliases *data.UnsortedKV, newcb APNewCallback, los } func (w *WiFi) MarshalJSON() ([]byte, error) { + w.RLock() + defer w.RUnlock() + doc := wifiJSON{ - AccessPoints: make([]*AccessPoint, 0), + // we know the length so preallocate to reduce memory allocations + AccessPoints: make([]*AccessPoint, len(w.aps)), } for _, ap := range w.aps { @@ -88,10 +92,11 @@ func (w *WiFi) EachAccessPoint(cb func(mac string, ap *AccessPoint)) { } func (w *WiFi) Stations() (list []*Station) { - w.Lock() - defer w.Unlock() + w.RLock() + defer w.RUnlock() + + list = make([]*Station, len(w.aps)) - list = make([]*Station, 0) for _, ap := range w.aps { list = append(list, ap.Station) } @@ -99,10 +104,11 @@ func (w *WiFi) Stations() (list []*Station) { } func (w *WiFi) List() (list []*AccessPoint) { - w.Lock() - defer w.Unlock() + w.RLock() + defer w.RUnlock() + + list = make([]*AccessPoint, len(w.aps)) - list = make([]*AccessPoint, 0) for _, ap := range w.aps { list = append(list, ap) } @@ -167,8 +173,8 @@ func (w *WiFi) AddIfNew(ssid, mac string, frequency int, rssi int8) (*AccessPoin } func (w *WiFi) Get(mac string) (*AccessPoint, bool) { - w.Lock() - defer w.Unlock() + w.RLock() + defer w.RUnlock() mac = NormalizeMac(mac) ap, found := w.aps[mac] @@ -176,8 +182,8 @@ func (w *WiFi) Get(mac string) (*AccessPoint, bool) { } func (w *WiFi) GetClient(mac string) (*Station, bool) { - w.Lock() - defer w.Unlock() + w.RLock() + defer w.RUnlock() mac = NormalizeMac(mac) for _, ap := range w.aps { @@ -196,8 +202,8 @@ func (w *WiFi) Clear() { } func (w *WiFi) NumHandshakes() int { - w.Lock() - defer w.Unlock() + w.RLock() + defer w.RUnlock() sum := 0 for _, ap := range w.aps { @@ -212,8 +218,8 @@ func (w *WiFi) NumHandshakes() int { } func (w *WiFi) SaveHandshakesTo(fileName string, linkType layers.LinkType) error { - w.Lock() - defer w.Unlock() + w.RLock() + defer w.RUnlock() // check if folder exists first dirName := filepath.Dir(fileName) diff --git a/network/wifi_ap.go b/network/wifi_ap.go index 58143cc0..fc889036 100644 --- a/network/wifi_ap.go +++ b/network/wifi_ap.go @@ -10,7 +10,7 @@ import ( type AccessPoint struct { *Station - sync.Mutex + sync.RWMutex aliases *data.UnsortedKV clients map[string]*Station @@ -32,12 +32,12 @@ func NewAccessPoint(essid, bssid string, frequency int, rssi int8, aliases *data } func (ap *AccessPoint) MarshalJSON() ([]byte, error) { - ap.Lock() - defer ap.Unlock() + ap.RLock() + defer ap.RUnlock() doc := apJSON{ Station: ap.Station, - Clients: make([]*Station, 0), + Clients: make([]*Station, len(ap.clients)), Handshake: ap.withKeyMaterial, } @@ -49,8 +49,8 @@ func (ap *AccessPoint) MarshalJSON() ([]byte, error) { } func (ap *AccessPoint) Get(bssid string) (*Station, bool) { - ap.Lock() - defer ap.Unlock() + ap.RLock() + defer ap.RUnlock() bssid = NormalizeMac(bssid) if s, found := ap.clients[bssid]; found { @@ -97,16 +97,16 @@ func (ap *AccessPoint) AddClientIfNew(bssid string, frequency int, rssi int8) (* } func (ap *AccessPoint) NumClients() int { - ap.Lock() - defer ap.Unlock() + ap.RLock() + defer ap.RUnlock() return len(ap.clients) } func (ap *AccessPoint) Clients() (list []*Station) { - ap.Lock() - defer ap.Unlock() + ap.RLock() + defer ap.RUnlock() - list = make([]*Station, 0) + list = make([]*Station, len(ap.clients)) for _, c := range ap.clients { list = append(list, c) } @@ -130,15 +130,15 @@ func (ap *AccessPoint) WithKeyMaterial(state bool) { } func (ap *AccessPoint) HasKeyMaterial() bool { - ap.Lock() - defer ap.Unlock() + ap.RLock() + defer ap.RUnlock() return ap.withKeyMaterial } func (ap *AccessPoint) NumHandshakes() int { - ap.Lock() - defer ap.Unlock() + ap.RLock() + defer ap.RUnlock() sum := 0 @@ -156,8 +156,8 @@ func (ap *AccessPoint) HasHandshakes() bool { } func (ap *AccessPoint) HasPMKID() bool { - ap.Lock() - defer ap.Unlock() + ap.RLock() + defer ap.RUnlock() for _, c := range ap.clients { if c.Handshake.HasPMKID() { diff --git a/network/wifi_handshake.go b/network/wifi_handshake.go index 341a4f9e..7181baec 100644 --- a/network/wifi_handshake.go +++ b/network/wifi_handshake.go @@ -8,7 +8,7 @@ import ( ) type Handshake struct { - sync.Mutex + sync.RWMutex Beacon gopacket.Packet Challenges []gopacket.Packet @@ -80,8 +80,8 @@ func (h *Handshake) AddFrame(n int, pkt gopacket.Packet) { } func (h *Handshake) Complete() bool { - h.Lock() - defer h.Unlock() + h.RLock() + defer h.RUnlock() nChal := len(h.Challenges) nResp := len(h.Responses) @@ -91,8 +91,8 @@ func (h *Handshake) Complete() bool { } func (h *Handshake) Half() bool { - h.Lock() - defer h.Unlock() + h.RLock() + defer h.RUnlock() /* * You can use every combination of the handshake to crack the net: @@ -110,14 +110,14 @@ func (h *Handshake) Half() bool { } func (h *Handshake) HasPMKID() bool { - h.Lock() - defer h.Unlock() + h.RLock() + defer h.RUnlock() return h.hasPMKID } func (h *Handshake) NumUnsaved() int { - h.Lock() - defer h.Unlock() + h.RLock() + defer h.RUnlock() return len(h.unsaved) } From 08cad808ef8da10f023cb67f63f714721a2122a3 Mon Sep 17 00:00:00 2001 From: bonedaddy Date: Tue, 22 Dec 2020 17:20:58 -0800 Subject: [PATCH 2/3] fix slice memory allocation optimization --- network/wifi.go | 6 +++--- network/wifi_ap.go | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/network/wifi.go b/network/wifi.go index 0a452d20..797e9009 100644 --- a/network/wifi.go +++ b/network/wifi.go @@ -72,7 +72,7 @@ func (w *WiFi) MarshalJSON() ([]byte, error) { doc := wifiJSON{ // we know the length so preallocate to reduce memory allocations - AccessPoints: make([]*AccessPoint, len(w.aps)), + AccessPoints: make([]*AccessPoint, 0, len(w.aps)), } for _, ap := range w.aps { @@ -95,7 +95,7 @@ func (w *WiFi) Stations() (list []*Station) { w.RLock() defer w.RUnlock() - list = make([]*Station, len(w.aps)) + list = make([]*Station, 0, len(w.aps)) for _, ap := range w.aps { list = append(list, ap.Station) @@ -107,7 +107,7 @@ func (w *WiFi) List() (list []*AccessPoint) { w.RLock() defer w.RUnlock() - list = make([]*AccessPoint, len(w.aps)) + list = make([]*AccessPoint, 0, len(w.aps)) for _, ap := range w.aps { list = append(list, ap) diff --git a/network/wifi_ap.go b/network/wifi_ap.go index fc889036..ed17f681 100644 --- a/network/wifi_ap.go +++ b/network/wifi_ap.go @@ -37,7 +37,7 @@ func (ap *AccessPoint) MarshalJSON() ([]byte, error) { doc := apJSON{ Station: ap.Station, - Clients: make([]*Station, len(ap.clients)), + Clients: make([]*Station, 0, len(ap.clients)), Handshake: ap.withKeyMaterial, } @@ -106,7 +106,7 @@ func (ap *AccessPoint) Clients() (list []*Station) { ap.RLock() defer ap.RUnlock() - list = make([]*Station, len(ap.clients)) + list = make([]*Station, 0, len(ap.clients)) for _, c := range ap.clients { list = append(list, c) } From 10817d574cdc29d92b4a60ddddb20cabe716db5e Mon Sep 17 00:00:00 2001 From: bonedaddy Date: Wed, 23 Dec 2020 13:22:48 -0800 Subject: [PATCH 3/3] wifi.go: dont claim read lock until it is needed --- network/wifi.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/network/wifi.go b/network/wifi.go index 797e9009..21e3b634 100644 --- a/network/wifi.go +++ b/network/wifi.go @@ -218,9 +218,6 @@ func (w *WiFi) NumHandshakes() int { } func (w *WiFi) SaveHandshakesTo(fileName string, linkType layers.LinkType) error { - w.RLock() - defer w.RUnlock() - // check if folder exists first dirName := filepath.Dir(fileName) if _, err := os.Stat(dirName); err != nil { @@ -244,6 +241,9 @@ func (w *WiFi) SaveHandshakesTo(fileName string, linkType layers.LinkType) error } } + w.RLock() + defer w.RUnlock() + for _, ap := range w.aps { for _, station := range ap.Clients() { // if half (which includes also complete) or has pmkid