Skip to content

Commit e351f8d

Browse files
ziyanli-amazonziyanli-amazon
authored
Merge pull request #120 from evangilo/issue-7-v1
[Backport] add support to set a prefix for the S3 key #7
2 parents b5cde9b + bf8e971 commit e351f8d

File tree

5 files changed

+218
-12
lines changed

5 files changed

+218
-12
lines changed

‎src/main/java/com/amazon/sqs/javamessaging/AmazonSQSExtendedClient.java

+11-4
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import com.amazonaws.services.sqs.model.SendMessageRequest;
5050
import com.amazonaws.services.sqs.model.SendMessageResult;
5151
import com.amazonaws.services.sqs.model.TooManyEntriesInBatchRequestException;
52+
import com.amazonaws.util.StringUtils;
5253
import org.apache.commons.logging.Log;
5354
import org.apache.commons.logging.LogFactory;
5455
import software.amazon.payloadoffloading.PayloadS3Pointer;
@@ -64,6 +65,7 @@
6465
import java.util.Map;
6566
import java.util.Map.Entry;
6667
import java.util.Optional;
68+
import java.util.UUID;
6769

6870

6971
/**
@@ -1263,8 +1265,7 @@ private SendMessageBatchRequestEntry storeMessageInS3(SendMessageBatchRequestEnt
12631265
}
12641266

12651267
// Store the message content in S3.
1266-
String largeMessagePointer = payloadStore.storeOriginalPayload(messageContentStr,
1267-
messageContentSize);
1268+
String largeMessagePointer = storeOriginalPayload(messageContentStr, messageContentSize);
12681269
batchEntry.setMessageBody(largeMessagePointer);
12691270

12701271
return batchEntry;
@@ -1291,11 +1292,17 @@ private SendMessageRequest storeMessageInS3(SendMessageRequest sendMessageReques
12911292
}
12921293

12931294
// Store the message content in S3.
1294-
String largeMessagePointer = payloadStore.storeOriginalPayload(messageContentStr,
1295-
messageContentSize);
1295+
String largeMessagePointer = storeOriginalPayload(messageContentStr, messageContentSize);
12961296
sendMessageRequest.setMessageBody(largeMessagePointer);
12971297

12981298
return sendMessageRequest;
12991299
}
13001300

1301+
private String storeOriginalPayload(String messageContentStr, Long messageContentSize) {
1302+
String s3KeyPrefix = clientConfiguration.getS3KeyPrefix();
1303+
if (StringUtils.isNullOrEmpty(s3KeyPrefix)) {
1304+
return payloadStore.storeOriginalPayload(messageContentStr, messageContentSize);
1305+
}
1306+
return payloadStore.storeOriginalPayload(messageContentStr, messageContentSize, s3KeyPrefix + UUID.randomUUID());
1307+
}
13011308
}

‎src/main/java/com/amazon/sqs/javamessaging/ExtendedClientConfiguration.java

+76
Original file line numberDiff line numberDiff line change
@@ -15,23 +15,36 @@
1515

1616
package com.amazon.sqs.javamessaging;
1717

18+
import com.amazonaws.AmazonClientException;
1819
import com.amazonaws.services.s3.AmazonS3;
1920
import com.amazonaws.services.s3.model.CannedAccessControlList;
2021
import com.amazonaws.services.s3.model.SSEAwsKeyManagementParams;
2122
import com.amazonaws.annotation.NotThreadSafe;
23+
import com.amazonaws.util.StringUtils;
24+
import org.apache.commons.logging.Log;
25+
import org.apache.commons.logging.LogFactory;
2226
import software.amazon.payloadoffloading.PayloadStorageConfiguration;
2327

28+
import java.util.regex.Pattern;
29+
2430

2531
/**
2632
* Amazon SQS extended client configuration options such as Amazon S3 client,
2733
* bucket name, and message size threshold for large-payload messages.
2834
*/
2935
@NotThreadSafe
3036
public class ExtendedClientConfiguration extends PayloadStorageConfiguration {
37+
private static final Log LOG = LogFactory.getLog(ExtendedClientConfiguration.class);
38+
39+
private static final int UUID_LENGTH = 36;
40+
private static final int MAX_S3_KEY_LENGTH = 1024;
41+
private static final int MAX_S3_KEY_PREFIX_LENGTH = MAX_S3_KEY_LENGTH - UUID_LENGTH;
42+
private static final Pattern INVALID_S3_PREFIX_KEY_CHARACTERS_PATTERN = Pattern.compile("[^a-zA-Z0-9./_-]");
3143

3244
private boolean cleanupS3Payload = true;
3345
private boolean useLegacyReservedAttributeName = true;
3446
private boolean ignorePayloadNotFound = false;
47+
private String s3KeyPrefix = "";
3548

3649
public ExtendedClientConfiguration() {
3750
super();
@@ -43,6 +56,7 @@ public ExtendedClientConfiguration(ExtendedClientConfiguration other) {
4356
this.cleanupS3Payload = other.doesCleanupS3Payload();
4457
this.useLegacyReservedAttributeName = other.usesLegacyReservedAttributeName();
4558
this.ignorePayloadNotFound = other.ignoresPayloadNotFound();
59+
this.s3KeyPrefix = other.getS3KeyPrefix();
4660
}
4761

4862
/**
@@ -160,6 +174,68 @@ public boolean ignoresPayloadNotFound() {
160174
return ignorePayloadNotFound;
161175
}
162176

177+
/**
178+
* Sets a string that will be used as prefix of the S3 Key.
179+
*
180+
* @param s3KeyPrefix
181+
* A S3 key prefix value
182+
*/
183+
public void setS3KeyPrefix(String s3KeyPrefix) {
184+
String trimmedPrefix = StringUtils.trim(s3KeyPrefix);
185+
186+
if (trimmedPrefix == null) {
187+
this.s3KeyPrefix = "";
188+
return;
189+
}
190+
191+
if (trimmedPrefix.length() > MAX_S3_KEY_PREFIX_LENGTH) {
192+
String errorMessage = "The S3 key prefix length must not be greater than " + MAX_S3_KEY_PREFIX_LENGTH;
193+
LOG.error(errorMessage);
194+
throw new AmazonClientException(errorMessage);
195+
}
196+
197+
if (trimmedPrefix.startsWith(".") || trimmedPrefix.startsWith("/")) {
198+
String errorMessage = "The S3 key prefix must not starts with '.' or '/'";
199+
LOG.error(errorMessage);
200+
throw new AmazonClientException(errorMessage);
201+
}
202+
203+
if (trimmedPrefix.contains("..")) {
204+
String errorMessage = "The S3 key prefix must not contains the string '..'";
205+
LOG.error(errorMessage);
206+
throw new AmazonClientException(errorMessage);
207+
}
208+
209+
if (INVALID_S3_PREFIX_KEY_CHARACTERS_PATTERN.matcher(trimmedPrefix).find()) {
210+
String errorMessage = "The S3 key prefix contain invalid characters. The allowed characters are: letters, digits, '/', '_', '-', and '.'";
211+
LOG.error(errorMessage);
212+
throw new AmazonClientException(errorMessage);
213+
}
214+
215+
this.s3KeyPrefix = trimmedPrefix;
216+
}
217+
218+
/**
219+
* Sets a string that will be used as prefix of the S3 Key.
220+
*
221+
* @param s3KeyPrefix
222+
* A S3 key prefix value
223+
*
224+
* @return the updated ExtendedClientConfiguration object.
225+
*/
226+
public ExtendedClientConfiguration withS3KeyPrefix(String s3KeyPrefix) {
227+
setS3KeyPrefix(s3KeyPrefix);
228+
return this;
229+
}
230+
231+
/**
232+
* Gets the S3 key prefix
233+
* @return the prefix value which is being used for compose the S3 key.
234+
*/
235+
public String getS3KeyPrefix() {
236+
return this.s3KeyPrefix;
237+
}
238+
163239
@Override
164240
public ExtendedClientConfiguration withAlwaysThroughS3(boolean alwaysThroughS3) {
165241
setAlwaysThroughS3(alwaysThroughS3);

‎src/test/java/com/amazon/sqs/javamessaging/AmazonSQSExtendedClientTest.java

+50-8
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515

1616
package com.amazon.sqs.javamessaging;
1717

18+
import static com.amazon.sqs.javamessaging.StringTestUtil.generateStringWithLength;
19+
1820
import com.amazonaws.AmazonServiceException;
1921
import com.amazonaws.services.s3.AmazonS3;
2022
import com.amazonaws.services.s3.model.CannedAccessControlList;
@@ -35,13 +37,14 @@
3537
import com.amazonaws.services.sqs.model.SendMessageBatchRequestEntry;
3638
import com.amazonaws.services.sqs.model.SendMessageRequest;
3739
import com.amazonaws.util.StringInputStream;
40+
import org.junit.jupiter.api.AfterEach;
3841
import org.junit.jupiter.api.BeforeEach;
3942
import org.junit.jupiter.api.Test;
4043
import org.mockito.ArgumentCaptor;
44+
import org.mockito.MockedStatic;
4145
import software.amazon.payloadoffloading.PayloadS3Pointer;
4246

4347
import java.util.ArrayList;
44-
import java.util.Arrays;
4548
import java.util.List;
4649
import java.util.Map;
4750
import java.util.UUID;
@@ -55,10 +58,12 @@
5558
import static org.junit.jupiter.api.Assertions.fail;
5659
import static org.mockito.ArgumentMatchers.any;
5760
import static org.mockito.ArgumentMatchers.anyString;
61+
import static org.mockito.ArgumentMatchers.argThat;
5862
import static org.mockito.ArgumentMatchers.eq;
5963
import static org.mockito.Mockito.doThrow;
6064
import static org.mockito.Mockito.isA;
6165
import static org.mockito.Mockito.mock;
66+
import static org.mockito.Mockito.mockStatic;
6267
import static org.mockito.Mockito.never;
6368
import static org.mockito.Mockito.spy;
6469
import static org.mockito.Mockito.times;
@@ -76,11 +81,16 @@ public class AmazonSQSExtendedClientTest {
7681
private AmazonSQS extendedSqsWithDefaultKMS;
7782
private AmazonSQS extendedSqsWithGenericReservedAttributeName;
7883
private AmazonSQS extendedSqsWithDeprecatedMethods;
84+
private AmazonSQS extendedSqsWithS3KeyPrefix;
7985
private AmazonSQS mockSqsBackend;
8086
private AmazonS3 mockS3;
87+
private MockedStatic<UUID> uuidMockStatic;
88+
8189
private static final String S3_BUCKET_NAME = "test-bucket-name";
8290
private static final String SQS_QUEUE_URL = "test-queue-url";
8391
private static final String S3_SERVER_SIDE_ENCRYPTION_KMS_KEY_ID = "test-customer-managed-kms-key-id";
92+
private static final String S3_KEY_PREFIX = "test-s3-key-prefix";
93+
private static final String S3_KEY_UUID = "test-s3-key-uuid";
8494

8595
private static final int LESS_THAN_SQS_SIZE_LIMIT = 3;
8696
private static final int SQS_SIZE_LIMIT = 262144;
@@ -91,6 +101,7 @@ public class AmazonSQSExtendedClientTest {
91101

92102
@BeforeEach
93103
public void setupClients() {
104+
uuidMockStatic = mockStatic(UUID.class);
94105
mockS3 = mock(AmazonS3.class);
95106
mockSqsBackend = mock(AmazonSQS.class);
96107
when(mockS3.putObject(isA(PutObjectRequest.class))).thenReturn(null);
@@ -111,11 +122,25 @@ public void setupClients() {
111122

112123
ExtendedClientConfiguration extendedClientConfigurationDeprecated = new ExtendedClientConfiguration().withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME);
113124

125+
ExtendedClientConfiguration extendedClientConfigurationWithS3KeyPrefix = new ExtendedClientConfiguration()
126+
.withPayloadSupportEnabled(mockS3, S3_BUCKET_NAME)
127+
.withS3KeyPrefix(S3_KEY_PREFIX);
128+
129+
UUID uuidMock = mock(UUID.class);
130+
when(uuidMock.toString()).thenReturn(S3_KEY_UUID);
131+
uuidMockStatic.when(UUID::randomUUID).thenReturn(uuidMock);
132+
114133
extendedSqsWithDefaultConfig = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfiguration));
115134
extendedSqsWithCustomKMS = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithCustomKMS));
116135
extendedSqsWithDefaultKMS = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithDefaultKMS));
117136
extendedSqsWithGenericReservedAttributeName = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithGenericReservedAttributeName));
118137
extendedSqsWithDeprecatedMethods = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationDeprecated));
138+
extendedSqsWithS3KeyPrefix = spy(new AmazonSQSExtendedClient(mockSqsBackend, extendedClientConfigurationWithS3KeyPrefix));
139+
}
140+
141+
@AfterEach
142+
public void tearDown() {
143+
uuidMockStatic.close();
119144
}
120145

121146
@Test
@@ -571,6 +596,30 @@ public void testWhenSendMessageWIthCannedAccessControlListDefined() {
571596
assertEquals(expected, captor.getValue().getCannedAcl());
572597
}
573598

599+
@Test
600+
public void testWhenSendLargeMessageWithS3PrefixKeyDefined() {
601+
String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
602+
603+
SendMessageRequest messageRequest = new SendMessageRequest(SQS_QUEUE_URL, messageBody);
604+
605+
extendedSqsWithS3KeyPrefix.sendMessage(messageRequest);
606+
607+
verify(mockS3, times(1)).putObject(
608+
argThat((PutObjectRequest obj) -> obj.getKey().equals(S3_KEY_PREFIX + S3_KEY_UUID)));
609+
}
610+
611+
@Test
612+
public void testWhenSendLargeMessageWithUndefinedS3PrefixKey() {
613+
String messageBody = generateStringWithLength(MORE_THAN_SQS_SIZE_LIMIT);
614+
615+
SendMessageRequest messageRequest = new SendMessageRequest(SQS_QUEUE_URL, messageBody);
616+
617+
extendedSqsWithDefaultConfig.sendMessage(messageRequest);
618+
619+
verify(mockS3, times(1)).putObject(
620+
argThat((PutObjectRequest obj) -> obj.getKey().equals(S3_KEY_UUID)));
621+
}
622+
574623
private void testReceiveMessage_when_MessageIsLarge(String reservedAttributeName) throws Exception {
575624
Message message = new Message().addMessageAttributesEntry(reservedAttributeName, mock(MessageAttributeValue.class));
576625
String pointer = new PayloadS3Pointer(S3_BUCKET_NAME, "S3Key").toJson();
@@ -607,11 +656,4 @@ private String getLargeReceiptHandle(String s3Key, String originalReceiptHandle)
607656
private String getSampleLargeReceiptHandle() {
608657
return getLargeReceiptHandle(UUID.randomUUID().toString(), UUID.randomUUID().toString());
609658
}
610-
611-
private String generateStringWithLength(int messageLength) {
612-
char[] charArray = new char[messageLength];
613-
Arrays.fill(charArray, 'x');
614-
return new String(charArray);
615-
}
616-
617659
}

‎src/test/java/com/amazon/sqs/javamessaging/ExtendedClientConfigurationTest.java

+70
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,22 @@
1515

1616
package com.amazon.sqs.javamessaging;
1717

18+
import static com.amazon.sqs.javamessaging.StringTestUtil.generateStringWithLength;
19+
20+
import com.amazonaws.AmazonClientException;
1821
import com.amazonaws.services.s3.AmazonS3;
1922
import com.amazonaws.services.s3.model.PutObjectRequest;
2023
import com.amazonaws.services.s3.model.SSEAwsKeyManagementParams;
2124
import org.junit.jupiter.api.Test;
25+
import org.junit.jupiter.params.ParameterizedTest;
26+
import org.junit.jupiter.params.provider.ValueSource;
2227

2328
import static org.junit.jupiter.api.Assertions.assertEquals;
2429
import static org.junit.jupiter.api.Assertions.assertFalse;
2530
import static org.junit.jupiter.api.Assertions.assertNotNull;
2631
import static org.junit.jupiter.api.Assertions.assertNotSame;
2732
import static org.junit.jupiter.api.Assertions.assertNull;
33+
import static org.junit.jupiter.api.Assertions.assertThrows;
2834
import static org.junit.jupiter.api.Assertions.assertTrue;
2935
import static org.mockito.ArgumentMatchers.isA;
3036
import static org.mockito.Mockito.mock;
@@ -173,4 +179,68 @@ public void testMessageSizeThreshold() {
173179
assertEquals(messageLength, extendedClientConfiguration.getPayloadSizeThreshold());
174180

175181
}
182+
183+
@ParameterizedTest
184+
@ValueSource(strings = {
185+
"test-s3-key-prefix",
186+
"TEST-S3-KEY-PREFIX",
187+
"test.s3.key.prefix",
188+
"test_s3_key_prefix",
189+
"test/s3/key/prefix/"
190+
})
191+
public void testS3keyPrefix(String s3KeyPrefix) {
192+
ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration();
193+
194+
extendedClientConfiguration.withS3KeyPrefix(s3KeyPrefix);
195+
196+
assertEquals(s3KeyPrefix, extendedClientConfiguration.getS3KeyPrefix());
197+
}
198+
199+
@Test
200+
public void testTrimS3keyPrefix() {
201+
String s3KeyPrefix = "test-s3-key-prefix";
202+
ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration();
203+
204+
extendedClientConfiguration.withS3KeyPrefix(String.format(" %s ", s3KeyPrefix));
205+
206+
assertEquals(s3KeyPrefix, extendedClientConfiguration.getS3KeyPrefix());
207+
}
208+
209+
@Test
210+
public void testSetS3keyPrefixWithNullValue() {
211+
String s3KeyPrefix = "test-s3-key-prefix";
212+
ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration();
213+
214+
extendedClientConfiguration.withS3KeyPrefix(null);
215+
216+
assertEquals("", extendedClientConfiguration.getS3KeyPrefix());
217+
}
218+
219+
@ParameterizedTest
220+
@ValueSource(strings = {
221+
".test-s3-key-prefix",
222+
"./test-s3-key-prefix",
223+
"../test-s3-key-prefix",
224+
"/test-s3-key-prefix",
225+
"test..s3..key..prefix",
226+
"test-s3-key-prefix@",
227+
"test s3 key prefix"
228+
})
229+
public void testS3KeyPrefixWithInvalidCharacters(String s3KeyPrefix) {
230+
ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration();
231+
232+
assertThrows(AmazonClientException.class, () -> extendedClientConfiguration.withS3KeyPrefix(s3KeyPrefix));
233+
}
234+
235+
@Test
236+
public void testS3keyPrefixWithALargeString() {
237+
int maxS3KeyLength = 1024;
238+
int uuidLength = 36;
239+
int maxS3KeyPrefixLength = maxS3KeyLength - uuidLength;
240+
String s3KeyPrefix = generateStringWithLength(maxS3KeyPrefixLength + 1);
241+
242+
ExtendedClientConfiguration extendedClientConfiguration = new ExtendedClientConfiguration();
243+
244+
assertThrows(AmazonClientException.class, () -> extendedClientConfiguration.withS3KeyPrefix(s3KeyPrefix));
245+
}
176246
}

‎src/test/java/com/amazon/sqs/javamessaging/StringTestUtil.java

+11
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package com.amazon.sqs.javamessaging;
2+
3+
import java.util.Arrays;
4+
5+
public class StringTestUtil {
6+
public static String generateStringWithLength(int messageLength) {
7+
char[] charArray = new char[messageLength];
8+
Arrays.fill(charArray, 'x');
9+
return new String(charArray);
10+
}
11+
}

0 commit comments

Comments
 (0)