device: fix TOCTOU race during session state update (#77)

API introduced in a927a66e has two cases of state determination
happening out of critical section for the state value:

(1) expiredSession loads sessionExpiresNano, then releases all locks
and calls noteSessionState(Expired). So a concurrent refresh that
lands in that gap gets clobbered by a stale Expired -- and sticks
until the next re-key.

(2) Likewise in noteSessionHandshakeStopped, hasKeyMaterial check
happens out of the session state lock and races with ZeroAndFlushAll.

Both lead to a wrong state emitted via the device.sessionState.fn,
but are otherwise benign.

This moves the expiry timestamp under a lock to address the former,
and provides a noteSessionStateLocked helper for the latter.
Also changes API semantics to serialize events per-peer, to avoid
sharing a single lock for all timestamps.

Updates tailscale/corp#42874

Signed-off-by: Alex Valiushko <alexvaliushko@tailscale.com>
Change-Id: Iee2cdf135375519e58a8e84362349d966a6a6964

device/device.go

@@ -64,10 +64,7 @@ type Device struct {
 		lookupFunc   PeerLookupFunc // or nil if unused
 	}
 
-	sessionState struct {
+	peerStateFn atomic.Pointer[PeerSessionStateFunc] // observes peer session state changes, nil if unset
-		sync.Mutex // serializes PeerSessionStateFunc calls and protects peer.sessionState
-		fn         PeerSessionStateFunc
-	}
 
 	rate struct {
 		underLoadUntil atomic.Int64
@@ -519,7 +516,7 @@ const (
 
 // PeerSessionStateFunc is called when a peer's WireGuard session state changes.
 //
-// Calls are serialized per Device and delivered in transition order. The
+// Calls are serialized per peer and delivered in that peer's transition order. The
 // callback must be cheap and must not call back into Device.
 type PeerSessionStateFunc func(peer NoisePublicKey, state PeerSessionState)
 
@@ -546,10 +543,14 @@ func (device *Device) SetPeerByIPPacketFunc(f PeerByIPPacketFunc) {
 // It does not replay current state. Callers that need a complete view should set
 // it before peers are started or lazily created, and maintain any snapshots,
 // sequence numbers, and pubsub state outside wireguard-go.
+//
+// The callback must be concurrent-safe and must not call back into Device.
 func (device *Device) SetSessionStateFunc(f PeerSessionStateFunc) {
-	device.sessionState.Lock()
+	if f == nil {
-	defer device.sessionState.Unlock()
+		device.peerStateFn.Store(nil)
-	device.sessionState.fn = f
+		return
+	}
+	device.peerStateFn.Store(&f)
 }
 
 func (device *Device) Close() {

device/peer.go

@@ -26,8 +26,12 @@ type Peer struct {
 	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
+	sessionState struct {
+		sync.Mutex
+		current        PeerSessionState
+		sessionExpires time.Time
+	}
 
 	queuedOutboundPackets atomic.Int32 // packets in staged+outbound queues, for input backpressure
 
@@ -266,7 +270,6 @@ func (peer *Peer) Start() {
 
 func (peer *Peer) ZeroAndFlushAll() {
 	device := peer.device
-	peer.sessionExpiresNano.Store(0)
 	if peer.timers.sessionExpired != nil {
 		peer.timers.sessionExpired.Del()
 	}
@@ -292,7 +295,11 @@ func (peer *Peer) ZeroAndFlushAll() {
 	handshake.mutex.Unlock()
 
 	peer.FlushStagedPackets()
-	peer.noteSessionState(PeerSessionNone)
+
+	peer.sessionState.Lock()
+	peer.sessionState.sessionExpires = time.Time{}
+	peer.noteSessionStateLocked(PeerSessionNone)
+	peer.sessionState.Unlock()
 }
 
 func (peer *Peer) ExpireCurrentKeypairs() {
@@ -313,8 +320,10 @@ func (peer *Peer) ExpireCurrentKeypairs() {
 	}
 	keypairs.Unlock()
 
-	peer.sessionExpiresNano.Store(0)
+	peer.sessionState.Lock()
-	peer.noteSessionState(PeerSessionExpired)
+	peer.sessionState.sessionExpires = time.Time{}
+	peer.noteSessionStateLocked(PeerSessionExpired)
+	peer.sessionState.Unlock()
 }
 
 func (peer *Peer) Stop() {
@@ -338,42 +347,41 @@ func (peer *Peer) Stop() {
 }
 
 func (peer *Peer) noteSessionState(state PeerSessionState) {
-	device := peer.device
+	peer.sessionState.Lock()
-	device.sessionState.Lock()
+	defer peer.sessionState.Unlock()
-	defer device.sessionState.Unlock()
+	peer.noteSessionStateLocked(state)
+}
 
-	if peer.sessionState == state {
+// noteSessionStateLocked records a session state transition and delivers the
+// callback. The caller must hold peer.sessionState.Mutex during the
+// state determination and transition.
+func (peer *Peer) noteSessionStateLocked(state PeerSessionState) {
+	if peer.sessionState.current == state {
 		return
 	}
-	peer.sessionState = state
+	peer.sessionState.current = state
-	if f := device.sessionState.fn; f != nil {
+	if f := peer.device.peerStateFn.Load(); f != nil {
-		f(peer.handshake.remoteStatic, state)
+		(*f)(peer.handshake.remoteStatic, state)
 	}
 }
 
 func (peer *Peer) noteSessionHandshakeStarted() {
-	device := peer.device
+	peer.sessionState.Lock()
-	device.sessionState.Lock()
+	defer peer.sessionState.Unlock()
-	defer device.sessionState.Unlock()
+	if peer.sessionState.current == PeerSessionEstablished {
-
-	switch peer.sessionState {
-	case PeerSessionEstablished:
-		return
-	case PeerSessionHandshake:
 		return
 	}
-	peer.sessionState = PeerSessionHandshake
+	peer.noteSessionStateLocked(PeerSessionHandshake)
-	if f := device.sessionState.fn; f != nil {
-		f(peer.handshake.remoteStatic, PeerSessionHandshake)
-	}
 }
 
 func (peer *Peer) noteSessionHandshakeStopped() {
+	peer.sessionState.Lock()
+	defer peer.sessionState.Unlock()
 	state := PeerSessionNone
 	if peer.hasKeyMaterial() {
 		state = PeerSessionExpired
 	}
-	peer.noteSessionState(state)
+	peer.noteSessionStateLocked(state)
 }
 
 func (peer *Peer) hasKeyMaterial() bool {

device/timers.go

@@ -141,12 +141,13 @@ func expiredZeroKeyMaterial(peer *Peer) {
 }
 
 func expiredSession(peer *Peer) {
-	expires := peer.sessionExpiresNano.Load()
+	peer.sessionState.Lock()
-	if expires == 0 || time.Now().UnixNano() < expires {
+	defer peer.sessionState.Unlock()
+	if peer.sessionState.sessionExpires.IsZero() || time.Now().Before(peer.sessionState.sessionExpires) {
 		return
 	}
 	peer.device.log.Verbosef("%s - Session expired after %d seconds", peer, int(RejectAfterTime.Seconds()))
-	peer.noteSessionState(PeerSessionExpired)
+	peer.noteSessionStateLocked(PeerSessionExpired)
 }
 
 func expiredPersistentKeepalive(peer *Peer) {
@@ -208,11 +209,15 @@ func (peer *Peer) timersHandshakeComplete() {
 /* Should be called after an ephemeral key is created, which is before sending a handshake response or after receiving a handshake response. */
 func (peer *Peer) timersSessionDerived() {
 	if peer.timersActive() {
-		peer.sessionExpiresNano.Store(time.Now().Add(RejectAfterTime).UnixNano())
+		peer.sessionState.Lock()
+		peer.sessionState.sessionExpires = time.Now().Add(RejectAfterTime)
+		peer.noteSessionStateLocked(PeerSessionEstablished)
+		peer.sessionState.Unlock()
 		peer.timers.sessionExpired.Mod(RejectAfterTime)
 		peer.timers.zeroKeyMaterial.Mod(RejectAfterTime * 3)
-	}
+	} else {
 		peer.noteSessionState(PeerSessionEstablished)
+	}
 }
 
 /* Should be called before a packet with authentication -- keepalive, data, or handshake -- is sent, or after one is received. */