From 3371f8dac6fe6bbd7522a8316b50f6473012e302 Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Tue, 11 Jun 2019 18:13:52 +0200
Subject: device: update transfer counters correctly

The rule is to always update them to the full packet size minus UDP/IP
encapsulation for all authenticated packet types.
---
 device/peer.go    | 7 ++++++-
 device/receive.go | 4 +++-
 device/send.go    | 2 --
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/device/peer.go b/device/peer.go
index 815dff4..4e7f2da 100644
--- a/device/peer.go
+++ b/device/peer.go
@@ -10,6 +10,7 @@ import (
 	"errors"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
@@ -140,7 +141,11 @@ func (peer *Peer) SendBuffer(buffer []byte) error {
 		return errors.New("no known endpoint for peer")
 	}
 
-	return peer.device.net.bind.Send(buffer, peer.endpoint)
+	err := peer.device.net.bind.Send(buffer, peer.endpoint)
+	if err == nil {
+		atomic.AddUint64(&peer.stats.txBytes, uint64(len(buffer)))
+	}
+	return err
 }
 
 func (peer *Peer) String() string {
diff --git a/device/receive.go b/device/receive.go
index 32d632b..62b5ef4 100644
--- a/device/receive.go
+++ b/device/receive.go
@@ -427,6 +427,7 @@ func (device *Device) RoutineHandshake() {
 			peer.SetEndpointFromPacket(elem.endpoint)
 
 			logDebug.Println(peer, "- Received handshake initiation")
+			atomic.AddUint64(&peer.stats.rxBytes, uint64(len(elem.packet)))
 
 			peer.SendHandshakeResponse()
 
@@ -457,6 +458,7 @@ func (device *Device) RoutineHandshake() {
 			peer.SetEndpointFromPacket(elem.endpoint)
 
 			logDebug.Println(peer, "- Received handshake response")
+			atomic.AddUint64(&peer.stats.rxBytes, uint64(len(elem.packet)))
 
 			// update timers
 
@@ -581,6 +583,7 @@ func (peer *Peer) RoutineSequentialReceiver() {
 		peer.keepKeyFreshReceiving()
 		peer.timersAnyAuthenticatedPacketTraversal()
 		peer.timersAnyAuthenticatedPacketReceived()
+		atomic.AddUint64(&peer.stats.rxBytes, uint64(len(elem.packet)+MinMessageSize))
 
 		// check for keepalive
 
@@ -656,7 +659,6 @@ func (peer *Peer) RoutineSequentialReceiver() {
 		// write to tun device
 
 		offset := MessageTransportOffsetContent
-		atomic.AddUint64(&peer.stats.rxBytes, uint64(len(elem.packet)))
 		_, err := device.tun.device.Write(elem.buffer[:offset+len(elem.packet)], offset)
 		if err == nil {
 			shouldFlush = true
diff --git a/device/send.go b/device/send.go
index 2198427..c4aa5b9 100644
--- a/device/send.go
+++ b/device/send.go
@@ -600,7 +600,6 @@ func (peer *Peer) RoutineSequentialSender() {
 
 			// send message and return buffer to pool
 
-			length := uint64(len(elem.packet))
 			err := peer.SendBuffer(elem.packet)
 			if len(elem.packet) != MessageKeepaliveSize {
 				peer.timersDataSent()
@@ -611,7 +610,6 @@ func (peer *Peer) RoutineSequentialSender() {
 				logError.Println(peer, "- Failed to send data packet", err)
 				continue
 			}
-			atomic.AddUint64(&peer.stats.txBytes, length)
 
 			peer.keepKeyFreshSending()
 		}
-- 
cgit v1.2.3-59-g8ed1b