diff --git a/services/commtest/tests/identity_onetime_key_tests.rs b/services/commtest/tests/identity_onetime_key_tests.rs new file mode 100644 --- /dev/null +++ b/services/commtest/tests/identity_onetime_key_tests.rs @@ -0,0 +1,35 @@ +mod client { + tonic::include_proto!("identity.client"); +} +mod auth_proto { + tonic::include_proto!("identity.authenticated"); +} +use client::identity_client_service_client::IdentityClientServiceClient; +use client::UploadOneTimeKeysRequest; +use commtest::identity::device::create_device; + +#[tokio::test] +async fn verify_access_token() { + let device_info = create_device().await; + + let mut identity_client = + IdentityClientServiceClient::connect("http://127.0.0.1:50054") + .await + .expect("Couldn't connect to identitiy service"); + + let upload_request = UploadOneTimeKeysRequest { + user_id: device_info.user_id, + device_id: device_info.device_id, + access_token: device_info.access_token, + content_one_time_pre_keys: vec![ + "content1".to_string(), + "content2".to_string(), + ], + notif_one_time_pre_keys: vec!["notif1".to_string(), "notif2".to_string()], + }; + + identity_client + .upload_one_time_keys(upload_request) + .await + .unwrap(); +} diff --git a/services/identity/src/client_service.rs b/services/identity/src/client_service.rs --- a/services/identity/src/client_service.rs +++ b/services/identity/src/client_service.rs @@ -791,7 +791,6 @@ self .client .append_one_time_prekeys( - message.user_id, message.device_id, message.content_one_time_pre_keys, message.notif_one_time_pre_keys, diff --git a/services/identity/src/constants.rs b/services/identity/src/constants.rs --- a/services/identity/src/constants.rs +++ b/services/identity/src/constants.rs @@ -88,7 +88,7 @@ // One time keys table, which need to exist in their own table to ensure // atomicity of additions and removals pub mod content_one_time_keys_table { - pub const NAME: &'static str = "identity-content-one-time-keys"; + pub const NAME: &'static str = "identity-content-onetime-keys"; pub const PARTITION_KEY: &'static str = "deviceID"; pub const DEVICE_ID: &'static str = PARTITION_KEY; pub const SORT_KEY: &'static str = "oneTimeKey"; @@ -96,7 +96,7 @@ } pub mod notif_one_time_keys_table { - pub const NAME: &'static str = "identity-notif-one-time-keys"; + pub const NAME: &'static str = "identity-notif-onetime-keys"; pub const PARTITION_KEY: &'static str = "deviceID"; pub const DEVICE_ID: &'static str = PARTITION_KEY; pub const SORT_KEY: &'static str = "oneTimeKey"; diff --git a/services/identity/src/database.rs b/services/identity/src/database.rs --- a/services/identity/src/database.rs +++ b/services/identity/src/database.rs @@ -4,6 +4,7 @@ use std::str::FromStr; use std::sync::Arc; +use crate::ddb_utils::into_onetime_put_requests; use crate::error::{DBItemAttributeError, DBItemError, Error}; use aws_config::SdkConfig; use aws_sdk_dynamodb::model::{AttributeValue, PutRequest, WriteRequest}; @@ -25,13 +26,11 @@ NONCE_TABLE_CREATED_ATTRIBUTE, NONCE_TABLE_PARTITION_KEY, RESERVED_USERNAMES_TABLE, RESERVED_USERNAMES_TABLE_PARTITION_KEY, USERS_TABLE, USERS_TABLE_DEVICES_ATTRIBUTE, - USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME, - USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME, @@ -295,42 +294,24 @@ pub async fn append_one_time_prekeys( &self, - user_id: String, device_id: String, - content_one_time_keys: Vec, - notif_one_time_keys: Vec, + content_onetime_keys: Vec, + notif_onetime_keys: Vec, ) -> Result<(), Error> { - let notif_keys_av: Vec = notif_one_time_keys - .into_iter() - .map(AttributeValue::S) - .collect(); - let content_keys_av: Vec = content_one_time_keys - .into_iter() - .map(AttributeValue::S) - .collect(); + use crate::constants::{ + content_one_time_keys_table, notif_one_time_keys_table, + }; - let update_expression = - format!("SET {0}.#{1}.{2} = list_append({0}.#{1}.{2}, :n), {0}.#{1}.{3} = list_append({0}.#{1}.{3}, :i)", - USERS_TABLE_DEVICES_ATTRIBUTE, - "deviceID", - USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME, - USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME - ); - let expression_attribute_names = - HashMap::from([(format!("#{}", "deviceID"), device_id)]); - let expression_attribute_values = HashMap::from([ - (":n".to_string(), AttributeValue::L(notif_keys_av)), - (":i".to_string(), AttributeValue::L(content_keys_av)), - ]); + let content_otk_requests = + into_onetime_put_requests(device_id.clone(), content_onetime_keys); + let notif_otk_requests = + into_onetime_put_requests(device_id.clone(), notif_onetime_keys); self .client - .update_item() - .table_name(USERS_TABLE) - .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id)) - .update_expression(update_expression) - .set_expression_attribute_names(Some(expression_attribute_names)) - .set_expression_attribute_values(Some(expression_attribute_values)) + .batch_write_item() + .request_items(content_one_time_keys_table::NAME, content_otk_requests) + .request_items(notif_one_time_keys_table::NAME, notif_otk_requests) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; @@ -344,14 +325,19 @@ flattened_device_key_upload: FlattenedDeviceKeyUpload, social_proof: Option, ) -> Result<(), Error> { + // Avoid borrowing from lifetime of flattened_device_key_upload + let device_id = flattened_device_key_upload.device_id_key.clone(); + let content_onetime_keys = + flattened_device_key_upload.content_onetime_keys.clone(); + let notif_onetime_keys = + flattened_device_key_upload.notif_onetime_keys.clone(); + let device_info = - create_device_info(flattened_device_key_upload.clone(), social_proof); + create_device_info(flattened_device_key_upload, social_proof); let update_expression = format!("SET {}.#{} = :v", USERS_TABLE_DEVICES_ATTRIBUTE, "deviceID",); - let expression_attribute_names = HashMap::from([( - format!("#{}", "deviceID"), - flattened_device_key_upload.device_id_key, - )]); + let expression_attribute_names = + HashMap::from([(format!("#{}", "deviceID"), device_id.clone())]); let expression_attribute_values = HashMap::from([(":v".to_string(), AttributeValue::M(device_info))]); @@ -367,6 +353,26 @@ .await .map_err(|e| Error::AwsSdk(e.into()))?; + let content_otk_requests = + into_onetime_put_requests(device_id.clone(), content_onetime_keys); + let notif_otk_requests = + into_onetime_put_requests(device_id.clone(), notif_onetime_keys); + + self + .client + .batch_write_item() + .request_items( + crate::constants::content_one_time_keys_table::NAME, + content_otk_requests, + ) + .request_items( + crate::constants::notif_one_time_keys_table::NAME, + notif_otk_requests, + ) + .send() + .await + .map_err(|e| Error::AwsSdk(e.into()))?; + Ok(()) } @@ -1042,16 +1048,6 @@ .to_string(), AttributeValue::S(flattened_device_key_upload.content_prekey_signature), ), - ( - USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(), - AttributeValue::L( - flattened_device_key_upload - .content_onetime_keys - .into_iter() - .map(AttributeValue::S) - .collect(), - ), - ), ( USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME.to_string(), AttributeValue::S(flattened_device_key_upload.notif_prekey), @@ -1060,16 +1056,6 @@ USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME.to_string(), AttributeValue::S(flattened_device_key_upload.notif_prekey_signature), ), - ( - USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(), - AttributeValue::L( - flattened_device_key_upload - .notif_onetime_keys - .into_iter() - .map(AttributeValue::S) - .collect(), - ), - ), ]); if let Some(social_proof) = social_proof { diff --git a/services/identity/src/ddb_utils.rs b/services/identity/src/ddb_utils.rs new file mode 100644 --- /dev/null +++ b/services/identity/src/ddb_utils.rs @@ -0,0 +1,38 @@ +use aws_sdk_dynamodb::model::{AttributeValue, PutRequest, WriteRequest}; +use std::collections::HashMap; +use std::iter::IntoIterator; + +fn create_onetime_key_put_request( + device_id: String, + onetime_key: String, +) -> WriteRequest { + // The content and notif onetime key tables use the same partition and sort + // key + use crate::constants::content_one_time_keys_table::*; + + let builder = PutRequest::builder(); + let attrs = HashMap::from([ + (PARTITION_KEY.to_string(), AttributeValue::S(device_id)), + (SORT_KEY.to_string(), AttributeValue::S(onetime_key)), + ]); + + let put_request = builder.set_item(Some(attrs)).build(); + + WriteRequest::builder().put_request(put_request).build() +} + +pub fn into_onetime_put_requests( + device_id: String, + onetime_keys: T, +) -> Vec +where + T: IntoIterator, + ::Item: ToString, +{ + onetime_keys + .into_iter() + .map(|otk| { + create_onetime_key_put_request(device_id.clone(), otk.to_string()) + }) + .collect() +} diff --git a/services/identity/src/main.rs b/services/identity/src/main.rs --- a/services/identity/src/main.rs +++ b/services/identity/src/main.rs @@ -9,6 +9,7 @@ mod config; pub mod constants; mod database; +pub mod ddb_utils; pub mod error; mod grpc_services; mod id;