From 412e04b9a214d10579eb51cf5e82961277bfaf8e Mon Sep 17 00:00:00 2001
From: full-bars <45684698+full-bars@users.noreply.github.com>
Date: Mon, 11 May 2026 17:08:02 -0700
Subject: [PATCH] device: expire peer sessions on drain and initiate handshakes
 on startup

Without a way to notify clients of a server restart, clients wait up to
RekeyAfterTime (120 s) before re-handshaking with the new instance.

Two changes to close that gap:

- upLocked: call SendHandshakeInitiation for every peer when the device
  comes up. The new server proactively reaches out to all configured
  peers so they can re-establish in under RekeyTimeout (5 s) rather
  than waiting for natural session expiry.

- DrainPeers / Config.Drain: add an explicit drain signal. Calling
  DrainPeers() (or setting Drain: true in IpcSet2) expires all current
  keypairs so the next send from each client triggers an immediate
  re-handshake instead of silently using a session the restarted server
  no longer knows about.
---
 device/device.go | 15 +++++++++++++++
 device/uapi.go   | 11 +++++++++++
 2 files changed, 26 insertions(+)

diff --git a/device/device.go b/device/device.go
index 40eb98f22..c91cbbc3c 100644
--- a/device/device.go
+++ b/device/device.go
@@ -183,6 +183,7 @@ func (device *Device) upLocked() error {
 	device.peers.RLock()
 	for _, peer := range device.peers.keyMap {
 		peer.Start()
+		peer.SendHandshakeInitiation(false)
 		if peer.persistentKeepaliveInterval.Load() > 0 {
 			peer.SendKeepalive()
 		}
@@ -191,6 +192,20 @@ func (device *Device) upLocked() error {
 	return nil
 }
 
+// DrainPeers expires the current keypairs for all peers, causing each peer to
+// immediately initiate a new handshake on the next send attempt. Call this
+// before gracefully shutting down or restarting the device so that clients
+// re-establish sessions with the new instance quickly rather than waiting up
+// to RekeyAfterTime (120 s) for natural expiry.
+func (device *Device) DrainPeers() {
+	device.peers.RLock()
+	defer device.peers.RUnlock()
+	for _, peer := range device.peers.keyMap {
+		peer.ExpireCurrentKeypairs()
+	}
+	device.log.Verbosef("DrainPeers: expired keypairs for %d peer(s)", len(device.peers.keyMap))
+}
+
 // downLocked attempts to bring the device down.
 // The caller must hold device.state.mu and is responsible for updating device.state.state.
 func (device *Device) downLocked() error {
diff --git a/device/uapi.go b/device/uapi.go
index b3d323c9d..d0be315ec 100644
--- a/device/uapi.go
+++ b/device/uapi.go
@@ -24,6 +24,10 @@ const (
 type Config struct {
 	BindIpv4 *string
 	BindIpv6 *string
+	// Drain, when true, expires all peer keypairs before applying any other
+	// config changes. Useful for signalling an imminent server restart so that
+	// clients re-handshake quickly rather than waiting for natural session expiry.
+	Drain bool
 
 	wgtypes.Config
 }
@@ -198,6 +202,13 @@ func (device *Device) IpcSet2(deviceConfig *Config) (err error) {
 		}
 	}()
 
+	// Drain peers before applying any other changes if requested.
+
+	if deviceConfig.Drain {
+		device.log.Verbosef("UAPI: Draining all peer sessions")
+		device.DrainPeers()
+	}
+
 	// Configuring Device //
 
 	// private key