From 7736319f2ed49768b9d9eb832465a1317dd0b106 Mon Sep 17 00:00:00 2001
From: Jeremy THERIN <jtherin@scaleway.com>
Date: Tue, 26 Oct 2021 17:41:11 +0200
Subject: [PATCH] fix: paginate through s3 multipart uploads

Signed-off-by: Jeremy THERIN <jtherin@scaleway.com>
---
 registry/storage/driver/s3-aws/s3.go | 68 ++++++++++++++++++----------
 1 file changed, 45 insertions(+), 23 deletions(-)

diff --git a/registry/storage/driver/s3-aws/s3.go b/registry/storage/driver/s3-aws/s3.go
index 75f80b7e..39cb372c 100644
--- a/registry/storage/driver/s3-aws/s3.go
+++ b/registry/storage/driver/s3-aws/s3.go
@@ -558,40 +558,62 @@ func (d *driver) Writer(ctx context.Context, path string, appendParam bool) (sto
 		}
 		return d.newWriter(key, *resp.UploadId, nil), nil
 	}
-	resp, err := d.S3.ListMultipartUploads(&s3.ListMultipartUploadsInput{
+
+	listMultipartUploadsInput := &s3.ListMultipartUploadsInput{
 		Bucket: aws.String(d.Bucket),
 		Prefix: aws.String(key),
-	})
-	if err != nil {
-		return nil, parseError(path, err)
 	}
-	var allParts []*s3.Part
-	for _, multi := range resp.Uploads {
-		if key != *multi.Key {
-			continue
-		}
-		resp, err := d.S3.ListParts(&s3.ListPartsInput{
-			Bucket:   aws.String(d.Bucket),
-			Key:      aws.String(key),
-			UploadId: multi.UploadId,
-		})
+	for {
+		resp, err := d.S3.ListMultipartUploads(listMultipartUploadsInput)
 		if err != nil {
 			return nil, parseError(path, err)
 		}
-		allParts = append(allParts, resp.Parts...)
-		for *resp.IsTruncated {
-			resp, err = d.S3.ListParts(&s3.ListPartsInput{
-				Bucket:           aws.String(d.Bucket),
-				Key:              aws.String(key),
-				UploadId:         multi.UploadId,
-				PartNumberMarker: resp.NextPartNumberMarker,
+
+		// resp.Uploads can only be empty on the first call
+		// if there were no more results to return after the first call, resp.IsTruncated would have been false
+		// and the loop would be exited without recalling ListMultipartUploads
+		if len(resp.Uploads) == 0 {
+			break
+		}
+
+		var allParts []*s3.Part
+		for _, multi := range resp.Uploads {
+			if key != *multi.Key {
+				continue
+			}
+
+			partsList, err := d.S3.ListParts(&s3.ListPartsInput{
+				Bucket:   aws.String(d.Bucket),
+				Key:      aws.String(key),
+				UploadId: multi.UploadId,
 			})
 			if err != nil {
 				return nil, parseError(path, err)
 			}
-			allParts = append(allParts, resp.Parts...)
+			allParts = append(allParts, partsList.Parts...)
+			for *resp.IsTruncated {
+				partsList, err = d.S3.ListParts(&s3.ListPartsInput{
+					Bucket:           aws.String(d.Bucket),
+					Key:              aws.String(key),
+					UploadId:         multi.UploadId,
+					PartNumberMarker: partsList.NextPartNumberMarker,
+				})
+				if err != nil {
+					return nil, parseError(path, err)
+				}
+				allParts = append(allParts, partsList.Parts...)
+			}
+			return d.newWriter(key, *multi.UploadId, allParts), nil
+		}
+
+		// resp.NextUploadIdMarker must have at least one element or we would have returned not found
+		listMultipartUploadsInput.UploadIdMarker = resp.NextUploadIdMarker
+
+		// from the s3 api docs, IsTruncated "specifies whether (true) or not (false) all of the results were returned"
+		// if everything has been returned, break
+		if resp.IsTruncated == nil || !*resp.IsTruncated {
+			break
 		}
-		return d.newWriter(key, *multi.UploadId, allParts), nil
 	}
 	return nil, storagedriver.PathNotFoundError{Path: path}
 }