protect against time of check vs use race condition

Author: skarim

cmd/utils.go

@@ -138,14 +138,19 @@ func loadStack(cfg *config.Config, branch string) (*loadStackResult, error) {
 }
 
 // handleSaveError translates a stack.Save error into the appropriate user
-// message and exit error.  Lock contention returns ErrLockFailed (exit 8);
-// other write failures return ErrSilent (exit 1).
+// message and exit error.  Lock contention and stale-file detection both
+// return ErrLockFailed (exit 8); other write failures return ErrSilent (exit 1).
 func handleSaveError(cfg *config.Config, err error) error {
 	var lockErr *stack.LockError
 	if errors.As(err, &lockErr) {
 		cfg.Errorf("another process is currently editing the stack — try again later")
 		return ErrLockFailed
 	}
+	var staleErr *stack.StaleError
+	if errors.As(err, &staleErr) {
+		cfg.Errorf("stack file was modified by another process — please re-run the command")
+		return ErrLockFailed
+	}
 	cfg.Errorf("failed to save stack state: %s", err)
 	return ErrSilent
 }

internal/stack/lock.go

@@ -19,10 +19,20 @@ type LockError struct {
 func (e *LockError) Error() string { return e.Err.Error() }
 func (e *LockError) Unwrap() error { return e.Err }
 
+// StaleError is returned when the stack file was modified on disk since it
+// was loaded.  This indicates another process wrote to the file concurrently.
+// Callers can check for this with errors.As.
+type StaleError struct {
+	Err error
+}
+
+func (e *StaleError) Error() string { return e.Err.Error() }
+func (e *StaleError) Unwrap() error { return e.Err }
+
 // LockTimeout is how long Lock() will wait for the exclusive lock before
 // giving up.  With the lock held only during file writes (milliseconds),
 // this timeout primarily guards against a hung process holding the lock.
-const LockTimeout = 5 * time.Second
+var LockTimeout = 5 * time.Second
 
 // lockRetryInterval is the sleep between non-blocking lock attempts.
 const lockRetryInterval = 100 * time.Millisecond

internal/stack/lock_test.go

@@ -1,6 +1,7 @@
 package stack
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -77,7 +78,7 @@ func TestLock_SerializesConcurrentAccess(t *testing.T) {
 	require.NoError(t, Save(dir, sf))
 
 	// Run 10 concurrent goroutines, each adding a stack under lock.
-	// Uses Lock + Load + SaveLocked for atomic read-modify-write.
+	// Uses Lock + Load + writeStackFile for atomic read-modify-write.
 	errCh := make(chan error, 10)
 	var wg sync.WaitGroup
 	for i := 0; i < 10; i++ {
@@ -99,8 +100,8 @@ func TestLock_SerializesConcurrentAccess(t *testing.T) {
 			}
 
 			loaded.AddStack(makeStack("main", "branch"))
-			if err := SaveLocked(dir, loaded); err != nil {
-				errCh <- fmt.Errorf("goroutine %d SaveLocked: %w", idx, err)
+			if err := writeStackFile(dir, loaded); err != nil {
+				errCh <- fmt.Errorf("goroutine %d writeStackFile: %w", idx, err)
 			}
 		}(i)
 	}
@@ -132,3 +133,112 @@ func TestLock_FileLeftOnDisk(t *testing.T) {
 	require.NoError(t, err, "should be able to re-lock after unlock")
 	lock2.Unlock()
 }
+
+func TestLock_TimesOut(t *testing.T) {
+	dir := t.TempDir()
+
+	// Hold the lock so the second attempt can never acquire it.
+	lock1, err := Lock(dir)
+	require.NoError(t, err)
+	defer lock1.Unlock()
+
+	// Save original timeout and set a short one for the test.
+	origTimeout := LockTimeout
+	LockTimeout = 200 * time.Millisecond
+	defer func() { LockTimeout = origTimeout }()
+
+	start := time.Now()
+	lock2, err := Lock(dir)
+	elapsed := time.Since(start)
+
+	assert.Nil(t, lock2, "should not acquire lock")
+	require.Error(t, err)
+
+	var lockErr *LockError
+	require.True(t, errors.As(err, &lockErr), "error should be *LockError, got %T", err)
+	assert.Contains(t, lockErr.Error(), "timed out")
+
+	// Should have waited roughly LockTimeout before giving up.
+	assert.GreaterOrEqual(t, elapsed, 150*time.Millisecond, "should wait near the timeout")
+}
+
+func TestSave_DetectsStaleFile(t *testing.T) {
+	dir := t.TempDir()
+
+	// Write an initial stack file.
+	sf := &StackFile{SchemaVersion: 1, Stacks: []Stack{}}
+	require.NoError(t, Save(dir, sf))
+
+	// Load — captures the on-disk checksum.
+	loaded, err := Load(dir)
+	require.NoError(t, err)
+
+	// Simulate another process: load, modify, save.
+	other, err := Load(dir)
+	require.NoError(t, err)
+	other.AddStack(makeStack("main", "sneaky"))
+	require.NoError(t, Save(dir, other))
+
+	// Our loaded copy tries to save — should detect staleness.
+	loaded.AddStack(makeStack("main", "my-branch"))
+	err = Save(dir, loaded)
+	require.Error(t, err)
+
+	var staleErr *StaleError
+	require.True(t, errors.As(err, &staleErr), "error should be *StaleError, got %T", err)
+	assert.Contains(t, staleErr.Error(), "modified by another process")
+}
+
+func TestSave_AllowsWriteWhenFileUnchanged(t *testing.T) {
+	dir := t.TempDir()
+
+	// Write, load, modify, save — no concurrent changes.
+	sf := &StackFile{SchemaVersion: 1, Stacks: []Stack{}}
+	require.NoError(t, Save(dir, sf))
+
+	loaded, err := Load(dir)
+	require.NoError(t, err)
+
+	loaded.AddStack(makeStack("main", "feature"))
+	require.NoError(t, Save(dir, loaded))
+
+	// Verify the write actually persisted.
+	final, err := Load(dir)
+	require.NoError(t, err)
+	assert.Len(t, final.Stacks, 1)
+}
+
+func TestSave_AllowsFirstWrite(t *testing.T) {
+	dir := t.TempDir()
+
+	// File doesn't exist — Load returns nil checksum, Save should succeed.
+	sf, err := Load(dir)
+	require.NoError(t, err)
+	assert.Empty(t, sf.Stacks)
+
+	sf.AddStack(makeStack("main", "first"))
+	require.NoError(t, Save(dir, sf), "first save to a new file should succeed")
+
+	final, err := Load(dir)
+	require.NoError(t, err)
+	assert.Len(t, final.Stacks, 1)
+}
+
+func TestSave_DoubleSaveSucceeds(t *testing.T) {
+	dir := t.TempDir()
+
+	sf, err := Load(dir)
+	require.NoError(t, err)
+
+	sf.AddStack(makeStack("main", "first"))
+	require.NoError(t, Save(dir, sf), "first save should succeed")
+
+	// A second Save on the same instance must not spuriously fail —
+	// writeStackFile refreshes loadChecksum after writing.
+	sf.AddStack(makeStack("main", "second"))
+	require.NoError(t, Save(dir, sf), "second save on same instance should succeed")
+
+	final, err := Load(dir)
+	require.NoError(t, err)
+	assert.Len(t, final.Stacks, 2)
+}

internal/stack/stack.go

@@ -1,6 +1,8 @@
 package stack
 
 import (
+	"bytes"
+	"crypto/sha256"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -168,6 +170,11 @@ type StackFile struct {
 	SchemaVersion int     `json:"schemaVersion"`
 	Repository    string  `json:"repository"`
 	Stacks        []Stack `json:"stacks"`
+
+	// loadChecksum is the SHA-256 of the raw file bytes at Load time.
+	// Save uses it to detect concurrent modifications (optimistic concurrency).
+	// nil means the file did not exist when loaded.
+	loadChecksum []byte
 }
 
 // FindAllStacksForBranch returns all stacks that contain the given branch.
@@ -233,11 +240,14 @@ func stackFilePath(gitDir string) string {
 
 // Load reads the stack file from the given git directory.
 // Returns an empty StackFile if the file does not exist.
+// The returned StackFile records a checksum of the on-disk content so that
+// Save can detect concurrent modifications.
 func Load(gitDir string) (*StackFile, error) {
 	path := stackFilePath(gitDir)
 	data, err := os.ReadFile(path)
 	if err != nil {
 		if errors.Is(err, os.ErrNotExist) {
+			// loadChecksum stays nil — sentinel for "file absent at load time".
 			return &StackFile{
 				SchemaVersion: schemaVersion,
 				Stacks:        []Stack{},
@@ -255,24 +265,32 @@ func Load(gitDir string) (*StackFile, error) {
 		return nil, fmt.Errorf("stack file has schema version %d, but this version of gh-stack only supports up to version %d — please upgrade gh-stack", sf.SchemaVersion, schemaVersion)
 	}
 
+	sum := sha256.Sum256(data)
+	sf.loadChecksum = sum[:]
 	return &sf, nil
 }
 
-// Save acquires an exclusive lock on the stack file, writes sf as JSON, and
-// releases the lock.  The lock is held only for the duration of the write.
-// Returns *LockError if the lock times out due to contention.
+// Save acquires an exclusive lock on the stack file, verifies the file hasn't
+// been modified since Load (optimistic concurrency), writes sf as JSON, and
+// releases the lock.  The lock is held only for the read-compare-write window.
+// Returns *LockError if the lock times out, or *StaleError if another process
+// modified the file since it was loaded.
 func Save(gitDir string, sf *StackFile) error {
 	lock, err := Lock(gitDir)
 	if err != nil {
 		return err // *LockError for contention, plain error for I/O failures
 	}
 	defer lock.Unlock()
+
+	if err := checkStale(gitDir, sf); err != nil {
+		return err
+	}
 	return writeStackFile(gitDir, sf)
 }
 
 // SaveNonBlocking attempts to save without blocking.  If another process holds
-// the lock, the save is silently skipped.  Use this for best-effort metadata
-// persistence (e.g. syncing PR state in view).
+// the lock or the file was modified since Load, the save is silently skipped.
+// Use this for best-effort metadata persistence (e.g. syncing PR state in view).
 func SaveNonBlocking(gitDir string, sf *StackFile) {
 	path := filepath.Join(gitDir, lockFileName)
 	f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0644)
@@ -285,14 +303,46 @@ func SaveNonBlocking(gitDir string, sf *StackFile) {
 	}
 	lock := &FileLock{f: f}
 	defer lock.Unlock()
+
+	if checkStale(gitDir, sf) != nil {
+		return
+	}
 	_ = writeStackFile(gitDir, sf)
 }
 
-// SaveLocked writes the stack file without acquiring the lock.  The caller
-// must already hold the lock (via Lock) to protect the write.  Use this when
-// you need an atomic Load-Modify-Save sequence.
-func SaveLocked(gitDir string, sf *StackFile) error {
-	return writeStackFile(gitDir, sf)
+// checkStale compares the current on-disk content against the checksum
+// captured at Load time.  Returns *StaleError if the file was modified
+// by another process.  The caller must hold the lock.
+func checkStale(gitDir string, sf *StackFile) error {
+	path := stackFilePath(gitDir)
+	data, err := os.ReadFile(path)
+
+	if errors.Is(err, os.ErrNotExist) {
+		// File absent on disk.
+		if sf.loadChecksum == nil {
+			return nil // was absent at Load time too — no conflict
+		}
+		// File existed at Load but is now gone.  Allow the write to
+		// recreate it rather than erroring; this is not a lost-update.
+		return nil
+	}
+	if err != nil {
+		return fmt.Errorf("reading stack file for staleness check: %w", err)
+	}
+
+	// File exists on disk.
+	if sf.loadChecksum == nil {
+		// File was absent at Load but another process created it.
+		return &StaleError{Err: fmt.Errorf(
+			"stack file was created by another process since it was loaded")}
+	}
+
+	sum := sha256.Sum256(data)
+	if !bytes.Equal(sf.loadChecksum, sum[:]) {
+		return &StaleError{Err: fmt.Errorf(
+			"stack file was modified by another process since it was loaded")}
+	}
+	return nil
 }
 
 func writeStackFile(gitDir string, sf *StackFile) error {
@@ -308,5 +358,9 @@ func writeStackFile(gitDir string, sf *StackFile) error {
 	if err := os.WriteFile(path, data, 0644); err != nil {
 		return fmt.Errorf("writing stack file: %w", err)
 	}
+	// Refresh checksum so a second Save on the same StackFile doesn't
+	// spuriously fail the staleness check.
+	sum := sha256.Sum256(data)
+	sf.loadChecksum = sum[:]
 	return nil
 }