From 9fe0f095170d80ff66b6a26052a3cf8a5fb19cb4 Mon Sep 17 00:00:00 2001
From: Dmitrii Neeman <dneeman@aerospike.com>
Date: Sun, 6 Oct 2024 10:09:36 +0300
Subject: [PATCH] FMWK-556-restore-folder-validation - added validator - added
 nested dirs

---
 io/aws/s3/reader.go                  | 15 +++++++++---
 io/azure/blob/reader.go              | 15 +++++++++---
 io/gcp/storage/reader.go             | 15 +++++++++---
 io/gcp/storage/reader_writer_test.go |  1 +
 io/local/reader.go                   | 36 ++++++++++++++++++++++------
 io/local/reader_test.go              |  5 ++--
 6 files changed, 69 insertions(+), 18 deletions(-)

diff --git a/io/aws/s3/reader.go b/io/aws/s3/reader.go
index 73b849e8..58201370 100644
--- a/io/aws/s3/reader.go
+++ b/io/aws/s3/reader.go
@@ -212,9 +212,18 @@ func (r *Reader) checkRestoreDirectory(ctx context.Context) error {
 			if p.Key == nil || isDirectory(r.prefix, *p.Key) && !r.withNestedDir {
 				continue
 			}
-			// If we found anything, then folder is not empty.
-			if p.Key != nil {
-				return nil
+
+			switch {
+			case r.validator != nil:
+				// If we found a valid file, return.
+				if err = r.validator.Run(*p.Key); err == nil {
+					return nil
+				}
+			default:
+				// If we found anything, then folder is not empty.
+				if p.Key != nil {
+					return nil
+				}
 			}
 		}
 
diff --git a/io/azure/blob/reader.go b/io/azure/blob/reader.go
index 5f2c6ad2..df326e74 100644
--- a/io/azure/blob/reader.go
+++ b/io/azure/blob/reader.go
@@ -183,9 +183,18 @@ func (r *Reader) checkRestoreDirectory(ctx context.Context) error {
 			if isDirectory(r.prefix, *blob.Name) && !r.withNestedDir {
 				continue
 			}
-			// If we found anything, then folder is not empty.
-			if blob.Name != nil {
-				return nil
+
+			switch {
+			case r.validator != nil:
+				// If we found a valid file, return.
+				if err = r.validator.Run(*blob.Name); err == nil {
+					return nil
+				}
+			default:
+				// If we found anything, then folder is not empty.
+				if blob.Name != nil {
+					return nil
+				}
 			}
 		}
 	}
diff --git a/io/gcp/storage/reader.go b/io/gcp/storage/reader.go
index 7a36dfc4..b3ad85c3 100644
--- a/io/gcp/storage/reader.go
+++ b/io/gcp/storage/reader.go
@@ -204,9 +204,18 @@ func (r *Reader) checkRestoreDirectory(ctx context.Context) error {
 		if isDirectory(r.prefix, objAttrs.Name) && !r.withNestedDir {
 			continue
 		}
-		// If we found anything, then folder is not empty.
-		if objAttrs.Name != "" {
-			return nil
+
+		switch {
+		case r.validator != nil:
+			// If we found a valid file, return.
+			if err = r.validator.Run(objAttrs.Name); err == nil {
+				return nil
+			}
+		default:
+			// If we found anything, then folder is not empty.
+			if objAttrs.Name != "" {
+				return nil
+			}
 		}
 	}
 
diff --git a/io/gcp/storage/reader_writer_test.go b/io/gcp/storage/reader_writer_test.go
index 02068574..c1d351bc 100644
--- a/io/gcp/storage/reader_writer_test.go
+++ b/io/gcp/storage/reader_writer_test.go
@@ -268,6 +268,7 @@ func (s *GCPSuite) TestReader_StreamFilesEmpty() {
 		testBucketName,
 		WithDir(testReadFolderEmpty),
 		WithValidator(validatorMock{}),
+		WithNestedDir(),
 	)
 	s.Require().NoError(err)
 
diff --git a/io/local/reader.go b/io/local/reader.go
index 148eeae8..ae1c9831 100644
--- a/io/local/reader.go
+++ b/io/local/reader.go
@@ -112,7 +112,7 @@ func (r *Reader) StreamFiles(
 	defer close(readersCh)
 	// If it is a folder, open and return.
 	if r.isDir {
-		err := r.checkRestoreDirectory()
+		err := r.checkRestoreDirectory(r.path)
 		if err != nil {
 			errorsCh <- err
 			return
@@ -201,9 +201,7 @@ func (r *Reader) streamFile(
 
 // checkRestoreDirectory checks that the restore directory exists,
 // is a readable directory, and contains backup files of the correct format.
-func (r *Reader) checkRestoreDirectory() error {
-	dir := r.path
-
+func (r *Reader) checkRestoreDirectory(dir string) error {
 	dirInfo, err := os.Stat(dir)
 	if err != nil {
 		// Handle the error
@@ -220,9 +218,33 @@ func (r *Reader) checkRestoreDirectory() error {
 		return fmt.Errorf("failed to read path %s: %w", dir, err)
 	}
 
-	// Check if the directory is empty
-	if len(fileInfo) == 0 {
-		return fmt.Errorf("%s is empty", dir)
+	switch {
+	case r.validator != nil:
+		for _, file := range fileInfo {
+			if file.IsDir() {
+				// Iterate over nested dirs recursively.
+				if r.withNestedDir {
+					nestedDir := filepath.Join(r.path, file.Name())
+					// If the nested folder is ok, then return nil.
+					err = r.checkRestoreDirectory(nestedDir)
+					if err != nil {
+						return err
+					}
+				}
+
+				continue
+			}
+
+			// If we found a valid file, return.
+			if err = r.validator.Run(file.Name()); err == nil {
+				return nil
+			}
+		}
+	default:
+		// Check if the directory is empty
+		if len(fileInfo) == 0 {
+			return fmt.Errorf("%s is empty", dir)
+		}
 	}
 
 	return nil
diff --git a/io/local/reader_test.go b/io/local/reader_test.go
index c7082a67..51bf856d 100644
--- a/io/local/reader_test.go
+++ b/io/local/reader_test.go
@@ -43,8 +43,9 @@ func (s *checkRestoreDirectoryTestSuite) TestCheckRestoreDirectory_Negative_Empt
 		return fmt.Errorf("invalid file extension")
 	})
 
-	reader, _ := NewReader(WithValidator(mockValidator), WithDir(dir))
-	err := reader.checkRestoreDirectory()
+	reader, err := NewReader(WithValidator(mockValidator), WithDir(dir))
+	s.NoError(err)
+	err = reader.checkRestoreDirectory(dir)
 	s.Error(err)
 }