device: add peer session state callback

Add a minimal callback API for observing WireGuard peer session state
changes.

Updates tailscale/corp#42874

device/peer.go

@@ -18,14 +18,16 @@ import (
 )
 
 type Peer struct {
-	isRunning         atomic.Bool
-	keypairs          Keypairs
-	handshake         Handshake
-	device            *Device
-	stopping          sync.WaitGroup // routines pending stop
-	txBytes           atomic.Uint64  // bytes send to peer (endpoint)
-	rxBytes           atomic.Uint64  // bytes received from peer
-	lastHandshakeNano atomic.Int64   // nano seconds since epoch
+	isRunning          atomic.Bool
+	keypairs           Keypairs
+	handshake          Handshake
+	device             *Device
+	stopping           sync.WaitGroup   // routines pending stop
+	txBytes            atomic.Uint64    // bytes send to peer (endpoint)
+	rxBytes            atomic.Uint64    // bytes received from peer
+	lastHandshakeNano  atomic.Int64     // nano seconds since epoch
+	sessionExpiresNano atomic.Int64     // nano seconds since epoch
+	sessionState       PeerSessionState // guarded by device.sessionState.Mutex
 
 	queuedOutboundPackets atomic.Int32 // packets in staged+outbound queues, for input backpressure
 
@@ -46,6 +48,7 @@ type Peer struct {
 		retransmitHandshake     *Timer
 		sendKeepalive           *Timer
 		newHandshake            *Timer
+		sessionExpired          *Timer
 		zeroKeyMaterial         *Timer
 		persistentKeepalive     *Timer
 		handshakeAttempts       atomic.Uint32
@@ -255,6 +258,10 @@ func (peer *Peer) Start() {
 
 func (peer *Peer) ZeroAndFlushAll() {
 	device := peer.device
+	peer.sessionExpiresNano.Store(0)
+	if peer.timers.sessionExpired != nil {
+		peer.timers.sessionExpired.Del()
+	}
 
 	// clear key pairs
 
@@ -277,6 +284,7 @@ func (peer *Peer) ZeroAndFlushAll() {
 	handshake.mutex.Unlock()
 
 	peer.FlushStagedPackets()
+	peer.noteSessionState(PeerSessionNone)
 }
 
 func (peer *Peer) ExpireCurrentKeypairs() {
@@ -296,6 +304,9 @@ func (peer *Peer) ExpireCurrentKeypairs() {
 		next.sendNonce.Store(RejectAfterMessages)
 	}
 	keypairs.Unlock()
+
+	peer.sessionExpiresNano.Store(0)
+	peer.noteSessionState(PeerSessionExpired)
 }
 
 func (peer *Peer) Stop() {
@@ -318,6 +329,52 @@ func (peer *Peer) Stop() {
 	peer.ZeroAndFlushAll()
 }
 
+func (peer *Peer) noteSessionState(state PeerSessionState) {
+	device := peer.device
+	device.sessionState.Lock()
+	defer device.sessionState.Unlock()
+
+	if peer.sessionState == state {
+		return
+	}
+	peer.sessionState = state
+	if f := device.sessionState.fn; f != nil {
+		f(peer.handshake.remoteStatic, state)
+	}
+}
+
+func (peer *Peer) noteSessionHandshakeStarted() {
+	device := peer.device
+	device.sessionState.Lock()
+	defer device.sessionState.Unlock()
+
+	switch peer.sessionState {
+	case PeerSessionEstablished:
+		return
+	case PeerSessionHandshake:
+		return
+	}
+	peer.sessionState = PeerSessionHandshake
+	if f := device.sessionState.fn; f != nil {
+		f(peer.handshake.remoteStatic, PeerSessionHandshake)
+	}
+}
+
+func (peer *Peer) noteSessionHandshakeStopped() {
+	state := PeerSessionNone
+	if peer.hasKeyMaterial() {
+		state = PeerSessionExpired
+	}
+	peer.noteSessionState(state)
+}
+
+func (peer *Peer) hasKeyMaterial() bool {
+	keypairs := &peer.keypairs
+	keypairs.RLock()
+	defer keypairs.RUnlock()
+	return keypairs.previous != nil || keypairs.current != nil || keypairs.next.Load() != nil
+}
+
 func (peer *Peer) SetEndpointFromPacket(endpoint conn.Endpoint) {
 	peer.endpoint.Lock()
 	defer peer.endpoint.Unlock()