wallet: strictly enforce CAS in lifecycle state transitions

Author: yyforyongyu

wallet/state.go

@@ -136,8 +136,15 @@ func (s *walletState) toStarting() error {
 
 // toStarted marks the wallet as fully started. This should be called only
 // after all resource initialization is complete.
-func (s *walletState) toStarted() {
-	s.lifecycle.Store(uint32(lifecycleStarted))
+func (s *walletState) toStarted() error {
+	if !s.lifecycle.CompareAndSwap(
+		uint32(lifecycleStarting), uint32(lifecycleStarted)) {
+
+		return fmt.Errorf("%w: cannot transition to started from %v",
+			ErrStateForbidden, lifecycle(s.lifecycle.Load()))
+	}
+
+	return nil
 }
 
 // toStopping transitions the wallet from Started to Stopping.
@@ -160,11 +167,40 @@ func (s *walletState) toStopping() error {
 }
 
 // toStopped marks the wallet as fully stopped.
-func (s *walletState) toStopped() {
-	s.lifecycle.Store(uint32(lifecycleStopped))
+func (s *walletState) toStopped() error {
+	// We allow transition from Stopping (normal shutdown) or Starting
+	// (failure during startup).
+	//
+	// We use a CAS loop here to handle potential races where the state
+	// might change between Load and CompareAndSwap.
+	//
+	// This loop is guaranteed to terminate because:
+	// 1. If CAS succeeds, we break.
+	// 2. If CAS fails, it means the state changed. We reload the new state.
+	// 3. If the new state is not Stopping or Starting (e.g. it became
+	//    Started or already Stopped), the validation check fails and we
+	//    return an error.
+	for {
+		current := s.lifecycle.Load()
+		lc := lifecycle(current)
+
+		if lc != lifecycleStopping && lc != lifecycleStarting {
+			return fmt.Errorf("%w: cannot transition to stopped "+
+				"from %v", ErrStateForbidden, lc)
+		}
+
+		if s.lifecycle.CompareAndSwap(
+			current, uint32(lifecycleStopped),
+		) {
+
+			break
+		}
+	}
 
 	// Force lock the wallet on shutdown for security.
 	s.unlocked.Store(false)
+
+	return nil
 }
 
 // toUnlocked marks the wallet as unlocked.

wallet/state_test.go

@@ -25,7 +25,8 @@ func TestStateSecureByDefault(t *testing.T) {
 	require.NoError(t, err)
 
 	// Act: Transition to Started.
-	s.toStarted()
+	err = s.toStarted()
+	require.NoError(t, err)
 	require.True(t, s.isStarted())
 	require.True(t, s.isRunning())
 
@@ -38,7 +39,8 @@ func TestStateSecureByDefault(t *testing.T) {
 	require.False(t, s.isRunning())
 
 	// Act: Transition to Stopped.
-	s.toStopped()
+	err = s.toStopped()
+	require.NoError(t, err)
 	require.False(t, s.isRunning())
 
 	// Assert: Invalid transition (Stop when already Stopped).
@@ -54,22 +56,26 @@ func TestStateAuthentication(t *testing.T) {
 	s := newWalletState(syncer)
 
 	// Arrange: Start the wallet (must be started to be useful).
-	s.toStarted()
+	require.NoError(t, s.toStarting())
+	err := s.toStarted()
+	require.NoError(t, err)
 
 	// Assert: Default is Locked.
 	require.False(t, s.isUnlocked())
 
 	// Act: Unlock.
 	s.toUnlocked()
+	require.NoError(t, err)
 	require.True(t, s.isUnlocked())
 
 	// Act: Lock.
 	s.toLocked()
+	require.NoError(t, err)
 	require.False(t, s.isUnlocked())
 
 	// Act: Verify canSign checks.
 	// Case 1: Locked -> Error.
-	err := s.canSign()
+	err = s.canSign()
 	require.ErrorIs(t, err, ErrStateForbidden)
 	require.ErrorContains(t, err, "wallet locked")
 
@@ -80,7 +86,9 @@ func TestStateAuthentication(t *testing.T) {
 
 	// Case 3: Stopped -> Error (even if unlocked, though stopped forces
 	// lock).
-	s.toStopped()
+	require.NoError(t, s.toStopping())
+	err = s.toStopped()
+	require.NoError(t, err)
 	// Note: toStopped forces lock, so we must check that logic too.
 	require.False(t, s.isUnlocked())
 
@@ -98,7 +106,8 @@ func TestStateSynchronization(t *testing.T) {
 
 	syncer := &mockChainSyncer{}
 	s := newWalletState(syncer)
-	s.toStarted()
+	require.NoError(t, s.toStarting())
+	require.NoError(t, s.toStarted())
 
 	// Arrange: Mock syncer to return Synced.
 	syncer.On("syncState").Return(syncStateSynced)
@@ -150,8 +159,18 @@ func TestStateThreadSafety(t *testing.T) {
 			defer wg.Done()
 
 			<-start
+
+			// NOTE: We ignore errors here because we are
+			// purposefully hammering the state machine from
+			// multiple goroutines. Many of these transitions will
+			// fail (e.g., trying to start an already starting
+			// wallet), which is expected behavior. We are
+			// primarily verifying that no data races or panics
+			// occur.
+			//
 			// Try to start.
 			_ = s.toStarting()
+
 			// Try to stop.
 			_ = s.toStopping()
 		}()
@@ -178,7 +197,8 @@ func TestValidateSynced(t *testing.T) {
 	require.ErrorIs(t, err, ErrStateForbidden)
 
 	// Case 2: Started but not synced.
-	s.toStarted()
+	require.NoError(t, s.toStarting())
+	require.NoError(t, s.toStarted())
 	syncer.On("syncState").Return(syncStateSyncing)
 
 	err = s.validateSynced()
@@ -273,7 +293,8 @@ func TestStateStartStop(t *testing.T) {
 		require.False(t, state.unlocked.Load())
 
 		// Now mark as started.
-		state.toStarted()
+		err = state.toStarted()
+		require.NoError(t, err)
 		require.Equal(t, uint32(lifecycleStarted),
 			state.lifecycle.Load())
 	})
@@ -410,7 +431,8 @@ func TestStateAuxiliaryMethods(t *testing.T) {
 	require.ErrorIs(t, s.canChangePassphrase(), ErrStateForbidden)
 
 	// Case 2: Started -> All allowed.
-	s.toStarted()
+	require.NoError(t, s.toStarting())
+	require.NoError(t, s.toStarted())
 	require.NoError(t, s.canUnlock())
 	require.NoError(t, s.canLock())
 	require.NoError(t, s.canChangePassphrase())