Skip to content

Commit c60ad59

Browse files
WesleyRosenblumWesleyRosenblum
authored
Making client suppliers composable (#163)
* Make client suppliers composable * Refactor to move suppliers to StandardAwsKmsClientSuppliers class * Using ConcurrentHashMap for the client cache to be thread safe
1 parent 2055728 commit c60ad59

File tree

7 files changed

+316
-246
lines changed

7 files changed

+316
-246
lines changed

‎src/examples/java/com/amazonaws/crypto/examples/datakeycaching/MultiRegionRecordPusherExample.java

+5-7
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313

1414
package com.amazonaws.crypto.examples.datakeycaching;
1515

16-
import com.amazonaws.ClientConfiguration;
1716
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
1817
import com.amazonaws.encryptionsdk.AwsCrypto;
1918
import com.amazonaws.encryptionsdk.AwsCryptoResult;
@@ -22,8 +21,8 @@
2221
import com.amazonaws.encryptionsdk.caching.LocalCryptoMaterialsCache;
2322
import com.amazonaws.encryptionsdk.keyrings.Keyring;
2423
import com.amazonaws.encryptionsdk.keyrings.StandardKeyrings;
25-
import com.amazonaws.encryptionsdk.kms.AwsKmsClientSupplier;
2624
import com.amazonaws.encryptionsdk.kms.AwsKmsCmkId;
25+
import com.amazonaws.encryptionsdk.kms.StandardAwsKmsClientSuppliers;
2726
import com.amazonaws.regions.Region;
2827
import com.amazonaws.services.kinesis.AmazonKinesis;
2928
import com.amazonaws.services.kinesis.AmazonKinesisClientBuilder;
@@ -32,7 +31,6 @@
3231
import java.nio.ByteBuffer;
3332
import java.util.ArrayList;
3433
import java.util.Collections;
35-
import java.util.HashMap;
3634
import java.util.List;
3735
import java.util.Map;
3836
import java.util.UUID;
@@ -73,10 +71,10 @@ public MultiRegionRecordPusherExample(final Region[] regions, final String kmsAl
7371
.build());
7472

7573
keyrings.add(StandardKeyrings.awsKmsBuilder()
76-
.awsKmsClientSupplier(AwsKmsClientSupplier.builder()
77-
.credentialsProvider(credentialsProvider)
78-
.allowedRegions(Collections.singleton(region.getName()))
79-
.build())
74+
.awsKmsClientSupplier(StandardAwsKmsClientSuppliers
75+
.allowRegionsBuilder(Collections.singleton(region.getName()))
76+
.baseClientSupplier(StandardAwsKmsClientSuppliers.defaultBuilder()
77+
.credentialsProvider(credentialsProvider).build()).build())
8078
.generatorKeyId(AwsKmsCmkId.fromString(kmsAliasName)).build());
8179
}
8280

‎src/main/java/com/amazonaws/encryptionsdk/keyrings/AwsKmsKeyringBuilder.java

+2-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import com.amazonaws.encryptionsdk.kms.AwsKmsClientSupplier;
1717
import com.amazonaws.encryptionsdk.kms.AwsKmsCmkId;
1818
import com.amazonaws.encryptionsdk.kms.DataKeyEncryptionDao;
19+
import com.amazonaws.encryptionsdk.kms.StandardAwsKmsClientSuppliers;
1920

2021
import java.util.List;
2122

@@ -111,7 +112,7 @@ public AwsKmsKeyringBuilder generatorKeyId(AwsKmsCmkId generatorKeyId) {
111112
*/
112113
public Keyring build() {
113114
if (awsKmsClientSupplier == null) {
114-
awsKmsClientSupplier = AwsKmsClientSupplier.builder().build();
115+
awsKmsClientSupplier = StandardAwsKmsClientSuppliers.defaultBuilder().build();
115116
}
116117

117118
return new AwsKmsKeyring(DataKeyEncryptionDao.awsKms(awsKmsClientSupplier, grantTokens),

‎src/main/java/com/amazonaws/encryptionsdk/keyrings/StandardKeyrings.java

+4-5
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@
1313

1414
package com.amazonaws.encryptionsdk.keyrings;
1515

16-
import com.amazonaws.encryptionsdk.kms.AwsKmsClientSupplier;
1716
import com.amazonaws.encryptionsdk.kms.AwsKmsCmkId;
17+
import com.amazonaws.encryptionsdk.kms.StandardAwsKmsClientSuppliers;
1818

1919
import java.util.Arrays;
2020
import java.util.List;
@@ -80,15 +80,14 @@ public static AwsKmsKeyringBuilder awsKmsBuilder() {
8080
* AWS KMS Discovery keyrings do not specify any CMKs to decrypt with, and thus will attempt to decrypt
8181
* using any encrypted data key in an encrypted message. AWS KMS Discovery keyrings do not perform encryption.
8282
* <p></p>
83-
* To create an AWS KMS Regional Discovery Keyring, construct an {@link AwsKmsClientSupplier} using
84-
* {@link AwsKmsClientSupplier#builder()} to specify which regions to include/exclude.
83+
* To create an AWS KMS Regional Discovery Keyring, use {@link StandardAwsKmsClientSuppliers#allowRegionsBuilder} or
84+
* {@link StandardAwsKmsClientSuppliers#denyRegionsBuilder} to specify which regions to include/exclude.
8585
* <p></p>
8686
* For example, to include only CMKs in the us-east-1 region:
8787
* <pre>
8888
* StandardKeyrings.awsKmsDiscovery()
8989
* .awsKmsClientSupplier(
90-
* AwsKmsClientSupplier.builder()
91-
* .allowedRegions(Collections.singleton("us-east-1")).build())
90+
* StandardAwsKmsClientSuppliers.allowRegionsBuilder(Collections.singleton("us-east-1")).build()
9291
* .build();
9392
* </pre>
9493
*

‎src/main/java/com/amazonaws/encryptionsdk/kms/AwsKmsClientSupplier.java

-174
Original file line numberDiff line numberDiff line change
@@ -13,26 +13,13 @@
1313

1414
package com.amazonaws.encryptionsdk.kms;
1515

16-
import com.amazonaws.ClientConfiguration;
1716
import com.amazonaws.arn.Arn;
18-
import com.amazonaws.auth.AWSCredentialsProvider;
1917
import com.amazonaws.encryptionsdk.exception.UnsupportedRegionException;
2018
import com.amazonaws.services.kms.AWSKMS;
21-
import com.amazonaws.services.kms.AWSKMSClientBuilder;
22-
import com.amazonaws.services.kms.model.AWSKMSException;
2319

2420
import javax.annotation.Nullable;
25-
import java.lang.reflect.InvocationTargetException;
26-
import java.lang.reflect.Proxy;
27-
import java.util.Collections;
28-
import java.util.HashMap;
29-
import java.util.HashSet;
30-
import java.util.Map;
31-
import java.util.Set;
3221

3322
import static java.util.Objects.requireNonNull;
34-
import static org.apache.commons.lang3.Validate.isTrue;
35-
import static org.apache.commons.lang3.Validate.notEmpty;
3623

3724
/**
3825
* Represents a function that accepts an AWS region and returns an {@code AWSKMS} client for that region. The
@@ -51,15 +38,6 @@ public interface AwsKmsClientSupplier {
5138
*/
5239
AWSKMS getClient(@Nullable String regionId) throws UnsupportedRegionException;
5340

54-
/**
55-
* Gets a Builder for constructing an AwsKmsClientSupplier
56-
*
57-
* @return The builder
58-
*/
59-
static Builder builder() {
60-
return new Builder(AWSKMSClientBuilder.standard());
61-
}
62-
6341
/**
6442
* Parses region from the given key id (if possible) and passes that region to the
6543
* given clientSupplier to produce an {@code AWSKMS} client.
@@ -78,156 +56,4 @@ static AWSKMS getClientByKeyId(AwsKmsCmkId keyId, AwsKmsClientSupplier clientSup
7856

7957
return clientSupplier.getClient(null);
8058
}
81-
82-
/**
83-
* Builder to construct an AwsKmsClientSupplier given various
84-
* optional settings.
85-
*/
86-
class Builder {
87-
88-
private AWSCredentialsProvider credentialsProvider;
89-
private ClientConfiguration clientConfiguration;
90-
private Set<String> allowedRegions = Collections.emptySet();
91-
private Set<String> excludedRegions = Collections.emptySet();
92-
private boolean clientCachingEnabled = true;
93-
private final Map<String, AWSKMS> clientsCache = new HashMap<>();
94-
private static final Set<String> AWSKMS_METHODS = new HashSet<>();
95-
private AWSKMSClientBuilder awsKmsClientBuilder;
96-
97-
static {
98-
AWSKMS_METHODS.add("generateDataKey");
99-
AWSKMS_METHODS.add("encrypt");
100-
AWSKMS_METHODS.add("decrypt");
101-
}
102-
103-
Builder(AWSKMSClientBuilder awsKmsClientBuilder) {
104-
this.awsKmsClientBuilder = awsKmsClientBuilder;
105-
}
106-
107-
public AwsKmsClientSupplier build() {
108-
isTrue(allowedRegions.isEmpty() || excludedRegions.isEmpty(),
109-
"Either allowed regions or excluded regions may be set, not both.");
110-
111-
return regionId -> {
112-
if (!allowedRegions.isEmpty() && !allowedRegions.contains(regionId)) {
113-
throw new UnsupportedRegionException(String.format("Region %s is not in the list of allowed regions %s",
114-
regionId, allowedRegions));
115-
}
116-
117-
if (excludedRegions.contains(regionId)) {
118-
throw new UnsupportedRegionException(String.format("Region %s is in the list of excluded regions %s",
119-
regionId, excludedRegions));
120-
}
121-
122-
if (clientsCache.containsKey(regionId)) {
123-
return clientsCache.get(regionId);
124-
}
125-
126-
if (credentialsProvider != null) {
127-
awsKmsClientBuilder = awsKmsClientBuilder.withCredentials(credentialsProvider);
128-
}
129-
130-
if (clientConfiguration != null) {
131-
awsKmsClientBuilder = awsKmsClientBuilder.withClientConfiguration(clientConfiguration);
132-
}
133-
134-
if (regionId != null) {
135-
awsKmsClientBuilder = awsKmsClientBuilder.withRegion(regionId);
136-
}
137-
138-
AWSKMS client = awsKmsClientBuilder.build();
139-
140-
if (clientCachingEnabled) {
141-
client = newCachingProxy(client, regionId);
142-
}
143-
144-
return client;
145-
};
146-
}
147-
148-
/**
149-
* Sets the AWSCredentialsProvider used by the client.
150-
*
151-
* @param credentialsProvider New AWSCredentialsProvider to use.
152-
*/
153-
public Builder credentialsProvider(AWSCredentialsProvider credentialsProvider) {
154-
this.credentialsProvider = credentialsProvider;
155-
return this;
156-
}
157-
158-
/**
159-
* Sets the ClientConfiguration to be used by the client.
160-
*
161-
* @param clientConfiguration Custom configuration to use.
162-
*/
163-
public Builder clientConfiguration(ClientConfiguration clientConfiguration) {
164-
this.clientConfiguration = clientConfiguration;
165-
return this;
166-
}
167-
168-
/**
169-
* Sets the AWS regions that the client supplier is permitted to use.
170-
*
171-
* @param regions The set of regions.
172-
*/
173-
public Builder allowedRegions(Set<String> regions) {
174-
notEmpty(regions, "At least one region is required");
175-
this.allowedRegions = Collections.unmodifiableSet(new HashSet<>(regions));
176-
return this;
177-
}
178-
179-
/**
180-
* Sets the AWS regions that the client supplier is not permitted to use.
181-
*
182-
* @param regions The set of regions.
183-
*/
184-
public Builder excludedRegions(Set<String> regions) {
185-
requireNonNull(regions, "regions is required");
186-
this.excludedRegions = Collections.unmodifiableSet(new HashSet<>(regions));
187-
return this;
188-
}
189-
190-
/**
191-
* When set to false, disables the AWSKMS client for each region from being cached and reused.
192-
* By default, client caching is enabled.
193-
*
194-
* @param enabled Whether or not caching is enabled.
195-
*/
196-
public Builder clientCaching(boolean enabled) {
197-
this.clientCachingEnabled = enabled;
198-
return this;
199-
}
200-
201-
/**
202-
* Creates a proxy for the AWSKMS client that will populate the client into the client cache
203-
* after an AWS KMS method successfully completes or an AWS KMS exception occurs. This is to prevent a
204-
* a malicious user from causing a local resource DOS by sending ciphertext with a large number
205-
* of spurious regions, thereby filling the cache with regions and exhausting resources.
206-
*
207-
* @param client The client to proxy
208-
* @param regionId The region the client is associated with
209-
* @return The proxy
210-
*/
211-
private AWSKMS newCachingProxy(AWSKMS client, String regionId) {
212-
return (AWSKMS) Proxy.newProxyInstance(
213-
AWSKMS.class.getClassLoader(),
214-
new Class[]{AWSKMS.class},
215-
(proxy, method, methodArgs) -> {
216-
try {
217-
final Object result = method.invoke(client, methodArgs);
218-
if (AWSKMS_METHODS.contains(method.getName())) {
219-
clientsCache.put(regionId, client);
220-
}
221-
return result;
222-
} catch (InvocationTargetException e) {
223-
if (e.getTargetException() instanceof AWSKMSException &&
224-
AWSKMS_METHODS.contains(method.getName())) {
225-
clientsCache.put(regionId, client);
226-
}
227-
228-
throw e.getTargetException();
229-
}
230-
});
231-
}
232-
}
23359
}

0 commit comments

Comments
 (0)