Skip to content

Commit a9c26bb

Browse files
zoewanggzoewangg
authored
Add validations for upload in s3 mulitpart client (#6282)
* Add validations for upload in s3 mulitpart client * Improve workflow and add more functional tests * Add test scope
1 parent 769cf01 commit a9c26bb

File tree

14 files changed

+591
-102
lines changed

14 files changed

+591
-102
lines changed

‎.changes/next-release/bugfix-AmazonS3-6522f77.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "bugfix",
3+
"category": "Amazon S3",
4+
"contributor": "",
5+
"description": "Add additional validations for multipart upload operations in the Java multipart S3 client."
6+
}

‎bom-internal/pom.xml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,12 @@
235235
<version>${rxjava.version}</version>
236236
<scope>test</scope>
237237
</dependency>
238+
<dependency>
239+
<groupId>io.reactivex.rxjava3</groupId>
240+
<artifactId>rxjava</artifactId>
241+
<version>${rxjava3.version}</version>
242+
<scope>test</scope>
243+
</dependency>
238244
<dependency>
239245
<artifactId>commons-lang3</artifactId>
240246
<groupId>org.apache.commons</groupId>

‎pom.xml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@
125125
<org.eclipse.jdt.version>3.10.0</org.eclipse.jdt.version>
126126
<org.eclipse.text.version>3.5.101</org.eclipse.text.version>
127127
<rxjava.version>2.2.21</rxjava.version>
128+
<rxjava3.version>3.1.5</rxjava3.version>
128129
<commons-codec.verion>1.17.1</commons-codec.verion>
129130
<jmh.version>1.37</jmh.version>
130131
<awscrt.version>0.38.1</awscrt.version>

‎services/s3/pom.xml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,5 +230,10 @@
230230
<artifactId>jimfs</artifactId>
231231
<scope>test</scope>
232232
</dependency>
233+
<dependency>
234+
<groupId>io.reactivex.rxjava3</groupId>
235+
<artifactId>rxjava</artifactId>
236+
<scope>test</scope>
237+
</dependency>
233238
</dependencies>
234239
</project>

‎services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java

Lines changed: 74 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,14 @@
1515

1616
package software.amazon.awssdk.services.s3.internal.multipart;
1717

18+
import static software.amazon.awssdk.services.s3.internal.multipart.MultipartUploadHelper.contentLengthMismatchForPart;
19+
import static software.amazon.awssdk.services.s3.internal.multipart.MultipartUploadHelper.partNumMismatch;
1820
import static software.amazon.awssdk.services.s3.multipart.S3MultipartExecutionAttribute.JAVA_PROGRESS_LISTENER;
1921

2022
import java.util.Collection;
2123
import java.util.HashMap;
2224
import java.util.Map;
25+
import java.util.Optional;
2326
import java.util.concurrent.CompletableFuture;
2427
import java.util.concurrent.ConcurrentLinkedQueue;
2528
import java.util.concurrent.atomic.AtomicBoolean;
@@ -32,6 +35,7 @@
3235
import software.amazon.awssdk.annotations.SdkInternalApi;
3336
import software.amazon.awssdk.core.async.AsyncRequestBody;
3437
import software.amazon.awssdk.core.async.listener.PublisherListener;
38+
import software.amazon.awssdk.core.exception.SdkClientException;
3539
import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
3640
import software.amazon.awssdk.services.s3.model.CompletedPart;
3741
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
@@ -54,10 +58,10 @@ public class KnownContentLengthAsyncRequestBodySubscriber implements Subscriber<
5458
private final AtomicBoolean failureActionInitiated = new AtomicBoolean(false);
5559
private final AtomicInteger partNumber = new AtomicInteger(1);
5660
private final MultipartUploadHelper multipartUploadHelper;
57-
private final long contentLength;
61+
private final long totalSize;
5862
private final long partSize;
59-
private final int partCount;
60-
private final int numExistingParts;
63+
private final int expectedNumParts;
64+
private final int existingNumParts;
6165
private final String uploadId;
6266
private final Collection<CompletableFuture<CompletedPart>> futures = new ConcurrentLinkedQueue<>();
6367
private final PutObjectRequest putObjectRequest;
@@ -77,25 +81,21 @@ public class KnownContentLengthAsyncRequestBodySubscriber implements Subscriber<
7781
KnownContentLengthAsyncRequestBodySubscriber(MpuRequestContext mpuRequestContext,
7882
CompletableFuture<PutObjectResponse> returnFuture,
7983
MultipartUploadHelper multipartUploadHelper) {
80-
this.contentLength = mpuRequestContext.contentLength();
84+
this.totalSize = mpuRequestContext.contentLength();
8185
this.partSize = mpuRequestContext.partSize();
82-
this.partCount = determinePartCount(contentLength, partSize);
86+
this.expectedNumParts = mpuRequestContext.expectedNumParts();
8387
this.putObjectRequest = mpuRequestContext.request().left();
8488
this.returnFuture = returnFuture;
8589
this.uploadId = mpuRequestContext.uploadId();
8690
this.existingParts = mpuRequestContext.existingParts() == null ? new HashMap<>() : mpuRequestContext.existingParts();
87-
this.numExistingParts = NumericUtils.saturatedCast(mpuRequestContext.numPartsCompleted());
88-
this.completedParts = new AtomicReferenceArray<>(partCount);
91+
this.existingNumParts = NumericUtils.saturatedCast(mpuRequestContext.numPartsCompleted());
92+
this.completedParts = new AtomicReferenceArray<>(expectedNumParts);
8993
this.multipartUploadHelper = multipartUploadHelper;
9094
this.progressListener = putObjectRequest.overrideConfiguration().map(c -> c.executionAttributes()
9195
.getAttribute(JAVA_PROGRESS_LISTENER))
9296
.orElseGet(PublisherListener::noOp);
9397
}
9498

95-
private int determinePartCount(long contentLength, long partSize) {
96-
return (int) Math.ceil(contentLength / (double) partSize);
97-
}
98-
9999
public S3ResumeToken pause() {
100100
isPaused = true;
101101

@@ -119,8 +119,8 @@ public S3ResumeToken pause() {
119119
return S3ResumeToken.builder()
120120
.uploadId(uploadId)
121121
.partSize(partSize)
122-
.totalNumParts((long) partCount)
123-
.numPartsCompleted(numPartsCompleted + numExistingParts)
122+
.totalNumParts((long) expectedNumParts)
123+
.numPartsCompleted(numPartsCompleted + existingNumParts)
124124
.build();
125125
}
126126

@@ -145,21 +145,32 @@ public void onSubscribe(Subscription s) {
145145

146146
@Override
147147
public void onNext(AsyncRequestBody asyncRequestBody) {
148-
if (isPaused) {
148+
if (isPaused || isDone) {
149149
return;
150150
}
151151

152-
if (existingParts.containsKey(partNumber.get())) {
153-
partNumber.getAndIncrement();
152+
int currentPartNum = partNumber.getAndIncrement();
153+
if (existingParts.containsKey(currentPartNum)) {
154154
asyncRequestBody.subscribe(new CancelledSubscriber<>());
155155
subscription.request(1);
156156
asyncRequestBody.contentLength().ifPresent(progressListener::subscriberOnNext);
157157
return;
158158
}
159159

160+
Optional<SdkClientException> sdkClientException = validatePart(asyncRequestBody, currentPartNum);
161+
if (sdkClientException.isPresent()) {
162+
multipartUploadHelper.failRequestsElegantly(futures,
163+
sdkClientException.get(),
164+
uploadId,
165+
returnFuture,
166+
putObjectRequest);
167+
subscription.cancel();
168+
return;
169+
}
170+
160171
asyncRequestBodyInFlight.incrementAndGet();
161172
UploadPartRequest uploadRequest = SdkPojoConversionUtils.toUploadPartRequest(putObjectRequest,
162-
partNumber.getAndIncrement(),
173+
currentPartNum,
163174
uploadId);
164175

165176
Consumer<CompletedPart> completedPartConsumer = completedPart -> completedParts.set(completedPart.partNumber() - 1,
@@ -179,6 +190,39 @@ public void onNext(AsyncRequestBody asyncRequestBody) {
179190
subscription.request(1);
180191
}
181192

193+
private Optional<SdkClientException> validatePart(AsyncRequestBody asyncRequestBody, int currentPartNum) {
194+
if (!asyncRequestBody.contentLength().isPresent()) {
195+
return Optional.of(MultipartUploadHelper.contentLengthMissingForPart(currentPartNum));
196+
}
197+
198+
Long currentPartSize = asyncRequestBody.contentLength().get();
199+
200+
if (currentPartNum > expectedNumParts) {
201+
return Optional.of(partNumMismatch(expectedNumParts, currentPartNum));
202+
}
203+
204+
if (currentPartNum == expectedNumParts) {
205+
return validateLastPartSize(currentPartSize);
206+
}
207+
208+
if (currentPartSize != partSize) {
209+
return Optional.of(contentLengthMismatchForPart(partSize, currentPartSize));
210+
}
211+
return Optional.empty();
212+
}
213+
214+
private Optional<SdkClientException> validateLastPartSize(Long currentPartSize) {
215+
long remainder = totalSize % partSize;
216+
long expectedLastPartSize = remainder == 0 ? partSize : remainder;
217+
if (currentPartSize != expectedLastPartSize) {
218+
return Optional.of(
219+
SdkClientException.create("Content length of the last part must be equal to the "
220+
+ "expected last part size. Expected: " + expectedLastPartSize
221+
+ ", Actual: " + currentPartSize));
222+
}
223+
return Optional.empty();
224+
}
225+
182226
private boolean shouldFailRequest() {
183227
return failureActionInitiated.compareAndSet(false, true) && !isPaused;
184228
}
@@ -187,6 +231,7 @@ private boolean shouldFailRequest() {
187231
public void onError(Throwable t) {
188232
log.debug(() -> "Received onError ", t);
189233
if (failureActionInitiated.compareAndSet(false, true)) {
234+
isDone = true;
190235
multipartUploadHelper.failRequestsElegantly(futures, t, uploadId, returnFuture, putObjectRequest);
191236
}
192237
}
@@ -203,6 +248,7 @@ public void onComplete() {
203248
private void completeMultipartUploadIfFinished(int requestsInFlight) {
204249
if (isDone && requestsInFlight == 0 && completedMultipartInitiated.compareAndSet(false, true)) {
205250
CompletedPart[] parts;
251+
206252
if (existingParts.isEmpty()) {
207253
parts =
208254
IntStream.range(0, completedParts.length())
@@ -212,15 +258,23 @@ private void completeMultipartUploadIfFinished(int requestsInFlight) {
212258
// List of CompletedParts needs to be in ascending order
213259
parts = mergeCompletedParts();
214260
}
261+
262+
int actualNumParts = partNumber.get() - 1;
263+
if (actualNumParts != expectedNumParts) {
264+
SdkClientException exception = partNumMismatch(expectedNumParts, actualNumParts);
265+
multipartUploadHelper.failRequestsElegantly(futures, exception, uploadId, returnFuture, putObjectRequest);
266+
return;
267+
}
268+
215269
completeMpuFuture = multipartUploadHelper.completeMultipartUpload(returnFuture, uploadId, parts, putObjectRequest,
216-
contentLength);
270+
totalSize);
217271
}
218272
}
219273

220274
private CompletedPart[] mergeCompletedParts() {
221-
CompletedPart[] merged = new CompletedPart[partCount];
275+
CompletedPart[] merged = new CompletedPart[expectedNumParts];
222276
int currPart = 1;
223-
while (currPart < partCount + 1) {
277+
while (currPart < expectedNumParts + 1) {
224278
CompletedPart completedPart = existingParts.containsKey(currPart) ? existingParts.get(currPart) :
225279
completedParts.get(currPart - 1);
226280
merged[currPart - 1] = completedPart;

‎services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/MpuRequestContext.java

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import software.amazon.awssdk.services.s3.model.CompletedPart;
2323
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
2424
import software.amazon.awssdk.utils.Pair;
25+
import software.amazon.awssdk.utils.Validate;
2526

2627
@SdkInternalApi
2728
public class MpuRequestContext {
@@ -32,6 +33,7 @@ public class MpuRequestContext {
3233
private final Long numPartsCompleted;
3334
private final String uploadId;
3435
private final Map<Integer, CompletedPart> existingParts;
36+
private final int expectedNumParts;
3537

3638
protected MpuRequestContext(Builder builder) {
3739
this.request = builder.request;
@@ -40,6 +42,8 @@ protected MpuRequestContext(Builder builder) {
4042
this.uploadId = builder.uploadId;
4143
this.existingParts = builder.existingParts;
4244
this.numPartsCompleted = builder.numPartsCompleted;
45+
this.expectedNumParts = Validate.paramNotNull(builder.expectedNumParts,
46+
"expectedNumParts");
4347
}
4448

4549
public static Builder builder() {
@@ -56,9 +60,13 @@ public boolean equals(Object o) {
5660
}
5761
MpuRequestContext that = (MpuRequestContext) o;
5862

59-
return Objects.equals(request, that.request) && Objects.equals(contentLength, that.contentLength)
60-
&& Objects.equals(partSize, that.partSize) && Objects.equals(numPartsCompleted, that.numPartsCompleted)
61-
&& Objects.equals(uploadId, that.uploadId) && Objects.equals(existingParts, that.existingParts);
63+
return expectedNumParts == that.expectedNumParts
64+
&& Objects.equals(request, that.request)
65+
&& Objects.equals(contentLength, that.contentLength)
66+
&& Objects.equals(partSize, that.partSize)
67+
&& Objects.equals(numPartsCompleted, that.numPartsCompleted)
68+
&& Objects.equals(uploadId, that.uploadId)
69+
&& Objects.equals(existingParts, that.existingParts);
6270
}
6371

6472
@Override
@@ -69,6 +77,7 @@ public int hashCode() {
6977
result = 31 * result + (contentLength != null ? contentLength.hashCode() : 0);
7078
result = 31 * result + (partSize != null ? partSize.hashCode() : 0);
7179
result = 31 * result + (numPartsCompleted != null ? numPartsCompleted.hashCode() : 0);
80+
result = 31 * result + expectedNumParts;
7281
return result;
7382
}
7483

@@ -92,6 +101,10 @@ public String uploadId() {
92101
return uploadId;
93102
}
94103

104+
public int expectedNumParts() {
105+
return expectedNumParts;
106+
}
107+
95108
public Map<Integer, CompletedPart> existingParts() {
96109
return existingParts;
97110
}
@@ -103,6 +116,7 @@ public static final class Builder {
103116
private Long numPartsCompleted;
104117
private String uploadId;
105118
private Map<Integer, CompletedPart> existingParts;
119+
private Integer expectedNumParts;
106120

107121
private Builder() {
108122
}
@@ -137,6 +151,11 @@ public Builder existingParts(Map<Integer, CompletedPart> existingParts) {
137151
return this;
138152
}
139153

154+
public Builder expectedNumParts(Integer expectedNumParts) {
155+
this.expectedNumParts = expectedNumParts;
156+
return this;
157+
}
158+
140159
public MpuRequestContext build() {
141160
return new MpuRequestContext(this);
142161
}

‎services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartUploadHelper.java

Lines changed: 30 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import software.amazon.awssdk.annotations.SdkInternalApi;
2626
import software.amazon.awssdk.core.async.AsyncRequestBody;
2727
import software.amazon.awssdk.core.async.listener.PublisherListener;
28+
import software.amazon.awssdk.core.exception.SdkClientException;
2829
import software.amazon.awssdk.services.s3.S3AsyncClient;
2930
import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
3031
import software.amazon.awssdk.services.s3.model.CompletedPart;
@@ -47,18 +48,15 @@ public final class MultipartUploadHelper {
4748
private static final Logger log = Logger.loggerFor(MultipartUploadHelper.class);
4849

4950
private final S3AsyncClient s3AsyncClient;
50-
private final long partSizeInBytes;
5151
private final GenericMultipartHelper<PutObjectRequest, PutObjectResponse> genericMultipartHelper;
5252

5353
private final long maxMemoryUsageInBytes;
5454
private final long multipartUploadThresholdInBytes;
5555

5656
public MultipartUploadHelper(S3AsyncClient s3AsyncClient,
57-
long partSizeInBytes,
5857
long multipartUploadThresholdInBytes,
5958
long maxMemoryUsageInBytes) {
6059
this.s3AsyncClient = s3AsyncClient;
61-
this.partSizeInBytes = partSizeInBytes;
6260
this.genericMultipartHelper = new GenericMultipartHelper<>(s3AsyncClient,
6361
SdkPojoConversionUtils::toAbortMultipartUploadRequest,
6462
SdkPojoConversionUtils::toPutObjectResponse);
@@ -123,11 +121,18 @@ void failRequestsElegantly(Collection<CompletableFuture<CompletedPart>> futures,
123121
String uploadId,
124122
CompletableFuture<PutObjectResponse> returnFuture,
125123
PutObjectRequest putObjectRequest) {
126-
genericMultipartHelper.handleException(returnFuture, () -> "Failed to send multipart upload requests", t);
127-
if (uploadId != null) {
128-
genericMultipartHelper.cleanUpParts(uploadId, toAbortMultipartUploadRequest(putObjectRequest));
124+
125+
try {
126+
genericMultipartHelper.handleException(returnFuture, () -> "Failed to send multipart upload requests", t);
127+
if (uploadId != null) {
128+
genericMultipartHelper.cleanUpParts(uploadId, toAbortMultipartUploadRequest(putObjectRequest));
129+
}
130+
cancelingOtherOngoingRequests(futures, t);
131+
} catch (Throwable throwable) {
132+
returnFuture.completeExceptionally(SdkClientException.create("Unexpected error occurred while handling the upstream "
133+
+ "exception.", throwable));
129134
}
130-
cancelingOtherOngoingRequests(futures, t);
135+
131136
}
132137

133138
static void cancelingOtherOngoingRequests(Collection<CompletableFuture<CompletedPart>> futures, Throwable t) {
@@ -152,4 +157,22 @@ void uploadInOneChunk(PutObjectRequest putObjectRequest,
152157
CompletableFutureUtils.forwardExceptionTo(returnFuture, putObjectResponseCompletableFuture);
153158
CompletableFutureUtils.forwardResultTo(putObjectResponseCompletableFuture, returnFuture);
154159
}
160+
161+
static SdkClientException contentLengthMissingForPart(int currentPartNum) {
162+
return SdkClientException.create("Content length is missing on the AsyncRequestBody for part number " + currentPartNum);
163+
}
164+
165+
static SdkClientException contentLengthMismatchForPart(long expected, long actual) {
166+
return SdkClientException.create(String.format("Content length must not be greater than "
167+
+ "part size. Expected: %d, Actual: %d",
168+
expected,
169+
actual));
170+
}
171+
172+
static SdkClientException partNumMismatch(int expectedNumParts, int actualNumParts) {
173+
return SdkClientException.create(String.format("The number of parts divided is "
174+
+ "not equal to the expected number of "
175+
+ "parts. Expected: %d, Actual: %d",
176+
expectedNumParts, actualNumParts));
177+
}
155178
}

0 commit comments

Comments
 (0)