Skip to content

Commit 33a8c85

Browse files
lucasmcdonald3lucasmcdonald3
authored
feat(Examples): Example using client supplier (#186)
1 parent c617880 commit 33a8c85

File tree

7 files changed

+393
-2
lines changed

7 files changed

+393
-2
lines changed

‎Examples/runtimes/java/DynamoDbEncryption/build.gradle.kts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,9 +61,11 @@ dependencies {
6161

6262
implementation(platform("software.amazon.awssdk:bom:2.19.1"))
6363
implementation("software.amazon.awssdk:arns")
64+
implementation("software.amazon.awssdk:auth")
6465
implementation("software.amazon.awssdk:dynamodb")
6566
implementation("software.amazon.awssdk:dynamodb-enhanced")
6667
implementation("software.amazon.awssdk:kms")
68+
implementation("software.amazon.awssdk:sts")
6769

6870
implementation("org.bouncycastle:bcprov-jdk18on:1.72")
6971

‎Examples/runtimes/java/DynamoDbEncryption/src/main/java/software/aws/cryptography/examples/clientsupplier/ClientSupplierExample.java

Lines changed: 280 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,280 @@
1+
package software.aws.cryptography.examples.clientsupplier;
2+
3+
import java.util.ArrayList;
4+
import java.util.HashMap;
5+
import java.util.List;
6+
import java.util.Map;
7+
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
8+
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
9+
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
10+
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
11+
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
12+
import software.amazon.awssdk.services.dynamodb.model.PutItemRequest;
13+
import software.amazon.awssdk.services.dynamodb.model.PutItemResponse;
14+
import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.DynamoDbTableEncryptionConfig;
15+
import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.DynamoDbTablesEncryptionConfig;
16+
import software.amazon.cryptography.dbencryptionsdk.structuredencryption.model.CryptoAction;
17+
import software.amazon.cryptography.materialproviders.IKeyring;
18+
import software.amazon.cryptography.materialproviders.MaterialProviders;
19+
import software.amazon.cryptography.materialproviders.model.CreateAwsKmsMrkDiscoveryMultiKeyringInput;
20+
import software.amazon.cryptography.materialproviders.model.CreateAwsKmsMrkMultiKeyringInput;
21+
import software.amazon.cryptography.materialproviders.model.DiscoveryFilter;
22+
import software.amazon.cryptography.materialproviders.model.MaterialProvidersConfig;
23+
import software.aws.cryptography.dbencryptionsdk.dynamodb.DynamoDbEncryptionInterceptor;
24+
25+
/*
26+
This example sets up an MRK multi-keyring and an MRK discovery
27+
multi-keyring using a custom client supplier.
28+
A custom client supplier grants users access to more granular
29+
configuration aspects of their authentication details and KMS
30+
client. In this example, we create a simple custom client supplier
31+
that authenticates with a different IAM role based on the
32+
region of the KMS key.
33+
34+
This example creates a MRK multi-keyring configured with a custom
35+
client supplier using a single MRK and puts an encrypted item to the
36+
table. Then, it creates a MRK discovery multi-keyring to decrypt the item
37+
and retrieves the item from the table.
38+
39+
Running this example requires access to the DDB Table whose name
40+
is provided in CLI arguments.
41+
This table must be configured with the following
42+
primary key configuration:
43+
- Partition key is named "partition_key" with type (S)
44+
- Sort key is named "sort_key" with type (S)
45+
*/
46+
public class ClientSupplierExample {
47+
48+
public static void ClientSupplierPutItemGetItem(String ddbTableName, String keyArn,
49+
List<String> accountIds, List<String> regions) {
50+
// 1. Create a single MRK multi-keyring.
51+
// This can be either a single-region KMS key or an MRK.
52+
// For this example to succeed, the key's region must either
53+
// 1) be in the regions list, or
54+
// 2) the key must be an MRK with a replica defined
55+
// in a region in the regions list, and the client
56+
// must have the correct permissions to access the replica.
57+
final MaterialProviders matProv = MaterialProviders.builder()
58+
.MaterialProvidersConfig(MaterialProvidersConfig.builder().build())
59+
.build();
60+
// Create the multi-keyring using our custom client supplier
61+
// defined in the RegionalRoleClientSupplier class in this directory.
62+
final CreateAwsKmsMrkMultiKeyringInput createAwsKmsMrkMultiKeyringInput =
63+
CreateAwsKmsMrkMultiKeyringInput.builder()
64+
// Note: RegionalRoleClientSupplier will internally use the keyArn's region
65+
// to retrieve the correct IAM role.
66+
.clientSupplier(new RegionalRoleClientSupplier())
67+
.generator(keyArn)
68+
.build();
69+
IKeyring mrkKeyringWithClientSupplier = matProv.CreateAwsKmsMrkMultiKeyring(createAwsKmsMrkMultiKeyringInput);
70+
71+
// 2. Configure which attributes are encrypted and/or signed when writing new items.
72+
// For each attribute that may exist on the items we plan to write to our DynamoDbTable,
73+
// we must explicitly configure how they should be treated during item encryption:
74+
// - ENCRYPT_AND_SIGN: The attribute is encrypted and included in the signature
75+
// - SIGN_ONLY: The attribute is not encrypted, but is still included in the signature
76+
// - DO_NOTHING: The attribute is not encrypted and not included in the signature
77+
final Map<String, CryptoAction> attributeActions = new HashMap<>();
78+
attributeActions.put("partition_key", CryptoAction.SIGN_ONLY); // Our partition attribute must be SIGN_ONLY
79+
attributeActions.put("sort_key", CryptoAction.SIGN_ONLY); // Our sort attribute must be SIGN_ONLY
80+
attributeActions.put("sensitive_data", CryptoAction.ENCRYPT_AND_SIGN);
81+
82+
// 3. Configure which attributes we expect to be included in the signature
83+
// when reading items. There are two options for configuring this:
84+
//
85+
// - (Recommended) Configure `allowedUnauthenticatedAttributesPrefix`:
86+
// When defining your DynamoDb schema and deciding on attribute names,
87+
// choose a distinguishing prefix (such as ":") for all attributes that
88+
// you do not want to include in the signature.
89+
// This has two main benefits:
90+
// - It is easier to reason about the security and authenticity of data within your item
91+
// when all unauthenticated data is easily distinguishable by their attribute name.
92+
// - If you need to add new unauthenticated attributes in the future,
93+
// you can easily make the corresponding update to your `attributeActions`
94+
// and immediately start writing to that new attribute, without
95+
// any other configuration update needed.
96+
// Once you configure this field, it is not safe to update it.
97+
//
98+
// - Configure `allowedUnauthenticatedAttributes`: You may also explicitly list
99+
// a set of attributes that should be considered unauthenticated when encountered
100+
// on read. Be careful if you use this configuration. Do not remove an attribute
101+
// name from this configuration, even if you are no longer writing with that attribute,
102+
// as old items may still include this attribute, and our configuration needs to know
103+
// to continue to exclude this attribute from the signature scope.
104+
// If you add new attribute names to this field, you must first deploy the update to this
105+
// field to all readers in your host fleet before deploying the update to start writing
106+
// with that new attribute.
107+
//
108+
// For this example, we currently authenticate all attributes. To make it easier to
109+
// add unauthenticated attributes in the future, we define a prefix ":" for such attributes.
110+
final String unauthAttrPrefix = ":";
111+
112+
// 4. Create the DynamoDb Encryption configuration for the table we will be writing to.
113+
final Map<String, DynamoDbTableEncryptionConfig> tableConfigs = new HashMap<>();
114+
final DynamoDbTableEncryptionConfig config = DynamoDbTableEncryptionConfig.builder()
115+
.logicalTableName(ddbTableName)
116+
.partitionKeyName("partition_key")
117+
.sortKeyName("sort_key")
118+
.attributeActions(attributeActions)
119+
.keyring(mrkKeyringWithClientSupplier)
120+
.allowedUnauthenticatedAttributePrefix(unauthAttrPrefix)
121+
.build();
122+
tableConfigs.put(ddbTableName, config);
123+
124+
// 5. Create the DynamoDb Encryption Interceptor
125+
DynamoDbEncryptionInterceptor encryptionInterceptor = DynamoDbEncryptionInterceptor.builder()
126+
.config(DynamoDbTablesEncryptionConfig.builder()
127+
.tableEncryptionConfigs(tableConfigs)
128+
.build())
129+
.build();
130+
131+
// 6. Create a new AWS SDK DynamoDb client using the DynamoDb Encryption Interceptor above
132+
final DynamoDbClient ddbClient = DynamoDbClient.builder()
133+
.overrideConfiguration(
134+
ClientOverrideConfiguration.builder()
135+
.addExecutionInterceptor(encryptionInterceptor)
136+
.build())
137+
.build();
138+
139+
// 7. Put an item into our table using the above client.
140+
// Before the item gets sent to DynamoDb, it will be encrypted
141+
// client-side using the MRK multi-keyring.
142+
// The data key protecting this item will be encrypted
143+
// with all the KMS Keys in this keyring, so that it can be
144+
// decrypted with any one of those KMS Keys.
145+
final HashMap<String, AttributeValue> item = new HashMap<>();
146+
item.put("partition_key", AttributeValue.builder().s("clientSupplierItem").build());
147+
item.put("sort_key", AttributeValue.builder().n("0").build());
148+
item.put("sensitive_data", AttributeValue.builder().s("encrypt and sign me!").build());
149+
150+
final PutItemRequest putRequest = PutItemRequest.builder()
151+
.tableName(ddbTableName)
152+
.item(item)
153+
.build();
154+
155+
final PutItemResponse putResponse = ddbClient.putItem(putRequest);
156+
157+
// Demonstrate that PutItem succeeded
158+
assert 200 == putResponse.sdkHttpResponse().statusCode();
159+
160+
// 8. Get the item back from our table using the same keyring.
161+
// The client will decrypt the item client-side using the MRK
162+
// and return the original item.
163+
final HashMap<String, AttributeValue> keyToGet = new HashMap<>();
164+
keyToGet.put("partition_key", AttributeValue.builder().s("clientSupplierItem").build());
165+
keyToGet.put("sort_key", AttributeValue.builder().n("0").build());
166+
167+
final GetItemRequest getRequest = GetItemRequest.builder()
168+
.key(keyToGet)
169+
.tableName(ddbTableName)
170+
.build();
171+
172+
final GetItemResponse getResponse = ddbClient.getItem(getRequest);
173+
174+
// Demonstrate that GetItem succeeded and returned the decrypted item
175+
assert 200 == getResponse.sdkHttpResponse().statusCode();
176+
final Map<String, AttributeValue> returnedItem = getResponse.item();
177+
assert returnedItem.get("sensitive_data").s().equals("encrypt and sign me!");
178+
179+
// 9. Create a MRK discovery multi-keyring with a custom client supplier.
180+
// A discovery MRK multi-keyring will be composed of
181+
// multiple discovery MRK keyrings, one for each region.
182+
// Each component keyring has its own KMS client in a particular region.
183+
// When we provide a client supplier to the multi-keyring, all component
184+
// keyrings will use that client supplier configuration.
185+
// In our tests, we make `keyArn` an MRK with a replica, and
186+
// provide only the replica region in our discovery filter.
187+
DiscoveryFilter discoveryFilter = DiscoveryFilter.builder()
188+
.partition("aws")
189+
.accountIds(accountIds)
190+
.build();
191+
192+
final CreateAwsKmsMrkDiscoveryMultiKeyringInput mrkDiscoveryClientSupplierInput =
193+
CreateAwsKmsMrkDiscoveryMultiKeyringInput.builder()
194+
.clientSupplier(new RegionalRoleClientSupplier())
195+
.discoveryFilter(discoveryFilter)
196+
.regions(regions)
197+
.build();
198+
IKeyring mrkDiscoveryClientSupplierKeyring = matProv.CreateAwsKmsMrkDiscoveryMultiKeyring(
199+
mrkDiscoveryClientSupplierInput);
200+
201+
// 10. Create a new config and client using the discovery keyring.
202+
// This is the same setup as above, except we provide the discovery keyring to the config.
203+
final Map<String, DynamoDbTableEncryptionConfig> onlyReplicaKeyTableConfigs = new HashMap<>();
204+
final DynamoDbTableEncryptionConfig onlyReplicaKeyConfig = DynamoDbTableEncryptionConfig.builder()
205+
.logicalTableName(ddbTableName)
206+
.partitionKeyName("partition_key")
207+
.sortKeyName("sort_key")
208+
.attributeActions(attributeActions)
209+
// Provide discovery keyring here
210+
.keyring(mrkDiscoveryClientSupplierKeyring)
211+
.allowedUnauthenticatedAttributePrefix(unauthAttrPrefix)
212+
.build();
213+
onlyReplicaKeyTableConfigs.put(ddbTableName, onlyReplicaKeyConfig);
214+
215+
DynamoDbEncryptionInterceptor onlyReplicaKeyEncryptionInterceptor = DynamoDbEncryptionInterceptor.builder()
216+
.config(DynamoDbTablesEncryptionConfig.builder()
217+
.tableEncryptionConfigs(onlyReplicaKeyTableConfigs)
218+
.build())
219+
.build();
220+
221+
final DynamoDbClient onlyReplicaKeyDdbClient = DynamoDbClient.builder()
222+
.overrideConfiguration(
223+
ClientOverrideConfiguration.builder()
224+
.addExecutionInterceptor(onlyReplicaKeyEncryptionInterceptor)
225+
.build())
226+
.build();
227+
228+
// 11. Get the item back from our table using the discovery keyring client.
229+
// The client will decrypt the item client-side using the keyring,
230+
// and return the original item.
231+
// The discovery keyring will only use KMS keys in the provided regions and
232+
// AWS accounts. Since we have provided it with a custom client supplier
233+
// which uses different IAM roles based on the key region,
234+
// the discovery keyring will use a particular IAM role to decrypt
235+
// based on the region of the KMS key it uses to decrypt.
236+
final HashMap<String, AttributeValue> onlyReplicaKeyKeyToGet = new HashMap<>();
237+
onlyReplicaKeyKeyToGet.put("partition_key", AttributeValue.builder().s("awsKmsMrkMultiKeyringItem").build());
238+
onlyReplicaKeyKeyToGet.put("sort_key", AttributeValue.builder().n("0").build());
239+
240+
final GetItemRequest onlyReplicaKeyGetRequest = GetItemRequest.builder()
241+
.key(onlyReplicaKeyKeyToGet)
242+
.tableName(ddbTableName)
243+
.build();
244+
245+
final GetItemResponse onlyReplicaKeyGetResponse = onlyReplicaKeyDdbClient.getItem(onlyReplicaKeyGetRequest);
246+
247+
// Demonstrate that GetItem succeeded and returned the decrypted item
248+
assert 200 == onlyReplicaKeyGetResponse.sdkHttpResponse().statusCode();
249+
final Map<String, AttributeValue> onlyReplicaKeyReturnedItem = onlyReplicaKeyGetResponse.item();
250+
assert onlyReplicaKeyReturnedItem.get("sensitive_data").s().equals("encrypt and sign me!");
251+
252+
// TODO: After adding MissingRegionException, give an example with a fake region
253+
// demonstrating that MissingRegionException extends AwsCryptographicMaterialProvidersException
254+
}
255+
256+
public static void main(final String[] args) {
257+
if (args.length <= 1) {
258+
throw new IllegalArgumentException("To run this example, include the ddbTable, keyArn, AWS accounts, and regions in args");
259+
}
260+
final String ddbTableName = args[0];
261+
final String keyArn = args[1];
262+
263+
// We will assume only 1 AWS account and 1 region will be passed into args.
264+
// To add more of either, change this number, then pass them into args.
265+
int numberOfAwsAccounts = 1;
266+
int numberOfRegions = 1;
267+
268+
List<String> accounts = new ArrayList<>();
269+
int firstAccountIndex = 2; // 2 because args[0] is ddbTableName and args[1] is mrkArn
270+
for (int i = firstAccountIndex; i < firstAccountIndex + numberOfAwsAccounts; i++) {
271+
accounts.add(args[i]);
272+
}
273+
List<String> regions = new ArrayList<>();
274+
int firstRegionIndex = firstAccountIndex + numberOfAwsAccounts;
275+
for (int i = firstRegionIndex; i < firstRegionIndex + numberOfRegions; i++) {
276+
regions.add(args[i]);
277+
}
278+
ClientSupplierPutItemGetItem(ddbTableName, keyArn, accounts, regions);
279+
}
280+
}

‎Examples/runtimes/java/DynamoDbEncryption/src/main/java/software/aws/cryptography/examples/clientsupplier/RegionalRoleClientSupplier.java

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
package software.aws.cryptography.examples.clientsupplier;
2+
3+
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
4+
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
5+
import software.amazon.awssdk.regions.Region;
6+
import software.amazon.awssdk.services.kms.KmsClient;
7+
import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
8+
import software.amazon.awssdk.services.sts.model.Credentials;
9+
import software.amazon.cryptography.materialproviders.IClientSupplier;
10+
import software.amazon.awssdk.services.sts.StsClient;
11+
import software.amazon.cryptography.materialproviders.model.GetClientInput;
12+
13+
/*
14+
Example class demonstrating an implementation of a custom client supplier.
15+
This particular implementation will create KMS clients with different IAM roles,
16+
depending on the region passed.
17+
*/
18+
public class RegionalRoleClientSupplier implements IClientSupplier {
19+
20+
private static StsClient stsClient;
21+
private static RegionalRoleClientSupplierConfig config;
22+
23+
public RegionalRoleClientSupplier() {
24+
stsClient = StsClient.create();
25+
config = new RegionalRoleClientSupplierConfig();
26+
}
27+
28+
@Override
29+
public KmsClient GetClient(GetClientInput getClientInput) {
30+
if (!config.regionIamRoleMap.containsKey(getClientInput.region())) {
31+
// TODO: Create a MissingRegionException that extends AwsCryptographicMaterialProvidersException.
32+
// The generated code for AwsCryptographicMaterialProvidersException cannot be extended as-is,
33+
// as its constructor requires access to a class private to itself.
34+
throw new RuntimeException("Missing region");
35+
}
36+
37+
String arn = config.regionIamRoleMap.get(getClientInput.region());
38+
Credentials creds = stsClient.assumeRole(AssumeRoleRequest.builder()
39+
.roleArn(arn)
40+
.durationSeconds(900) // 15 minutes is the minimum value
41+
.roleSessionName("Java-Client-Supplier-Example-Session")
42+
.build()).credentials();
43+
44+
return KmsClient.builder()
45+
.credentialsProvider(
46+
StaticCredentialsProvider.create(
47+
AwsSessionCredentials.create(
48+
creds.accessKeyId(),
49+
creds.secretAccessKey(),
50+
creds.sessionToken()
51+
)
52+
)
53+
)
54+
.region(Region.of(getClientInput.region()))
55+
.build();
56+
}
57+
}

‎Examples/runtimes/java/DynamoDbEncryption/src/main/java/software/aws/cryptography/examples/clientsupplier/RegionalRoleClientSupplierConfig.java

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package software.aws.cryptography.examples.clientsupplier;
2+
3+
import java.util.HashMap;
4+
import java.util.Map;
5+
6+
import software.amazon.awssdk.regions.Region;
7+
8+
/*
9+
Class containing config for the RegionalRoleClientSupplier.
10+
In your own code, this might be hardcoded, or reference
11+
an external source, e.g. environment variables or AWS AppConfig.
12+
*/
13+
14+
public class RegionalRoleClientSupplierConfig {
15+
16+
private static final String US_EAST_1_IAM_ROLE = "arn:aws:iam::370957321024:role/GitHub-CI-DDBEC-Dafny-Role-only-us-east-1-KMS-keys";
17+
private static final String EU_WEST_1_IAM_ROLE = "arn:aws:iam::370957321024:role/GitHub-CI-DDBEC-Dafny-Role-only-eu-west-1-KMS-keys";
18+
19+
public Map<String, String> regionIamRoleMap;
20+
21+
public RegionalRoleClientSupplierConfig() {
22+
regionIamRoleMap = new HashMap<>();
23+
regionIamRoleMap.put(Region.US_EAST_1.id(), US_EAST_1_IAM_ROLE);
24+
regionIamRoleMap.put(Region.EU_WEST_1.id(), EU_WEST_1_IAM_ROLE);
25+
}
26+
27+
}

0 commit comments

Comments
 (0)