diff --git a/services/commtest/tests/identity_access_tokens_tests.rs b/services/commtest/tests/identity_access_tokens_tests.rs --- a/services/commtest/tests/identity_access_tokens_tests.rs +++ b/services/commtest/tests/identity_access_tokens_tests.rs @@ -5,7 +5,8 @@ use proto::identity_client_service_client::IdentityClientServiceClient; use proto::{ DeviceKeyUpload, IdentityKeyInfo, PreKey, RegistrationFinishRequest, - RegistrationStartRequest, VerifyUserAccessTokenRequest, + RegistrationStartRequest, UploadOneTimeKeysRequest, + VerifyUserAccessTokenRequest, }; use rand::{distributions::Alphanumeric, Rng}; @@ -113,3 +114,27 @@ assert_eq!(response.into_inner().token_valid, true); } + +#[tokio::test] +async fn upload_one_time_keys() { + 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!["a".to_string(), "b".to_string()], + notif_one_time_pre_keys: vec!["c".to_string(), "d".to_string()], + }; + + // This send will fail if the one-time keys weren't successfully added + 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 @@ -35,7 +35,7 @@ use moka::future::Cache; use rand::rngs::OsRng; use tonic::Response; -use tracing::error; +use tracing::{debug, error}; use self::client_proto::ReservedRegistrationStartRequest; @@ -749,9 +749,41 @@ async fn upload_one_time_keys( &self, - _request: tonic::Request, + request: tonic::Request, ) -> Result, tonic::Status> { - unimplemented!(); + let message = request.into_inner(); + + debug!("Validating token: {:?}", message); + let token_valid = self + .client + .verify_access_token( + message.user_id.clone(), + message.device_id.clone(), + message.access_token, + ) + .await + .map_err(handle_db_error)?; + + if !token_valid { + return Err(tonic::Status::unauthenticated("Invalid token")); + } + + debug!( + "Attempting to update one time keys for user: {}", + message.user_id + ); + self + .client + .append_one_time_prekeys( + message.user_id, + message.device_id, + message.content_one_time_pre_keys, + message.notif_one_time_pre_keys, + ) + .await + .map_err(handle_db_error)?; + + Ok(tonic::Response::new(Empty {})) } async fn refresh_user_pre_keys( 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 @@ -218,6 +218,51 @@ .await } + pub async fn append_one_time_prekeys( + &self, + user_id: String, + device_id: String, + content_one_time_keys: Vec, + notif_one_time_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(); + + 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)), + ]); + + 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)) + .send() + .await + .map_err(|e| Error::AwsSdk(e.into()))?; + + Ok(()) + } + async fn add_device_to_users_table( &self, user_id: String,