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 @@ -12,19 +12,19 @@ // Workspace crate imports use crate::client_service::client_proto::{ - inbound_keys_for_user_request::Identifier, AddReservedUsernamesRequest, - DeleteUserRequest, Empty, GenerateNonceResponse, InboundKeyInfo, - InboundKeysForUserRequest, InboundKeysForUserResponse, LogoutRequest, - OpaqueLoginFinishRequest, OpaqueLoginFinishResponse, OpaqueLoginStartRequest, - OpaqueLoginStartResponse, OutboundKeysForUserRequest, - OutboundKeysForUserResponse, RefreshUserPreKeysRequest, - RegistrationFinishRequest, RegistrationFinishResponse, - RegistrationStartRequest, RegistrationStartResponse, - RemoveReservedUsernameRequest, ReservedRegistrationStartRequest, - UpdateUserPasswordFinishRequest, UpdateUserPasswordStartRequest, - UpdateUserPasswordStartResponse, UploadOneTimeKeysRequest, - VerifyUserAccessTokenRequest, VerifyUserAccessTokenResponse, - WalletLoginRequest, WalletLoginResponse, + inbound_keys_for_user_request, outbound_keys_for_user_request, + AddReservedUsernamesRequest, DeleteUserRequest, Empty, GenerateNonceResponse, + InboundKeyInfo, InboundKeysForUserRequest, InboundKeysForUserResponse, + LogoutRequest, OpaqueLoginFinishRequest, OpaqueLoginFinishResponse, + OpaqueLoginStartRequest, OpaqueLoginStartResponse, OutboundKeyInfo, + OutboundKeysForUserRequest, OutboundKeysForUserResponse, + RefreshUserPreKeysRequest, RegistrationFinishRequest, + RegistrationFinishResponse, RegistrationStartRequest, + RegistrationStartResponse, RemoveReservedUsernameRequest, + ReservedRegistrationStartRequest, UpdateUserPasswordFinishRequest, + UpdateUserPasswordStartRequest, UpdateUserPasswordStartResponse, + UploadOneTimeKeysRequest, VerifyUserAccessTokenRequest, + VerifyUserAccessTokenResponse, WalletLoginRequest, WalletLoginResponse, }; use crate::config::CONFIG; use crate::database::{DatabaseClient, Device, KeyPayload}; @@ -786,9 +786,51 @@ async fn get_outbound_keys_for_user( &self, - _request: tonic::Request, + request: tonic::Request, ) -> Result, tonic::Status> { - unimplemented!(); + let message = request.into_inner(); + + use outbound_keys_for_user_request::Identifier; + let (user_ident, auth_type) = match message.identifier { + None => { + return Err(tonic::Status::invalid_argument("no identifier provided")) + } + Some(Identifier::Username(username)) => (username, AuthType::Password), + Some(Identifier::WalletAddress(address)) => (address, AuthType::Wallet), + }; + + let devices_map = self + .client + .get_keys_for_user(user_ident, &auth_type, true) + .await + .map_err(handle_db_error)? + .ok_or_else(|| match auth_type { + AuthType::Password => tonic::Status::not_found("username not found"), + AuthType::Wallet => { + tonic::Status::not_found("wallet address not found") + } + })?; + + let transformed_devices = devices_map + .into_iter() + .filter_map(|(key, device_info)| { + let device_info_with_auth = DeviceInfoWithAuth { + device_info, + auth_type: &auth_type, + }; + match OutboundKeyInfo::try_from(device_info_with_auth) { + Ok(key_info) => Some((key, key_info)), + Err(_) => { + error!("Failed to transform device info for key {}", key); + None + } + } + }) + .collect::>(); + + Ok(tonic::Response::new(OutboundKeysForUserResponse { + devices: transformed_devices, + })) } async fn get_inbound_keys_for_user( @@ -797,6 +839,7 @@ ) -> Result, tonic::Status> { let message = request.into_inner(); + use inbound_keys_for_user_request::Identifier; let (user_ident, auth_type) = match message.identifier { None => { return Err(tonic::Status::invalid_argument("no identifier provided")) @@ -807,7 +850,7 @@ let devices_map = self .client - .get_keys_for_user(user_ident, &auth_type) + .get_keys_for_user(user_ident, &auth_type, false) .await .map_err(handle_db_error)? .ok_or_else(|| match auth_type { 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 @@ -100,6 +100,10 @@ pub const ONE_TIME_KEY: &str = SORT_KEY; } +// One-time key constants for device info map +pub const CONTENT_ONE_TIME_KEY: &str = "contentOneTimeKey"; +pub const NOTIF_ONE_TIME_KEY: &str = "notifOneTimeKey"; + // Tokio pub const MPSC_CHANNEL_BUFFER_CAPACITY: usize = 1; 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 @@ -24,11 +24,12 @@ ACCESS_TOKEN_SORT_KEY, ACCESS_TOKEN_TABLE, ACCESS_TOKEN_TABLE_AUTH_TYPE_ATTRIBUTE, ACCESS_TOKEN_TABLE_CREATED_ATTRIBUTE, ACCESS_TOKEN_TABLE_PARTITION_KEY, ACCESS_TOKEN_TABLE_TOKEN_ATTRIBUTE, - ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE, NONCE_TABLE, + ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE, CONTENT_ONE_TIME_KEY, NONCE_TABLE, NONCE_TABLE_CREATED_ATTRIBUTE, NONCE_TABLE_EXPIRATION_TIME_ATTRIBUTE, NONCE_TABLE_EXPIRATION_TIME_UNIX_ATTRIBUTE, NONCE_TABLE_PARTITION_KEY, - RESERVED_USERNAMES_TABLE, RESERVED_USERNAMES_TABLE_PARTITION_KEY, - USERS_TABLE, USERS_TABLE_DEVICES_ATTRIBUTE, + NOTIF_ONE_TIME_KEY, RESERVED_USERNAMES_TABLE, + RESERVED_USERNAMES_TABLE_PARTITION_KEY, USERS_TABLE, + USERS_TABLE_DEVICES_ATTRIBUTE, USERS_TABLE_DEVICES_MAP_CONTENT_ONE_TIME_KEYS_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME, @@ -325,12 +326,12 @@ .await?; debug!( - "Able to get notif key for keyserver {}: {}", + "Able to get notif one-time key for keyserver {}: {}", keyserver_id, notif_one_time_key.is_some() ); debug!( - "Able to get content key for keyserver {}: {}", + "Able to get content one-time key for keyserver {}: {}", keyserver_id, content_one_time_key.is_some() ); @@ -382,8 +383,8 @@ Ok(Some(outbound_payload)) } - /// Will "mint" a single one time key by attempting to successfully deleting - /// a key + /// Will "mint" a single one-time key by attempting to successfully delete a + /// key pub async fn get_one_time_key( &self, device_id: &str, @@ -395,14 +396,8 @@ let query_result = self.get_one_time_keys(device_id, account_type).await?; let items = query_result.items(); - // If no one time keys exists, return none early - let Some(item_vec) = items else { - debug!("Unable to find {:?} one time key", account_type); - return Ok(None); - }; - - if item_vec.len() < ONE_TIME_KEY_MINIMUM_THRESHOLD { - // Avoid device_id being moved out-of-scope by "move" + fn spawn_refresh_keys_task(device_id: &str) { + // Clone the string slice to move into the async block let device_id = device_id.to_string(); tokio::spawn(async move { debug!("Attempting to request more keys for device: {}", &device_id); @@ -412,9 +407,22 @@ }); } + // If no one-time keys exist, or if there aren't enough, request more. + // Additionally, if no one-time keys exist, return early. + let item_vec = if let Some(items_list) = items { + if items_list.len() < ONE_TIME_KEY_MINIMUM_THRESHOLD { + spawn_refresh_keys_task(device_id); + } + items_list + } else { + debug!("Unable to find {:?} one-time key", account_type); + spawn_refresh_keys_task(device_id); + return Ok(None); + }; + let mut result = None; - // Attempt to delete the one time keys individually, a successful delete - // mints the one time key to the requester + // Attempt to delete the one-time keys individually, a successful delete + // mints the one-time key to the requester for item in item_vec { let pk = item.get_string(otk_table::PARTITION_KEY)?; let otk = item.get_string(otk_table::SORT_KEY)?; @@ -956,6 +964,7 @@ &self, user_info: String, auth_type: &AuthType, + get_one_time_keys: bool, ) -> Result, Error> { let Some(mut user) = self.get_user_from_user_info(user_info, auth_type).await? @@ -989,6 +998,23 @@ device_info_string_map.insert(attribute_name, attribute_value_str); } + if get_one_time_keys { + if let Some(notif_one_time_key) = self + .get_one_time_key(&device_id_key, OlmAccountType::Notification) + .await? + { + device_info_string_map + .insert(NOTIF_ONE_TIME_KEY.to_string(), notif_one_time_key); + } + if let Some(content_one_time_key) = self + .get_one_time_key(&device_id_key, OlmAccountType::Content) + .await? + { + device_info_string_map + .insert(CONTENT_ONE_TIME_KEY.to_string(), content_one_time_key); + } + } + devices_response.insert(device_id_key, device_info_string_map); } diff --git a/services/identity/src/grpc_utils.rs b/services/identity/src/grpc_utils.rs --- a/services/identity/src/grpc_utils.rs +++ b/services/identity/src/grpc_utils.rs @@ -4,8 +4,11 @@ use tracing::error; use crate::{ - client_service::client_proto::{IdentityKeyInfo, InboundKeyInfo, PreKey}, + client_service::client_proto::{ + IdentityKeyInfo, InboundKeyInfo, OutboundKeyInfo, PreKey, + }, constants::{ + CONTENT_ONE_TIME_KEY, NOTIF_ONE_TIME_KEY, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME, @@ -24,53 +27,56 @@ } impl TryFrom> for InboundKeyInfo { - type Error = tonic::Status; + type Error = Status; fn try_from(data: DeviceInfoWithAuth) -> Result { let mut device_info = data.device_info; - let payload = extract_key( - &mut device_info, - USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME, - )?; - let payload_signature = extract_key( - &mut device_info, - USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME, - )?; - - let social_proof = - device_info.remove(USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME); - if social_proof.is_none() && data.auth_type == &AuthType::Wallet { - error!("Social proof missing for wallet user"); - return Err(tonic::Status::failed_precondition( - "Database item malformed", - )); - } - - let identity_info = IdentityKeyInfo { - payload, - payload_signature, - social_proof, - }; - - let mut create_prekey = - |key_attr, signature_attr| -> Result { - Ok(PreKey { - pre_key: extract_key(&mut device_info, key_attr)?, - pre_key_signature: extract_key(&mut device_info, signature_attr)?, - }) - }; + let identity_info = + extract_identity_info(&mut device_info, &data.auth_type)?; Ok(InboundKeyInfo { identity_info: Some(identity_info), content_prekey: Some(create_prekey( + &mut device_info, + USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME, + USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME, + )?), + notif_prekey: Some(create_prekey( + &mut device_info, + USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME, + USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME, + )?), + }) + } +} + +impl TryFrom> for OutboundKeyInfo { + type Error = Status; + + fn try_from(data: DeviceInfoWithAuth) -> Result { + let mut device_info = data.device_info; + + let identity_info = + extract_identity_info(&mut device_info, &data.auth_type)?; + + let content_one_time_key = device_info.remove(CONTENT_ONE_TIME_KEY); + let notif_one_time_key = device_info.remove(NOTIF_ONE_TIME_KEY); + + Ok(OutboundKeyInfo { + identity_info: Some(identity_info), + content_prekey: Some(create_prekey( + &mut device_info, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME, )?), notif_prekey: Some(create_prekey( + &mut device_info, USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME, )?), + one_time_content_prekey: content_one_time_key, + one_time_notif_prekey: notif_one_time_key, }) } } @@ -84,3 +90,40 @@ Status::failed_precondition("Database item malformed") }) } + +fn extract_identity_info( + device_info: &mut HashMap, + auth_type: &AuthType, +) -> Result { + let payload = extract_key( + device_info, + USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME, + )?; + let payload_signature = extract_key( + device_info, + USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME, + )?; + let social_proof = + device_info.remove(USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME); + if social_proof.is_none() && auth_type == &AuthType::Wallet { + error!("Social proof missing for wallet user"); + return Err(Status::failed_precondition("Database item malformed")); + } + + Ok(IdentityKeyInfo { + payload, + payload_signature, + social_proof, + }) +} + +fn create_prekey( + device_info: &mut HashMap, + key_attr: &str, + signature_attr: &str, +) -> Result { + Ok(PreKey { + pre_key: extract_key(device_info, key_attr)?, + pre_key_signature: extract_key(device_info, signature_attr)?, + }) +}