From cdc9e91750036fc370db65a44618f3139db11ae1 Mon Sep 17 00:00:00 2001
From: FuXiaoHei <fuxiaohei@vip.qq.com>
Date: Mon, 13 Mar 2023 18:23:51 +0800
Subject: [PATCH] add path prefix to ObjectStorage.Iterator (#23332)

Support to iterator subdirectory in ObjectStorage for
ObjectStorage.Iterator method.

It's required for https://github.com/go-gitea/gitea/pull/22738 to make
artifact files cleanable.

---------

Co-authored-by: Jason Song <i@wolfogre.com>
Co-authored-by: Lunny Xiao <xiaolunwen@gmail.com>
---
 cmd/dump.go                    |  6 ++---
 modules/doctor/storage.go      |  2 +-
 modules/storage/helper.go      |  2 +-
 modules/storage/helper_test.go |  2 +-
 modules/storage/local.go       |  8 +++++--
 modules/storage/local_test.go  | 42 ++++++++++++++++++++++++++++++++++
 modules/storage/minio.go       | 12 +++++++---
 modules/storage/storage.go     |  4 ++--
 tests/test_utils.go            |  4 ++--
 9 files changed, 67 insertions(+), 15 deletions(-)

diff --git a/cmd/dump.go b/cmd/dump.go
index c802849f8e..00d279b991 100644
--- a/cmd/dump.go
+++ b/cmd/dump.go
@@ -250,7 +250,7 @@ func runDump(ctx *cli.Context) error {
 
 		if ctx.IsSet("skip-lfs-data") && ctx.Bool("skip-lfs-data") {
 			log.Info("Skip dumping LFS data")
-		} else if err := storage.LFS.IterateObjects(func(objPath string, object storage.Object) error {
+		} else if err := storage.LFS.IterateObjects("", func(objPath string, object storage.Object) error {
 			info, err := object.Stat()
 			if err != nil {
 				return err
@@ -351,7 +351,7 @@ func runDump(ctx *cli.Context) error {
 
 	if ctx.IsSet("skip-attachment-data") && ctx.Bool("skip-attachment-data") {
 		log.Info("Skip dumping attachment data")
-	} else if err := storage.Attachments.IterateObjects(func(objPath string, object storage.Object) error {
+	} else if err := storage.Attachments.IterateObjects("", func(objPath string, object storage.Object) error {
 		info, err := object.Stat()
 		if err != nil {
 			return err
@@ -364,7 +364,7 @@ func runDump(ctx *cli.Context) error {
 
 	if ctx.IsSet("skip-package-data") && ctx.Bool("skip-package-data") {
 		log.Info("Skip dumping package data")
-	} else if err := storage.Packages.IterateObjects(func(objPath string, object storage.Object) error {
+	} else if err := storage.Packages.IterateObjects("", func(objPath string, object storage.Object) error {
 		info, err := object.Stat()
 		if err != nil {
 			return err
diff --git a/modules/doctor/storage.go b/modules/doctor/storage.go
index aa987de447..c20566d675 100644
--- a/modules/doctor/storage.go
+++ b/modules/doctor/storage.go
@@ -31,7 +31,7 @@ func commonCheckStorage(ctx context.Context, logger log.Logger, autofix bool, op
 	totalSize, orphanedSize := int64(0), int64(0)
 
 	var pathsToDelete []string
-	if err := opts.storer.IterateObjects(func(p string, obj storage.Object) error {
+	if err := opts.storer.IterateObjects("", func(p string, obj storage.Object) error {
 		defer obj.Close()
 
 		totalCount++
diff --git a/modules/storage/helper.go b/modules/storage/helper.go
index 1ab99d98b3..d1959830b9 100644
--- a/modules/storage/helper.go
+++ b/modules/storage/helper.go
@@ -90,6 +90,6 @@ func (s discardStorage) URL(_, _ string) (*url.URL, error) {
 	return nil, fmt.Errorf("%s", s)
 }
 
-func (s discardStorage) IterateObjects(_ func(string, Object) error) error {
+func (s discardStorage) IterateObjects(_ string, _ func(string, Object) error) error {
 	return fmt.Errorf("%s", s)
 }
diff --git a/modules/storage/helper_test.go b/modules/storage/helper_test.go
index 7d74671c54..f4c2d0467f 100644
--- a/modules/storage/helper_test.go
+++ b/modules/storage/helper_test.go
@@ -42,7 +42,7 @@ func Test_discardStorage(t *testing.T) {
 				assert.Errorf(t, err, string(tt))
 			}
 			{
-				err := tt.IterateObjects(func(_ string, _ Object) error { return nil })
+				err := tt.IterateObjects("", func(_ string, _ Object) error { return nil })
 				assert.Error(t, err, string(tt))
 			}
 		})
diff --git a/modules/storage/local.go b/modules/storage/local.go
index 05bf1fb28a..15f5761e8f 100644
--- a/modules/storage/local.go
+++ b/modules/storage/local.go
@@ -127,8 +127,12 @@ func (l *LocalStorage) URL(path, name string) (*url.URL, error) {
 }
 
 // IterateObjects iterates across the objects in the local storage
-func (l *LocalStorage) IterateObjects(fn func(path string, obj Object) error) error {
-	return filepath.WalkDir(l.dir, func(path string, d os.DirEntry, err error) error {
+func (l *LocalStorage) IterateObjects(prefix string, fn func(path string, obj Object) error) error {
+	dir := l.dir
+	if prefix != "" {
+		dir = filepath.Join(l.dir, util.CleanPath(prefix))
+	}
+	return filepath.WalkDir(dir, func(path string, d os.DirEntry, err error) error {
 		if err != nil {
 			return err
 		}
diff --git a/modules/storage/local_test.go b/modules/storage/local_test.go
index 994c54e859..2b112df8f1 100644
--- a/modules/storage/local_test.go
+++ b/modules/storage/local_test.go
@@ -4,6 +4,10 @@
 package storage
 
 import (
+	"bytes"
+	"context"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -50,3 +54,41 @@ func TestBuildLocalPath(t *testing.T) {
 		})
 	}
 }
+
+func TestLocalStorageIterator(t *testing.T) {
+	dir := filepath.Join(os.TempDir(), "TestLocalStorageIteratorTestDir")
+	l, err := NewLocalStorage(context.Background(), LocalStorageConfig{Path: dir})
+	assert.NoError(t, err)
+
+	testFiles := [][]string{
+		{"a/1.txt", "a1"},
+		{"/a/1.txt", "aa1"}, // same as above, but with leading slash that will be trim
+		{"b/1.txt", "b1"},
+		{"b/2.txt", "b2"},
+		{"b/3.txt", "b3"},
+		{"b/x 4.txt", "bx4"},
+	}
+	for _, f := range testFiles {
+		_, err = l.Save(f[0], bytes.NewBufferString(f[1]), -1)
+		assert.NoError(t, err)
+	}
+
+	expectedList := map[string][]string{
+		"a":           {"a/1.txt"},
+		"b":           {"b/1.txt", "b/2.txt", "b/3.txt", "b/x 4.txt"},
+		"":            {"a/1.txt", "b/1.txt", "b/2.txt", "b/3.txt", "b/x 4.txt"},
+		"/":           {"a/1.txt", "b/1.txt", "b/2.txt", "b/3.txt", "b/x 4.txt"},
+		"a/b/../../a": {"a/1.txt"},
+	}
+	for dir, expected := range expectedList {
+		count := 0
+		err = l.IterateObjects(dir, func(path string, f Object) error {
+			defer f.Close()
+			assert.Contains(t, expected, path)
+			count++
+			return nil
+		})
+		assert.NoError(t, err)
+		assert.Equal(t, count, len(expected))
+	}
+}
diff --git a/modules/storage/minio.go b/modules/storage/minio.go
index 24da14b634..8cc06bcdd3 100644
--- a/modules/storage/minio.go
+++ b/modules/storage/minio.go
@@ -209,12 +209,18 @@ func (m *MinioStorage) URL(path, name string) (*url.URL, error) {
 }
 
 // IterateObjects iterates across the objects in the miniostorage
-func (m *MinioStorage) IterateObjects(fn func(path string, obj Object) error) error {
+func (m *MinioStorage) IterateObjects(prefix string, fn func(path string, obj Object) error) error {
 	opts := minio.GetObjectOptions{}
 	lobjectCtx, cancel := context.WithCancel(m.ctx)
 	defer cancel()
+
+	basePath := m.basePath
+	if prefix != "" {
+		basePath = m.buildMinioPath(prefix)
+	}
+
 	for mObjInfo := range m.client.ListObjects(lobjectCtx, m.bucket, minio.ListObjectsOptions{
-		Prefix:    m.basePath,
+		Prefix:    basePath,
 		Recursive: true,
 	}) {
 		object, err := m.client.GetObject(lobjectCtx, m.bucket, mObjInfo.Key, opts)
@@ -223,7 +229,7 @@ func (m *MinioStorage) IterateObjects(fn func(path string, obj Object) error) er
 		}
 		if err := func(object *minio.Object, fn func(path string, obj Object) error) error {
 			defer object.Close()
-			return fn(strings.TrimPrefix(mObjInfo.Key, m.basePath), &minioObject{object})
+			return fn(strings.TrimPrefix(mObjInfo.Key, basePath), &minioObject{object})
 		}(object, fn); err != nil {
 			return convertMinioErr(err)
 		}
diff --git a/modules/storage/storage.go b/modules/storage/storage.go
index d8998b1922..caecab306e 100644
--- a/modules/storage/storage.go
+++ b/modules/storage/storage.go
@@ -65,7 +65,7 @@ type ObjectStorage interface {
 	Stat(path string) (os.FileInfo, error)
 	Delete(path string) error
 	URL(path, name string) (*url.URL, error)
-	IterateObjects(func(path string, obj Object) error) error
+	IterateObjects(path string, iterator func(path string, obj Object) error) error
 }
 
 // Copy copies a file from source ObjectStorage to dest ObjectStorage
@@ -87,7 +87,7 @@ func Copy(dstStorage ObjectStorage, dstPath string, srcStorage ObjectStorage, sr
 
 // Clean delete all the objects in this storage
 func Clean(storage ObjectStorage) error {
-	return storage.IterateObjects(func(path string, obj Object) error {
+	return storage.IterateObjects("", func(path string, obj Object) error {
 		_ = obj.Close()
 		return storage.Delete(path)
 	})
diff --git a/tests/test_utils.go b/tests/test_utils.go
index e3e5becfbe..102dd3d298 100644
--- a/tests/test_utils.go
+++ b/tests/test_utils.go
@@ -201,7 +201,7 @@ func PrepareTestEnv(t testing.TB, skip ...int) func() {
 	lfsFixtures, err := storage.NewStorage("", storage.LocalStorageConfig{Path: path.Join(filepath.Dir(setting.AppPath), "tests/gitea-lfs-meta")})
 	assert.NoError(t, err)
 	assert.NoError(t, storage.Clean(storage.LFS))
-	assert.NoError(t, lfsFixtures.IterateObjects(func(path string, _ storage.Object) error {
+	assert.NoError(t, lfsFixtures.IterateObjects("", func(path string, _ storage.Object) error {
 		_, err := storage.Copy(storage.LFS, path, lfsFixtures, path)
 		return err
 	}))
@@ -258,7 +258,7 @@ func ResetFixtures(t *testing.T) {
 	lfsFixtures, err := storage.NewStorage("", storage.LocalStorageConfig{Path: path.Join(filepath.Dir(setting.AppPath), "tests/gitea-lfs-meta")})
 	assert.NoError(t, err)
 	assert.NoError(t, storage.Clean(storage.LFS))
-	assert.NoError(t, lfsFixtures.IterateObjects(func(path string, _ storage.Object) error {
+	assert.NoError(t, lfsFixtures.IterateObjects("", func(path string, _ storage.Object) error {
 		_, err := storage.Copy(storage.LFS, path, lfsFixtures, path)
 		return err
 	}))