From f1352dfe13ecccf2abf09ec7b3ee4301f98dbf79 Mon Sep 17 00:00:00 2001
From: Giuseppe Scrivano <gscrivan@redhat.com>
Date: Tue, 30 Apr 2024 15:15:45 +0200
Subject: [PATCH] overlay: lock staging directories

lock any staging directory while it is being used so that another
process cannot delete it.

Now the Cleanup() function deletes only the staging directories that
are not locked by any other user.

Closes: https://github.com/containers/storage/issues/1915

Signed-off-by: Giuseppe Scrivano <gscrivan@redhat.com>
---
 drivers/overlay/overlay.go | 76 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 73 insertions(+), 3 deletions(-)

diff --git a/drivers/overlay/overlay.go b/drivers/overlay/overlay.go
index 605fae8775..91f457ce20 100644
--- a/drivers/overlay/overlay.go
+++ b/drivers/overlay/overlay.go
@@ -28,6 +28,7 @@ import (
 	"github.com/containers/storage/pkg/fsutils"
 	"github.com/containers/storage/pkg/idmap"
 	"github.com/containers/storage/pkg/idtools"
+	"github.com/containers/storage/pkg/lockfile"
 	"github.com/containers/storage/pkg/mount"
 	"github.com/containers/storage/pkg/parsers"
 	"github.com/containers/storage/pkg/system"
@@ -83,6 +84,8 @@ const (
 	lowerFile  = "lower"
 	maxDepth   = 500
 
+	stagingLockFile = "staging.lock"
+
 	tocArtifact             = "toc"
 	fsVerityDigestsArtifact = "fs-verity-digests"
 
@@ -127,6 +130,8 @@ type Driver struct {
 	usingMetacopy    bool
 	usingComposefs   bool
 
+	stagingDirsLocks map[string]*lockfile.LockFile
+
 	supportsIDMappedMounts *bool
 }
 
@@ -460,6 +465,7 @@ func Init(home string, options graphdriver.Options) (graphdriver.Driver, error)
 		supportsVolatile: supportsVolatile,
 		usingComposefs:   opts.useComposefs,
 		options:          *opts,
+		stagingDirsLocks: make(map[string]*lockfile.LockFile),
 	}
 
 	d.naiveDiff = graphdriver.NewNaiveDiffDriver(d, graphdriver.NewNaiveLayerIDMapUpdater(d))
@@ -880,10 +886,44 @@ func (d *Driver) Metadata(id string) (map[string]string, error) {
 // the storage is being shutdown.  The only state created by the driver
 // is the bind mount on the home directory.
 func (d *Driver) Cleanup() error {
-	_ = os.RemoveAll(filepath.Join(d.home, stagingDir))
+	anyPresent := d.pruneStagingDirectories()
+	if anyPresent {
+		return nil
+	}
 	return mount.Unmount(d.home)
 }
 
+// pruneStagingDirectories cleans up any staging directory that was leaked.
+// It returns whether any staging directory is still present.
+func (d *Driver) pruneStagingDirectories() bool {
+	for _, lock := range d.stagingDirsLocks {
+		lock.Unlock()
+	}
+	d.stagingDirsLocks = make(map[string]*lockfile.LockFile)
+
+	anyPresent := false
+
+	homeStagingDir := filepath.Join(d.home, stagingDir)
+	dirs, err := os.ReadDir(homeStagingDir)
+	if err == nil {
+		for _, dir := range dirs {
+			stagingDirToRemove := filepath.Join(homeStagingDir, dir.Name())
+			lock, err := lockfile.GetLockFile(filepath.Join(stagingDirToRemove, stagingLockFile))
+			if err != nil {
+				anyPresent = true
+				continue
+			}
+			if err := lock.TryLock(); err != nil {
+				anyPresent = true
+				continue
+			}
+			_ = os.RemoveAll(stagingDirToRemove)
+			lock.Unlock()
+		}
+	}
+	return anyPresent
+}
+
 // LookupAdditionalLayer looks up additional layer store by the specified
 // digest and ref and returns an object representing that layer.
 // This API is experimental and can be changed without bumping the major version number.
@@ -2029,7 +2069,14 @@ func (d *Driver) DiffGetter(id string) (graphdriver.FileGetCloser, error) {
 
 // CleanupStagingDirectory cleanups the staging directory.
 func (d *Driver) CleanupStagingDirectory(stagingDirectory string) error {
-	return os.RemoveAll(filepath.Dir(stagingDirectory))
+	parentStagingDir := filepath.Dir(stagingDirectory)
+
+	if lock, ok := d.stagingDirsLocks[parentStagingDir]; ok {
+		delete(d.stagingDirsLocks, parentStagingDir)
+		lock.Unlock()
+	}
+
+	return os.RemoveAll(parentStagingDir)
 }
 
 func supportsDataOnlyLayersCached(home, runhome string) (bool, error) {
@@ -2080,6 +2127,19 @@ func (d *Driver) ApplyDiffWithDiffer(id, parent string, options *graphdriver.App
 		if err := os.Mkdir(applyDir, perms); err != nil {
 			return graphdriver.DriverWithDifferOutput{}, err
 		}
+
+		lock, err := lockfile.GetLockFile(filepath.Join(layerDir, stagingLockFile))
+		if err != nil {
+			return graphdriver.DriverWithDifferOutput{}, err
+		}
+		defer func() {
+			if errRet != nil {
+				delete(d.stagingDirsLocks, layerDir)
+				lock.Unlock()
+			}
+		}()
+		d.stagingDirsLocks[layerDir] = lock
+		lock.Lock()
 	} else {
 		var err error
 		applyDir, err = d.getDiffPath(id)
@@ -2113,9 +2173,19 @@ func (d *Driver) ApplyDiffWithDiffer(id, parent string, options *graphdriver.App
 // ApplyDiffFromStagingDirectory applies the changes using the specified staging directory.
 func (d *Driver) ApplyDiffFromStagingDirectory(id, parent string, diffOutput *graphdriver.DriverWithDifferOutput, options *graphdriver.ApplyDiffWithDifferOpts) error {
 	stagingDirectory := diffOutput.Target
-	if filepath.Dir(filepath.Dir(stagingDirectory)) != d.getStagingDir(id) {
+	parentStagingDir := filepath.Dir(stagingDirectory)
+
+	defer func() {
+		if lock, ok := d.stagingDirsLocks[parentStagingDir]; ok {
+			delete(d.stagingDirsLocks, parentStagingDir)
+			lock.Unlock()
+		}
+	}()
+
+	if filepath.Dir(parentStagingDir) != d.getStagingDir(id) {
 		return fmt.Errorf("%q is not a staging directory", stagingDirectory)
 	}
+
 	diffPath, err := d.getDiffPath(id)
 	if err != nil {
 		return err