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 @@ -335,9 +335,10 @@ let login_time = chrono::Utc::now(); self .client - .add_password_user_device_to_users_table( + .add_user_device( state.user_id.clone(), state.flattened_device_key_upload.clone(), + None, code_version, login_time, ) @@ -417,10 +418,10 @@ // User already exists, so we should update the DDB item self .client - .add_wallet_user_device_to_users_table( + .add_user_device( id.clone(), flattened_device_key_upload.clone(), - social_proof, + Some(social_proof), code_version, chrono::Utc::now(), ) 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 @@ -16,12 +16,15 @@ use std::str::FromStr; use std::sync::Arc; -use crate::ddb_utils::{ - create_one_time_key_partition_key, into_one_time_put_requests, Identifier, - OlmAccountType, -}; use crate::error::{consume_error, Error}; use crate::reserved_users::UserDetail; +use crate::{ + ddb_utils::{ + create_one_time_key_partition_key, into_one_time_put_requests, Identifier, + OlmAccountType, + }, + grpc_services::protos, +}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use tracing::{debug, error, info, warn}; @@ -52,6 +55,8 @@ mod device_list; pub use device_list::DeviceListRow; +use self::device_list::PreKey; + #[derive(Serialize, Deserialize)] pub struct OlmKeys { pub curve25519: String, @@ -94,12 +99,6 @@ } } -// This is very similar to the protobuf definitions, however, -// coupling the protobuf schema to the database API should be avoided. -pub struct PreKey { - pub prekey: String, - pub prekey_signature: String, -} pub struct OutboundKeys { pub key_payload: String, pub key_payload_signature: String, @@ -110,6 +109,23 @@ pub notif_one_time_key: Option, } +impl From for protos::auth::OutboundKeyInfo { + fn from(db_keys: OutboundKeys) -> Self { + use protos::unauth::IdentityKeyInfo; + Self { + identity_info: Some(IdentityKeyInfo { + payload: db_keys.key_payload, + payload_signature: db_keys.key_payload_signature, + social_proof: db_keys.social_proof, + }), + content_prekey: Some(db_keys.content_prekey.into()), + notif_prekey: Some(db_keys.notif_prekey.into()), + one_time_content_prekey: db_keys.content_one_time_key, + one_time_notif_prekey: db_keys.notif_one_time_key, + } + } +} + #[derive(Clone)] pub struct DatabaseClient { client: Arc, @@ -149,7 +165,6 @@ device_key_upload.clone(), Some((registration_state.username, Blob::new(password_file))), None, - None, registration_state.user_id, ) .await?; @@ -182,7 +197,6 @@ flattened_device_key_upload.clone(), None, Some(wallet_address), - social_proof.clone(), user_id, ) .await?; @@ -205,7 +219,6 @@ flattened_device_key_upload: FlattenedDeviceKeyUpload, username_and_password_file: Option<(String, Blob)>, wallet_address: Option, - social_proof: Option, user_id: Option, ) -> Result { let user_id = user_id.unwrap_or_else(generate_uuid); @@ -255,61 +268,11 @@ Ok(user_id) } - pub async fn add_password_user_device_to_users_table( - &self, - user_id: String, - flattened_device_key_upload: FlattenedDeviceKeyUpload, - code_version: u64, - access_token_creation_time: DateTime, - ) -> Result<(), Error> { - let content_one_time_keys = - flattened_device_key_upload.content_one_time_keys.clone(); - let notif_one_time_keys = - flattened_device_key_upload.notif_one_time_keys.clone(); - - // add device to the device list if not exists - let device_id = flattened_device_key_upload.device_id_key.clone(); - let device_exists = self - .device_exists(user_id.clone(), device_id.clone()) - .await?; - - if device_exists { - self - .update_device_login_time( - user_id.clone(), - device_id, - access_token_creation_time, - ) - .await?; - return Ok(()); - } - - self - .add_device( - user_id, - flattened_device_key_upload, - None, - code_version, - access_token_creation_time, - ) - .await?; - - self - .append_one_time_prekeys( - device_id, - content_one_time_keys, - notif_one_time_keys, - ) - .await?; - - Ok(()) - } - - pub async fn add_wallet_user_device_to_users_table( + pub async fn add_user_device( &self, user_id: String, flattened_device_key_upload: FlattenedDeviceKeyUpload, - social_proof: String, + social_proof: Option, code_version: u64, access_token_creation_time: DateTime, ) -> Result<(), Error> { @@ -340,7 +303,7 @@ .add_device( user_id, flattened_device_key_upload, - Some(social_proof), + social_proof, code_version, access_token_creation_time, ) @@ -408,14 +371,8 @@ key_payload: keyserver.device_key_info.key_payload, key_payload_signature: keyserver.device_key_info.key_payload_signature, social_proof: keyserver.device_key_info.social_proof, - content_prekey: PreKey { - prekey: keyserver.content_prekey.pre_key, - prekey_signature: keyserver.content_prekey.pre_key_signature, - }, - notif_prekey: PreKey { - prekey: keyserver.notif_prekey.pre_key, - prekey_signature: keyserver.notif_prekey.pre_key_signature, - }, + content_prekey: keyserver.content_prekey, + notif_prekey: keyserver.notif_prekey, content_one_time_key, notif_one_time_key, }; @@ -534,27 +491,6 @@ }) } - 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> { - self - .update_device_prekeys( - user_id, - device_id, - content_prekey, - content_prekey_signature, - notif_prekey, - notif_prekey_signature, - ) - .await - } - pub async fn append_one_time_prekeys( &self, device_id: String, @@ -589,14 +525,6 @@ Ok(()) } - pub async fn remove_device_from_users_table( - &self, - user_id: String, - device_id_key: String, - ) -> Result<(), Error> { - self.remove_device(&user_id, &device_id_key).await - } - pub async fn update_user_password( &self, user_id: String, @@ -916,25 +844,16 @@ } } - pub async fn get_keys_for_user_id( + pub async fn get_keys_for_user( &self, user_id: &str, get_one_time_keys: bool, ) -> Result, Error> { - let Some(user) = self.get_item_from_users_table(user_id).await?.item else { - return Ok(None); - }; - - self.get_keys_for_user(user, get_one_time_keys).await - } - - async fn get_keys_for_user( - &self, - mut user: AttributeMap, - get_one_time_keys: bool, - ) -> Result, Error> { - let user_id: String = user.take_attr(USERS_TABLE_PARTITION_KEY)?; let mut devices_response = self.get_keys_for_user_devices(user_id).await?; + if devices_response.is_empty() { + debug!("No devices found for user {}", user_id); + return Ok(None); + } if get_one_time_keys { for (device_id_key, device_info_map) in devices_response.iter_mut() { diff --git a/services/identity/src/database/device_list.rs b/services/identity/src/database/device_list.rs --- a/services/identity/src/database/device_list.rs +++ b/services/identity/src/database/device_list.rs @@ -32,7 +32,7 @@ }, ddb_utils::AttributesOptionExt, error::{DeviceListError, Error, FromAttributeValue}, - grpc_services::protos::unauth::DeviceType, + grpc_services::protos::{self, unauth::DeviceType}, }; use super::DatabaseClient; @@ -324,6 +324,24 @@ } } +impl From for protos::unauth::Prekey { + fn from(value: PreKey) -> Self { + Self { + prekey: value.pre_key, + prekey_signature: value.pre_key_signature, + } + } +} + +impl From for PreKey { + fn from(value: protos::unauth::Prekey) -> Self { + Self { + pre_key: value.prekey, + pre_key_signature: value.prekey_signature, + } + } +} + impl TryFrom for PreKey { type Error = DBItemError; fn try_from(mut attrs: AttributeMap) -> Result { @@ -536,20 +554,9 @@ &self, user_id: impl Into, device_id: impl Into, - content_prekey: String, - content_prekey_signature: String, - notif_prekey: String, - notif_prekey_signature: String, + content_prekey: PreKey, + notif_prekey: PreKey, ) -> Result<(), Error> { - let content_prekey = PreKey { - pre_key: content_prekey, - pre_key_signature: content_prekey_signature, - }; - let notif_prekey = PreKey { - pre_key: notif_prekey, - pre_key_signature: notif_prekey_signature, - }; - self .client .update_item() 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 @@ -29,7 +29,7 @@ UpdateUserPasswordFinishRequest, UpdateUserPasswordStartRequest, UpdateUserPasswordStartResponse, UploadOneTimeKeysRequest, }; -use super::protos::unauth::{Empty, IdentityKeyInfo, Prekey}; +use super::protos::unauth::Empty; #[derive(derive_more::Constructor)] pub struct AuthenticatedService { @@ -109,13 +109,11 @@ self .db_client - .set_prekey( + .update_device_prekeys( user_id, device_id, - content_keys.prekey, - content_keys.prekey_signature, - notif_keys.prekey, - notif_keys.prekey_signature, + content_keys.into(), + notif_keys.into(), ) .await .map_err(handle_db_error)?; @@ -132,7 +130,7 @@ let devices_map = self .db_client - .get_keys_for_user_id(&message.user_id, true) + .get_keys_for_user(&message.user_id, true) .await .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; @@ -169,7 +167,7 @@ let devices_map = self .db_client - .get_keys_for_user_id(&message.user_id, false) + .get_keys_for_user(&message.user_id, false) .await .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; @@ -220,23 +218,7 @@ .get_keyserver_keys_for_user(&message.user_id) .await .map_err(handle_db_error)? - .map(|db_keys| OutboundKeyInfo { - identity_info: Some(IdentityKeyInfo { - payload: db_keys.key_payload, - payload_signature: db_keys.key_payload_signature, - social_proof: db_keys.social_proof, - }), - content_prekey: Some(Prekey { - prekey: db_keys.content_prekey.prekey, - prekey_signature: db_keys.content_prekey.prekey_signature, - }), - notif_prekey: Some(Prekey { - prekey: db_keys.notif_prekey.prekey, - prekey_signature: db_keys.notif_prekey.prekey_signature, - }), - one_time_content_prekey: db_keys.content_one_time_key, - one_time_notif_prekey: db_keys.notif_one_time_key, - }); + .map(OutboundKeyInfo::from); let identifier = self .db_client @@ -378,7 +360,7 @@ self .db_client - .remove_device_from_users_table(user_id.clone(), device_id.clone()) + .remove_device(&user_id, &device_id) .await .map_err(handle_db_error)?;