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() {

View file

@ -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.noteSessionState(PeerSessionExpired)
peer.sessionState.Lock()
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
device.sessionState.Lock()
defer device.sessionState.Unlock()
peer.sessionState.Lock()
defer peer.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
if f := device.sessionState.fn; f != nil {
f(peer.handshake.remoteStatic, state)
peer.sessionState.current = state
if f := peer.device.peerStateFn.Load(); 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:
peer.sessionState.Lock()
defer peer.sessionState.Unlock()
if peer.sessionState.current == PeerSessionEstablished {
return
}
peer.sessionState = PeerSessionHandshake
if f := device.sessionState.fn; f != nil {
f(peer.handshake.remoteStatic, PeerSessionHandshake)
}
peer.noteSessionStateLocked(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 {

View file

@ -141,12 +141,13 @@ func expiredZeroKeyMaterial(peer *Peer) {
}
func expiredSession(peer *Peer) {
expires := peer.sessionExpiresNano.Load()
if expires == 0 || time.Now().UnixNano() < expires {
peer.sessionState.Lock()
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,12 +209,16 @@ 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. */
func (peer *Peer) timersAnyAuthenticatedPacketTraversal() {