Skip to content

Commit e6f904f

Browse files
philasmarGarrettBeatty
authored andcommitted
Add additional validation to UploadPartRequests in Transfer Utility (#4083)
1 parent 423231e commit e6f904f

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
if ( (this._fileTransporterRequest.InputStream != null && !this._fileTransporterRequest.InputStream.CanSeek) || this._fileTransporterRequest.ContentLength == -1)
@@ -57,6 +59,29 @@ public override async Task ExecuteAsync(CancellationToken cancellationToken)
5759
cancellationToken.ThrowIfCancellationRequested();
5860

5961
var uploadRequest = ConstructUploadPartRequest(i, filePosition, initResponse);
62+
63+
var expectedFileOffset = (i - 1) * this._partSize;
64+
// Calculating how many bytes are remaining to be uploaded from the current part.
65+
// This is mainly used for the last part scenario.
66+
var remainingBytes = this._contentLength - expectedFileOffset;
67+
// We then check based on the remaining bytes and the content length if this is the last part.
68+
var isLastPart = calculateIsLastPart(remainingBytes);
69+
// To maintain the same behavior as the ConstructUploadPartRequest.
70+
// We are setting the remainingBytes/partSize when using the IAmazonS3Encryption client to 0.
71+
if (isLastPart
72+
&& _s3Client is Amazon.S3.Internal.IAmazonS3Encryption)
73+
{
74+
remainingBytes = 0;
75+
}
76+
this._expectedUploadParts.Add(i, new ExpectedUploadPart {
77+
PartNumber = i,
78+
ExpectedContentLength =
79+
isLastPart ?
80+
remainingBytes :
81+
this._partSize,
82+
ExpectedFileOffset = expectedFileOffset,
83+
IsLastPart = isLastPart
84+
});
6085
this._partsToUpload.Enqueue(uploadRequest);
6186
filePosition += this._partSize;
6287
}
@@ -133,8 +158,50 @@ private async Task<UploadPartResponse> UploadPartAsync(UploadPartRequest uploadR
133158
{
134159
try
135160
{
136-
return await _s3Client.UploadPartAsync(uploadRequest, internalCts.Token)
161+
var response = await _s3Client.UploadPartAsync(uploadRequest, internalCts.Token)
137162
.ConfigureAwait(continueOnCapturedContext: false);
163+
164+
if (response.PartNumber is null)
165+
{
166+
throw new ArgumentNullException(nameof(response.PartNumber));
167+
}
168+
else
169+
{
170+
if (this._expectedUploadParts.TryGetValue((int) response.PartNumber, out var expectedUploadPart))
171+
{
172+
var actualContentLength = uploadRequest.PartSize;
173+
if (actualContentLength != expectedUploadPart.ExpectedContentLength)
174+
{
175+
throw new InvalidOperationException($"Cannot complete multipart upload request. The expected content length of part {expectedUploadPart.PartNumber} " +
176+
$"does not equal the actual content length.");
177+
}
178+
179+
if (expectedUploadPart.IsLastPart)
180+
{
181+
if (actualContentLength < 0 ||
182+
actualContentLength > expectedUploadPart.ExpectedContentLength)
183+
{
184+
throw new InvalidOperationException($"Cannot complete multipart upload request. The last part " +
185+
$"has an invalid content length.");
186+
}
187+
}
188+
189+
var actualFileOsset = uploadRequest.FilePosition;
190+
if (uploadRequest.IsSetFilePath() &&
191+
actualFileOsset != expectedUploadPart.ExpectedFileOffset)
192+
{
193+
throw new InvalidOperationException($"Cannot complete multipart upload request. The expected file offset of part {expectedUploadPart.PartNumber} " +
194+
$"does not equal the actual file offset.");
195+
}
196+
}
197+
else
198+
{
199+
throw new InvalidOperationException("Multipart upload request part was unexpected.");
200+
}
201+
}
202+
203+
204+
return response;
138205
}
139206
catch (Exception exception)
140207
{
@@ -293,5 +360,13 @@ await _s3Client.AbortMultipartUploadAsync(new AbortMultipartUploadRequest()
293360
throw;
294361
}
295362
}
363+
364+
private class ExpectedUploadPart
365+
{
366+
public int PartNumber { get; set; }
367+
public long? ExpectedContentLength { get; set; }
368+
public long? ExpectedFileOffset { get; set; }
369+
public bool IsLastPart { get; set; }
370+
}
296371
}
297372
}
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)