From 9d930f5d183da66f0859a2c21cdd5e9919b84db4 Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Thu, 19 Oct 2017 04:41:13 +0200
Subject: stats: more robust accounting

---
 src/receive.c      | 12 +++++++-----
 src/send.c         |  1 +
 src/socket.c       |  2 ++
 src/tests/netns.sh |  9 ++++++++-
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/receive.c b/src/receive.c
index 0f896ee..b27876c 100644
--- a/src/receive.c
+++ b/src/receive.c
@@ -19,11 +19,11 @@ static inline void rx_stats(struct wireguard_peer *peer, size_t len)
 	struct pcpu_sw_netstats *tstats = get_cpu_ptr(peer->device->dev->tstats);
 
 	u64_stats_update_begin(&tstats->syncp);
-	tstats->rx_bytes += len;
 	++tstats->rx_packets;
+	tstats->rx_bytes += len;
+	peer->rx_bytes += len;
 	u64_stats_update_end(&tstats->syncp);
 	put_cpu_ptr(tstats);
-	peer->rx_bytes += len;
 }
 
 #define SKB_TYPE_LE32(skb) ((struct message_header *)(skb)->data)->type
@@ -277,7 +277,7 @@ static void packet_consume_data_done(struct sk_buff *skb, struct endpoint *endpo
 {
 	struct wireguard_peer *peer = PACKET_PEER(skb), *routed_peer;
 	struct net_device *dev = peer->device->dev;
-	unsigned int len;
+	unsigned int len, len_before_trim;
 
 	socket_set_peer_endpoint(peer, endpoint);
 
@@ -290,6 +290,7 @@ static void packet_consume_data_done(struct sk_buff *skb, struct endpoint *endpo
 
 	/* A packet with length 0 is a keepalive packet */
 	if (unlikely(!skb->len)) {
+		rx_stats(peer, message_data_len(0));
 		net_dbg_ratelimited("%s: Receiving keepalive packet from peer %Lu (%pISpfsc)\n", dev->name, peer->internal_id, &peer->endpoint.addr);
 		goto packet_processed;
 	}
@@ -317,6 +318,7 @@ static void packet_consume_data_done(struct sk_buff *skb, struct endpoint *endpo
 
 	if (unlikely(len > skb->len))
 		goto dishonest_packet_size;
+	len_before_trim = skb->len;
 	if (unlikely(pskb_trim(skb, len)))
 		goto packet_processed;
 
@@ -328,12 +330,11 @@ static void packet_consume_data_done(struct sk_buff *skb, struct endpoint *endpo
 	if (unlikely(routed_peer != peer))
 		goto dishonest_packet_peer;
 
-	len = skb->len;
 	if (unlikely(netif_receive_skb(skb) == NET_RX_DROP)) {
 		++dev->stats.rx_dropped;
 		net_dbg_ratelimited("%s: Failed to give packet to userspace from peer %Lu (%pISpfsc)\n", dev->name, peer->internal_id, &peer->endpoint.addr);
 	} else
-		rx_stats(peer, len);
+		rx_stats(peer, message_data_len(len_before_trim));
 	goto continue_processing;
 
 dishonest_packet_peer:
@@ -451,6 +452,7 @@ void packet_receive(struct wireguard_device *wg, struct sk_buff *skb)
 	case cpu_to_le32(MESSAGE_HANDSHAKE_RESPONSE):
 	case cpu_to_le32(MESSAGE_HANDSHAKE_COOKIE): {
 		int cpu;
+
 		if (skb_queue_len(&wg->incoming_handshakes) > MAX_QUEUED_INCOMING_HANDSHAKES) {
 			net_dbg_skb_ratelimited("%s: Too many handshakes queued, dropping packet from %pISpfsc\n", wg->dev->name, skb);
 			goto err;
diff --git a/src/send.c b/src/send.c
index 05a8472..d9fefc6 100644
--- a/src/send.c
+++ b/src/send.c
@@ -149,6 +149,7 @@ static inline bool skb_encrypt(struct sk_buff *skb, struct noise_keypair *keypai
 		skb_checksum_help(skb);
 
 	/* Only after checksumming can we safely add on the padding at the end and the header. */
+	skb_set_inner_network_header(skb, 0);
 	header = (struct message_data *)skb_push(skb, sizeof(struct message_data));
 	header->header.type = cpu_to_le32(MESSAGE_DATA);
 	header->key_idx = keypair->remote_index;
diff --git a/src/socket.c b/src/socket.c
index 8d0e9ca..913ad0a 100644
--- a/src/socket.c
+++ b/src/socket.c
@@ -173,6 +173,7 @@ int socket_send_buffer_to_peer(struct wireguard_peer *peer, void *buffer, size_t
 		return -ENOMEM;
 
 	skb_reserve(skb, SKB_HEADER_LEN);
+	skb_set_inner_network_header(skb, 0);
 	memcpy(skb_put(skb, len), buffer, len);
 	return socket_send_skb_to_peer(peer, skb, ds);
 }
@@ -193,6 +194,7 @@ int socket_send_buffer_as_reply_to_skb(struct wireguard_device *wg, struct sk_bu
 	if (unlikely(!skb))
 		return -ENOMEM;
 	skb_reserve(skb, SKB_HEADER_LEN);
+	skb_set_inner_network_header(skb, 0);
 	memcpy(skb_put(skb, len), out_buffer, len);
 
 	if (endpoint.addr.sa_family == AF_INET)
diff --git a/src/tests/netns.sh b/src/tests/netns.sh
index 94666f3..9a12a19 100755
--- a/src/tests/netns.sh
+++ b/src/tests/netns.sh
@@ -144,7 +144,14 @@ n2 wg set wg0 peer "$pub1" endpoint 127.0.0.1:1
 # Before calling tests, we first make sure that the stats counters are working
 n2 ping -c 10 -f -W 1 192.168.241.1
 { read _; read _; read _; read rx_bytes _; read _; read tx_bytes _; } < <(ip2 -stats link show dev wg0)
-[[ $rx_bytes -ge 932 && $tx_bytes -ge 1516 && $rx_bytes -lt 2500 && $rx_bytes -lt 2500 ]]
+(( rx_bytes == 1372 && (tx_bytes == 1428 || tx_bytes == 1460) ))
+{ read _; read _; read _; read rx_bytes _; read _; read tx_bytes _; } < <(ip1 -stats link show dev wg0)
+(( tx_bytes == 1372 && (rx_bytes == 1428 || rx_bytes == 1460) ))
+read _ rx_bytes tx_bytes < <(n2 wg show wg0 transfer)
+(( rx_bytes == 1372 && (tx_bytes == 1428 || tx_bytes == 1460) ))
+read _ rx_bytes tx_bytes < <(n1 wg show wg0 transfer)
+(( tx_bytes == 1372 && (rx_bytes == 1428 || rx_bytes == 1460) ))
+
 tests
 ip1 link set wg0 mtu $big_mtu
 ip2 link set wg0 mtu $big_mtu
-- 
cgit v1.2.3-59-g8ed1b