Fix TOCTOU race condition in Docker image download logic

pkg/cli/docker_images.go

@@ -34,18 +34,18 @@ var pullState = &dockerPullState{
 	mockAvailable: make(map[string]bool),
 }
 
-// IsDockerImageAvailable checks if a Docker image is available locally
-func IsDockerImageAvailable(image string) bool {
+// isDockerImageAvailableUnlocked checks if a Docker image is available locally
+// This function must be called with pullState.mu held (either RLock or Lock)
+func isDockerImageAvailableUnlocked(image string) bool {
 	// Check if we're in mock mode (for testing)
-	pullState.mu.RLock()
 	if pullState.mockAvailableInUse {
 		available := pullState.mockAvailable[image]
-		pullState.mu.RUnlock()
 		dockerImagesLog.Printf("Mock: Checking if image %s is available: %v", image, available)
 		return available
 	}
-	pullState.mu.RUnlock()
 
+	// For non-mock mode, we need to execute docker command
+	// This is safe to do under lock since it's just a subprocess call
 	cmd := exec.Command("docker", "image", "inspect", image)
 	// Suppress output - we only care about exit code
 	cmd.Stdout = nil
@@ -56,6 +56,13 @@ func IsDockerImageAvailable(image string) bool {
 	return available
 }
 
+// IsDockerImageAvailable checks if a Docker image is available locally
+func IsDockerImageAvailable(image string) bool {
+	pullState.mu.RLock()
+	defer pullState.mu.RUnlock()
+	return isDockerImageAvailableUnlocked(image)
+}
+
 // IsDockerImageDownloading checks if a Docker image is currently being downloaded
 func IsDockerImageDownloading(image string) bool {
 	pullState.mu.RLock()
@@ -66,21 +73,23 @@ func IsDockerImageDownloading(image string) bool {
 // StartDockerImageDownload starts downloading a Docker image in the background
 // Returns true if download was started, false if already downloading or available
 func StartDockerImageDownload(image string) bool {
-	// First check if already available
-	if IsDockerImageAvailable(image) {
+	// Check availability and downloading status atomically under lock
+	pullState.mu.Lock()
+	defer pullState.mu.Unlock()
+
+	// Check if already available (inside lock for atomicity)
+	if isDockerImageAvailableUnlocked(image) {
 		dockerImagesLog.Printf("Image %s is already available", image)
 		return false
 	}
 
 	// Check if already downloading
-	pullState.mu.Lock()
 	if pullState.downloading[image] {
-		pullState.mu.Unlock()
 		dockerImagesLog.Printf("Image %s is already downloading", image)
 		return false
 	}
+
 	pullState.downloading[image] = true
-	pullState.mu.Unlock()
 
 	// Start the download in a goroutine with retry logic
 	go func() {

pkg/cli/docker_images_test.go

@@ -261,3 +261,153 @@ func TestMockImageAvailability(t *testing.T) {
 	// Clean up
 	ResetDockerPullState()
 }
+
+func TestStartDockerImageDownload_ConcurrentCalls(t *testing.T) {
+	// Reset state before test
+	ResetDockerPullState()
+
+	testImage := "test/concurrent-image:v1.0.0"
+
+	// Mock the image as not available
+	SetMockImageAvailable(testImage, false)
+
+	// Track how many times StartDockerImageDownload returns true (indicating it started a download)
+	const numGoroutines = 10
+	started := make([]bool, numGoroutines)
+
+	// Use a channel to synchronize all goroutines to start at roughly the same time
+	startChan := make(chan struct{})
+	doneChan := make(chan int, numGoroutines)
+
+	// Launch multiple goroutines that all try to start downloading the same image
+	for i := 0; i < numGoroutines; i++ {
+		go func(index int) {
+			<-startChan // Wait for the signal to start
+			started[index] = StartDockerImageDownload(testImage)
+			doneChan <- index
+		}(i)
+	}
+
+	// Signal all goroutines to start simultaneously
+	close(startChan)
+
+	// Wait for all goroutines to finish
+	for i := 0; i < numGoroutines; i++ {
+		<-doneChan
+	}
+
+	// Count how many goroutines successfully started a download
+	downloadCount := 0
+	for _, didStart := range started {
+		if didStart {
+			downloadCount++
+		}
+	}
+
+	// Only ONE goroutine should have successfully started the download
+	if downloadCount != 1 {
+		t.Errorf("Expected exactly 1 goroutine to start download, but %d did", downloadCount)
+	}
+
+	// Verify the image is marked as downloading
+	if !IsDockerImageDownloading(testImage) {
+		t.Error("Expected image to be marked as downloading")
+	}
+
+	// Clean up
+	ResetDockerPullState()
+}
+
+func TestStartDockerImageDownload_ConcurrentCallsWithAvailableImage(t *testing.T) {
+	// Reset state before test
+	ResetDockerPullState()
+
+	testImage := "test/concurrent-available-image:v1.0.0"
+
+	// Mock the image as available
+	SetMockImageAvailable(testImage, true)
+
+	// Track how many times StartDockerImageDownload returns true
+	const numGoroutines = 10
+	started := make([]bool, numGoroutines)
+
+	// Use a channel to synchronize all goroutines
+	startChan := make(chan struct{})
+	doneChan := make(chan int, numGoroutines)
+
+	// Launch multiple goroutines
+	for i := 0; i < numGoroutines; i++ {
+		go func(index int) {
+			<-startChan
+			started[index] = StartDockerImageDownload(testImage)
+			doneChan <- index
+		}(i)
+	}
+
+	// Signal all goroutines to start
+	close(startChan)
+
+	// Wait for all to finish
+	for i := 0; i < numGoroutines; i++ {
+		<-doneChan
+	}
+
+	// Count successful starts
+	downloadCount := 0
+	for _, didStart := range started {
+		if didStart {
+			downloadCount++
+		}
+	}
+
+	// NO goroutine should have started a download since image is available
+	if downloadCount != 0 {
+		t.Errorf("Expected 0 goroutines to start download for available image, but %d did", downloadCount)
+	}
+
+	// Verify the image is NOT marked as downloading
+	if IsDockerImageDownloading(testImage) {
+		t.Error("Expected image to not be marked as downloading since it's available")
+	}
+
+	// Clean up
+	ResetDockerPullState()
+}
+
+func TestStartDockerImageDownload_RaceWithExternalDownload(t *testing.T) {
+	// This test simulates the scenario where an image becomes available
+	// (e.g., externally downloaded) between when we check availability
+	// and when we mark it as downloading
+	ResetDockerPullState()
+
+	testImage := "test/race-image:v1.0.0"
+
+	// Initially not available
+	SetMockImageAvailable(testImage, false)
+
+	// Start multiple goroutines attempting to download
+	const numGoroutines = 5
+	results := make(chan bool, numGoroutines)
+
+	for i := 0; i < numGoroutines; i++ {
+		go func() {
+			results <- StartDockerImageDownload(testImage)
+		}()
+	}
+
+	// Collect results
+	downloadStarts := 0
+	for i := 0; i < numGoroutines; i++ {
+		if <-results {
+			downloadStarts++
+		}
+	}
+
+	// Should only have one successful start
+	if downloadStarts != 1 {
+		t.Errorf("Expected exactly 1 download to start, got %d", downloadStarts)
+	}
+
+	// Clean up
+	ResetDockerPullState()
+}