Skip to content

Commit e843444

Browse files
zoewanggzoewangg
committed
Add validations for upload in s3 mulitpart client
1 parent 8724d3a commit e843444

File tree

9 files changed

+416
-84
lines changed

9 files changed

+416
-84
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+
}

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

Lines changed: 61 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import software.amazon.awssdk.annotations.SdkInternalApi;
3333
import software.amazon.awssdk.core.async.AsyncRequestBody;
3434
import software.amazon.awssdk.core.async.listener.PublisherListener;
35+
import software.amazon.awssdk.core.exception.SdkClientException;
3536
import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
3637
import software.amazon.awssdk.services.s3.model.CompletedPart;
3738
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
@@ -54,9 +55,9 @@ public class KnownContentLengthAsyncRequestBodySubscriber implements Subscriber<
5455
private final AtomicBoolean failureActionInitiated = new AtomicBoolean(false);
5556
private final AtomicInteger partNumber = new AtomicInteger(1);
5657
private final MultipartUploadHelper multipartUploadHelper;
57-
private final long contentLength;
58+
private final long totalSize;
5859
private final long partSize;
59-
private final int partCount;
60+
private final int expectedNumPart;
6061
private final int numExistingParts;
6162
private final String uploadId;
6263
private final Collection<CompletableFuture<CompletedPart>> futures = new ConcurrentLinkedQueue<>();
@@ -77,25 +78,21 @@ public class KnownContentLengthAsyncRequestBodySubscriber implements Subscriber<
7778
KnownContentLengthAsyncRequestBodySubscriber(MpuRequestContext mpuRequestContext,
7879
CompletableFuture<PutObjectResponse> returnFuture,
7980
MultipartUploadHelper multipartUploadHelper) {
80-
this.contentLength = mpuRequestContext.contentLength();
81+
this.totalSize = mpuRequestContext.contentLength();
8182
this.partSize = mpuRequestContext.partSize();
82-
this.partCount = determinePartCount(contentLength, partSize);
83+
this.expectedNumPart = mpuRequestContext.expectedNumParts();
8384
this.putObjectRequest = mpuRequestContext.request().left();
8485
this.returnFuture = returnFuture;
8586
this.uploadId = mpuRequestContext.uploadId();
8687
this.existingParts = mpuRequestContext.existingParts() == null ? new HashMap<>() : mpuRequestContext.existingParts();
8788
this.numExistingParts = NumericUtils.saturatedCast(mpuRequestContext.numPartsCompleted());
88-
this.completedParts = new AtomicReferenceArray<>(partCount);
89+
this.completedParts = new AtomicReferenceArray<>(expectedNumPart);
8990
this.multipartUploadHelper = multipartUploadHelper;
9091
this.progressListener = putObjectRequest.overrideConfiguration().map(c -> c.executionAttributes()
9192
.getAttribute(JAVA_PROGRESS_LISTENER))
9293
.orElseGet(PublisherListener::noOp);
9394
}
9495

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

@@ -119,7 +116,7 @@ public S3ResumeToken pause() {
119116
return S3ResumeToken.builder()
120117
.uploadId(uploadId)
121118
.partSize(partSize)
122-
.totalNumParts((long) partCount)
119+
.totalNumParts((long) expectedNumPart)
123120
.numPartsCompleted(numPartsCompleted + numExistingParts)
124121
.build();
125122
}
@@ -145,21 +142,23 @@ public void onSubscribe(Subscription s) {
145142

146143
@Override
147144
public void onNext(AsyncRequestBody asyncRequestBody) {
148-
if (isPaused) {
145+
if (isPaused || isDone) {
149146
return;
150147
}
151148

152-
if (existingParts.containsKey(partNumber.get())) {
153-
partNumber.getAndIncrement();
149+
int currentPartNum = partNumber.getAndIncrement();
150+
if (existingParts.containsKey(currentPartNum)) {
154151
asyncRequestBody.subscribe(new CancelledSubscriber<>());
155152
subscription.request(1);
156153
asyncRequestBody.contentLength().ifPresent(progressListener::subscriberOnNext);
157154
return;
158155
}
159156

157+
validatePart(asyncRequestBody, currentPartNum);
158+
160159
asyncRequestBodyInFlight.incrementAndGet();
161160
UploadPartRequest uploadRequest = SdkPojoConversionUtils.toUploadPartRequest(putObjectRequest,
162-
partNumber.getAndIncrement(),
161+
currentPartNum,
163162
uploadId);
164163

165164
Consumer<CompletedPart> completedPartConsumer = completedPart -> completedParts.set(completedPart.partNumber() - 1,
@@ -179,6 +178,49 @@ public void onNext(AsyncRequestBody asyncRequestBody) {
179178
subscription.request(1);
180179
}
181180

181+
private void validatePart(AsyncRequestBody asyncRequestBody, int currentPartNum) {
182+
if (!asyncRequestBody.contentLength().isPresent()) {
183+
SdkClientException e = SdkClientException.create("Content length must be present on the AsyncRequestBody");
184+
multipartUploadHelper.failRequestsElegantly(futures, e, uploadId, returnFuture, putObjectRequest);
185+
return;
186+
}
187+
188+
Long currentPartSize = asyncRequestBody.contentLength().get();
189+
if (currentPartNum > expectedNumPart) {
190+
SdkClientException exception = SdkClientException.create(String.format("The number of parts divided is "
191+
+ "not equal to the expected number of "
192+
+ "parts. Expected: %d, Actual: %d",
193+
expectedNumPart, currentPartNum));
194+
multipartUploadHelper.failRequestsElegantly(futures, exception, uploadId, returnFuture, putObjectRequest);
195+
return;
196+
}
197+
198+
if (currentPartNum == expectedNumPart) {
199+
validateLastPartSize(currentPartSize);
200+
return;
201+
}
202+
203+
if (currentPartSize != partSize) {
204+
SdkClientException e = SdkClientException.create(String.format("Content length must not be greater than the "
205+
+ "part size. Expected: %d, Actual: %d",
206+
partSize,
207+
currentPartSize));
208+
multipartUploadHelper.failRequestsElegantly(futures, e, uploadId, returnFuture, putObjectRequest);
209+
}
210+
}
211+
212+
private void validateLastPartSize(Long currentPartSize) {
213+
long remainder = totalSize % partSize;
214+
long expectedLastPartSize = remainder == 0 ? partSize : remainder;
215+
if (currentPartSize != expectedLastPartSize) {
216+
SdkClientException exception =
217+
SdkClientException.create("Content length of the last part must be equal to the "
218+
+ "expected last part size. Expected: " + expectedLastPartSize
219+
+ ", Actual: " + currentPartSize);
220+
multipartUploadHelper.failRequestsElegantly(futures, exception, uploadId, returnFuture, putObjectRequest);
221+
}
222+
}
223+
182224
private boolean shouldFailRequest() {
183225
return failureActionInitiated.compareAndSet(false, true) && !isPaused;
184226
}
@@ -187,6 +229,7 @@ private boolean shouldFailRequest() {
187229
public void onError(Throwable t) {
188230
log.debug(() -> "Received onError ", t);
189231
if (failureActionInitiated.compareAndSet(false, true)) {
232+
isDone = true;
190233
multipartUploadHelper.failRequestsElegantly(futures, t, uploadId, returnFuture, putObjectRequest);
191234
}
192235
}
@@ -203,6 +246,7 @@ public void onComplete() {
203246
private void completeMultipartUploadIfFinished(int requestsInFlight) {
204247
if (isDone && requestsInFlight == 0 && completedMultipartInitiated.compareAndSet(false, true)) {
205248
CompletedPart[] parts;
249+
206250
if (existingParts.isEmpty()) {
207251
parts =
208252
IntStream.range(0, completedParts.length())
@@ -213,14 +257,14 @@ private void completeMultipartUploadIfFinished(int requestsInFlight) {
213257
parts = mergeCompletedParts();
214258
}
215259
completeMpuFuture = multipartUploadHelper.completeMultipartUpload(returnFuture, uploadId, parts, putObjectRequest,
216-
contentLength);
260+
totalSize);
217261
}
218262
}
219263

220264
private CompletedPart[] mergeCompletedParts() {
221-
CompletedPart[] merged = new CompletedPart[partCount];
265+
CompletedPart[] merged = new CompletedPart[expectedNumPart];
222266
int currPart = 1;
223-
while (currPart < partCount + 1) {
267+
while (currPart < expectedNumPart + 1) {
224268
CompletedPart completedPart = existingParts.containsKey(currPart) ? existingParts.get(currPart) :
225269
completedParts.get(currPart - 1);
226270
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: 13 additions & 4 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;
@@ -123,11 +124,19 @@ void failRequestsElegantly(Collection<CompletableFuture<CompletedPart>> futures,
123124
String uploadId,
124125
CompletableFuture<PutObjectResponse> returnFuture,
125126
PutObjectRequest putObjectRequest) {
126-
genericMultipartHelper.handleException(returnFuture, () -> "Failed to send multipart upload requests", t);
127-
if (uploadId != null) {
128-
genericMultipartHelper.cleanUpParts(uploadId, toAbortMultipartUploadRequest(putObjectRequest));
127+
128+
try {
129+
genericMultipartHelper.handleException(returnFuture, () -> "Failed to send multipart upload requests", t);
130+
if (uploadId != null) {
131+
genericMultipartHelper.cleanUpParts(uploadId, toAbortMultipartUploadRequest(putObjectRequest));
132+
}
133+
cancelingOtherOngoingRequests(futures, t);
134+
} catch (Throwable throwable) {
135+
log.debug(() -> "Unexpected error occurred.", throwable);
136+
returnFuture.completeExceptionally(SdkClientException.create("Unexpected error occurred while handling the upstream "
137+
+ "exception.", t));
129138
}
130-
cancelingOtherOngoingRequests(futures, t);
139+
131140
}
132141

133142
static void cancelingOtherOngoingRequests(Collection<CompletableFuture<CompletedPart>> futures, Throwable t) {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ private void uploadFromBeginning(Pair<PutObjectRequest, AsyncRequestBody> reques
137137
.partSize(partSize)
138138
.uploadId(uploadId)
139139
.numPartsCompleted(numPartsCompleted)
140+
.expectedNumParts(partCount)
140141
.build();
141142

142143
splitAndSubscribe(mpuRequestContext, returnFuture);
@@ -170,6 +171,7 @@ private void resumePausedUpload(ResumeRequestContext resumeContext) {
170171
.partSize(resumeToken.partSize())
171172
.uploadId(uploadId)
172173
.existingParts(existingParts)
174+
.expectedNumParts(Math.toIntExact(resumeToken.totalNumParts()))
173175
.numPartsCompleted(resumeToken.numPartsCompleted())
174176
.build();
175177

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

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,9 @@ public void onSubscribe(Subscription s) {
160160

161161
@Override
162162
public void onNext(AsyncRequestBody asyncRequestBody) {
163+
if (isDone) {
164+
return;
165+
}
163166
int currentPartNum = partNumber.incrementAndGet();
164167
log.trace(() -> "Received asyncRequestBody " + asyncRequestBody.contentLength());
165168
asyncRequestBodyInFlight.incrementAndGet();
@@ -211,7 +214,17 @@ private void sendUploadPartRequest(String uploadId,
211214
SdkClientException e = SdkClientException.create("Content length must be present on the AsyncRequestBody");
212215
multipartUploadHelper.failRequestsElegantly(futures, e, uploadId, returnFuture, putObjectRequest);
213216
}
214-
this.contentLength.getAndAdd(contentLength.get());
217+
218+
Long contentLengthCurrentPart = contentLength.get();
219+
if (contentLengthCurrentPart > partSizeInBytes) {
220+
SdkClientException e = SdkClientException.create(String.format("Content length must not be greater than the "
221+
+ "part size. Expected: %d, Actual: %d",
222+
partSizeInBytes,
223+
contentLengthCurrentPart));
224+
multipartUploadHelper.failRequestsElegantly(futures, e, uploadId, returnFuture, putObjectRequest);
225+
}
226+
227+
this.contentLength.getAndAdd(contentLengthCurrentPart);
215228

216229
multipartUploadHelper
217230
.sendIndividualUploadPartRequest(uploadId, completedParts::add, futures,
@@ -235,13 +248,15 @@ private Pair<UploadPartRequest, AsyncRequestBody> uploadPart(AsyncRequestBody as
235248
SdkPojoConversionUtils.toUploadPartRequest(putObjectRequest,
236249
partNum,
237250
uploadId);
251+
238252
return Pair.of(uploadRequest, asyncRequestBody);
239253
}
240254

241255
@Override
242256
public void onError(Throwable t) {
243257
log.debug(() -> "Received onError() ", t);
244258
if (failureActionInitiated.compareAndSet(false, true)) {
259+
isDone = true;
245260
multipartUploadHelper.failRequestsElegantly(futures, t, uploadId, returnFuture, putObjectRequest);
246261
}
247262
}
@@ -264,8 +279,19 @@ private void completeMultipartUploadIfFinish(int requestsInFlight) {
264279
CompletedPart[] parts = completedParts.stream()
265280
.sorted(Comparator.comparingInt(CompletedPart::partNumber))
266281
.toArray(CompletedPart[]::new);
282+
283+
long totalLength = contentLength.get();
284+
int expectedNumPart = genericMultipartHelper.determinePartCount(totalLength, partSizeInBytes);
285+
if (parts.length != expectedNumPart) {
286+
SdkClientException exception = SdkClientException.create(
287+
String.format("The number of UploadParts requests is not equal to the expected number of parts. "
288+
+ "Expected: %d, Actual: %d", expectedNumPart, parts.length));
289+
multipartUploadHelper.failRequestsElegantly(futures, exception, uploadId, returnFuture, putObjectRequest);
290+
return;
291+
}
292+
267293
multipartUploadHelper.completeMultipartUpload(returnFuture, uploadId, parts, putObjectRequest,
268-
this.contentLength.get());
294+
totalLength);
269295
}
270296
}
271297
}

0 commit comments

Comments
 (0)