From 227fdaa00b226e107c21ce0c6d6aa4772e2c8efe Mon Sep 17 00:00:00 2001
From: jonaustin09 <jonaustin460@gmail.com>
Date: Wed, 28 Aug 2024 11:07:44 -0400
Subject: [PATCH] fix: Fixed the pagination for common prefixes in ListObjects
 & ListObjectsV2 actions

---
 backend/walk.go                  | 11 +++++--
 tests/integration/group-tests.go |  2 ++
 tests/integration/tests.go       | 55 ++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/backend/walk.go b/backend/walk.go
index d2f1cea0..bb1f0b57 100644
--- a/backend/walk.go
+++ b/backend/walk.go
@@ -192,9 +192,16 @@ func Walk(ctx context.Context, fileSystem fs.FS, prefix, delimiter, marker strin
 		// Common prefixes are a set, so should not have duplicates.
 		// These are abstractly a "directory", so need to include the
 		// delimiter at the end.
-		cpmap[prefix+before+delimiter] = struct{}{}
+		cpref := prefix + before + delimiter
+		if cpref == marker {
+			pastMarker = true
+			return nil
+		}
+		cpmap[cpref] = struct{}{}
 		if (len(objects) + len(cpmap)) == int(max) {
-			pastMax = true
+			newMarker = cpref
+			truncated = true
+			return fs.SkipAll
 		}
 
 		return nil
diff --git a/tests/integration/group-tests.go b/tests/integration/group-tests.go
index 06e0ecd1..292c0407 100644
--- a/tests/integration/group-tests.go
+++ b/tests/integration/group-tests.go
@@ -184,6 +184,7 @@ func TestListObjectsV2(s *S3Conf) {
 	ListObjectsV2_start_after_empty_result(s)
 	ListObjectsV2_both_delimiter_and_prefix(s)
 	ListObjectsV2_single_dir_object_with_delim_and_prefix(s)
+	ListObjectsV2_truncated_common_prefixes(s)
 }
 
 func TestDeleteObject(s *S3Conf) {
@@ -620,6 +621,7 @@ func GetIntTests() IntTests {
 		"ListObjectsV2_start_after_empty_result":                              ListObjectsV2_start_after_empty_result,
 		"ListObjectsV2_both_delimiter_and_prefix":                             ListObjectsV2_both_delimiter_and_prefix,
 		"ListObjectsV2_single_dir_object_with_delim_and_prefix":               ListObjectsV2_single_dir_object_with_delim_and_prefix,
+		"ListObjectsV2_truncated_common_prefixes":                             ListObjectsV2_truncated_common_prefixes,
 		"DeleteObject_non_existing_object":                                    DeleteObject_non_existing_object,
 		"DeleteObject_name_too_long":                                          DeleteObject_name_too_long,
 		"DeleteObject_non_existing_dir_object":                                DeleteObject_non_existing_dir_object,
diff --git a/tests/integration/tests.go b/tests/integration/tests.go
index 7a971e04..dd6f02af 100644
--- a/tests/integration/tests.go
+++ b/tests/integration/tests.go
@@ -4108,6 +4108,61 @@ func ListObjectsV2_single_dir_object_with_delim_and_prefix(s *S3Conf) error {
 	})
 }
 
+func ListObjectsV2_truncated_common_prefixes(s *S3Conf) error {
+	testName := "ListObjectsV2_truncated_common_prefixes"
+	return actionHandler(s, testName, func(s3client *s3.Client, bucket string) error {
+		if err := putObjects(s3client, []string{"d1/f1", "d2/f2", "d3/f3", "d4/f4"}, bucket); err != nil {
+			return err
+		}
+
+		delim, maxKeys := "/", int32(3)
+
+		ctx, cancel := context.WithTimeout(context.Background(), shortTimeout)
+		out, err := s3client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
+			Bucket:    &bucket,
+			Delimiter: &delim,
+			MaxKeys:   &maxKeys,
+		})
+		cancel()
+		if err != nil {
+			return err
+		}
+
+		if !comparePrefixes([]string{"d1/", "d2/", "d3/"}, out.CommonPrefixes) {
+			return fmt.Errorf("expected the common prefixes to be %v, instead got %v", []string{"d1/", "d2/", "d3/"}, out.CommonPrefixes)
+		}
+		if *out.MaxKeys != maxKeys {
+			return fmt.Errorf("expected the max-keys to be %v, instead got %v", maxKeys, *out.MaxKeys)
+		}
+		if *out.NextContinuationToken != "d3/" {
+			return fmt.Errorf("expected the NextContinuationToken to be d3/, instead got %v", *out.NextContinuationToken)
+		}
+		if *out.Delimiter != delim {
+			return fmt.Errorf("expected the delimiter to be %v, instead got %v", delim, *out.Delimiter)
+		}
+
+		ctx, cancel = context.WithTimeout(context.Background(), shortTimeout)
+		out, err = s3client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
+			Bucket:            &bucket,
+			Delimiter:         &delim,
+			ContinuationToken: out.NextContinuationToken,
+		})
+		cancel()
+		if err != nil {
+			return err
+		}
+
+		if !comparePrefixes([]string{"d4/"}, out.CommonPrefixes) {
+			return fmt.Errorf("expected the common prefixes to be %v, instead got %v", []string{"d4/"}, out.CommonPrefixes)
+		}
+		if *out.Delimiter != delim {
+			return fmt.Errorf("expected the delimiter to be %v, instead got %v", delim, *out.Delimiter)
+		}
+
+		return nil
+	})
+}
+
 func DeleteObject_non_existing_object(s *S3Conf) error {
 	testName := "DeleteObject_non_existing_object"
 	return actionHandler(s, testName, func(s3client *s3.Client, bucket string) error {