Skip to content

Commit a789a13

Browse files
robin-awsrobin-aws
authored
Merge pull request #17 from awslabs/robin-aws/cross-account
Whitelisting most queue attributes for temporary queues
2 parents 4ac18df + d0ff760 commit a789a13

5 files changed

+283
-6
lines changed

‎src/main/java/com/amazonaws/services/sqs/AmazonSQSTemporaryQueuesClient.java

+31-2
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,18 @@
11
package com.amazonaws.services.sqs;
22

3+
import java.util.Arrays;
34
import java.util.HashMap;
5+
import java.util.HashSet;
46
import java.util.Map;
7+
import java.util.Set;
58
import java.util.UUID;
69
import java.util.concurrent.ConcurrentHashMap;
710
import java.util.concurrent.ConcurrentMap;
811
import java.util.concurrent.TimeUnit;
912

1013
import com.amazonaws.services.sqs.model.CreateQueueRequest;
1114
import com.amazonaws.services.sqs.model.CreateQueueResult;
15+
import com.amazonaws.services.sqs.model.QueueAttributeName;
1216
import com.amazonaws.services.sqs.util.AbstractAmazonSQSClientWrapper;
1317
import com.amazonaws.services.sqs.util.SQSQueueUtils;
1418

@@ -30,7 +34,20 @@ class AmazonSQSTemporaryQueuesClient extends AbstractAmazonSQSClientWrapper {
3034

3135
// TODO-RS: Expose configuration
3236
private final static String QUEUE_RETENTION_PERIOD_SECONDS = Long.toString(TimeUnit.MINUTES.toSeconds(5));
33-
37+
38+
// We don't necessary support all queue attributes - some will behave differently on a virtual queue
39+
// In particular, a virtual FIFO queue will deduplicate at the scope of its host queue!
40+
private final static Set<String> SUPPORTED_QUEUE_ATTRIBUTES = new HashSet<>(Arrays.asList(
41+
QueueAttributeName.DelaySeconds.name(),
42+
QueueAttributeName.MaximumMessageSize.name(),
43+
QueueAttributeName.MessageRetentionPeriod.name(),
44+
QueueAttributeName.Policy.name(),
45+
QueueAttributeName.ReceiveMessageWaitTimeSeconds.name(),
46+
QueueAttributeName.RedrivePolicy.name(),
47+
QueueAttributeName.VisibilityTimeout.name(),
48+
QueueAttributeName.KmsMasterKeyId.name(),
49+
QueueAttributeName.KmsDataKeyReusePeriodSeconds.name()));
50+
3451
// These clients are owned by this one, and need to be shutdown when this client is.
3552
private final AmazonSQSIdleQueueDeletingClient deleter;
3653
private final AmazonSQS virtualizer;
@@ -85,6 +102,14 @@ AmazonSQSRequester getRequester() {
85102

86103
@Override
87104
public CreateQueueResult createQueue(CreateQueueRequest request) {
105+
// Check for unsupported queue attributes first
106+
Set<String> unsupportedQueueAttributes = new HashSet<>(request.getAttributes().keySet());
107+
unsupportedQueueAttributes.removeAll(SUPPORTED_QUEUE_ATTRIBUTES);
108+
if (!unsupportedQueueAttributes.isEmpty()) {
109+
throw new IllegalArgumentException("Cannot create a temporary queue with the following attributes: "
110+
+ String.join(", ", unsupportedQueueAttributes));
111+
}
112+
88113
Map<String, String> extraQueueAttributes = new HashMap<>();
89114
// Add the retention period to both the host queue and each virtual queue
90115
extraQueueAttributes.put(AmazonSQSIdleQueueDeletingClient.IDLE_QUEUE_RETENTION_PERIOD, QUEUE_RETENTION_PERIOD_SECONDS);
@@ -95,7 +120,11 @@ public CreateQueueResult createQueue(CreateQueueRequest request) {
95120
});
96121

97122
extraQueueAttributes.put(AmazonSQSVirtualQueuesClient.VIRTUAL_QUEUE_HOST_QUEUE_ATTRIBUTE, hostQueueUrl);
98-
CreateQueueRequest createVirtualQueueRequest = SQSQueueUtils.copyWithExtraAttributes(request, extraQueueAttributes);
123+
// The host queue takes care of all the other queue attributes, so don't specify them when creating the virtual
124+
// queue or else the client may think we're trying to set them independently!
125+
CreateQueueRequest createVirtualQueueRequest = new CreateQueueRequest()
126+
.withQueueName(request.getQueueName())
127+
.withAttributes(extraQueueAttributes);
99128
return amazonSqsToBeExtended.createQueue(createVirtualQueueRequest);
100129
}
101130

‎src/test/java/com/amazonaws/services/sqs/AmazonSQSResponsesClientCrossAccountIT.java

+69
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package com.amazonaws.services.sqs;
2+
3+
import com.amazonaws.services.sqs.model.Message;
4+
import com.amazonaws.services.sqs.model.QueueAttributeName;
5+
import com.amazonaws.services.sqs.model.SendMessageRequest;
6+
import com.amazonaws.services.sqs.util.IntegrationTest;
7+
import com.amazonaws.services.sqs.util.SQSMessageConsumer;
8+
import com.amazonaws.services.sqs.util.SQSMessageConsumerBuilder;
9+
import org.junit.After;
10+
import org.junit.Before;
11+
import org.junit.Test;
12+
13+
import java.util.Collections;
14+
import java.util.UUID;
15+
import java.util.concurrent.TimeUnit;
16+
17+
import static org.junit.Assert.assertEquals;
18+
19+
public class AmazonSQSResponsesClientCrossAccountIT extends IntegrationTest {
20+
private static AmazonSQSRequester sqsRequester;
21+
private static AmazonSQSResponder sqsResponder;
22+
private static String requestQueueUrl;
23+
24+
@Override
25+
protected String testSuiteName() {
26+
return "SQSXAccountResponsesClientIT";
27+
}
28+
29+
@Before
30+
public void setup() {
31+
String policyString = allowSendMessagePolicy().toJson();
32+
sqsRequester = new AmazonSQSRequesterClient(sqs, queueNamePrefix,
33+
Collections.singletonMap(QueueAttributeName.Policy.toString(), policyString),
34+
exceptionHandler);
35+
// Use the second account for the responder
36+
sqsResponder = new AmazonSQSResponderClient(getBuddyPrincipalClient());
37+
requestQueueUrl = sqs.createQueue("RequestQueue-" + UUID.randomUUID().toString()).getQueueUrl();
38+
}
39+
40+
@After
41+
public void teardown() {
42+
sqs.deleteQueue(requestQueueUrl);
43+
sqsResponder.shutdown();
44+
sqsRequester.shutdown();
45+
}
46+
47+
@Test
48+
public void test() throws Exception {
49+
SQSMessageConsumer consumer = SQSMessageConsumerBuilder.standard()
50+
.withAmazonSQS(sqs)
51+
.withQueueUrl(requestQueueUrl)
52+
.withConsumer(message -> {
53+
sqsResponder.sendResponseMessage(MessageContent.fromMessage(message),
54+
new MessageContent("Right back atcha buddy!"));
55+
})
56+
.build();
57+
consumer.start();
58+
try {
59+
SendMessageRequest request = new SendMessageRequest()
60+
.withMessageBody("Hi there!")
61+
.withQueueUrl(requestQueueUrl);
62+
Message replyMessage = sqsRequester.sendMessageAndGetResponse(request, 5, TimeUnit.SECONDS);
63+
64+
assertEquals("Right back atcha buddy!", replyMessage.getBody());
65+
} finally {
66+
consumer.terminate();
67+
}
68+
}
69+
}

‎src/test/java/com/amazonaws/services/sqs/AmazonSQSTemporaryQueuesClientCrossAccountIT.java

+102
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
package com.amazonaws.services.sqs;
2+
3+
import com.amazonaws.services.sqs.model.AmazonSQSException;
4+
import com.amazonaws.services.sqs.model.CreateQueueRequest;
5+
import com.amazonaws.services.sqs.model.QueueAttributeName;
6+
import com.amazonaws.services.sqs.util.AbstractAmazonSQSClientWrapper;
7+
import com.amazonaws.services.sqs.util.IntegrationTest;
8+
import org.junit.After;
9+
import org.junit.Assert;
10+
import org.junit.Before;
11+
import org.junit.Test;
12+
import org.junit.runner.RunWith;
13+
import org.junit.runners.Parameterized;
14+
import org.junit.runners.Parameterized.Parameters;
15+
16+
import java.util.Arrays;
17+
import java.util.Collections;
18+
import java.util.concurrent.TimeUnit;
19+
20+
import static org.junit.Assert.assertEquals;
21+
22+
@RunWith(Parameterized.class)
23+
public class AmazonSQSTemporaryQueuesClientCrossAccountIT extends IntegrationTest {
24+
25+
@Parameters(name= "With temporary queues = {0}")
26+
public static Iterable<Object[]> data() {
27+
return Arrays.asList(new Object[][] { { false }, { true } } );
28+
}
29+
30+
// Parameterized to emphasize that the same client code for physical SQS queues
31+
// should work for temporary queues as well, even thought they are virtual.
32+
private final boolean withTemporaryQueues;
33+
private AmazonSQS client;
34+
private AmazonSQS otherAccountClient;
35+
36+
public AmazonSQSTemporaryQueuesClientCrossAccountIT(boolean withTemporaryQueues) {
37+
this.withTemporaryQueues = withTemporaryQueues;
38+
}
39+
40+
@Before
41+
public void setup() {
42+
client = makeTemporaryQueueClient(sqs);
43+
otherAccountClient = makeTemporaryQueueClient(getBuddyPrincipalClient());
44+
}
45+
46+
@After
47+
public void teardown() {
48+
client.shutdown();
49+
}
50+
51+
private AmazonSQS makeTemporaryQueueClient(AmazonSQS sqs) {
52+
if (withTemporaryQueues) {
53+
AmazonSQSRequesterClientBuilder requesterBuilder =
54+
AmazonSQSRequesterClientBuilder.standard()
55+
.withAmazonSQS(sqs)
56+
.withInternalQueuePrefix(queueNamePrefix)
57+
.withIdleQueueSweepingPeriod(0, TimeUnit.SECONDS);
58+
return AmazonSQSTemporaryQueuesClient.make(requesterBuilder);
59+
} else {
60+
// Use a wrapper just to avoid shutting down the client from
61+
// the base class too early.
62+
return new AbstractAmazonSQSClientWrapper(sqs);
63+
}
64+
}
65+
66+
@Override
67+
protected String testSuiteName() {
68+
return "SQSXAccountTempQueueIT";
69+
}
70+
71+
@Test
72+
public void accessDenied() {
73+
// Assert that a different principal is not permitted to
74+
// send to virtual queues
75+
String virtualQueueUrl = client.createQueue(queueNamePrefix + "TestQueueWithoutAccess").getQueueUrl();
76+
try {
77+
otherAccountClient.sendMessage(virtualQueueUrl, "Haxxors!!");
78+
Assert.fail("Should not have been able to send a message");
79+
} catch (AmazonSQSException e) {
80+
// Access Denied
81+
assertEquals(403, e.getStatusCode());
82+
} finally {
83+
client.deleteQueue(virtualQueueUrl);
84+
}
85+
}
86+
87+
@Test
88+
public void withAccess() {
89+
String policyString = allowSendMessagePolicy().toJson();
90+
CreateQueueRequest createQueueRequest = new CreateQueueRequest()
91+
.withQueueName(queueNamePrefix + "TestQueueWithAccess")
92+
.withAttributes(Collections.singletonMap(QueueAttributeName.Policy.toString(), policyString));
93+
94+
String queueUrl = client.createQueue(createQueueRequest).getQueueUrl();
95+
try {
96+
otherAccountClient.sendMessage(queueUrl, "Hi there!");
97+
} finally {
98+
client.deleteQueue(queueUrl);
99+
}
100+
}
101+
102+
}

‎src/test/java/com/amazonaws/services/sqs/AmazonSQSTemporaryQueuesClientIT.java

+18-2
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
import java.util.List;
99
import java.util.Map;
1010

11+
import com.amazonaws.services.sqs.model.CreateQueueRequest;
12+
import com.amazonaws.services.sqs.model.QueueAttributeName;
1113
import org.junit.After;
1214
import org.junit.Assert;
1315
import org.junit.Before;
@@ -31,13 +33,15 @@ public void setup() {
3133

3234
@After
3335
public void teardown() {
34-
client.deleteQueue(queueUrl);
36+
if (queueUrl != null) {
37+
client.deleteQueue(queueUrl);
38+
}
3539
client.shutdown();
3640
}
3741

3842
@Test
3943
public void createQueueAddsAttributes() {
40-
queueUrl = client.createQueue("TestQueue").getQueueUrl();
44+
queueUrl = client.createQueue(queueNamePrefix + "TestQueue").getQueueUrl();
4145
Map<String, String> attributes = client.getQueueAttributes(queueUrl, Collections.singletonList("All")).getAttributes();
4246
String hostQueueUrl = attributes.get(VIRTUAL_QUEUE_HOST_QUEUE_ATTRIBUTE);
4347
assertNotNull(hostQueueUrl);
@@ -46,4 +50,16 @@ public void createQueueAddsAttributes() {
4650
Map<String, String> hostQueueAttributes = client.getQueueAttributes(queueUrl, Collections.singletonList("All")).getAttributes();
4751
Assert.assertEquals("300", hostQueueAttributes.get(AmazonSQSIdleQueueDeletingClient.IDLE_QUEUE_RETENTION_PERIOD));
4852
}
53+
54+
@Test
55+
public void createQueueWithUnsupportedAttributes() {
56+
try {
57+
client.createQueue(new CreateQueueRequest()
58+
.withQueueName(queueNamePrefix + "InvalidQueue")
59+
.withAttributes(Collections.singletonMap(QueueAttributeName.FifoQueue.name(), "true")));
60+
Assert.fail("Shouldn't be able to create a FIFO temporary queue");
61+
} catch (IllegalArgumentException e) {
62+
Assert.assertEquals("Cannot create a temporary queue with the following attributes: FifoQueue", e.getMessage());
63+
}
64+
}
4965
}

‎src/test/java/com/amazonaws/services/sqs/util/IntegrationTest.java

+63-2
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,45 @@
11
package com.amazonaws.services.sqs.util;
22

3+
import java.util.Collections;
34
import java.util.concurrent.ThreadLocalRandom;
45

6+
import com.amazonaws.auth.AWSCredentialsProvider;
7+
import com.amazonaws.auth.policy.Policy;
8+
import com.amazonaws.auth.policy.Principal;
9+
import com.amazonaws.auth.policy.Resource;
10+
import com.amazonaws.auth.policy.Statement;
11+
import com.amazonaws.auth.policy.actions.SQSActions;
12+
import com.amazonaws.auth.profile.ProfileCredentialsProvider;
13+
import com.amazonaws.services.sqs.model.AmazonSQSException;
514
import org.junit.After;
615
import org.junit.Before;
716

817
import com.amazonaws.services.sqs.AmazonSQS;
918
import com.amazonaws.services.sqs.AmazonSQSClientBuilder;
1019
import com.amazonaws.services.sqs.model.QueueDoesNotExistException;
1120

21+
import static org.hamcrest.CoreMatchers.equalTo;
22+
import static org.junit.Assume.assumeNoException;
23+
import static org.junit.Assume.assumeThat;
24+
import static org.junit.Assume.assumeTrue;
25+
1226
/**
1327
* Base class for integration tests
1428
*/
1529
public abstract class IntegrationTest {
1630

1731
protected AmazonSQS sqs;
1832
// UUIDs are too long for this
19-
protected String queueNamePrefix = "__" + getClass().getSimpleName() + "-" + ThreadLocalRandom.current().nextInt(1000000);
33+
protected String queueNamePrefix = "__" + testSuiteName() + "-" + ThreadLocalRandom.current().nextInt(1000000);
2034
protected ExceptionAsserter exceptionHandler = new ExceptionAsserter();
21-
35+
36+
/**
37+
* Customizable to ensure queue names stay under 80 characters
38+
*/
39+
protected String testSuiteName() {
40+
return getClass().getSimpleName();
41+
}
42+
2243
@Before
2344
public void setupSQSClient() {
2445
sqs = AmazonSQSClientBuilder.defaultClient();
@@ -40,4 +61,44 @@ public void teardownSQSClient() {
4061
}
4162
exceptionHandler.assertNothingThrown();
4263
}
64+
65+
protected AmazonSQS getBuddyPrincipalClient() {
66+
AWSCredentialsProvider credentialsProvider = new ProfileCredentialsProvider("buddy");
67+
try {
68+
credentialsProvider.getCredentials();
69+
} catch (Exception e) {
70+
assumeNoException("This test requires a second 'buddy' credential profile.", e);
71+
}
72+
73+
AmazonSQS client = AmazonSQSClientBuilder.standard()
74+
.withRegion("us-west-2")
75+
.withCredentials(credentialsProvider)
76+
.build();
77+
78+
// Assume that the principal is not able to send messages to arbitrary queues
79+
String queueUrl = sqs.createQueue(queueNamePrefix + "TestQueue").getQueueUrl();
80+
try {
81+
client.sendMessage(queueUrl, "Haxxors!!");
82+
assumeTrue("The buddy credentials should not authorize sending to arbitrary queues", false);
83+
} catch (AmazonSQSException e) {
84+
// Access Denied
85+
assumeThat(e.getStatusCode(), equalTo(403));
86+
} finally {
87+
sqs.deleteQueue(queueUrl);
88+
}
89+
90+
return client;
91+
}
92+
93+
protected Policy allowSendMessagePolicy() {
94+
Policy policy = new Policy();
95+
Statement statement = new Statement(Statement.Effect.Allow);
96+
statement.setActions(Collections.singletonList(SQSActions.SendMessage));
97+
// Ideally we would only allow the principal we're testing with, but we
98+
// only have access to the credentials and not necessarily the account number.
99+
statement.setPrincipals(Principal.All);
100+
statement.setResources(Collections.singletonList(new Resource("arn:aws:sqs:*:*:*")));
101+
policy.setStatements(Collections.singletonList(statement));
102+
return policy;
103+
}
43104
}

0 commit comments

Comments
 (0)