From d60648f0dbadf2cc35e51458db12ccfb8baf918b Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Fri, 23 Feb 2018 10:15:51 +0100
Subject: [PATCH] fix: using normal mutex instead of atomic increments ( fixes
 #69 )

---
 modules/net_probe_udp.go  |  5 +++--
 modules/net_recon_show.go | 11 ++++++-----
 packets/queue.go          | 19 ++++++++++++++-----
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/modules/net_probe_udp.go b/modules/net_probe_udp.go
index 7e8db8fc..14e299a9 100644
--- a/modules/net_probe_udp.go
+++ b/modules/net_probe_udp.go
@@ -3,7 +3,6 @@ package modules
 import (
 	"fmt"
 	"net"
-	"sync/atomic"
 
 	"github.com/bettercap/bettercap/log"
 )
@@ -20,7 +19,9 @@ func (p *Prober) sendProbeUDP(from net.IP, from_hw net.HardwareAddr, ip net.IP)
 		wrote, _ := con.Write([]byte{0x00})
 
 		if wrote > 0 {
-			atomic.AddUint64(&p.Session.Queue.Stats.Sent, uint64(wrote))
+			p.Session.Queue.Stats.Lock()
+			p.Session.Queue.Stats.Sent += uint64(wrote)
+			p.Session.Queue.Stats.Unlock()
 		}
 	}
 }
diff --git a/modules/net_recon_show.go b/modules/net_recon_show.go
index ca5bc56e..0ed426b5 100644
--- a/modules/net_recon_show.go
+++ b/modules/net_recon_show.go
@@ -5,7 +5,6 @@ import (
 	"os"
 	"sort"
 	"strings"
-	"sync/atomic"
 	"time"
 
 	"github.com/bettercap/bettercap/core"
@@ -169,13 +168,15 @@ func (d *Discovery) Show(by string) error {
 
 	d.showTable(colNames, rows)
 
+	d.Session.Queue.Stats.RLock()
 	fmt.Printf("\n%s %s / %s %s / %d pkts / %d errs\n\n",
 		core.Red("↑"),
-		humanize.Bytes(atomic.LoadUint64(&d.Session.Queue.Stats.Sent)),
+		humanize.Bytes(d.Session.Queue.Stats.Sent),
 		core.Green("↓"),
-		humanize.Bytes(atomic.LoadUint64(&d.Session.Queue.Stats.Received)),
-		atomic.LoadUint64(&d.Session.Queue.Stats.PktReceived),
-		atomic.LoadUint64(&d.Session.Queue.Stats.Errors))
+		humanize.Bytes(d.Session.Queue.Stats.Received),
+		d.Session.Queue.Stats.PktReceived,
+		d.Session.Queue.Stats.Errors)
+	d.Session.Queue.Stats.RUnlock()
 
 	s := EventsStream{}
 	events := d.Session.Events.Sorted()
diff --git a/packets/queue.go b/packets/queue.go
index 9bfb1048..38dd853d 100644
--- a/packets/queue.go
+++ b/packets/queue.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"net"
 	"sync"
-	"sync/atomic"
 
 	"github.com/bettercap/bettercap/network"
 
@@ -26,6 +25,8 @@ type Traffic struct {
 }
 
 type Stats struct {
+	sync.RWMutex
+
 	Sent        uint64
 	Received    uint64
 	PktReceived uint64
@@ -146,8 +147,12 @@ func (q *Queue) worker() {
 
 		pktSize := uint64(len(pkt.Data()))
 
-		atomic.AddUint64(&q.Stats.PktReceived, 1)
-		atomic.AddUint64(&q.Stats.Received, pktSize)
+		q.Stats.Lock()
+
+		q.Stats.PktReceived++
+		q.Stats.Received += pktSize
+
+		q.Stats.Unlock()
 
 		q.onPacketCallback(pkt)
 
@@ -179,10 +184,14 @@ func (q *Queue) Send(raw []byte) error {
 	}
 
 	if err := q.handle.WritePacketData(raw); err != nil {
-		atomic.AddUint64(&q.Stats.Errors, 1)
+		q.Stats.Lock()
+		q.Stats.Errors++
+		q.Stats.Unlock()
 		return err
 	} else {
-		atomic.AddUint64(&q.Stats.Sent, uint64(len(raw)))
+		q.Stats.Lock()
+		q.Stats.Sent += uint64(len(raw))
+		q.Stats.Unlock()
 	}
 
 	return nil