Skip to content

Commit f790f4f

Browse files
committed
m
1 parent a2662a6 commit f790f4f

35 files changed

+430
-58
lines changed

releases/rust/db_esdk/Cargo.toml

+3
Original file line numberDiff line numberDiff line change
@@ -24,3 +24,6 @@ uuid = { version = "1.10.0", features = ["v4"] }
2424

2525
[lib]
2626
path = "src/implementation_from_dafny.rs"
27+
28+
[dev-dependencies]
29+
aws-sdk-sts = "1.43.0"
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,257 @@
1+
use super::regional_role_client_supplier::RegionalRoleClientSupplier;
2+
use crate::test_utils;
3+
use aws_db_esdk::aws_cryptography_dbEncryptionSdk_dynamoDb::types::DynamoDbTableEncryptionConfig;
4+
use aws_db_esdk::aws_cryptography_dbEncryptionSdk_structuredEncryption::types::CryptoAction;
5+
use aws_db_esdk::aws_cryptography_materialProviders::client as mpl_client;
6+
use aws_db_esdk::aws_cryptography_materialProviders::types::client_supplier::ClientSupplierRef;
7+
use aws_db_esdk::aws_cryptography_materialProviders::types::material_providers_config::MaterialProvidersConfig;
8+
use aws_db_esdk::aws_cryptography_materialProviders::types::DiscoveryFilter;
9+
use aws_db_esdk::intercept::DbEsdkInterceptor;
10+
use aws_db_esdk::DynamoDbTablesEncryptionConfig;
11+
use aws_sdk_dynamodb::types::AttributeValue;
12+
use std::collections::HashMap;
13+
14+
/*
15+
This example sets up an MRK multi-keyring and an MRK discovery
16+
multi-keyring using a custom client supplier.
17+
A custom client supplier grants users access to more granular
18+
configuration aspects of their authentication details and KMS
19+
client. In this example, we create a simple custom client supplier
20+
that authenticates with a different IAM role based on the
21+
region of the KMS key.
22+
23+
This example creates a MRK multi-keyring configured with a custom
24+
client supplier using a single MRK and puts an encrypted item to the
25+
table. Then, it creates a MRK discovery multi-keyring to decrypt the item
26+
and retrieves the item from the table.
27+
28+
Running this example requires access to the DDB Table whose name
29+
is provided in CLI arguments.
30+
This table must be configured with the following
31+
primary key configuration:
32+
- Partition key is named "partition_key" with type (S)
33+
- Sort key is named "sort_key" with type (S)
34+
*/
35+
pub async fn put_item_get_item() {
36+
let ddb_table_name = test_utils::TEST_DDB_TABLE_NAME;
37+
// Note that we pass in an MRK in us-east-1...
38+
let key_arn = test_utils::TEST_MRK_REPLICA_KEY_ID_US_EAST_1.to_string();
39+
let account_ids = vec![test_utils::TEST_AWS_ACCOUNT_ID.to_string()];
40+
// ...and access its replica in eu-west-1
41+
let regions = vec!["eu-west-1".to_string()];
42+
43+
// 1. Create a single MRK multi-keyring.
44+
// This can be either a single-region KMS key or an MRK.
45+
// For this example to succeed, the key's region must either
46+
// 1) be in the regions list, or
47+
// 2) the key must be an MRK with a replica defined
48+
// in a region in the regions list, and the client
49+
// must have the correct permissions to access the replica.
50+
let mpl_config = MaterialProvidersConfig::builder().build().unwrap();
51+
let mpl = mpl_client::Client::from_conf(mpl_config).unwrap();
52+
53+
// Create the multi-keyring using our custom client supplier
54+
// defined in the RegionalRoleClientSupplier class in this directory.
55+
// Note: RegionalRoleClientSupplier will internally use the key_arn's region
56+
// to retrieve the correct IAM role.
57+
let supplier_ref = ClientSupplierRef {
58+
inner: std::rc::Rc::new(std::cell::RefCell::new(RegionalRoleClientSupplier::new())),
59+
};
60+
61+
let mrk_keyring_with_client_supplier = mpl
62+
.create_aws_kms_mrk_multi_keyring()
63+
.client_supplier(supplier_ref.clone())
64+
.generator(key_arn)
65+
.send()
66+
.await
67+
.unwrap();
68+
69+
// 2. Configure which attributes are encrypted and/or signed when writing new items.
70+
// For each attribute that may exist on the items we plan to write to our DynamoDbTable,
71+
// we must explicitly configure how they should be treated during item encryption:
72+
// - ENCRYPT_AND_SIGN: The attribute is encrypted and included in the signature
73+
// - SIGN_ONLY: The attribute is not encrypted, but is still included in the signature
74+
// - DO_NOTHING: The attribute is not encrypted and not included in the signature
75+
let attribute_actions_on_encrypt = HashMap::from([
76+
("partition_key".to_string(), CryptoAction::SignOnly), // Our partition attribute must be SIGN_ONLY
77+
("sort_key".to_string(), CryptoAction::SignOnly), // Our sort attribute must be SIGN_ONLY
78+
("sensitive_data".to_string(), CryptoAction::EncryptAndSign),
79+
]);
80+
81+
// 3. Configure which attributes we expect to be included in the signature
82+
// when reading items. There are two options for configuring this:
83+
//
84+
// - (Recommended) Configure `allowedUnsignedAttributesPrefix`:
85+
// When defining your DynamoDb schema and deciding on attribute names,
86+
// choose a distinguishing prefix (such as ":") for all attributes that
87+
// you do not want to include in the signature.
88+
// This has two main benefits:
89+
// - It is easier to reason about the security and authenticity of data within your item
90+
// when all unauthenticated data is easily distinguishable by their attribute name.
91+
// - If you need to add new unauthenticated attributes in the future,
92+
// you can easily make the corresponding update to your `attributeActionsOnEncrypt`
93+
// and immediately start writing to that new attribute, without
94+
// any other configuration update needed.
95+
// Once you configure this field, it is not safe to update it.
96+
//
97+
// - Configure `allowedUnsignedAttributes`: You may also explicitly list
98+
// a set of attributes that should be considered unauthenticated when encountered
99+
// on read. Be careful if you use this configuration. Do not remove an attribute
100+
// name from this configuration, even if you are no longer writing with that attribute,
101+
// as old items may still include this attribute, and our configuration needs to know
102+
// to continue to exclude this attribute from the signature scope.
103+
// If you add new attribute names to this field, you must first deploy the update to this
104+
// field to all readers in your host fleet before deploying the update to start writing
105+
// with that new attribute.
106+
//
107+
// For this example, we currently authenticate all attributes. To make it easier to
108+
// add unauthenticated attributes in the future, we define a prefix ":" for such attributes.
109+
const UNSIGNED_ATTR_PREFIX: &str = ":";
110+
111+
// 4. Create the DynamoDb Encryption configuration for the table we will be writing to.
112+
let table_config = DynamoDbTableEncryptionConfig::builder()
113+
.logical_table_name(ddb_table_name)
114+
.partition_key_name("partition_key")
115+
.sort_key_name("sort_key")
116+
.attribute_actions_on_encrypt(attribute_actions_on_encrypt.clone())
117+
.keyring(mrk_keyring_with_client_supplier)
118+
.allowed_unsigned_attribute_prefix(UNSIGNED_ATTR_PREFIX)
119+
.build()
120+
.unwrap();
121+
122+
let table_configs = DynamoDbTablesEncryptionConfig::builder()
123+
.table_encryption_configs(HashMap::from([(ddb_table_name.to_string(), table_config)]))
124+
.build()
125+
.unwrap();
126+
127+
// 5. Create a new AWS SDK DynamoDb client using the DynamoDb Config above
128+
let sdk_config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
129+
let dynamo_config = aws_sdk_dynamodb::config::Builder::from(&sdk_config)
130+
.interceptor(DbEsdkInterceptor::new(table_configs))
131+
.build();
132+
let ddb = aws_sdk_dynamodb::Client::from_conf(dynamo_config);
133+
134+
// 6. Put an item into our table using the above client.
135+
// Before the item gets sent to DynamoDb, it will be encrypted
136+
// client-side using the MRK multi-keyring.
137+
// The data key protecting this item will be encrypted
138+
// with all the KMS Keys in this keyring, so that it can be
139+
// decrypted with any one of those KMS Keys.
140+
let item = HashMap::from([
141+
(
142+
"partition_key".to_string(),
143+
AttributeValue::S("clientSupplierItem".to_string()),
144+
),
145+
("sort_key".to_string(), AttributeValue::N("0".to_string())),
146+
(
147+
"sensitive_data".to_string(),
148+
AttributeValue::S("encrypt and sign me!".to_string()),
149+
),
150+
]);
151+
152+
let _resp = ddb
153+
.put_item()
154+
.table_name(ddb_table_name)
155+
.set_item(Some(item.clone()))
156+
.send()
157+
.await
158+
.unwrap();
159+
160+
// 7. 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+
let key_to_get = HashMap::from([
164+
(
165+
"partition_key".to_string(),
166+
AttributeValue::S("clientSupplierItem".to_string()),
167+
),
168+
("sort_key".to_string(), AttributeValue::N("0".to_string())),
169+
]);
170+
171+
let resp = ddb
172+
.get_item()
173+
.table_name(ddb_table_name)
174+
.set_key(Some(key_to_get.clone()))
175+
.consistent_read(true)
176+
.send()
177+
.await
178+
.unwrap();
179+
180+
assert_eq!(
181+
resp.item.unwrap()["sensitive_data"],
182+
AttributeValue::S("encrypt and sign me!".to_string())
183+
);
184+
185+
// 7. Create a MRK discovery multi-keyring with a custom client supplier.
186+
// A discovery MRK multi-keyring will be composed of
187+
// multiple discovery MRK keyrings, one for each region.
188+
// Each component keyring has its own KMS client in a particular region.
189+
// When we provide a client supplier to the multi-keyring, all component
190+
// keyrings will use that client supplier configuration.
191+
// In our tests, we make `key_arn` an MRK with a replica, and
192+
// provide only the replica region in our discovery filter.
193+
let discovery_filter = DiscoveryFilter::builder()
194+
.partition("aws")
195+
.account_ids(account_ids)
196+
.build()
197+
.unwrap();
198+
199+
let mrk_discovery_client_supplier_keyring = mpl
200+
.create_aws_kms_mrk_discovery_multi_keyring()
201+
.client_supplier(supplier_ref.clone())
202+
.discovery_filter(discovery_filter)
203+
.regions(regions)
204+
.send()
205+
.await
206+
.unwrap();
207+
208+
// 9. Create a new config and client using the discovery keyring.
209+
// This is the same setup as above, except we provide the discovery keyring to the config.
210+
let only_replica_table_config = DynamoDbTableEncryptionConfig::builder()
211+
.logical_table_name(ddb_table_name)
212+
.partition_key_name("partition_key")
213+
.sort_key_name("sort_key")
214+
.attribute_actions_on_encrypt(attribute_actions_on_encrypt)
215+
.keyring(mrk_discovery_client_supplier_keyring)
216+
.allowed_unsigned_attribute_prefix(UNSIGNED_ATTR_PREFIX)
217+
.build()
218+
.unwrap();
219+
220+
let only_replica_table_configs = DynamoDbTablesEncryptionConfig::builder()
221+
.table_encryption_configs(HashMap::from([(
222+
ddb_table_name.to_string(),
223+
only_replica_table_config,
224+
)]))
225+
.build()
226+
.unwrap();
227+
228+
let only_replica_dynamo_config = aws_sdk_dynamodb::config::Builder::from(&sdk_config)
229+
.interceptor(DbEsdkInterceptor::new(only_replica_table_configs))
230+
.build();
231+
let only_replica_ddb = aws_sdk_dynamodb::Client::from_conf(only_replica_dynamo_config);
232+
233+
// 10. Get the item back from our table using the discovery keyring client.
234+
// The client will decrypt the item client-side using the keyring,
235+
// and return the original item.
236+
// The discovery keyring will only use KMS keys in the provided regions and
237+
// AWS accounts. Since we have provided it with a custom client supplier
238+
// which uses different IAM roles based on the key region,
239+
// the discovery keyring will use a particular IAM role to decrypt
240+
// based on the region of the KMS key it uses to decrypt.
241+
242+
let resp = only_replica_ddb
243+
.get_item()
244+
.table_name(ddb_table_name)
245+
.set_key(Some(key_to_get))
246+
.consistent_read(true)
247+
.send()
248+
.await
249+
.unwrap();
250+
251+
assert_eq!(
252+
resp.item.unwrap()["sensitive_data"],
253+
AttributeValue::S("encrypt and sign me!".to_string())
254+
);
255+
256+
println!("client_supplier_example successful.");
257+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
pub mod client_supplier_example;
2+
pub mod regional_role_client_supplier;
3+
pub mod regional_role_client_supplier_config;
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
use aws_config::Region;
2+
use aws_db_esdk::aws_cryptography_materialProviders::types::ClientSupplier;
3+
use aws_db_esdk::deps::aws_cryptography_materialProviders::operation::get_client::GetClientInput;
4+
use aws_db_esdk::deps::aws_cryptography_materialProviders::types::error::Error;
5+
use aws_db_esdk::deps::com_amazonaws_kms::client::Client as kms_client;
6+
use aws_sdk_sts::Client as sts_client;
7+
8+
/*
9+
Example class demonstrating an implementation of a custom client supplier.
10+
This particular implementation will create KMS clients with different IAM roles,
11+
depending on the region passed.
12+
*/
13+
14+
pub struct RegionalRoleClientSupplier {
15+
sts_client: sts_client, // private readonly AmazonSecurityTokenServiceClient _stsClient = new AmazonSecurityTokenServiceClient();
16+
}
17+
18+
impl RegionalRoleClientSupplier {
19+
pub fn new() -> Self {
20+
let sdk_config = tokio::task::block_in_place(|| {
21+
tokio::runtime::Handle::current().block_on(async {
22+
aws_config::load_defaults(aws_config::BehaviorVersion::v2024_03_28()).await
23+
})
24+
});
25+
Self {
26+
sts_client: sts_client::new(&sdk_config),
27+
}
28+
}
29+
}
30+
31+
impl ClientSupplier for RegionalRoleClientSupplier {
32+
fn get_client(&mut self, input: GetClientInput) -> Result<kms_client, Error> {
33+
let region = input.region.unwrap();
34+
let arn =
35+
super::regional_role_client_supplier_config::region_iam_role_map()[&region].clone();
36+
let creds = tokio::task::block_in_place(|| {
37+
tokio::runtime::Handle::current().block_on(async {
38+
self.sts_client
39+
.assume_role()
40+
.role_arn(arn)
41+
.duration_seconds(900)
42+
.role_session_name("Rust-Client-Supplier-Example-Session")
43+
.send()
44+
.await
45+
})
46+
})
47+
.unwrap();
48+
49+
let types_cred = creds.credentials.unwrap();
50+
let config_creds = aws_sdk_sts::config::Credentials::new(
51+
types_cred.access_key_id(),
52+
types_cred.secret_access_key(),
53+
Some(types_cred.session_token().to_string()),
54+
Some(
55+
std::time::SystemTime::UNIX_EPOCH
56+
+ std::time::Duration::from_secs(types_cred.expiration().secs() as u64),
57+
),
58+
"SomeProvider",
59+
);
60+
let cred_prov = aws_sdk_kms::config::SharedCredentialsProvider::new(config_creds);
61+
62+
let sdk_config = tokio::task::block_in_place(|| {
63+
tokio::runtime::Handle::current().block_on(async {
64+
aws_config::load_defaults(aws_config::BehaviorVersion::v2024_03_28()).await
65+
})
66+
});
67+
let kms_config = aws_sdk_kms::config::Builder::from(&sdk_config)
68+
.credentials_provider(cred_prov)
69+
.region(Region::new(region))
70+
.build();
71+
72+
let inner_client = aws_sdk_kms::Client::from_conf(kms_config);
73+
Ok(aws_db_esdk::deps::com_amazonaws_kms::client::Client {
74+
inner: inner_client,
75+
})
76+
}
77+
}
78+
// protected override IAmazonKeyManagementService _GetClient(GetClientInput getClientInput)
79+
// {
80+
// String arn = _config.regionIamRoleMap[getClientInput.Region];
81+
// Credentials creds = _stsClient.AssumeRoleAsync(new AssumeRoleRequest
82+
// {
83+
// RoleArn = arn,
84+
// DurationSeconds = 900, // 15 minutes is the minimum value
85+
// RoleSessionName = "Java-Client-Supplier-Example-Session"
86+
// }
87+
// ).Result.Credentials;
88+
89+
// return new AmazonKeyManagementServiceClient(creds, RegionEndpoint.GetBySystemName(getClientInput.Region));
90+
// }
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
use std::collections::HashMap;
2+
3+
/*
4+
Class containing config for the RegionalRoleClientSupplier.
5+
In your own code, this might be hardcoded, or reference
6+
an external source, e.g. environment variables or AWS AppConfig.
7+
*/
8+
9+
const US_EAST_1_IAM_ROLE: &str =
10+
"arn:aws:iam::370957321024:role/GitHub-CI-DDBEC-Dafny-Role-only-us-east-1-KMS-keys";
11+
12+
const EU_WEST_1_IAM_ROLE: &str =
13+
"arn:aws:iam::370957321024:role/GitHub-CI-DDBEC-Dafny-Role-only-eu-west-1-KMS-keys";
14+
15+
pub fn region_iam_role_map() -> HashMap<String, String> {
16+
HashMap::from([
17+
("us-east-1".to_string(), US_EAST_1_IAM_ROLE.to_string()),
18+
("eu-west-1".to_string(), EU_WEST_1_IAM_ROLE.to_string()),
19+
])
20+
}

0 commit comments

Comments
 (0)