From ccdc8dff236c00ec86223c852a3e86ca4e5756da Mon Sep 17 00:00:00 2001
From: glorv <glorvs@163.com>
Date: Wed, 3 Feb 2021 11:30:59 +0800
Subject: [PATCH] cherry pick #713 to release-4.0

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
---
 pkg/storage/s3.go      | 16 +++++++---
 pkg/storage/s3_test.go | 66 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/pkg/storage/s3.go b/pkg/storage/s3.go
index 4962ae3ba..69df47120 100644
--- a/pkg/storage/s3.go
+++ b/pkg/storage/s3.go
@@ -372,10 +372,16 @@ func (rs *S3Storage) WalkDir(ctx context.Context, opt *WalkOption, fn func(strin
 	if opt == nil {
 		opt = &WalkOption{}
 	}
-	prefix := rs.options.Prefix + opt.SubDir
-	if !strings.HasSuffix(prefix, "/") {
-		prefix += "/"
+
+	// NOTE: leave prefix empty if subDir is not set. Else, s3 will return empty result!
+	prefix := ""
+	if len(opt.SubDir) > 0 {
+		prefix = rs.options.Prefix + opt.SubDir
+		if len(prefix) > 0 && !strings.HasSuffix(prefix, "/") {
+			prefix += "/"
+		}
 	}
+
 	maxKeys := int64(1000)
 	if opt.ListCount > 0 {
 		maxKeys = opt.ListCount
@@ -383,9 +389,11 @@ func (rs *S3Storage) WalkDir(ctx context.Context, opt *WalkOption, fn func(strin
 
 	req := &s3.ListObjectsInput{
 		Bucket:  aws.String(rs.options.Bucket),
-		Prefix:  aws.String(prefix),
 		MaxKeys: aws.Int64(maxKeys),
 	}
+	if len(prefix) > 0 {
+		req.Prefix = aws.String(prefix)
+	}
 	for {
 		// FIXME: We can't use ListObjectsV2, it is not universally supported.
 		// (Ceph RGW supported ListObjectsV2 since v15.1.0, released 2020 Jan 30th)
diff --git a/pkg/storage/s3_test.go b/pkg/storage/s3_test.go
index 2017bdbfa..b9e9a04e8 100644
--- a/pkg/storage/s3_test.go
+++ b/pkg/storage/s3_test.go
@@ -29,7 +29,12 @@ type s3Suite struct {
 	storage    *S3Storage
 }
 
-var _ = Suite(&s3Suite{})
+type s3SuiteCustom struct{}
+
+var (
+	_ = Suite(&s3Suite{})
+	_ = Suite(&s3SuiteCustom{})
+)
 
 // FIXME: Cannot use the real SetUpTest/TearDownTest to set up the mock
 // otherwise the mock error will be ignored.
@@ -802,3 +807,62 @@ func (s *s3Suite) TestWalkDir(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(i, Equals, len(contents))
 }
+
+// TestWalkDirBucket checks WalkDir retrieves all directory content under a bucket.
+func (s *s3SuiteCustom) TestWalkDirWithEmptyPrefix(c *C) {
+	controller := gomock.NewController(c)
+	s3API := mock.NewMockS3API(controller)
+	storage := NewS3StorageForTest(
+		s3API,
+		&backup.S3{
+			Region:       "us-west-2",
+			Bucket:       "bucket",
+			Prefix:       "",
+			Acl:          "acl",
+			Sse:          "sse",
+			StorageClass: "sc",
+		},
+	)
+	defer controller.Finish()
+	ctx := aws.BackgroundContext()
+
+	contents := []*s3.Object{
+		{
+			Key:  aws.String("sp/.gitignore"),
+			Size: aws.Int64(437),
+		},
+		{
+			Key:  aws.String("prefix/sp/01.jpg"),
+			Size: aws.Int64(27499),
+		},
+	}
+	s3API.EXPECT().
+		ListObjectsWithContext(ctx, gomock.Any()).
+		DoAndReturn(func(_ context.Context, input *s3.ListObjectsInput) (*s3.ListObjectsOutput, error) {
+			c.Assert(aws.StringValue(input.Bucket), Equals, "bucket")
+			c.Assert(aws.StringValue(input.Prefix), Equals, "")
+			c.Assert(aws.StringValue(input.Marker), Equals, "")
+			c.Assert(aws.Int64Value(input.MaxKeys), Equals, int64(2))
+			c.Assert(aws.StringValue(input.Delimiter), Equals, "")
+			return &s3.ListObjectsOutput{
+				IsTruncated: aws.Bool(false),
+				Contents:    contents,
+			}, nil
+		})
+
+	// Ensure we receive the items in order.
+	i := 0
+	err := storage.WalkDir(
+		ctx,
+		&WalkOption{SubDir: "", ListCount: 2},
+		func(path string, size int64) error {
+			comment := Commentf("index = %d", i)
+			c.Assert(path, Equals, *contents[i].Key, comment)
+			c.Assert(size, Equals, *contents[i].Size, comment)
+			i++
+			return nil
+		},
+	)
+	c.Assert(err, IsNil)
+	c.Assert(i, Equals, len(contents))
+}