diff --git a/services/commtest/tests/identity_prekey_tests.rs b/services/commtest/tests/identity_prekey_tests.rs new file mode 100644 --- /dev/null +++ b/services/commtest/tests/identity_prekey_tests.rs @@ -0,0 +1,49 @@ +mod proto { + tonic::include_proto!("identity.client"); +} +use proto as client; +mod auth_proto { + tonic::include_proto!("identity.authenticated"); +} +use auth_proto::identity_client_service_client::IdentityClientServiceClient as AuthClient; +use auth_proto::RefreshUserPreKeysRequest; +use client::PreKey; +use commtest::identity::device::create_device; +use tonic::{transport::Endpoint, Request}; + +#[tokio::test] +async fn set_prekey() { + let device_info = create_device().await; + + let channel = Endpoint::from_static("http://[::1]:50054") + .connect() + .await + .unwrap(); + + let mut client = + AuthClient::with_interceptor(channel, |mut request: Request<()>| { + let metadata = request.metadata_mut(); + metadata.append("user_id", device_info.user_id.parse().unwrap()); + metadata.append("device_id", device_info.device_id.parse().unwrap()); + metadata + .append("access_token", device_info.access_token.parse().unwrap()); + Ok(request) + }); + + let upload_request = RefreshUserPreKeysRequest { + new_content_pre_keys: Some(PreKey { + pre_key: "content_prekey".to_string(), + pre_key_signature: "content_prekey_signature".to_string(), + }), + new_notif_pre_keys: Some(PreKey { + pre_key: "content_prekey".to_string(), + pre_key_signature: "content_prekey_signature".to_string(), + }), + }; + + // This send will fail if the one-time keys weren't successfully added + println!( + "Error: {:?}", + client.refresh_user_pre_keys(upload_request).await + ); +} 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 @@ -219,6 +219,60 @@ .await } + pub async fn set_prekey( + &self, + user_id: String, + device_id: String, + content_prekey: String, + content_prekey_signature: String, + notif_prekey: String, + notif_prekey_signature: String, + ) -> Result<(), Error> { + let notif_prekey_av = AttributeValue::S(notif_prekey); + let notif_prekey_signature_av = AttributeValue::S(notif_prekey_signature); + let content_prekey_av = AttributeValue::S(content_prekey); + let content_prekey_signature_av = + AttributeValue::S(content_prekey_signature); + + let update_expression = + format!("SET {0}.#{1}.{2} = :n, {0}.#{1}.{3} = :p, {0}.#{1}.{4} = :c, {0}.#{1}.{5} = :d", + USERS_TABLE_DEVICES_ATTRIBUTE, + "deviceID", + USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME, + USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME, + USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME, + USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME, + ); + let expression_attribute_names = HashMap::from([ + (format!("#{}", "deviceID"), device_id), + ( + "#user_id".to_string(), + USERS_TABLE_PARTITION_KEY.to_string(), + ), + ]); + let expression_attribute_values = HashMap::from([ + (":n".to_string(), notif_prekey_av), + (":p".to_string(), notif_prekey_signature_av), + (":c".to_string(), content_prekey_av), + (":d".to_string(), content_prekey_signature_av), + ]); + + self + .client + .update_item() + .table_name(USERS_TABLE) + .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id)) + .update_expression(update_expression) + .condition_expression("attribute_exists(#user_id)") + .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(()) + } + pub async fn append_one_time_prekeys( &self, user_id: String, diff --git a/services/identity/src/grpc_services/authenticated.rs b/services/identity/src/grpc_services/authenticated.rs --- a/services/identity/src/grpc_services/authenticated.rs +++ b/services/identity/src/grpc_services/authenticated.rs @@ -72,8 +72,37 @@ impl IdentityClientService for AuthenticatedService { async fn refresh_user_pre_keys( &self, - _request: Request, + request: Request, ) -> Result, Status> { - unimplemented!(); + let user_id = get_value(&request, "user_id") + .ok_or(Status::unauthenticated("Missing user_id field"))?; + let device_id = get_value(&request, "device_id") + .ok_or(Status::unauthenticated("Missing device_id field"))?; + let message = request.into_inner(); + + debug!("Refreshing prekeys for user: {}", user_id); + + let content_keys = message + .new_content_pre_keys + .ok_or(Status::invalid_argument("Missing content keys"))?; + let notif_keys = message + .new_notif_pre_keys + .ok_or(Status::invalid_argument("Missing notification keys"))?; + + self + .db_client + .set_prekey( + user_id, + device_id, + content_keys.pre_key, + content_keys.pre_key_signature, + notif_keys.pre_key, + notif_keys.pre_key_signature, + ) + .await + .map_err(handle_db_error)?; + + let response = Response::new(Empty {}); + Ok(response) } }