From 08b3f5f070738e88d4b9f5d2e5d74251bdd9bce5 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Mon, 29 Aug 2022 09:46:26 -0700
Subject: [PATCH] wgengine/wgint: add shady temporary package to get at
 wireguard internals

For #5451

Change-Id: I43482289e323ba9142a446d551ab7a94a467c43a
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 wgengine/wgint/wgint.go      | 59 ++++++++++++++++++++++++++++++++++++
 wgengine/wgint/wgint_test.go | 24 +++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 wgengine/wgint/wgint.go
 create mode 100644 wgengine/wgint/wgint_test.go

diff --git a/wgengine/wgint/wgint.go b/wgengine/wgint/wgint.go
new file mode 100644
index 000000000..16e5fd1f3
--- /dev/null
+++ b/wgengine/wgint/wgint.go
@@ -0,0 +1,59 @@
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package wgint provides somewhat shady access to wireguard-go
+// internals that don't (yet) have public APIs.
+package wgint
+
+import (
+	"reflect"
+	"sync/atomic"
+	"unsafe"
+
+	"golang.zx2c4.com/wireguard/device"
+)
+
+var (
+	offHandshake = getPeerStatsOffset("lastHandshakeNano")
+	offRxBytes   = getPeerStatsOffset("rxBytes")
+	offTxBytes   = getPeerStatsOffset("txBytes")
+)
+
+func getPeerStatsOffset(name string) uintptr {
+	peerType := reflect.TypeOf(device.Peer{})
+	sf, ok := peerType.FieldByName("stats")
+	if !ok {
+		panic("no stats field in device.Peer")
+	}
+	if sf.Type.Kind() != reflect.Struct {
+		panic("stats field is not a struct")
+	}
+	base := sf.Offset
+
+	st := sf.Type
+	field, ok := st.FieldByName(name)
+	if !ok {
+		panic("no " + name + " field in device.Peer.stats")
+	}
+	if field.Type.Kind() != reflect.Int64 && field.Type.Kind() != reflect.Uint64 {
+		panic("unexpected kind of " + name + " field in device.Peer.stats")
+	}
+	return base + field.Offset
+}
+
+// PeerLastHandshakeNano returns the last handshake time in nanoseconds since the
+// unix epoch.
+func PeerLastHandshakeNano(peer *device.Peer) int64 {
+	return atomic.LoadInt64((*int64)(unsafe.Add(unsafe.Pointer(peer), offHandshake)))
+}
+
+// PeerRxBytes returns the number of bytes received from this peer.
+func PeerRxBytes(peer *device.Peer) uint64 {
+	return atomic.LoadUint64((*uint64)(unsafe.Add(unsafe.Pointer(peer), offRxBytes)))
+}
+
+// PeerTxBytes returns the number of bytes sent to this peer.
+func PeerTxBytes(peer *device.Peer) uint64 {
+	return atomic.LoadUint64((*uint64)(unsafe.Add(unsafe.Pointer(peer), offTxBytes)))
+}
diff --git a/wgengine/wgint/wgint_test.go b/wgengine/wgint/wgint_test.go
new file mode 100644
index 000000000..bfa0770a5
--- /dev/null
+++ b/wgengine/wgint/wgint_test.go
@@ -0,0 +1,24 @@
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package wgint
+
+import (
+	"testing"
+
+	"golang.zx2c4.com/wireguard/device"
+)
+
+func TestPeerStats(t *testing.T) {
+	peer := new(device.Peer)
+	if got := PeerLastHandshakeNano(peer); got != 0 {
+		t.Errorf("PeerLastHandshakeNano = %v, want 0", got)
+	}
+	if got := PeerRxBytes(peer); got != 0 {
+		t.Errorf("PeerRxBytes = %v, want 0", got)
+	}
+	if got := PeerTxBytes(peer); got != 0 {
+		t.Errorf("PeerTxBytes = %v, want 0", got)
+	}
+}