From 6a68bd48257d1c1b1e20f34b1aba7b2323ddf9de Mon Sep 17 00:00:00 2001
From: Martin Brennan <mjrbrennan@gmail.com>
Date: Wed, 10 Nov 2021 08:01:28 +1000
Subject: [PATCH] DEV: Limit list multipart parts to 1 (#14853)

We are only using list_multipart_parts right now in the
uploads controller for multipart uploads to check if the
upload exists; thus we don't need up to 1000 parts.

Also adding a note for future explorers that list_multipart_parts
only gets 1000 parts max, and adding params for max parts
and starting parts.
---
 app/controllers/uploads_controller.rb    |  4 +++-
 lib/file_store/s3_store.rb               | 28 ++++++++++++++++++++----
 spec/requests/uploads_controller_spec.rb |  4 ++--
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/app/controllers/uploads_controller.rb b/app/controllers/uploads_controller.rb
index b95b33b5c04..ba1a9a239d5 100644
--- a/app/controllers/uploads_controller.rb
+++ b/app/controllers/uploads_controller.rb
@@ -372,7 +372,9 @@ class UploadsController < ApplicationController
   def multipart_upload_exists?(external_upload_stub)
     begin
       Discourse.store.list_multipart_parts(
-        upload_id: external_upload_stub.external_upload_identifier, key: external_upload_stub.key
+        upload_id: external_upload_stub.external_upload_identifier,
+        key: external_upload_stub.key,
+        max_parts: 1
       )
     rescue Aws::S3::Errors::NoSuchUpload => err
       debug_upload_error(err, "upload.external_upload_not_found", { additional_detail: "path: #{external_upload_stub.key}" })
diff --git a/lib/file_store/s3_store.rb b/lib/file_store/s3_store.rb
index 2e636818215..f67bedff408 100644
--- a/lib/file_store/s3_store.rb
+++ b/lib/file_store/s3_store.rb
@@ -347,12 +347,32 @@ module FileStore
       )
     end
 
-    def list_multipart_parts(upload_id:, key:)
-      s3_helper.s3_client.list_parts(
+    # Important note from the S3 documentation:
+    #
+    # This request returns a default and maximum of 1000 parts.
+    # You can restrict the number of parts returned by specifying the
+    # max_parts argument. If your multipart upload consists of more than 1,000
+    # parts, the response returns an IsTruncated field with the value of true,
+    # and a NextPartNumberMarker element.
+    #
+    # In subsequent ListParts requests you can include the part_number_marker arg
+    # using the NextPartNumberMarker the field value from the previous response to
+    # get more parts.
+    #
+    # See https://docs.aws.amazon.com/sdk-for-ruby/v3/api/Aws/S3/Client.html#list_parts-instance_method
+    def list_multipart_parts(upload_id:, key:, max_parts: 1000, start_from_part_number: nil)
+      options = {
         bucket: s3_bucket_name,
         key: key,
-        upload_id: upload_id
-      )
+        upload_id: upload_id,
+        max_parts: max_parts
+      }
+
+      if start_from_part_number.present?
+        options[:part_number_marker] = start_from_part_number
+      end
+
+      s3_helper.s3_client.list_parts(options)
     end
 
     def complete_multipart(upload_id:, key:, parts:)
diff --git a/spec/requests/uploads_controller_spec.rb b/spec/requests/uploads_controller_spec.rb
index def58bec7ed..2247af51049 100644
--- a/spec/requests/uploads_controller_spec.rb
+++ b/spec/requests/uploads_controller_spec.rb
@@ -974,7 +974,7 @@ describe UploadsController do
            <StorageClass>STANDARD</StorageClass>
         </ListPartsResult>
         BODY
-        stub_request(:get, "https://s3-upload-bucket.s3.us-west-1.amazonaws.com/#{external_upload_stub.key}?uploadId=#{mock_multipart_upload_id}").to_return({ status: 200, body: list_multipart_result })
+        stub_request(:get, "https://s3-upload-bucket.s3.us-west-1.amazonaws.com/#{external_upload_stub.key}?max-parts=1&uploadId=#{mock_multipart_upload_id}").to_return({ status: 200, body: list_multipart_result })
       end
 
       it "errors if the correct params are not provided" do
@@ -1118,7 +1118,7 @@ describe UploadsController do
            <StorageClass>STANDARD</StorageClass>
         </ListPartsResult>
         BODY
-        stub_request(:get, "#{upload_base_url}/#{external_upload_stub.key}?uploadId=#{mock_multipart_upload_id}").to_return({ status: 200, body: list_multipart_result })
+        stub_request(:get, "#{upload_base_url}/#{external_upload_stub.key}?max-parts=1&uploadId=#{mock_multipart_upload_id}").to_return({ status: 200, body: list_multipart_result })
       end
 
       it "errors if the correct params are not provided" do