Merge pull request #101 from buildkite/feat_clean_before_restore

feat: adds clean before restore to avoid merging data

Author: wolfeidau

restore.go

@@ -2,9 +2,13 @@ package zstash
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"io/fs"
+	"log/slog"
 	"os"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"time"
 
@@ -151,6 +155,25 @@ func (c *Cache) Restore(ctx context.Context, cacheID string) (RestoreResult, err
 		Concurrency:      transferInfo.Concurrency,
 	}
 
+	c.callProgress(cacheID, "cleaning", "Cleaning paths", 0, 0)
+
+	for _, path := range cacheConfig.Paths {
+		extractedPath, err := archive.ResolveHomeDir(path)
+		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, "failed to resolve home dir")
+			return result, fmt.Errorf("failed to resolve home dir for %q: %w", path, err)
+		}
+
+		slog.Debug("cleaning path", "path", path, "extractedPath", extractedPath)
+
+		if err := cleanPath(ctx, extractedPath); err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, "failed to clean path")
+			return result, fmt.Errorf("failed to clean path %q: %w", extractedPath, err)
+		}
+	}
+
 	c.callProgress(cacheID, "extracting", "Extracting files from cache", 0, int(transferInfo.BytesTransferred))
 
 	// Extract files
@@ -279,3 +302,70 @@ func (c *Cache) extractCache(ctx context.Context, archiveFile string, archiveSiz
 
 	return archiveInfo, nil
 }
+
+// cleanPath removes a directory tree for a configured cache path.
+// It handles Go module cache directories that have 0555 permissions by
+// making them writable before removal.
+func cleanPath(ctx context.Context, dir string) error {
+	if dir == "" {
+		return fmt.Errorf("cleanPath: empty directory path")
+	}
+
+	clean := filepath.Clean(dir)
+
+	// Refuse to delete root or current directory
+	if clean == "." || clean == string(os.PathSeparator) {
+		return fmt.Errorf("cleanPath: refusing to remove %q", clean)
+	}
+
+	// On Windows, also check for drive roots like "C:\"
+	if runtime.GOOS == "windows" && len(clean) == 3 && clean[1] == ':' && clean[2] == '\\' {
+		return fmt.Errorf("cleanPath: refusing to remove drive root %q", clean)
+	}
+
+	// Refuse to delete home directory
+	if home, err := os.UserHomeDir(); err == nil {
+		if clean == filepath.Clean(home) {
+			return fmt.Errorf("cleanPath: refusing to remove home directory %q", clean)
+		}
+	}
+
+	// Module cache has 0555 directories; make them writable in order to remove content.
+	err := filepath.WalkDir(clean, func(path string, info fs.DirEntry, walkErr error) error {
+		// Respect context cancellation for long directory trees
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
+		if walkErr != nil {
+			slog.Debug("cleanPath: error walking path", "path", path, "err", walkErr)
+			return nil
+		}
+
+		if info.IsDir() {
+			if chmodErr := os.Chmod(path, 0o755); chmodErr != nil {
+				return fmt.Errorf("chmod %q: %w", path, chmodErr)
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			return err
+		}
+		return fmt.Errorf("cleanPath: error preparing %q for removal: %w", clean, err)
+	}
+
+	// Check context again before potentially long RemoveAll
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
+	if err := os.RemoveAll(clean); err != nil {
+		return fmt.Errorf("cleanPath: failed to remove %q: %w", clean, err)
+	}
+
+	return nil
+}

restore_test.go

@@ -0,0 +1,99 @@
+package zstash
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCleanPath(t *testing.T) {
+	t.Run("removes directory and contents", func(t *testing.T) {
+		dir := t.TempDir()
+		testDir := filepath.Join(dir, "cache")
+		require.NoError(t, os.MkdirAll(filepath.Join(testDir, "subdir"), 0o755))
+		require.NoError(t, os.WriteFile(filepath.Join(testDir, "file.txt"), []byte("test"), 0o600))
+		require.NoError(t, os.WriteFile(filepath.Join(testDir, "subdir", "nested.txt"), []byte("nested"), 0o600))
+
+		err := cleanPath(context.Background(), testDir)
+		require.NoError(t, err)
+
+		_, err = os.Stat(testDir)
+		assert.True(t, os.IsNotExist(err), "directory should be removed")
+	})
+
+	t.Run("handles read-only directories (like go module cache)", func(t *testing.T) {
+		dir := t.TempDir()
+		testDir := filepath.Join(dir, "modcache")
+		subdir := filepath.Join(testDir, "pkg")
+		require.NoError(t, os.MkdirAll(subdir, 0o755))
+		require.NoError(t, os.WriteFile(filepath.Join(subdir, "mod.go"), []byte("package mod"), 0o400))
+		require.NoError(t, os.Chmod(subdir, 0o555))
+		require.NoError(t, os.Chmod(testDir, 0o555))
+
+		err := cleanPath(context.Background(), testDir)
+		require.NoError(t, err)
+
+		_, err = os.Stat(testDir)
+		assert.True(t, os.IsNotExist(err), "directory should be removed")
+	})
+
+	t.Run("succeeds on non-existent path", func(t *testing.T) {
+		err := cleanPath(context.Background(), "/nonexistent/path/that/does/not/exist")
+		require.NoError(t, err)
+	})
+
+	t.Run("rejects empty path", func(t *testing.T) {
+		err := cleanPath(context.Background(), "")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "empty directory path")
+	})
+
+	t.Run("rejects root path", func(t *testing.T) {
+		err := cleanPath(context.Background(), "/")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "refusing to remove")
+	})
+
+	t.Run("rejects current directory", func(t *testing.T) {
+		err := cleanPath(context.Background(), ".")
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "refusing to remove")
+	})
+
+	t.Run("rejects home directory", func(t *testing.T) {
+		home, err := os.UserHomeDir()
+		require.NoError(t, err)
+
+		err = cleanPath(context.Background(), home)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "refusing to remove home directory")
+	})
+
+	t.Run("respects context cancellation", func(t *testing.T) {
+		dir := t.TempDir()
+		testDir := filepath.Join(dir, "cache")
+		require.NoError(t, os.MkdirAll(testDir, 0o755))
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		err := cleanPath(ctx, testDir)
+		require.Error(t, err)
+		assert.ErrorIs(t, err, context.Canceled)
+	})
+}
+
+func TestCleanPathWindowsDriveRoot(t *testing.T) {
+	if runtime.GOOS != "windows" {
+		t.Skip("Windows-specific test")
+	}
+
+	err := cleanPath(context.Background(), "C:\\")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "refusing to remove drive root")
+}