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
This commit is contained in:
Alex Valiushko 2026-06-24 14:09:25 -07:00 committed by 世界
parent 35a60acb84
commit 15b912c1c0
No known key found for this signature in database
GPG key ID: CD109927C34A63C4
3 changed files with 61 additions and 47 deletions

View file

@ -64,10 +64,7 @@ type Device struct {
lookupFunc PeerLookupFunc // or nil if unused
}
sessionState struct {
sync.Mutex // serializes PeerSessionStateFunc calls and protects peer.sessionState
fn PeerSessionStateFunc
}
peerStateFn atomic.Pointer[PeerSessionStateFunc] // observes peer session state changes, nil if unset
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()
defer device.sessionState.Unlock()
device.sessionState.fn = f
if f == nil {
device.peerStateFn.Store(nil)
return
}
device.peerStateFn.Store(&f)
}
func (device *Device) Close() {