Skip to content

Commit fbdbe92

Browse files
author
Anirav Kareddy
committed
Added tests for S3EncryptonClient with Instruction File + Multipart enabled
1 parent ef353d7 commit fbdbe92

File tree

1 file changed

+165
-0
lines changed

1 file changed

+165
-0
lines changed

‎src/test/java/software/amazon/encryption/s3/S3EncryptionClientInstructionFileTest.java‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,26 +9,46 @@
99
import com.amazonaws.services.s3.model.EncryptionMaterialsProvider;
1010
import com.amazonaws.services.s3.model.KMSEncryptionMaterials;
1111
import com.amazonaws.services.s3.model.StaticEncryptionMaterialsProvider;
12+
import org.apache.commons.io.IOUtils;
1213
import org.junit.jupiter.api.Test;
1314
import software.amazon.awssdk.core.ResponseBytes;
15+
import software.amazon.awssdk.core.ResponseInputStream;
1416
import software.amazon.awssdk.core.sync.RequestBody;
17+
import software.amazon.awssdk.regions.Region;
1518
import software.amazon.awssdk.services.s3.S3Client;
19+
import software.amazon.awssdk.services.s3.model.CompletedPart;
20+
import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse;
1621
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
1722
import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
23+
import software.amazon.awssdk.services.s3.model.SdkPartType;
24+
import software.amazon.awssdk.services.s3.model.UploadPartRequest;
25+
import software.amazon.awssdk.services.s3.model.UploadPartResponse;
1826
import software.amazon.encryption.s3.internal.InstructionFileConfig;
27+
import software.amazon.encryption.s3.utils.BoundedInputStream;
1928

2029
import javax.crypto.KeyGenerator;
2130
import javax.crypto.SecretKey;
2231

32+
import java.io.ByteArrayInputStream;
33+
import java.io.ByteArrayOutputStream;
34+
import java.io.IOException;
35+
import java.io.InputStream;
2336
import java.security.KeyPair;
2437
import java.security.KeyPairGenerator;
2538
import java.security.NoSuchAlgorithmException;
39+
import java.util.ArrayList;
40+
import java.util.HashMap;
41+
import java.util.List;
42+
import java.util.Map;
2643

2744
import static org.junit.jupiter.api.Assertions.assertEquals;
45+
import static org.junit.jupiter.api.Assertions.assertNotNull;
2846
import static org.junit.jupiter.api.Assertions.assertTrue;
2947
import static org.junit.jupiter.api.Assertions.fail;
48+
import static software.amazon.encryption.s3.S3EncryptionClient.withAdditionalConfiguration;
3049
import static software.amazon.encryption.s3.utils.S3EncryptionClientTestResources.BUCKET;
3150
import static software.amazon.encryption.s3.utils.S3EncryptionClientTestResources.KMS_KEY_ID;
51+
import static software.amazon.encryption.s3.utils.S3EncryptionClientTestResources.S3_REGION;
3252
import static software.amazon.encryption.s3.utils.S3EncryptionClientTestResources.appendTestSuffix;
3353
import static software.amazon.encryption.s3.utils.S3EncryptionClientTestResources.deleteObject;
3454

@@ -290,4 +310,149 @@ public void testPutWithInstructionFileV3ToV2Rsa() throws NoSuchAlgorithmExceptio
290310
deleteObject(BUCKET, objectKey, s3Client);
291311
s3Client.close();
292312
}
313+
314+
@Test
315+
public void testMultipartPutWithInstructionFile() throws IOException, NoSuchAlgorithmException {
316+
final String object_key = appendTestSuffix("test-multipart-put-instruction-file");
317+
318+
final long fileSizeLimit = 1024 * 1024 * 50; //50 MB
319+
final InputStream inputStream = new BoundedInputStream(fileSizeLimit);
320+
final InputStream objectStreamForResult = new BoundedInputStream(fileSizeLimit);
321+
322+
S3Client wrappedClient = S3Client.create();
323+
S3Client s3Client = S3EncryptionClient.builder()
324+
.instructionFileConfig(InstructionFileConfig.builder()
325+
.instructionFileClient(wrappedClient)
326+
.enableInstructionFilePutObject(true)
327+
.build())
328+
.kmsKeyId(KMS_KEY_ID)
329+
.build();
330+
331+
Map<String, String> encryptionContext = new HashMap<>();
332+
encryptionContext.put("test-key", "test-value");
333+
334+
335+
s3Client.putObject(builder -> builder
336+
.bucket(BUCKET)
337+
.overrideConfiguration(withAdditionalConfiguration(encryptionContext))
338+
.key(object_key), RequestBody.fromInputStream(inputStream, fileSizeLimit));
339+
340+
S3Client defaultClient = S3Client.create();
341+
ResponseBytes<GetObjectResponse> directInstGetResponse = defaultClient.getObjectAsBytes(builder -> builder
342+
.bucket(BUCKET)
343+
.key(object_key + ".instruction")
344+
.build());
345+
assertTrue(directInstGetResponse.response().metadata().containsKey("x-amz-crypto-instr-file"));
346+
347+
ResponseInputStream<GetObjectResponse> getResponse = s3Client.getObject(builder -> builder
348+
.bucket(BUCKET)
349+
.overrideConfiguration(withAdditionalConfiguration(encryptionContext))
350+
.key(object_key));
351+
352+
assertTrue(IOUtils.contentEquals(objectStreamForResult, getResponse));
353+
354+
deleteObject(BUCKET, object_key, s3Client);
355+
s3Client.close();
356+
357+
}
358+
359+
@Test
360+
public void testLowLevelMultipartPutWithInstructionFile() throws NoSuchAlgorithmException, IOException {
361+
final String object_key = appendTestSuffix("test-low-level-multipart-put-instruction-file");
362+
363+
final long fileSizeLimit = 1024 * 1024 * 50;
364+
final int PART_SIZE = 10 * 1024 * 1024;
365+
final InputStream inputStream = new BoundedInputStream(fileSizeLimit);
366+
final InputStream objectStreamForResult = new BoundedInputStream(fileSizeLimit);
367+
368+
KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA");
369+
keyPairGen.initialize(2048);
370+
KeyPair rsaKey = keyPairGen.generateKeyPair();
371+
372+
S3Client wrappedClient = S3Client.create();
373+
374+
S3Client v3Client = S3EncryptionClient.builder()
375+
.rsaKeyPair(rsaKey)
376+
.instructionFileConfig(InstructionFileConfig.builder()
377+
.instructionFileClient(wrappedClient)
378+
.enableInstructionFilePutObject(true)
379+
.build())
380+
.enableDelayedAuthenticationMode(true)
381+
.build();
382+
383+
384+
CreateMultipartUploadResponse initiateResult = v3Client.createMultipartUpload(builder ->
385+
builder.bucket(BUCKET).key(object_key));
386+
387+
List<CompletedPart> partETags = new ArrayList<>();
388+
389+
int bytesRead, bytesSent = 0;
390+
byte[] partData = new byte[PART_SIZE];
391+
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
392+
int partsSent = 1;
393+
while ((bytesRead = inputStream.read(partData, 0, partData.length)) != -1) {
394+
outputStream.write(partData, 0, bytesRead);
395+
if (bytesSent < PART_SIZE) {
396+
bytesSent += bytesRead;
397+
continue;
398+
}
399+
UploadPartRequest uploadPartRequest = UploadPartRequest.builder()
400+
.bucket(BUCKET)
401+
.key(object_key)
402+
.uploadId(initiateResult.uploadId())
403+
.partNumber(partsSent)
404+
.build();
405+
406+
final InputStream partInputStream = new ByteArrayInputStream(outputStream.toByteArray());
407+
408+
UploadPartResponse uploadPartResult = v3Client.uploadPart(uploadPartRequest,
409+
RequestBody.fromInputStream(partInputStream, partInputStream.available()));
410+
411+
partETags.add(CompletedPart.builder()
412+
.partNumber(partsSent)
413+
.eTag(uploadPartResult.eTag())
414+
.build());
415+
outputStream.reset();
416+
bytesSent = 0;
417+
partsSent++;
418+
}
419+
inputStream.close();
420+
UploadPartRequest uploadPartRequest = UploadPartRequest.builder()
421+
.bucket(BUCKET)
422+
.key(object_key)
423+
.uploadId(initiateResult.uploadId())
424+
.partNumber(partsSent)
425+
.sdkPartType(SdkPartType.LAST)
426+
.build();
427+
final InputStream partInputStream = new ByteArrayInputStream(outputStream.toByteArray());
428+
UploadPartResponse uploadPartResult = v3Client.uploadPart(uploadPartRequest,
429+
RequestBody.fromInputStream(partInputStream, partInputStream.available()));
430+
partETags.add(CompletedPart.builder()
431+
.partNumber(partsSent)
432+
.eTag(uploadPartResult.eTag())
433+
.build());
434+
v3Client.completeMultipartUpload(builder -> builder
435+
.bucket(BUCKET)
436+
.key(object_key)
437+
.uploadId(initiateResult.uploadId())
438+
.multipartUpload(partBuilder -> partBuilder.parts(partETags)));
439+
440+
S3Client defaultClient = S3Client.create();
441+
ResponseBytes<GetObjectResponse> directInstGetResponse = defaultClient.getObjectAsBytes(builder -> builder
442+
.bucket(BUCKET)
443+
.key(object_key + ".instruction")
444+
.build());
445+
assertTrue(directInstGetResponse.response().metadata().containsKey("x-amz-crypto-instr-file"));
446+
447+
ResponseInputStream<GetObjectResponse> getResponse = v3Client.getObject(builder -> builder
448+
.bucket(BUCKET)
449+
.key(object_key));
450+
451+
assertTrue(IOUtils.contentEquals(objectStreamForResult, getResponse));
452+
453+
deleteObject(BUCKET, object_key, v3Client);
454+
v3Client.close();
455+
}
456+
293457
}
458+

0 commit comments

Comments
 (0)