From 97c97a668a3f7be3508fbf350caf2a762103e218 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Wed, 17 Jan 2018 17:31:51 +0100
Subject: [PATCH] fix: counting udp probe packets in sent stats

---
 modules/net_probe.go | 8 +++++++-
 packets/queue.go     | 4 ++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/modules/net_probe.go b/modules/net_probe.go
index d61f6039..be0c17b6 100644
--- a/modules/net_probe.go
+++ b/modules/net_probe.go
@@ -75,7 +75,13 @@ func (p *Prober) sendProbe(from net.IP, from_hw net.HardwareAddr, ip net.IP) {
 	} else {
 		// log.Debug("UDP connection to %s enstablished.", name)
 		defer con.Close()
-		con.Write([]byte{0xde, 0xad, 0xbe, 0xef})
+		wrote, _ := con.Write([]byte{0xde, 0xad, 0xbe, 0xef})
+
+		if wrote > 0 {
+			p.Session.Queue.Lock()
+			p.Session.Queue.Sent += uint64(wrote)
+			p.Session.Queue.Unlock()
+		}
 	}
 }
 
diff --git a/packets/queue.go b/packets/queue.go
index 0413e07f..88864be1 100644
--- a/packets/queue.go
+++ b/packets/queue.go
@@ -74,6 +74,8 @@ func (q *Queue) worker() {
 		if q.active == false {
 			return
 		}
+
+		q.Lock()
 		pktSize := uint64(len(pkt.Data()))
 
 		q.PktReceived++
@@ -131,6 +133,8 @@ func (q *Queue) worker() {
 				q.Traffic[addr].Received += pktSize
 			}
 		}
+
+		q.Unlock()
 	}
 }