From 96961f79aa1540ad02a949f21a6e07d947b01398 Mon Sep 17 00:00:00 2001
From: Willem Kaufmann <willem.kaufmann@redpanda.com>
Date: Fri, 3 May 2024 15:05:35 -0400
Subject: [PATCH] cloud_storage: correct `list_objects_resp` behavior

Fixes behavior with `max_keys` and `continuation_token` within
`list_objects_resp()`. Fixture tests that use the `s3_imposter`
for requests can now expect proper behavior around these parameters.
---
 src/v/cloud_storage/tests/s3_imposter.cc | 58 ++++++++++++++++++------
 src/v/cloud_storage/tests/s3_imposter.h  |  1 +
 2 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/src/v/cloud_storage/tests/s3_imposter.cc b/src/v/cloud_storage/tests/s3_imposter.cc
index 62dda7df9808..56c24c38e289 100644
--- a/src/v/cloud_storage/tests/s3_imposter.cc
+++ b/src/v/cloud_storage/tests/s3_imposter.cc
@@ -42,23 +42,39 @@ uint16_t unit_test_httpd_port_number() { return 4442; }
 
 namespace {
 
+using expectation_map_t
+  = std::map<ss::sstring, s3_imposter_fixture::expectation>;
+
 // Takes the input map of keys to expectations and returns a stringified XML
 // corresponding to the appropriate S3 response.
 ss::sstring list_objects_resp(
-  const std::map<ss::sstring, s3_imposter_fixture::expectation>& objects,
+  const expectation_map_t& objects,
   ss::sstring prefix,
   ss::sstring delimiter,
-  std::optional<size_t> max_keys) {
+  std::optional<size_t> max_keys_opt,
+  std::optional<ss::sstring> continuation_token_opt) {
     std::map<ss::sstring, size_t> content_key_to_size;
     std::set<ss::sstring> common_prefixes;
     // Filter by prefix and group by the substring between the prefix and first
     // delimiter.
-    for (const auto& [_, expectation] : objects) {
-        if (
-          max_keys.has_value()
-          && content_key_to_size.size() == max_keys.value()) {
+    auto max_keys = max_keys_opt.has_value()
+                      ? std::min(
+                        max_keys_opt.value(),
+                        s3_imposter_fixture::default_max_keys)
+                      : s3_imposter_fixture::default_max_keys;
+    auto it = (continuation_token_opt.has_value())
+                ? objects.find(continuation_token_opt.value())
+                : objects.begin();
+    auto end_it = objects.end();
+    ss::sstring next_continuation_token = "";
+    for (; it != end_it; ++it) {
+        const auto& expectation = it->second;
+
+        if (content_key_to_size.size() == max_keys) {
+            next_continuation_token = it->first;
             break;
         }
+
         auto key = expectation.url;
         if (!key.empty() && key[0] == '/') {
             // Remove / character that S3 client adds
@@ -95,6 +111,8 @@ ss::sstring list_objects_resp(
           prefix
           + key.substr(prefix.size(), delimiter_pos - prefix.size() + 1));
     }
+
+    const bool is_truncated = (it != end_it);
     // Populate the returned XML.
     ss::sstring ret;
     ret += fmt::format(
@@ -103,14 +121,17 @@ ss::sstring list_objects_resp(
   <Name>test-bucket</Name>
   <Prefix>{}</Prefix>
   <KeyCount>{}</KeyCount>
-  <MaxKeys>1000</MaxKeys>
+  <MaxKeys>{}</MaxKeys>
   <Delimiter>{}</Delimiter>
-  <IsTruncated>false</IsTruncated>
-  <NextContinuationToken>next</NextContinuationToken>
+  <IsTruncated>{}</IsTruncated>
+  <NextContinuationToken>{}</NextContinuationToken>
 )xml",
       prefix,
       content_key_to_size.size(),
-      delimiter);
+      max_keys,
+      delimiter,
+      is_truncated,
+      next_continuation_token);
     for (const auto& [key, size] : content_key_to_size) {
         ret += fmt::format(
           R"xml(
@@ -221,18 +242,29 @@ struct s3_imposter_fixture::content_handler {
                 auto prefix = request.get_query_param("prefix");
                 auto delimiter = request.get_query_param("delimiter");
                 auto max_keys_str = request.get_query_param("max-keys");
+                auto continuation_token_str = request.get_query_param(
+                  "continuation-token");
                 std::optional<size_t> max_keys = (max_keys_str.empty())
                                                    ? std::optional<size_t>{}
                                                    : std::stoi(max_keys_str);
+                std::optional<ss::sstring> continuation_token
+                  = (continuation_token_str.empty())
+                      ? std::optional<ss::sstring>{}
+                      : continuation_token_str;
                 vlog(
                   fixt_log.trace,
-                  "S3 imposter list request {} - {} - {} - {}",
+                  "S3 imposter list request {} - {} - {} - {} - {}",
                   prefix,
                   delimiter,
                   max_keys,
+                  continuation_token,
                   request._method);
                 return list_objects_resp(
-                  expectations, prefix, delimiter, max_keys);
+                  expectations,
+                  prefix,
+                  delimiter,
+                  max_keys,
+                  continuation_token);
             }
             if (
               expect_iter == expectations.end()
@@ -331,7 +363,7 @@ struct s3_imposter_fixture::content_handler {
         RPTEST_FAIL("Unexpected request");
         return "";
     }
-    std::map<ss::sstring, s3_imposter_fixture::expectation> expectations;
+    expectation_map_t expectations;
     s3_imposter_fixture& fixture;
     std::optional<absl::flat_hash_set<ss::sstring>> headers = std::nullopt;
 };
diff --git a/src/v/cloud_storage/tests/s3_imposter.h b/src/v/cloud_storage/tests/s3_imposter.h
index b654329751de..cac3eb53b8a2 100644
--- a/src/v/cloud_storage/tests/s3_imposter.h
+++ b/src/v/cloud_storage/tests/s3_imposter.h
@@ -38,6 +38,7 @@
 /// be retrieved using the GET request or deleted using the DELETE request.
 class s3_imposter_fixture {
 public:
+    static constexpr size_t default_max_keys = 100;
     uint16_t httpd_port_number();
     static constexpr const char* httpd_host_name = "127.0.0.1";