From 1cf85822d0a0daf42e3e26e5b646947bfabf21bd Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Sun, 25 Feb 2024 06:40:35 -0800
Subject: [PATCH] ipn/ipnstate, wgengine/wgint: add handshake attempts
 accessors

Not yet used. This is being made available so magicsock/wgengine can
use it to ignore certain sends (UDP + DERP) later on at least mobile,
letting wireguard-go think it's doing its full attempt schedule, but
we can cut it short conditionally based on what we know from the
control plane.

Updates #7617

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
Change-Id: Ia367cf6bd87b2aeedd3c6f4989528acdb6773ca7
---
 ipn/ipnstate/ipnstate.go     |  4 ++++
 wgengine/userspace.go        |  1 +
 wgengine/wgint/wgint.go      | 24 ++++++++++++++++++++++++
 wgengine/wgint/wgint_test.go |  5 ++++-
 4 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/ipn/ipnstate/ipnstate.go b/ipn/ipnstate/ipnstate.go
index 4a3c80106..5292bac47 100644
--- a/ipn/ipnstate/ipnstate.go
+++ b/ipn/ipnstate/ipnstate.go
@@ -196,6 +196,10 @@ type PeerStatusLite struct {
 	LastHandshake time.Time
 	// NodeKey is this peer's public node key.
 	NodeKey key.NodePublic
+	// HandshakeAttempts is how many failed attempts there have been at
+	// completing the current WireGuard handshake. This resets to zero on every
+	// successful handshake.
+	HandshakeAttempts uint32
 }
 
 // PeerStatus describes a peer node and its current state.
diff --git a/wgengine/userspace.go b/wgengine/userspace.go
index 93eb53bf2..b965d77e5 100644
--- a/wgengine/userspace.go
+++ b/wgengine/userspace.go
@@ -1026,6 +1026,7 @@ func (e *userspaceEngine) getPeerStatusLite(pk key.NodePublic) (status ipnstate.
 	status.RxBytes = int64(wgint.PeerRxBytes(peer))
 	status.TxBytes = int64(wgint.PeerTxBytes(peer))
 	status.LastHandshake = time.Unix(0, wgint.PeerLastHandshakeNano(peer))
+	status.HandshakeAttempts = wgint.PeerHandshakeAttempts(peer)
 	return status, true
 }
 
diff --git a/wgengine/wgint/wgint.go b/wgengine/wgint/wgint.go
index 09c6ccab4..ce8109d5f 100644
--- a/wgengine/wgint/wgint.go
+++ b/wgengine/wgint/wgint.go
@@ -17,6 +17,8 @@ var (
 	offHandshake = getPeerStatsOffset("lastHandshakeNano")
 	offRxBytes   = getPeerStatsOffset("rxBytes")
 	offTxBytes   = getPeerStatsOffset("txBytes")
+
+	offHandshakeAttempts = getPeerHandshakeAttemptsOffset()
 )
 
 func getPeerStatsOffset(name string) uintptr {
@@ -31,6 +33,22 @@ func getPeerStatsOffset(name string) uintptr {
 	return field.Offset
 }
 
+func getPeerHandshakeAttemptsOffset() uintptr {
+	peerType := reflect.TypeFor[device.Peer]()
+	field, ok := peerType.FieldByName("timers")
+	if !ok {
+		panic("no timers field in device.Peer")
+	}
+	field2, ok := field.Type.FieldByName("handshakeAttempts")
+	if !ok {
+		panic("no handshakeAttempts field in device.Peer.timers")
+	}
+	if g, w := field2.Type.String(), "atomic.Uint32"; g != w {
+		panic("unexpected type " + g + " of field handshakeAttempts in device.Peer.timers; want " + w)
+	}
+	return field.Offset + field2.Offset
+}
+
 // PeerLastHandshakeNano returns the last handshake time in nanoseconds since the
 // unix epoch.
 func PeerLastHandshakeNano(peer *device.Peer) int64 {
@@ -46,3 +64,9 @@ func PeerRxBytes(peer *device.Peer) uint64 {
 func PeerTxBytes(peer *device.Peer) uint64 {
 	return (*atomic.Uint64)(unsafe.Add(unsafe.Pointer(peer), offTxBytes)).Load()
 }
+
+// PeerHandshakeAttempts returns the number of WireGuard handshake attempts
+// made for the current handshake. It resets to zero before every new handshake.
+func PeerHandshakeAttempts(peer *device.Peer) uint32 {
+	return (*atomic.Uint32)(unsafe.Add(unsafe.Pointer(peer), offHandshakeAttempts)).Load()
+}
diff --git a/wgengine/wgint/wgint_test.go b/wgengine/wgint/wgint_test.go
index 9eae29c7f..3e9e85e92 100644
--- a/wgengine/wgint/wgint_test.go
+++ b/wgengine/wgint/wgint_test.go
@@ -9,7 +9,7 @@ import (
 	"github.com/tailscale/wireguard-go/device"
 )
 
-func TestPeerStats(t *testing.T) {
+func TestInternalOffsets(t *testing.T) {
 	peer := new(device.Peer)
 	if got := PeerLastHandshakeNano(peer); got != 0 {
 		t.Errorf("PeerLastHandshakeNano = %v, want 0", got)
@@ -20,4 +20,7 @@ func TestPeerStats(t *testing.T) {
 	if got := PeerTxBytes(peer); got != 0 {
 		t.Errorf("PeerTxBytes = %v, want 0", got)
 	}
+	if got := PeerHandshakeAttempts(peer); got != 0 {
+		t.Errorf("PeerHandshakeAttempts = %v, want 0", got)
+	}
 }