Skip to content

Commit 2b20531

Browse files
authored
Add additional validation to UploadPartRequests in Transfer Utility (#4083)
1 parent a62545e commit 2b20531

File tree

4 files changed

+237
-5
lines changed

4 files changed

+237
-5
lines changed
Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
{
2+
"services": [
3+
{
4+
"serviceName": "S3",
5+
"type": "patch",
6+
"changeLogMessages": [
7+
"Fixed issue where PartSize and IsLastPart fields were not properly set on Transfer Utility Upload Part Request.",
8+
"Add additional validations for Transfer Utility requests to ensure Upload Parts have the proper Content Length and File Offsets."
9+
]
10+
}
11+
]
12+
}

sdk/src/Services/S3/Custom/Transfer/Internal/MultipartUploadCommand.cs

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,6 @@ internal partial class MultipartUploadCommand : BaseCommand
4949
long _totalTransferredBytes;
5050
Queue<UploadPartRequest> _partsToUpload = new Queue<UploadPartRequest>();
5151

52-
5352
long _contentLength;
5453
private static Logger Logger
5554
{
@@ -211,17 +210,30 @@ internal CompleteMultipartUploadRequest ConstructCompleteMultipartUploadRequest(
211210
return compRequest;
212211
}
213212

213+
private bool calculateIsLastPart(long remainingBytes)
214+
{
215+
var isLastPart = false;
216+
if (remainingBytes <= this._partSize)
217+
isLastPart = true;
218+
return isLastPart;
219+
}
220+
214221
internal UploadPartRequest ConstructUploadPartRequest(int partNumber, long filePosition, InitiateMultipartUploadResponse initiateResponse)
215222
{
216223
UploadPartRequest uploadPartRequest = ConstructGenericUploadPartRequest(initiateResponse);
217224

225+
// Calculating how many bytes are remaining to be uploaded from the current part.
226+
// This is mainly used for the last part scenario.
227+
var remainingBytes = this._contentLength - filePosition;
228+
// We then check based on the remaining bytes and the content length if this is the last part.
229+
var isLastPart = calculateIsLastPart(remainingBytes);
218230
uploadPartRequest.PartNumber = partNumber;
219-
uploadPartRequest.PartSize = this._partSize;
231+
uploadPartRequest.PartSize = isLastPart ? remainingBytes : this._partSize;
232+
uploadPartRequest.IsLastPart = isLastPart;
220233

221-
if ((filePosition + this._partSize >= this._contentLength)
234+
if (isLastPart
222235
&& _s3Client is Amazon.S3.Internal.IAmazonS3Encryption)
223236
{
224-
uploadPartRequest.IsLastPart = true;
225237
uploadPartRequest.PartSize = 0;
226238
}
227239

sdk/src/Services/S3/Custom/Transfer/Internal/_async/MultipartUploadCommand.async.cs

Lines changed: 76 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ internal partial class MultipartUploadCommand : BaseCommand
3131
{
3232
public SemaphoreSlim AsyncThrottler { get; set; }
3333

34+
Dictionary<int, ExpectedUploadPart> _expectedUploadParts = new Dictionary<int, ExpectedUploadPart>();
35+
3436
public override async Task ExecuteAsync(CancellationToken cancellationToken)
3537
{
3638
// Fire transfer initiated event FIRST, before choosing path
@@ -69,6 +71,29 @@ public override async Task ExecuteAsync(CancellationToken cancellationToken)
6971
cancellationToken.ThrowIfCancellationRequested();
7072

7173
var uploadRequest = ConstructUploadPartRequest(i, filePosition, initResponse);
74+
75+
var expectedFileOffset = (i - 1) * this._partSize;
76+
// Calculating how many bytes are remaining to be uploaded from the current part.
77+
// This is mainly used for the last part scenario.
78+
var remainingBytes = this._contentLength - expectedFileOffset;
79+
// We then check based on the remaining bytes and the content length if this is the last part.
80+
var isLastPart = calculateIsLastPart(remainingBytes);
81+
// To maintain the same behavior as the ConstructUploadPartRequest.
82+
// We are setting the remainingBytes/partSize when using the IAmazonS3Encryption client to 0.
83+
if (isLastPart
84+
&& _s3Client is Amazon.S3.Internal.IAmazonS3Encryption)
85+
{
86+
remainingBytes = 0;
87+
}
88+
this._expectedUploadParts.Add(i, new ExpectedUploadPart {
89+
PartNumber = i,
90+
ExpectedContentLength =
91+
isLastPart ?
92+
remainingBytes :
93+
this._partSize,
94+
ExpectedFileOffset = expectedFileOffset,
95+
IsLastPart = isLastPart
96+
});
7297
this._partsToUpload.Enqueue(uploadRequest);
7398
filePosition += this._partSize;
7499
}
@@ -150,8 +175,50 @@ private async Task<UploadPartResponse> UploadPartAsync(UploadPartRequest uploadR
150175
{
151176
try
152177
{
153-
return await _s3Client.UploadPartAsync(uploadRequest, internalCts.Token)
178+
var response = await _s3Client.UploadPartAsync(uploadRequest, internalCts.Token)
154179
.ConfigureAwait(continueOnCapturedContext: false);
180+
181+
if (response.PartNumber is null)
182+
{
183+
throw new ArgumentNullException(nameof(response.PartNumber));
184+
}
185+
else
186+
{
187+
if (this._expectedUploadParts.TryGetValue((int) response.PartNumber, out var expectedUploadPart))
188+
{
189+
var actualContentLength = uploadRequest.PartSize;
190+
if (actualContentLength != expectedUploadPart.ExpectedContentLength)
191+
{
192+
throw new InvalidOperationException($"Cannot complete multipart upload request. The expected content length of part {expectedUploadPart.PartNumber} " +
193+
$"does not equal the actual content length.");
194+
}
195+
196+
if (expectedUploadPart.IsLastPart)
197+
{
198+
if (actualContentLength < 0 ||
199+
actualContentLength > expectedUploadPart.ExpectedContentLength)
200+
{
201+
throw new InvalidOperationException($"Cannot complete multipart upload request. The last part " +
202+
$"has an invalid content length.");
203+
}
204+
}
205+
206+
var actualFileOsset = uploadRequest.FilePosition;
207+
if (uploadRequest.IsSetFilePath() &&
208+
actualFileOsset != expectedUploadPart.ExpectedFileOffset)
209+
{
210+
throw new InvalidOperationException($"Cannot complete multipart upload request. The expected file offset of part {expectedUploadPart.PartNumber} " +
211+
$"does not equal the actual file offset.");
212+
}
213+
}
214+
else
215+
{
216+
throw new InvalidOperationException("Multipart upload request part was unexpected.");
217+
}
218+
}
219+
220+
221+
return response;
155222
}
156223
catch (Exception exception)
157224
{
@@ -326,5 +393,13 @@ await _s3Client.AbortMultipartUploadAsync(new AbortMultipartUploadRequest()
326393
throw;
327394
}
328395
}
396+
397+
private class ExpectedUploadPart
398+
{
399+
public int PartNumber { get; set; }
400+
public long? ExpectedContentLength { get; set; }
401+
public long? ExpectedFileOffset { get; set; }
402+
public bool IsLastPart { get; set; }
403+
}
329404
}
330405
}
Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
using Amazon.S3;
2+
using Amazon.S3.Model;
3+
using Amazon.S3.Transfer;
4+
using Amazon.S3.Transfer.Internal;
5+
using Amazon.S3.Util;
6+
using Microsoft.VisualStudio.TestTools.UnitTesting;
7+
using Moq;
8+
using System;
9+
using System.Collections.Generic;
10+
using System.IO;
11+
using System.Linq;
12+
using System.Text;
13+
using System.Threading;
14+
using System.Threading.Tasks;
15+
16+
namespace AWSSDK.UnitTests
17+
{
18+
[TestClass]
19+
public class MultipartUploadValidationTests
20+
{
21+
private static string _tempFilePath;
22+
private const long fileSizeInBytes = 40 * 1024 * 1024;
23+
24+
[ClassInitialize]
25+
public static void ClassInitialize(TestContext context)
26+
{
27+
_tempFilePath = Path.GetTempFileName();
28+
29+
CreateFileWithSpecificSize(_tempFilePath, fileSizeInBytes);
30+
}
31+
32+
[ClassCleanup]
33+
public static void ClassCleanup()
34+
{
35+
if (File.Exists(_tempFilePath))
36+
{
37+
File.Delete(_tempFilePath);
38+
}
39+
}
40+
41+
private static void CreateFileWithSpecificSize(string path, long size)
42+
{
43+
using (var fileStream = new FileStream(path, FileMode.Create, FileAccess.Write))
44+
{
45+
fileStream.SetLength(size);
46+
}
47+
}
48+
49+
[TestMethod]
50+
[TestCategory("S3")]
51+
public async Task Validation_HappyPath()
52+
{
53+
var initiateMultipartUploadResponse = new InitiateMultipartUploadResponse
54+
{
55+
UploadId = "test"
56+
};
57+
58+
var s3Client = new Mock<AmazonS3Client>();
59+
s3Client
60+
.Setup(x => x.InitiateMultipartUploadAsync(
61+
It.IsAny<InitiateMultipartUploadRequest>(),
62+
It.IsAny<CancellationToken>()))
63+
.ReturnsAsync(initiateMultipartUploadResponse);
64+
65+
s3Client
66+
.Setup(x => x.UploadPartAsync(It.IsAny<UploadPartRequest>(), It.IsAny<CancellationToken>()))
67+
.ReturnsAsync((UploadPartRequest request, CancellationToken cancellationToken) =>
68+
{
69+
return new UploadPartResponse { PartNumber = request.PartNumber };
70+
});
71+
72+
var uploadRequest = new TransferUtilityUploadRequest
73+
{
74+
FilePath = _tempFilePath,
75+
BucketName = "test-bucket",
76+
Key = "test"
77+
};
78+
var multipartUpload = new MultipartUploadCommand(s3Client.Object, new TransferUtilityConfig(), uploadRequest);
79+
await multipartUpload.ExecuteAsync(new CancellationToken());
80+
}
81+
82+
[TestMethod]
83+
[TestCategory("S3")]
84+
public void Validation_ConstructUploadPartRequest()
85+
{
86+
var initiateMultipartUploadResponse = new InitiateMultipartUploadResponse
87+
{
88+
UploadId = "test"
89+
};
90+
91+
var s3Client = new Mock<AmazonS3Client>();
92+
93+
s3Client
94+
.Setup(x => x.InitiateMultipartUploadAsync(
95+
It.IsAny<InitiateMultipartUploadRequest>(),
96+
It.IsAny<CancellationToken>()))
97+
.ReturnsAsync(initiateMultipartUploadResponse);
98+
99+
var uploadRequest = new TransferUtilityUploadRequest
100+
{
101+
FilePath = _tempFilePath,
102+
BucketName = "test-bucket",
103+
Key = "test"
104+
};
105+
106+
var multipartUpload = new MultipartUploadCommand(s3Client.Object, new TransferUtilityConfig(), uploadRequest);
107+
108+
var partSize = Math.Max(S3Constants.DefaultPartSize, uploadRequest.ContentLength / S3Constants.MaxNumberOfParts);
109+
110+
long filePosition = 0;
111+
for (int i = 1; filePosition < uploadRequest.ContentLength; i++)
112+
{
113+
var constructUploadPartRequest = multipartUpload.ConstructUploadPartRequest(i, filePosition, initiateMultipartUploadResponse);
114+
115+
var expectedFileOffset = (i - 1) * partSize;
116+
var remainingBytes = uploadRequest.ContentLength - expectedFileOffset;
117+
var isLastPart = false;
118+
if (remainingBytes <= partSize)
119+
isLastPart = true;
120+
121+
Assert.AreEqual(i, constructUploadPartRequest.PartNumber);
122+
Assert.AreEqual(isLastPart, constructUploadPartRequest.IsLastPart);
123+
Assert.AreEqual(
124+
isLastPart ? remainingBytes : partSize,
125+
constructUploadPartRequest.PartSize);
126+
Assert.AreEqual(expectedFileOffset, constructUploadPartRequest.FilePosition);
127+
128+
filePosition += partSize;
129+
}
130+
131+
}
132+
}
133+
}

0 commit comments

Comments
 (0)