diff --git a/services/identity/src/client_service.rs b/services/identity/src/client_service.rs index c98e801a7..a257f6e2b 100644 --- a/services/identity/src/client_service.rs +++ b/services/identity/src/client_service.rs @@ -1,1171 +1,1177 @@ // Standard library imports use std::str::FromStr; // External crate imports use comm_lib::aws::DynamoDBError; use comm_lib::shared::reserved_users::RESERVED_USERNAME_SET; use comm_opaque2::grpc::protocol_error_to_grpc_status; use rand::rngs::OsRng; use serde::{Deserialize, Serialize}; use siwe::eip55; use tonic::Response; use tracing::{debug, error, info, warn}; // Workspace crate imports use crate::config::CONFIG; -use crate::constants::request_metadata; +use crate::constants::{ + request_metadata, + error_types +}; use crate::database::{ DBDeviceTypeInt, DatabaseClient, DeviceType, KeyPayload, }; use crate::error::{DeviceListError, Error as DBError}; use crate::grpc_services::authenticated::DeletePasswordUserInfo; use crate::grpc_services::protos::unauth::{ find_user_id_request, AddReservedUsernamesRequest, AuthResponse, Empty, ExistingDeviceLoginRequest, FindUserIdRequest, FindUserIdResponse, GenerateNonceResponse, OpaqueLoginFinishRequest, OpaqueLoginStartRequest, OpaqueLoginStartResponse, RegistrationFinishRequest, RegistrationStartRequest, RegistrationStartResponse, RemoveReservedUsernameRequest, ReservedRegistrationStartRequest, ReservedWalletRegistrationRequest, SecondaryDeviceKeysUploadRequest, VerifyUserAccessTokenRequest, VerifyUserAccessTokenResponse, WalletAuthRequest, GetFarcasterUsersRequest, GetFarcasterUsersResponse }; use crate::grpc_services::shared::get_value; use crate::grpc_utils::{ SignedNonce, DeviceKeyUploadActions, }; use crate::nonce::generate_nonce_data; use crate::reserved_users::{ validate_account_ownership_message_and_get_user_id, validate_add_reserved_usernames_message, validate_remove_reserved_username_message, }; use crate::siwe::{ is_valid_ethereum_address, parse_and_verify_siwe_message, SocialProof, }; use crate::token::{AccessTokenData, AuthType}; pub use crate::grpc_services::protos::unauth::identity_client_service_server::{ IdentityClientService, IdentityClientServiceServer, }; use crate::regex::is_valid_username; #[derive(Clone, Serialize, Deserialize)] pub enum WorkflowInProgress { Registration(Box), Login(Box), Update(UpdateState), PasswordUserDeletion(Box), } #[derive(Clone, Serialize, Deserialize)] pub struct UserRegistrationInfo { pub username: String, pub flattened_device_key_upload: FlattenedDeviceKeyUpload, pub user_id: Option, pub farcaster_id: Option, } #[derive(Clone, Serialize, Deserialize)] pub struct UserLoginInfo { pub user_id: String, pub flattened_device_key_upload: FlattenedDeviceKeyUpload, pub opaque_server_login: comm_opaque2::server::Login, pub device_to_remove: Option, } #[derive(Clone, Serialize, Deserialize)] pub struct UpdateState { pub user_id: String, } #[derive(Clone, Serialize, Deserialize)] pub struct FlattenedDeviceKeyUpload { pub device_id_key: String, pub key_payload: String, pub key_payload_signature: String, pub content_prekey: String, pub content_prekey_signature: String, pub content_one_time_keys: Vec, pub notif_prekey: String, pub notif_prekey_signature: String, pub notif_one_time_keys: Vec, pub device_type: DeviceType, } #[derive(derive_more::Constructor)] pub struct ClientService { client: DatabaseClient, } #[tonic::async_trait] impl IdentityClientService for ClientService { #[tracing::instrument(skip_all)] async fn register_password_user_start( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); debug!("Received registration request for: {}", message.username); if !is_valid_username(&message.username) || is_valid_ethereum_address(&message.username) { return Err(tonic::Status::invalid_argument("invalid username")); } self.check_username_taken(&message.username).await?; let username_in_reserved_usernames_table = self .client .username_in_reserved_usernames_table(&message.username) .await .map_err(handle_db_error)?; if username_in_reserved_usernames_table { return Err(tonic::Status::already_exists("username already exists")); } if RESERVED_USERNAME_SET.contains(&message.username) { return Err(tonic::Status::invalid_argument("username reserved")); } if let Some(fid) = &message.farcaster_id { self.check_farcaster_id_taken(fid).await?; } let registration_state = construct_user_registration_info( &message, None, message.username.clone(), message.farcaster_id.clone(), )?; let server_registration = comm_opaque2::server::Registration::new(); let server_message = server_registration .start( &CONFIG.server_setup, &message.opaque_registration_request, message.username.as_bytes(), ) .map_err(protocol_error_to_grpc_status)?; let session_id = self .client .insert_workflow(WorkflowInProgress::Registration(Box::new( registration_state, ))) .await .map_err(handle_db_error)?; let response = RegistrationStartResponse { session_id, opaque_registration_response: server_message, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn register_reserved_password_user_start( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); self.check_username_taken(&message.username).await?; if RESERVED_USERNAME_SET.contains(&message.username) { return Err(tonic::Status::invalid_argument("username reserved")); } let username_in_reserved_usernames_table = self .client .username_in_reserved_usernames_table(&message.username) .await .map_err(handle_db_error)?; if !username_in_reserved_usernames_table { return Err(tonic::Status::permission_denied("username not reserved")); } let user_id = validate_account_ownership_message_and_get_user_id( &message.username, &message.keyserver_message, &message.keyserver_signature, )?; let registration_state = construct_user_registration_info( &message, Some(user_id), message.username.clone(), None, )?; let server_registration = comm_opaque2::server::Registration::new(); let server_message = server_registration .start( &CONFIG.server_setup, &message.opaque_registration_request, message.username.as_bytes(), ) .map_err(protocol_error_to_grpc_status)?; let session_id = self .client .insert_workflow(WorkflowInProgress::Registration(Box::new( registration_state, ))) .await .map_err(handle_db_error)?; let response = RegistrationStartResponse { session_id, opaque_registration_response: server_message, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn register_password_user_finish( &self, request: tonic::Request, ) -> Result, tonic::Status> { let code_version = get_code_version(&request); let message = request.into_inner(); if let Some(WorkflowInProgress::Registration(state)) = self .client .get_workflow(message.session_id) .await .map_err(handle_db_error)? { let server_registration = comm_opaque2::server::Registration::new(); let password_file = server_registration .finish(&message.opaque_registration_upload) .map_err(protocol_error_to_grpc_status)?; let login_time = chrono::Utc::now(); let device_id = state.flattened_device_key_upload.device_id_key.clone(); let user_id = self .client .add_password_user_to_users_table( *state, password_file, code_version, login_time, ) .await .map_err(handle_db_error)?; // Create access token let token = AccessTokenData::with_created_time( user_id.clone(), device_id, login_time, crate::token::AuthType::Password, &mut OsRng, ); let access_token = token.access_token.clone(); self .client .put_access_token_data(token) .await .map_err(handle_db_error)?; let response = AuthResponse { user_id, access_token, }; Ok(Response::new(response)) } else { Err(tonic::Status::not_found("session not found")) } } #[tracing::instrument(skip_all)] async fn log_in_password_user_start( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); debug!("Attempting to log in user: {:?}", &message.username); let user_id_and_password_file = self .client .get_user_id_and_password_file_from_username(&message.username) .await .map_err(handle_db_error)?; let (user_id, password_file_bytes) = if let Some(data) = user_id_and_password_file { data } else { // It's possible that the user attempting login is already registered // on Ashoat's keyserver. If they are, we should send back a gRPC status // code instructing them to get a signed message from Ashoat's keyserver // in order to claim their username and register with the Identity // service. let username_in_reserved_usernames_table = self .client .username_in_reserved_usernames_table(&message.username) .await .map_err(handle_db_error)?; if username_in_reserved_usernames_table { return Err(tonic::Status::permission_denied( "need keyserver message to claim username", )); } return Err(tonic::Status::not_found("user not found")); }; let flattened_device_key_upload = construct_flattened_device_key_upload(&message)?; let maybe_device_to_remove = self .get_keyserver_device_to_remove( &user_id, &flattened_device_key_upload.device_id_key, message.force.unwrap_or(false), &flattened_device_key_upload.device_type, ) .await?; let mut server_login = comm_opaque2::server::Login::new(); let server_response = server_login .start( &CONFIG.server_setup, &password_file_bytes, &message.opaque_login_request, message.username.as_bytes(), ) .map_err(protocol_error_to_grpc_status)?; let login_state = construct_user_login_info( user_id, server_login, flattened_device_key_upload, maybe_device_to_remove, )?; let session_id = self .client .insert_workflow(WorkflowInProgress::Login(Box::new(login_state))) .await .map_err(handle_db_error)?; let response = Response::new(OpaqueLoginStartResponse { session_id, opaque_login_response: server_response, }); Ok(response) } #[tracing::instrument(skip_all)] async fn log_in_password_user_finish( &self, request: tonic::Request, ) -> Result, tonic::Status> { let code_version = get_code_version(&request); let message = request.into_inner(); if let Some(WorkflowInProgress::Login(state)) = self .client .get_workflow(message.session_id) .await .map_err(handle_db_error)? { let mut server_login = state.opaque_server_login.clone(); server_login .finish(&message.opaque_login_upload) .map_err(protocol_error_to_grpc_status)?; if let Some(device_to_remove) = state.device_to_remove { self .client .remove_device(state.user_id.clone(), device_to_remove) .await .map_err(handle_db_error)?; } let login_time = chrono::Utc::now(); self .client .add_user_device( state.user_id.clone(), state.flattened_device_key_upload.clone(), code_version, login_time, ) .await .map_err(handle_db_error)?; // Create access token let token = AccessTokenData::with_created_time( state.user_id.clone(), state.flattened_device_key_upload.device_id_key, login_time, crate::token::AuthType::Password, &mut OsRng, ); let access_token = token.access_token.clone(); self .client .put_access_token_data(token) .await .map_err(handle_db_error)?; let response = AuthResponse { user_id: state.user_id, access_token, }; Ok(Response::new(response)) } else { Err(tonic::Status::not_found("session not found")) } } #[tracing::instrument(skip_all)] async fn log_in_wallet_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let code_version = get_code_version(&request); let message = request.into_inner(); let parsed_message = parse_and_verify_siwe_message( &message.siwe_message, &message.siwe_signature, )?; self.verify_and_remove_nonce(&parsed_message.nonce).await?; let wallet_address = eip55(&parsed_message.address); let flattened_device_key_upload = construct_flattened_device_key_upload(&message)?; let login_time = chrono::Utc::now(); let Some(user_id) = self .client .get_user_id_from_user_info(wallet_address.clone(), &AuthType::Wallet) .await .map_err(handle_db_error)? else { // It's possible that the user attempting login is already registered // on Ashoat's keyserver. If they are, we should send back a gRPC status // code instructing them to get a signed message from Ashoat's keyserver // in order to claim their wallet address and register with the Identity // service. let username_in_reserved_usernames_table = self .client .username_in_reserved_usernames_table(&wallet_address) .await .map_err(handle_db_error)?; if username_in_reserved_usernames_table { return Err(tonic::Status::permission_denied( "need keyserver message to claim username", )); } return Err(tonic::Status::not_found("user not found")); }; self .client .add_user_device( user_id.clone(), flattened_device_key_upload.clone(), code_version, chrono::Utc::now(), ) .await .map_err(handle_db_error)?; // Create access token let token = AccessTokenData::with_created_time( user_id.clone(), flattened_device_key_upload.device_id_key, login_time, crate::token::AuthType::Wallet, &mut OsRng, ); let access_token = token.access_token.clone(); self .client .put_access_token_data(token) .await .map_err(handle_db_error)?; let response = AuthResponse { user_id, access_token, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn register_wallet_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let code_version = get_code_version(&request); let message = request.into_inner(); let parsed_message = parse_and_verify_siwe_message( &message.siwe_message, &message.siwe_signature, )?; match self .client .get_nonce_from_nonces_table(&parsed_message.nonce) .await .map_err(handle_db_error)? { None => return Err(tonic::Status::invalid_argument("invalid nonce")), Some(nonce) if nonce.is_expired() => { // we don't need to remove the nonce from the table here // because the DynamoDB TTL will take care of it return Err(tonic::Status::aborted("nonce expired")); } Some(_) => self .client .remove_nonce_from_nonces_table(&parsed_message.nonce) .await .map_err(handle_db_error)?, }; let wallet_address = eip55(&parsed_message.address); self.check_wallet_address_taken(&wallet_address).await?; let username_in_reserved_usernames_table = self .client .username_in_reserved_usernames_table(&wallet_address) .await .map_err(handle_db_error)?; if username_in_reserved_usernames_table { return Err(tonic::Status::already_exists( "wallet address already exists", )); } if let Some(fid) = &message.farcaster_id { self.check_farcaster_id_taken(fid).await?; } let flattened_device_key_upload = construct_flattened_device_key_upload(&message)?; let login_time = chrono::Utc::now(); let social_proof = SocialProof::new(message.siwe_message, message.siwe_signature); let user_id = self .client .add_wallet_user_to_users_table( flattened_device_key_upload.clone(), wallet_address, social_proof, None, code_version, login_time, message.farcaster_id, ) .await .map_err(handle_db_error)?; // Create access token let token = AccessTokenData::with_created_time( user_id.clone(), flattened_device_key_upload.device_id_key, login_time, crate::token::AuthType::Wallet, &mut OsRng, ); let access_token = token.access_token.clone(); self .client .put_access_token_data(token) .await .map_err(handle_db_error)?; let response = AuthResponse { user_id, access_token, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn register_reserved_wallet_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let code_version = get_code_version(&request); let message = request.into_inner(); let parsed_message = parse_and_verify_siwe_message( &message.siwe_message, &message.siwe_signature, )?; self.verify_and_remove_nonce(&parsed_message.nonce).await?; let wallet_address = eip55(&parsed_message.address); self.check_wallet_address_taken(&wallet_address).await?; let wallet_address_in_reserved_usernames_table = self .client .username_in_reserved_usernames_table(&wallet_address) .await .map_err(handle_db_error)?; if !wallet_address_in_reserved_usernames_table { return Err(tonic::Status::permission_denied( "wallet address not reserved", )); } let user_id = validate_account_ownership_message_and_get_user_id( &wallet_address, &message.keyserver_message, &message.keyserver_signature, )?; let flattened_device_key_upload = construct_flattened_device_key_upload(&message)?; let social_proof = SocialProof::new(message.siwe_message, message.siwe_signature); let login_time = chrono::Utc::now(); self .client .add_wallet_user_to_users_table( flattened_device_key_upload.clone(), wallet_address, social_proof, Some(user_id.clone()), code_version, login_time, None, ) .await .map_err(handle_db_error)?; let token = AccessTokenData::with_created_time( user_id.clone(), flattened_device_key_upload.device_id_key, login_time, crate::token::AuthType::Wallet, &mut OsRng, ); let access_token = token.access_token.clone(); self .client .put_access_token_data(token) .await .map_err(handle_db_error)?; let response = AuthResponse { user_id, access_token, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn upload_keys_for_registered_device_and_log_in( &self, request: tonic::Request, ) -> Result, tonic::Status> { let code_version = get_code_version(&request); let message = request.into_inner(); let challenge_response = SignedNonce::try_from(&message)?; let flattened_device_key_upload = construct_flattened_device_key_upload(&message)?; let user_id = message.user_id; let device_id = flattened_device_key_upload.device_id_key.clone(); let nonce = challenge_response.verify_and_get_nonce(&device_id)?; self.verify_and_remove_nonce(&nonce).await?; let user_identifier = self .client .get_user_identifier(&user_id) .await .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; let Some(device_list) = self .client .get_current_device_list(&user_id) .await .map_err(handle_db_error)? else { warn!("User {} does not have valid device list. Secondary device auth impossible.", user_id); return Err(tonic::Status::aborted("device list error")); }; if !device_list.device_ids.contains(&device_id) { return Err(tonic::Status::permission_denied( "device not in device list", )); } let login_time = chrono::Utc::now(); let token = AccessTokenData::with_created_time( user_id.clone(), device_id, login_time, user_identifier.into(), &mut OsRng, ); let access_token = token.access_token.clone(); self .client .put_access_token_data(token) .await .map_err(handle_db_error)?; self .client .put_device_data( &user_id, flattened_device_key_upload, code_version, login_time, ) .await .map_err(handle_db_error)?; let response = AuthResponse { user_id, access_token, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn log_in_existing_device( &self, request: tonic::Request, ) -> std::result::Result, tonic::Status> { let message = request.into_inner(); let challenge_response = SignedNonce::try_from(&message)?; let ExistingDeviceLoginRequest { user_id, device_id, .. } = message; let nonce = challenge_response.verify_and_get_nonce(&device_id)?; self.verify_and_remove_nonce(&nonce).await?; let (identifier_response, device_list_response) = tokio::join!( self.client.get_user_identifier(&user_id), self.client.get_current_device_list(&user_id) ); let user_identifier = identifier_response .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; let device_list = device_list_response .map_err(handle_db_error)? .ok_or_else(|| { warn!("User {} does not have a valid device list.", user_id); tonic::Status::aborted("device list error") })?; if !device_list.device_ids.contains(&device_id) { return Err(tonic::Status::permission_denied( "device not in device list", )); } let login_time = chrono::Utc::now(); let token = AccessTokenData::with_created_time( user_id.clone(), device_id, login_time, user_identifier.into(), &mut OsRng, ); let access_token = token.access_token.clone(); self .client .put_access_token_data(token) .await .map_err(handle_db_error)?; let response = AuthResponse { user_id, access_token, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn generate_nonce( &self, _request: tonic::Request, ) -> Result, tonic::Status> { let nonce_data = generate_nonce_data(&mut OsRng); match self .client .add_nonce_to_nonces_table(nonce_data.clone()) .await { Ok(_) => Ok(Response::new(GenerateNonceResponse { nonce: nonce_data.nonce, })), Err(e) => Err(handle_db_error(e)), } } #[tracing::instrument(skip_all)] async fn verify_user_access_token( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); debug!("Verifying device: {}", &message.device_id); let token_valid = self .client .verify_access_token( message.user_id, message.device_id.clone(), message.access_token, ) .await .map_err(handle_db_error)?; let response = Response::new(VerifyUserAccessTokenResponse { token_valid }); debug!( "device {} was verified: {}", &message.device_id, token_valid ); Ok(response) } #[tracing::instrument(skip_all)] async fn add_reserved_usernames( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let user_details = validate_add_reserved_usernames_message( &message.message, &message.signature, )?; let filtered_user_details = self .client .filter_out_taken_usernames(user_details) .await .map_err(handle_db_error)?; self .client .add_usernames_to_reserved_usernames_table(filtered_user_details) .await .map_err(handle_db_error)?; let response = Response::new(Empty {}); Ok(response) } #[tracing::instrument(skip_all)] async fn remove_reserved_username( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let username = validate_remove_reserved_username_message( &message.message, &message.signature, )?; self .client .delete_username_from_reserved_usernames_table(username) .await .map_err(handle_db_error)?; let response = Response::new(Empty {}); Ok(response) } #[tracing::instrument(skip_all)] async fn ping( &self, _request: tonic::Request, ) -> Result, tonic::Status> { let response = Response::new(Empty {}); Ok(response) } #[tracing::instrument(skip_all)] async fn find_user_id( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); use find_user_id_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 (is_reserved_result, user_id_result) = tokio::join!( self .client .username_in_reserved_usernames_table(&user_ident), self .client .get_user_id_from_user_info(user_ident.clone(), &auth_type), ); let is_reserved = is_reserved_result.map_err(handle_db_error)?; let user_id = user_id_result.map_err(handle_db_error)?; Ok(Response::new(FindUserIdResponse { user_id, is_reserved, })) } #[tracing::instrument(skip_all)] async fn get_farcaster_users( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let farcaster_users = self .client .get_farcaster_users(message.farcaster_ids) .await .map_err(handle_db_error)? .into_iter() .map(|d| d.0) .collect(); Ok(Response::new(GetFarcasterUsersResponse { farcaster_users })) } } impl ClientService { async fn check_username_taken( &self, username: &str, ) -> Result<(), tonic::Status> { let username_taken = self .client .username_taken(username.to_string()) .await .map_err(handle_db_error)?; if username_taken { return Err(tonic::Status::already_exists("username already exists")); } Ok(()) } async fn check_wallet_address_taken( &self, wallet_address: &str, ) -> Result<(), tonic::Status> { let wallet_address_taken = self .client .wallet_address_taken(wallet_address.to_string()) .await .map_err(handle_db_error)?; if wallet_address_taken { return Err(tonic::Status::already_exists( "wallet address already exists", )); } Ok(()) } async fn check_farcaster_id_taken( &self, farcaster_id: &str, ) -> Result<(), tonic::Status> { let fid_already_registered = !self .client .get_farcaster_users(vec![farcaster_id.to_string()]) .await .map_err(handle_db_error)? .is_empty(); if fid_already_registered { return Err(tonic::Status::already_exists( "farcaster ID already associated with different user", )); } Ok(()) } async fn verify_and_remove_nonce( &self, nonce: &str, ) -> Result<(), tonic::Status> { match self .client .get_nonce_from_nonces_table(nonce) .await .map_err(handle_db_error)? { None => return Err(tonic::Status::invalid_argument("invalid nonce")), Some(nonce) if nonce.is_expired() => { // we don't need to remove the nonce from the table here // because the DynamoDB TTL will take care of it return Err(tonic::Status::aborted("nonce expired")); } Some(nonce_data) => self .client .remove_nonce_from_nonces_table(&nonce_data.nonce) .await .map_err(handle_db_error)?, }; Ok(()) } async fn get_keyserver_device_to_remove( &self, user_id: &str, new_keyserver_device_id: &str, force: bool, device_type: &DeviceType, ) -> Result, tonic::Status> { if device_type != &DeviceType::Keyserver { return Ok(None); } let maybe_keyserver_device_id = self .client .get_keyserver_device_id_for_user(user_id) .await .map_err(handle_db_error)?; let Some(existing_keyserver_device_id) = maybe_keyserver_device_id else { return Ok(None); }; if new_keyserver_device_id == existing_keyserver_device_id { return Ok(None); } if force { info!( "keyserver {} will be removed from the device list", existing_keyserver_device_id ); Ok(Some(existing_keyserver_device_id)) } else { Err(tonic::Status::already_exists( "user already has a keyserver", )) } } } #[tracing::instrument(skip_all)] pub fn handle_db_error(db_error: DBError) -> tonic::Status { match db_error { DBError::AwsSdk(DynamoDBError::InternalServerError(_)) | DBError::AwsSdk(DynamoDBError::ProvisionedThroughputExceededException( _, )) | DBError::AwsSdk(DynamoDBError::RequestLimitExceeded(_)) => { tonic::Status::unavailable("please retry") } DBError::DeviceList(DeviceListError::InvalidDeviceListUpdate) => { tonic::Status::invalid_argument("invalid device list update") } e => { - error!("Encountered an unexpected error: {}", e); + error!( + errorType = error_types::GENERIC_DB_LOG, + "Encountered an unexpected error: {}", e + ); tonic::Status::failed_precondition("unexpected error") } } } fn construct_user_registration_info( message: &impl DeviceKeyUploadActions, user_id: Option, username: String, farcaster_id: Option, ) -> Result { Ok(UserRegistrationInfo { username, flattened_device_key_upload: construct_flattened_device_key_upload( message, )?, user_id, farcaster_id, }) } fn construct_user_login_info( user_id: String, opaque_server_login: comm_opaque2::server::Login, flattened_device_key_upload: FlattenedDeviceKeyUpload, device_to_remove: Option, ) -> Result { Ok(UserLoginInfo { user_id, flattened_device_key_upload, opaque_server_login, device_to_remove, }) } fn construct_flattened_device_key_upload( message: &impl DeviceKeyUploadActions, ) -> Result { let key_info = KeyPayload::from_str(&message.payload()?) .map_err(|_| tonic::Status::invalid_argument("malformed payload"))?; let flattened_device_key_upload = FlattenedDeviceKeyUpload { device_id_key: key_info.primary_identity_public_keys.ed25519, key_payload: message.payload()?, key_payload_signature: message.payload_signature()?, content_prekey: message.content_prekey()?, content_prekey_signature: message.content_prekey_signature()?, content_one_time_keys: message.one_time_content_prekeys()?, notif_prekey: message.notif_prekey()?, notif_prekey_signature: message.notif_prekey_signature()?, notif_one_time_keys: message.one_time_notif_prekeys()?, device_type: DeviceType::try_from(DBDeviceTypeInt(message.device_type()?)) .map_err(handle_db_error)?, }; Ok(flattened_device_key_upload) } fn get_code_version(req: &tonic::Request) -> u64 { get_value(req, request_metadata::CODE_VERSION) .and_then(|version| version.parse().ok()) .unwrap_or_else(|| { warn!( "Could not retrieve code version from request: {:?}. Defaulting to 0", req ); Default::default() }) } diff --git a/services/identity/src/constants.rs b/services/identity/src/constants.rs index 268984977..418fbbd34 100644 --- a/services/identity/src/constants.rs +++ b/services/identity/src/constants.rs @@ -1,263 +1,279 @@ use tokio::time::Duration; // Secrets pub const SECRETS_DIRECTORY: &str = "secrets"; pub const SECRETS_SETUP_FILE: &str = "server_setup.txt"; // DynamoDB // User table information, supporting opaque_ke 2.0 and X3DH information // Users can sign in either through username+password or Eth wallet. // // This structure should be aligned with the messages defined in // shared/protos/identity_unauthenticated.proto // // Structure for a user should be: // { // userID: String, // opaqueRegistrationData: Option, // username: Option, // walletAddress: Option, // devices: HashMap // } // // A device is defined as: // { // deviceType: String, # client or keyserver // keyPayload: String, // keyPayloadSignature: String, // identityPreKey: String, // identityPreKeySignature: String, // identityOneTimeKeys: Vec, // notifPreKey: String, // notifPreKeySignature: String, // notifOneTimeKeys: Vec, // socialProof: Option // } // } // // Additional context: // "devices" uses the signing public identity key of the device as a key for the devices map // "keyPayload" is a JSON encoded string containing identity and notif keys (both signature and verification) // if "deviceType" == "keyserver", then the device will not have any notif key information pub const USERS_TABLE: &str = "identity-users"; pub const USERS_TABLE_PARTITION_KEY: &str = "userID"; pub const USERS_TABLE_REGISTRATION_ATTRIBUTE: &str = "opaqueRegistrationData"; pub const USERS_TABLE_USERNAME_ATTRIBUTE: &str = "username"; pub const USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME: &str = "deviceType"; pub const USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE: &str = "walletAddress"; pub const USERS_TABLE_SOCIAL_PROOF_ATTRIBUTE_NAME: &str = "socialProof"; pub const USERS_TABLE_DEVICELIST_TIMESTAMP_ATTRIBUTE_NAME: &str = "deviceListTimestamp"; pub const USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME: &str = "farcasterID"; pub const USERS_TABLE_USERNAME_INDEX: &str = "username-index"; pub const USERS_TABLE_WALLET_ADDRESS_INDEX: &str = "walletAddress-index"; pub const USERS_TABLE_FARCASTER_ID_INDEX: &str = "farcasterID-index"; pub mod token_table { pub const NAME: &str = "identity-tokens"; pub const PARTITION_KEY: &str = "userID"; pub const SORT_KEY: &str = "signingPublicKey"; pub const ATTR_CREATED: &str = "created"; pub const ATTR_AUTH_TYPE: &str = "authType"; pub const ATTR_VALID: &str = "valid"; pub const ATTR_TOKEN: &str = "token"; } pub const NONCE_TABLE: &str = "identity-nonces"; pub const NONCE_TABLE_PARTITION_KEY: &str = "nonce"; pub const NONCE_TABLE_CREATED_ATTRIBUTE: &str = "created"; pub const NONCE_TABLE_EXPIRATION_TIME_ATTRIBUTE: &str = "expirationTime"; pub const NONCE_TABLE_EXPIRATION_TIME_UNIX_ATTRIBUTE: &str = "expirationTimeUnix"; pub const WORKFLOWS_IN_PROGRESS_TABLE: &str = "identity-workflows-in-progress"; pub const WORKFLOWS_IN_PROGRESS_PARTITION_KEY: &str = "id"; pub const WORKFLOWS_IN_PROGRESS_WORKFLOW_ATTRIBUTE: &str = "workflow"; pub const WORKFLOWS_IN_PROGRESS_TABLE_EXPIRATION_TIME_UNIX_ATTRIBUTE: &str = "expirationTimeUnix"; // Usernames reserved because they exist in Ashoat's keyserver already pub const RESERVED_USERNAMES_TABLE: &str = "identity-reserved-usernames"; pub const RESERVED_USERNAMES_TABLE_PARTITION_KEY: &str = "username"; pub const RESERVED_USERNAMES_TABLE_USER_ID_ATTRIBUTE: &str = "userID"; // Users table social proof attribute pub const SOCIAL_PROOF_MESSAGE_ATTRIBUTE: &str = "siweMessage"; pub const SOCIAL_PROOF_SIGNATURE_ATTRIBUTE: &str = "siweSignature"; pub mod devices_table { /// table name pub const NAME: &str = "identity-devices"; pub const TIMESTAMP_INDEX_NAME: &str = "deviceList-timestamp-index"; /// partition key pub const ATTR_USER_ID: &str = "userID"; /// sort key pub const ATTR_ITEM_ID: &str = "itemID"; // itemID prefixes (one shouldn't be a prefix of the other) pub const DEVICE_ITEM_KEY_PREFIX: &str = "device-"; pub const DEVICE_LIST_KEY_PREFIX: &str = "devicelist-"; // device-specific attrs pub const ATTR_DEVICE_TYPE: &str = "deviceType"; pub const ATTR_DEVICE_KEY_INFO: &str = "deviceKeyInfo"; pub const ATTR_CONTENT_PREKEY: &str = "contentPreKey"; pub const ATTR_NOTIF_PREKEY: &str = "notifPreKey"; // IdentityKeyInfo constants pub const ATTR_KEY_PAYLOAD: &str = "keyPayload"; pub const ATTR_KEY_PAYLOAD_SIGNATURE: &str = "keyPayloadSignature"; // PreKey constants pub const ATTR_PREKEY: &str = "preKey"; pub const ATTR_PREKEY_SIGNATURE: &str = "preKeySignature"; // device-list-specific attrs pub const ATTR_TIMESTAMP: &str = "timestamp"; pub const ATTR_DEVICE_IDS: &str = "deviceIDs"; // migration-specific attrs pub const ATTR_CODE_VERSION: &str = "codeVersion"; pub const ATTR_LOGIN_TIME: &str = "loginTime"; // one-time key constants pub const ATTR_CONTENT_OTK_COUNT: &str = "contentOTKCount"; pub const ATTR_NOTIF_OTK_COUNT: &str = "notifOTKCount"; } // One time keys table, which need to exist in their own table to ensure // atomicity of additions and removals pub mod one_time_keys_table { pub const NAME: &str = "identity-one-time-keys"; pub const PARTITION_KEY: &str = "userID#deviceID#olmAccount"; pub const SORT_KEY: &str = "timestamp#keyNumber"; pub const ATTR_ONE_TIME_KEY: &str = "oneTimeKey"; } // Tokio pub const MPSC_CHANNEL_BUFFER_CAPACITY: usize = 1; pub const IDENTITY_SERVICE_SOCKET_ADDR: &str = "[::]:50054"; pub const IDENTITY_SERVICE_WEBSOCKET_ADDR: &str = "[::]:51004"; pub const SOCKET_HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(3); // Token pub const ACCESS_TOKEN_LENGTH: usize = 512; // Temporary config pub const AUTH_TOKEN: &str = "COMM_IDENTITY_SERVICE_AUTH_TOKEN"; pub const KEYSERVER_PUBLIC_KEY: &str = "KEYSERVER_PUBLIC_KEY"; // Nonce pub const NONCE_LENGTH: usize = 17; pub const NONCE_TTL_DURATION: Duration = Duration::from_secs(120); // seconds // Device list pub const DEVICE_LIST_TIMESTAMP_VALID_FOR: Duration = Duration::from_secs(300); // Workflows in progress pub const WORKFLOWS_IN_PROGRESS_TTL_DURATION: Duration = Duration::from_secs(120); // Identity pub const DEFAULT_IDENTITY_ENDPOINT: &str = "http://localhost:50054"; // LocalStack pub const LOCALSTACK_ENDPOINT: &str = "LOCALSTACK_ENDPOINT"; // OPAQUE Server Setup pub const OPAQUE_SERVER_SETUP: &str = "OPAQUE_SERVER_SETUP"; // Identity Search pub const OPENSEARCH_ENDPOINT: &str = "OPENSEARCH_ENDPOINT"; pub const DEFAULT_OPENSEARCH_ENDPOINT: &str = "identity-search-domain.us-east-2.opensearch.localhost.localstack.cloud:4566"; pub const IDENTITY_SEARCH_INDEX: &str = "users"; pub const IDENTITY_SEARCH_RESULT_SIZE: u32 = 20; +// Log Error Types + +pub mod error_types { + pub const GENERIC_DB_LOG: &str = "DB Error"; + pub const OTK_DB_LOG: &str = "One-time Key DB Error"; + pub const DEVICE_LIST_DB_LOG: &str = "Device List DB Error"; + pub const TOKEN_DB_LOG: &str = "Token DB Error"; + pub const FARCASTER_DB_LOG: &str = "Farcaster DB Error"; + + pub const SYNC_LOG: &str = "Sync Error"; + pub const SEARCH_LOG: &str = "Search Error"; + pub const SIWE_LOG: &str = "SIWE Error"; + pub const GRPC_SERVICES_LOG: &str = "gRPC Services Error"; + pub const TUNNELBROKER_LOG: &str = "Tunnelbroker Error"; +} + // Tunnelbroker pub const TUNNELBROKER_GRPC_ENDPOINT: &str = "TUNNELBROKER_GRPC_ENDPOINT"; pub const DEFAULT_TUNNELBROKER_ENDPOINT: &str = "http://localhost:50051"; // X3DH key management // Threshold for requesting more one_time keys pub const ONE_TIME_KEY_MINIMUM_THRESHOLD: usize = 5; // Number of keys to be refreshed when below the threshold pub const ONE_TIME_KEY_REFRESH_NUMBER: u32 = 5; // Minimum supported code versions pub const MIN_SUPPORTED_NATIVE_VERSION: u64 = 270; // Request metadata pub mod request_metadata { pub const CODE_VERSION: &str = "code_version"; pub const DEVICE_TYPE: &str = "device_type"; pub const USER_ID: &str = "user_id"; pub const DEVICE_ID: &str = "device_id"; pub const ACCESS_TOKEN: &str = "access_token"; } // CORS pub mod cors { use std::time::Duration; pub const DEFAULT_MAX_AGE: Duration = Duration::from_secs(24 * 60 * 60); pub const DEFAULT_EXPOSED_HEADERS: [&str; 3] = ["grpc-status", "grpc-message", "grpc-status-details-bin"]; pub const DEFAULT_ALLOW_HEADERS: [&str; 9] = [ "x-grpc-web", "content-type", "x-user-agent", "grpc-timeout", super::request_metadata::CODE_VERSION, super::request_metadata::DEVICE_TYPE, super::request_metadata::USER_ID, super::request_metadata::DEVICE_ID, super::request_metadata::ACCESS_TOKEN, ]; pub const ALLOW_ORIGIN_LIST: &str = "ALLOW_ORIGIN_LIST"; } // Tracing pub const COMM_SERVICES_USE_JSON_LOGS: &str = "COMM_SERVICES_USE_JSON_LOGS"; // Regex pub const VALID_USERNAME_REGEX_STRING: &str = r"^[a-zA-Z0-9][a-zA-Z0-9-_]{0,190}$"; // Retry // TODO: Replace this with `ExponentialBackoffConfig` from `comm-lib` pub mod retry { pub const MAX_ATTEMPTS: usize = 8; pub const CONDITIONAL_CHECK_FAILED: &str = "ConditionalCheckFailed"; pub const TRANSACTION_CONFLICT: &str = "TransactionConflict"; } // One-time keys pub const ONE_TIME_KEY_UPLOAD_LIMIT_PER_ACCOUNT: usize = 24; pub const ONE_TIME_KEY_SIZE: usize = 43; // as defined in olm pub const MAX_ONE_TIME_KEYS: usize = 100; // as defined in olm diff --git a/services/identity/src/database.rs b/services/identity/src/database.rs index 4a17157ec..50dfa940c 100644 --- a/services/identity/src/database.rs +++ b/services/identity/src/database.rs @@ -1,1129 +1,1142 @@ use comm_lib::aws::ddb::{ operation::{ delete_item::DeleteItemOutput, get_item::GetItemOutput, put_item::PutItemOutput, query::QueryOutput, }, primitives::Blob, types::{AttributeValue, PutRequest, WriteRequest}, }; use comm_lib::aws::{AwsConfig, DynamoDBClient}; use comm_lib::database::{ AttributeExtractor, AttributeMap, DBItemAttributeError, DBItemError, TryFromAttribute, }; use std::collections::{HashMap, HashSet}; use std::str::FromStr; use std::sync::Arc; pub use crate::database::device_list::DeviceIDAttribute; pub use crate::database::one_time_keys::OTKRow; use crate::{ - constants::USERS_TABLE_SOCIAL_PROOF_ATTRIBUTE_NAME, + constants::error_types, constants::USERS_TABLE_SOCIAL_PROOF_ATTRIBUTE_NAME, ddb_utils::EthereumIdentity, reserved_users::UserDetail, siwe::SocialProof, }; use crate::{ ddb_utils::{Identifier, OlmAccountType}, grpc_services::protos, }; use crate::{error::Error, grpc_utils::DeviceKeysInfo}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use tracing::{debug, error, info, warn}; use crate::client_service::{FlattenedDeviceKeyUpload, UserRegistrationInfo}; use crate::config::CONFIG; use crate::constants::{ 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, RESERVED_USERNAMES_TABLE_USER_ID_ATTRIBUTE, USERS_TABLE, USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME, USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME, USERS_TABLE_PARTITION_KEY, USERS_TABLE_REGISTRATION_ATTRIBUTE, USERS_TABLE_USERNAME_ATTRIBUTE, USERS_TABLE_USERNAME_INDEX, USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE, USERS_TABLE_WALLET_ADDRESS_INDEX, }; use crate::id::generate_uuid; use crate::nonce::NonceData; use crate::token::AuthType; pub use grpc_clients::identity::DeviceType; mod device_list; mod farcaster; mod one_time_keys; mod token; mod workflows; pub use device_list::{DeviceListRow, DeviceListUpdate, DeviceRow}; use self::device_list::Prekey; #[derive(Serialize, Deserialize)] pub struct OlmKeys { pub curve25519: String, pub ed25519: String, } #[derive(Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct KeyPayload { pub notification_identity_public_keys: OlmKeys, pub primary_identity_public_keys: OlmKeys, } impl FromStr for KeyPayload { type Err = serde_json::Error; // The payload is held in the database as an escaped JSON payload. // Escaped double quotes need to be trimmed before attempting to serialize fn from_str(payload: &str) -> Result { serde_json::from_str(&payload.replace(r#"\""#, r#"""#)) } } pub struct DBDeviceTypeInt(pub i32); impl TryFrom for DeviceType { type Error = crate::error::Error; fn try_from(value: DBDeviceTypeInt) -> Result { let device_result = DeviceType::try_from(value.0); device_result.map_err(|_| { Error::Attribute(DBItemError { attribute_name: USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME .to_string(), attribute_value: Some(AttributeValue::N(value.0.to_string())).into(), attribute_error: DBItemAttributeError::InvalidValue, }) }) } } pub struct OutboundKeys { pub key_payload: String, pub key_payload_signature: String, pub content_prekey: Prekey, pub notif_prekey: Prekey, pub content_one_time_key: Option, 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, }), 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, } impl DatabaseClient { pub fn new(aws_config: &AwsConfig) -> Self { let client = match &CONFIG.localstack_endpoint { Some(endpoint) => { info!( "Configuring DynamoDB client to use LocalStack endpoint: {}", endpoint ); let ddb_config_builder = comm_lib::aws::ddb::config::Builder::from(aws_config) .endpoint_url(endpoint); DynamoDBClient::from_conf(ddb_config_builder.build()) } None => DynamoDBClient::new(aws_config), }; DatabaseClient { client: Arc::new(client), } } pub async fn add_password_user_to_users_table( &self, registration_state: UserRegistrationInfo, password_file: Vec, code_version: u64, access_token_creation_time: DateTime, ) -> Result { let device_key_upload = registration_state.flattened_device_key_upload; let user_id = self .add_user_to_users_table( Some((registration_state.username, Blob::new(password_file))), None, registration_state.user_id, registration_state.farcaster_id, ) .await?; self .add_device( &user_id, device_key_upload.clone(), code_version, access_token_creation_time, ) .await?; self .append_one_time_prekeys( &user_id, &device_key_upload.device_id_key, &device_key_upload.content_one_time_keys, &device_key_upload.notif_one_time_keys, ) .await?; Ok(user_id) } pub async fn add_wallet_user_to_users_table( &self, flattened_device_key_upload: FlattenedDeviceKeyUpload, wallet_address: String, social_proof: SocialProof, user_id: Option, code_version: u64, access_token_creation_time: DateTime, farcaster_id: Option, ) -> Result { let wallet_identity = EthereumIdentity { wallet_address, social_proof, }; let user_id = self .add_user_to_users_table( None, Some(wallet_identity), user_id, farcaster_id, ) .await?; self .add_device( &user_id, flattened_device_key_upload.clone(), code_version, access_token_creation_time, ) .await?; self .append_one_time_prekeys( &user_id, &flattened_device_key_upload.device_id_key, &flattened_device_key_upload.content_one_time_keys, &flattened_device_key_upload.notif_one_time_keys, ) .await?; Ok(user_id) } async fn add_user_to_users_table( &self, username_and_password_file: Option<(String, Blob)>, wallet_identity: Option, user_id: Option, farcaster_id: Option, ) -> Result { let user_id = user_id.unwrap_or_else(generate_uuid); let mut user = HashMap::from([( USERS_TABLE_PARTITION_KEY.to_string(), AttributeValue::S(user_id.clone()), )]); if let Some((username, password_file)) = username_and_password_file { user.insert( USERS_TABLE_USERNAME_ATTRIBUTE.to_string(), AttributeValue::S(username), ); user.insert( USERS_TABLE_REGISTRATION_ATTRIBUTE.to_string(), AttributeValue::B(password_file), ); } if let Some(eth_identity) = wallet_identity { user.insert( USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE.to_string(), AttributeValue::S(eth_identity.wallet_address), ); user.insert( USERS_TABLE_SOCIAL_PROOF_ATTRIBUTE_NAME.to_string(), eth_identity.social_proof.into(), ); } if let Some(fid) = farcaster_id { user.insert( USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME.to_string(), AttributeValue::S(fid), ); } self .client .put_item() .table_name(USERS_TABLE) .set_item(Some(user)) // make sure we don't accidentaly overwrite existing row .condition_expression("attribute_not_exists(#pk)") .expression_attribute_names("#pk", USERS_TABLE_PARTITION_KEY) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; Ok(user_id) } pub async fn add_user_device( &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(()); } // add device to the new device list self .add_device( &user_id, flattened_device_key_upload, code_version, access_token_creation_time, ) .await?; self .append_one_time_prekeys( &user_id, &device_id, &content_one_time_keys, ¬if_one_time_keys, ) .await?; Ok(()) } pub async fn get_keyserver_keys_for_user( &self, user_id: &str, ) -> Result, Error> { use crate::grpc_services::protos::unauth::DeviceType as GrpcDeviceType; let user_devices = self.get_current_devices(user_id).await?; let maybe_keyserver_device = user_devices .into_iter() .find(|device| device.device_type == GrpcDeviceType::Keyserver); let Some(keyserver) = maybe_keyserver_device else { return Ok(None); }; debug!( "Found keyserver in devices table (ID={})", &keyserver.device_id ); let (notif_one_time_key, requested_more_keys) = self .get_one_time_key( user_id, &keyserver.device_id, OlmAccountType::Notification, true, ) .await?; let (content_one_time_key, _) = self .get_one_time_key( user_id, &keyserver.device_id, OlmAccountType::Content, !requested_more_keys, ) .await?; debug!( "Able to get notif one-time key for keyserver {}: {}", &keyserver.device_id, notif_one_time_key.is_some() ); debug!( "Able to get content one-time key for keyserver {}: {}", &keyserver.device_id, content_one_time_key.is_some() ); let outbound_payload = OutboundKeys { key_payload: keyserver.device_key_info.key_payload, key_payload_signature: keyserver.device_key_info.key_payload_signature, content_prekey: keyserver.content_prekey, notif_prekey: keyserver.notif_prekey, content_one_time_key, notif_one_time_key, }; Ok(Some(outbound_payload)) } pub async fn get_keyserver_device_id_for_user( &self, user_id: &str, ) -> Result, Error> { use crate::grpc_services::protos::unauth::DeviceType as GrpcDeviceType; let user_devices = self.get_current_devices(user_id).await?; let maybe_keyserver_device_id = user_devices .into_iter() .find(|device| device.device_type == GrpcDeviceType::Keyserver) .map(|device| device.device_id); Ok(maybe_keyserver_device_id) } pub async fn update_user_password( &self, user_id: String, password_file: Vec, ) -> Result<(), Error> { let update_expression = format!("SET {} = :p", USERS_TABLE_REGISTRATION_ATTRIBUTE); let expression_attribute_values = HashMap::from([( ":p".to_string(), AttributeValue::B(Blob::new(password_file)), )]); self .client .update_item() .table_name(USERS_TABLE) .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id)) .update_expression(update_expression) .set_expression_attribute_values(Some(expression_attribute_values)) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; Ok(()) } #[tracing::instrument(skip_all)] pub async fn delete_user( &self, user_id: String, ) -> Result { // We must delete the one-time keys first because doing so requires device // IDs from the devices table debug!(user_id, "Attempting to delete user's one-time keys"); self.delete_otks_table_rows_for_user(&user_id).await?; debug!(user_id, "Attempting to delete user's devices"); self.delete_devices_table_rows_for_user(&user_id).await?; debug!(user_id, "Attempting to delete user's access tokens"); self.delete_all_tokens_for_user(&user_id).await?; debug!(user_id, "Attempting to delete user"); match self .client .delete_item() .table_name(USERS_TABLE) .key( USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id.clone()), ) .send() .await { Ok(out) => { info!("User has been deleted {}", user_id); Ok(out) } Err(e) => { - error!("DynamoDB client failed to delete user {}", user_id); + error!( + errorType = error_types::GENERIC_DB_LOG, + "DynamoDB client failed to delete user {}", user_id + ); Err(Error::AwsSdk(e.into())) } } } pub async fn wallet_address_taken( &self, wallet_address: String, ) -> Result { let result = self .get_user_id_from_user_info(wallet_address, &AuthType::Wallet) .await?; Ok(result.is_some()) } pub async fn username_taken(&self, username: String) -> Result { let result = self .get_user_id_from_user_info(username, &AuthType::Password) .await?; Ok(result.is_some()) } pub async fn filter_out_taken_usernames( &self, user_details: Vec, ) -> Result, Error> { let db_usernames = self.get_all_usernames().await?; let db_usernames_set: HashSet = db_usernames.into_iter().collect(); let available_user_details: Vec = user_details .into_iter() .filter(|user_detail| !db_usernames_set.contains(&user_detail.username)) .collect(); Ok(available_user_details) } #[tracing::instrument(skip_all)] async fn get_user_from_user_info( &self, user_info: String, auth_type: &AuthType, ) -> Result>, Error> { let (index, attribute_name) = match auth_type { AuthType::Password => { (USERS_TABLE_USERNAME_INDEX, USERS_TABLE_USERNAME_ATTRIBUTE) } AuthType::Wallet => ( USERS_TABLE_WALLET_ADDRESS_INDEX, USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE, ), }; match self .client .query() .table_name(USERS_TABLE) .index_name(index) .key_condition_expression(format!("{} = :u", attribute_name)) .expression_attribute_values(":u", AttributeValue::S(user_info.clone())) .send() .await { Ok(QueryOutput { items: Some(items), .. }) => { let num_items = items.len(); if num_items == 0 { return Ok(None); } if num_items > 1 { warn!( "{} user IDs associated with {} {}: {:?}", num_items, attribute_name, user_info, items ); } let first_item = items[0].clone(); let user_id = first_item .get(USERS_TABLE_PARTITION_KEY) .ok_or(DBItemError { attribute_name: USERS_TABLE_PARTITION_KEY.to_string(), attribute_value: None.into(), attribute_error: DBItemAttributeError::Missing, })? .as_s() .map_err(|_| DBItemError { attribute_name: USERS_TABLE_PARTITION_KEY.to_string(), attribute_value: first_item .get(USERS_TABLE_PARTITION_KEY) .cloned() .into(), attribute_error: DBItemAttributeError::IncorrectType, })?; let result = self.get_item_from_users_table(user_id).await?; Ok(result.item) } Ok(_) => { info!( "No item found for {} {} in users table", attribute_name, user_info ); Ok(None) } Err(e) => { error!( + errorType = error_types::GENERIC_DB_LOG, "DynamoDB client failed to get user from {} {}: {}", - attribute_name, user_info, e + attribute_name, + user_info, + e ); Err(Error::AwsSdk(e.into())) } } } pub async fn get_keys_for_user( &self, user_id: &str, get_one_time_keys: bool, ) -> Result, Error> { 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_keys) in devices_response.iter_mut() { let requested_more_keys; (device_keys.notif_one_time_key, requested_more_keys) = self .get_one_time_key( user_id, device_id_key, OlmAccountType::Notification, true, ) .await?; (device_keys.content_one_time_key, _) = self .get_one_time_key( user_id, device_id_key, OlmAccountType::Content, !requested_more_keys, ) .await?; } } Ok(Some(devices_response)) } pub async fn get_user_id_from_user_info( &self, user_info: String, auth_type: &AuthType, ) -> Result, Error> { match self .get_user_from_user_info(user_info.clone(), auth_type) .await { Ok(Some(mut user)) => user .take_attr(USERS_TABLE_PARTITION_KEY) .map(Some) .map_err(Error::Attribute), Ok(_) => Ok(None), Err(e) => Err(e), } } #[tracing::instrument(skip_all)] pub async fn get_user_id_and_password_file_from_username( &self, username: &str, ) -> Result)>, Error> { match self .get_user_from_user_info(username.to_string(), &AuthType::Password) .await { Ok(Some(mut user)) => { let user_id = user.take_attr(USERS_TABLE_PARTITION_KEY)?; let password_file = parse_registration_data_attribute( user.remove(USERS_TABLE_REGISTRATION_ATTRIBUTE), )?; Ok(Some((user_id, password_file))) } Ok(_) => { info!( "No item found for user {} in PAKE registration table", username ); Ok(None) } Err(e) => { error!( + errorType = error_types::GENERIC_DB_LOG, "DynamoDB client failed to get registration data for user {}: {}", - username, e + username, + e ); Err(e) } } } pub async fn get_username_and_password_file( &self, user_id: &str, ) -> Result)>, Error> { let Some(mut user) = self.get_item_from_users_table(user_id).await?.item else { return Ok(None); }; let username = user.take_attr(USERS_TABLE_USERNAME_ATTRIBUTE)?; let password_file = parse_registration_data_attribute( user.remove(USERS_TABLE_REGISTRATION_ATTRIBUTE), )?; Ok(Some((username, password_file))) } async fn get_item_from_users_table( &self, user_id: &str, ) -> Result { let primary_key = create_simple_primary_key(( USERS_TABLE_PARTITION_KEY.to_string(), user_id.to_string(), )); self .client .get_item() .table_name(USERS_TABLE) .set_key(Some(primary_key)) .consistent_read(true) .send() .await .map_err(|e| Error::AwsSdk(e.into())) } /// Retrieves username for password users or wallet address for wallet users /// Returns `None` if user not found #[tracing::instrument(skip_all)] pub async fn get_user_identifier( &self, user_id: &str, ) -> Result, Error> { self .get_item_from_users_table(user_id) .await? .item .map(Identifier::try_from) .transpose() .map_err(|e| { - error!(user_id, "Database item is missing an identifier"); + error!( + user_id, + errorType = error_types::GENERIC_DB_LOG, + "Database item is missing an identifier" + ); e }) } async fn get_all_usernames(&self) -> Result, Error> { let scan_output = self .client .scan() .table_name(USERS_TABLE) .projection_expression(USERS_TABLE_USERNAME_ATTRIBUTE) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; let mut result = Vec::new(); if let Some(attributes) = scan_output.items { for mut attribute in attributes { if let Ok(username) = attribute.take_attr(USERS_TABLE_USERNAME_ATTRIBUTE) { result.push(username); } } } Ok(result) } pub async fn get_all_user_details(&self) -> Result, Error> { let scan_output = self .client .scan() .table_name(USERS_TABLE) .projection_expression(format!( "{USERS_TABLE_USERNAME_ATTRIBUTE}, {USERS_TABLE_PARTITION_KEY}" )) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; let mut result = Vec::new(); if let Some(attributes) = scan_output.items { for mut attribute in attributes { if let (Ok(username), Ok(user_id)) = ( attribute.take_attr(USERS_TABLE_USERNAME_ATTRIBUTE), attribute.take_attr(USERS_TABLE_PARTITION_KEY), ) { result.push(UserDetail { username, user_id }); } } } Ok(result) } pub async fn get_all_reserved_user_details( &self, ) -> Result, Error> { let scan_output = self .client .scan() .table_name(RESERVED_USERNAMES_TABLE) .projection_expression(format!( "{RESERVED_USERNAMES_TABLE_PARTITION_KEY},\ {RESERVED_USERNAMES_TABLE_USER_ID_ATTRIBUTE}" )) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; let mut result = Vec::new(); if let Some(attributes) = scan_output.items { for mut attribute in attributes { if let (Ok(username), Ok(user_id)) = ( attribute.take_attr(USERS_TABLE_USERNAME_ATTRIBUTE), attribute.take_attr(RESERVED_USERNAMES_TABLE_USER_ID_ATTRIBUTE), ) { result.push(UserDetail { username, user_id }); } } } Ok(result) } pub async fn add_nonce_to_nonces_table( &self, nonce_data: NonceData, ) -> Result { let item = HashMap::from([ ( NONCE_TABLE_PARTITION_KEY.to_string(), AttributeValue::S(nonce_data.nonce), ), ( NONCE_TABLE_CREATED_ATTRIBUTE.to_string(), AttributeValue::S(nonce_data.created.to_rfc3339()), ), ( NONCE_TABLE_EXPIRATION_TIME_ATTRIBUTE.to_string(), AttributeValue::S(nonce_data.expiration_time.to_rfc3339()), ), ( NONCE_TABLE_EXPIRATION_TIME_UNIX_ATTRIBUTE.to_string(), AttributeValue::N(nonce_data.expiration_time.timestamp().to_string()), ), ]); self .client .put_item() .table_name(NONCE_TABLE) .set_item(Some(item)) .send() .await .map_err(|e| Error::AwsSdk(e.into())) } pub async fn get_nonce_from_nonces_table( &self, nonce_value: impl Into, ) -> Result, Error> { let get_response = self .client .get_item() .table_name(NONCE_TABLE) .key( NONCE_TABLE_PARTITION_KEY, AttributeValue::S(nonce_value.into()), ) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; let Some(mut item) = get_response.item else { return Ok(None); }; let nonce = item.take_attr(NONCE_TABLE_PARTITION_KEY)?; let created = DateTime::::try_from_attr( NONCE_TABLE_CREATED_ATTRIBUTE, item.remove(NONCE_TABLE_CREATED_ATTRIBUTE), )?; let expiration_time = DateTime::::try_from_attr( NONCE_TABLE_EXPIRATION_TIME_ATTRIBUTE, item.remove(NONCE_TABLE_EXPIRATION_TIME_ATTRIBUTE), )?; Ok(Some(NonceData { nonce, created, expiration_time, })) } pub async fn remove_nonce_from_nonces_table( &self, nonce: impl Into, ) -> Result<(), Error> { self .client .delete_item() .table_name(NONCE_TABLE) .key(NONCE_TABLE_PARTITION_KEY, AttributeValue::S(nonce.into())) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; Ok(()) } pub async fn add_usernames_to_reserved_usernames_table( &self, user_details: Vec, ) -> Result<(), Error> { // A single call to BatchWriteItem can consist of up to 25 operations for user_chunk in user_details.chunks(25) { let write_requests = user_chunk .iter() .map(|user_detail| { let put_request = PutRequest::builder() .item( RESERVED_USERNAMES_TABLE_PARTITION_KEY, AttributeValue::S(user_detail.username.to_string()), ) .item( RESERVED_USERNAMES_TABLE_USER_ID_ATTRIBUTE, AttributeValue::S(user_detail.user_id.to_string()), ) .build(); WriteRequest::builder().put_request(put_request).build() }) .collect(); self .client .batch_write_item() .request_items(RESERVED_USERNAMES_TABLE, write_requests) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; } info!("Batch write item to reserved usernames table succeeded"); Ok(()) } #[tracing::instrument(skip_all)] pub async fn delete_username_from_reserved_usernames_table( &self, username: String, ) -> Result { debug!( "Attempting to delete username {} from reserved usernames table", username ); match self .client .delete_item() .table_name(RESERVED_USERNAMES_TABLE) .key( RESERVED_USERNAMES_TABLE_PARTITION_KEY, AttributeValue::S(username.clone()), ) .send() .await { Ok(out) => { info!( "Username {} has been deleted from reserved usernames table", username ); Ok(out) } Err(e) => { - error!("DynamoDB client failed to delete username {} from reserved usernames table", username); + error!(errorType = error_types::GENERIC_DB_LOG, "DynamoDB client failed to delete username {} from reserved usernames table", username); Err(Error::AwsSdk(e.into())) } } } pub async fn username_in_reserved_usernames_table( &self, username: &str, ) -> Result { match self .client .get_item() .table_name(RESERVED_USERNAMES_TABLE) .key( RESERVED_USERNAMES_TABLE_PARTITION_KEY.to_string(), AttributeValue::S(username.to_string()), ) .consistent_read(true) .send() .await { Ok(GetItemOutput { item: Some(_), .. }) => Ok(true), Ok(_) => Ok(false), Err(e) => Err(Error::AwsSdk(e.into())), } } } type AttributeName = String; type Devices = HashMap; fn create_simple_primary_key( partition_key: (AttributeName, String), ) -> HashMap { HashMap::from([(partition_key.0, AttributeValue::S(partition_key.1))]) } fn create_composite_primary_key( partition_key: (AttributeName, String), sort_key: (AttributeName, String), ) -> HashMap { let mut primary_key = create_simple_primary_key(partition_key); primary_key.insert(sort_key.0, AttributeValue::S(sort_key.1)); primary_key } fn parse_registration_data_attribute( attribute: Option, ) -> Result, DBItemError> { match attribute { Some(AttributeValue::B(server_registration_bytes)) => { Ok(server_registration_bytes.into_inner()) } Some(_) => Err(DBItemError::new( USERS_TABLE_REGISTRATION_ATTRIBUTE.to_string(), attribute.into(), DBItemAttributeError::IncorrectType, )), None => Err(DBItemError::new( USERS_TABLE_REGISTRATION_ATTRIBUTE.to_string(), attribute.into(), DBItemAttributeError::Missing, )), } } #[deprecated(note = "Use `comm_lib` counterpart instead")] #[allow(dead_code)] fn parse_map_attribute( attribute_name: &str, attribute_value: Option, ) -> Result { match attribute_value { Some(AttributeValue::M(map)) => Ok(map), Some(_) => { error!( attribute = attribute_name, value = ?attribute_value, error_type = "IncorrectType", - "Unexpected attribute type when parsing map attribute" + errorType = error_types::GENERIC_DB_LOG, "Unexpected attribute type when parsing map attribute" ); Err(DBItemError::new( attribute_name.to_string(), attribute_value.into(), DBItemAttributeError::IncorrectType, )) } None => { error!( attribute = attribute_name, error_type = "Missing", + errorType = error_types::GENERIC_DB_LOG, "Attribute is missing" ); Err(DBItemError::new( attribute_name.to_string(), attribute_value.into(), DBItemAttributeError::Missing, )) } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_create_simple_primary_key() { let partition_key_name = "userID".to_string(); let partition_key_value = "12345".to_string(); let partition_key = (partition_key_name.clone(), partition_key_value.clone()); let mut primary_key = create_simple_primary_key(partition_key); assert_eq!(primary_key.len(), 1); let attribute = primary_key.remove(&partition_key_name); assert!(attribute.is_some()); assert_eq!(attribute, Some(AttributeValue::S(partition_key_value))); } #[test] fn test_create_composite_primary_key() { let partition_key_name = "userID".to_string(); let partition_key_value = "12345".to_string(); let partition_key = (partition_key_name.clone(), partition_key_value.clone()); let sort_key_name = "deviceID".to_string(); let sort_key_value = "54321".to_string(); let sort_key = (sort_key_name.clone(), sort_key_value.clone()); let mut primary_key = create_composite_primary_key(partition_key, sort_key); assert_eq!(primary_key.len(), 2); let partition_key_attribute = primary_key.remove(&partition_key_name); assert!(partition_key_attribute.is_some()); assert_eq!( partition_key_attribute, Some(AttributeValue::S(partition_key_value)) ); let sort_key_attribute = primary_key.remove(&sort_key_name); assert!(sort_key_attribute.is_some()); assert_eq!(sort_key_attribute, Some(AttributeValue::S(sort_key_value))) } #[test] fn validate_keys() { // Taken from test user let example_payload = r#"{\"notificationIdentityPublicKeys\":{\"curve25519\":\"DYmV8VdkjwG/VtC8C53morogNJhpTPT/4jzW0/cxzQo\",\"ed25519\":\"D0BV2Y7Qm36VUtjwyQTJJWYAycN7aMSJmhEsRJpW2mk\"},\"primaryIdentityPublicKeys\":{\"curve25519\":\"Y4ZIqzpE1nv83kKGfvFP6rifya0itRg2hifqYtsISnk\",\"ed25519\":\"cSlL+VLLJDgtKSPlIwoCZg0h0EmHlQoJC08uV/O+jvg\"}}"#; let serialized_payload = KeyPayload::from_str(example_payload).unwrap(); assert_eq!( serialized_payload .notification_identity_public_keys .curve25519, "DYmV8VdkjwG/VtC8C53morogNJhpTPT/4jzW0/cxzQo" ); } #[test] fn test_int_to_device_type() { let valid_result = DeviceType::try_from(3); assert!(valid_result.is_ok()); assert_eq!(valid_result.unwrap(), DeviceType::Android); let invalid_result = DeviceType::try_from(6); assert!(invalid_result.is_err()); } } diff --git a/services/identity/src/database/device_list.rs b/services/identity/src/database/device_list.rs index 1077049d3..5f8a3911e 100644 --- a/services/identity/src/database/device_list.rs +++ b/services/identity/src/database/device_list.rs @@ -1,1441 +1,1496 @@ use std::collections::HashMap; use chrono::{DateTime, Utc}; use comm_lib::{ aws::ddb::{ operation::{get_item::GetItemOutput, query::builders::QueryFluentBuilder}, types::{ error::TransactionCanceledException, AttributeValue, Delete, DeleteRequest, Put, TransactWriteItem, Update, WriteRequest, }, }, database::{ AttributeExtractor, AttributeMap, DBItemAttributeError, DBItemError, DynamoDBError, TryFromAttribute, }, }; use tracing::{debug, error, warn}; use crate::{ client_service::FlattenedDeviceKeyUpload, constants::{ devices_table::{self, *}, - USERS_TABLE, USERS_TABLE_DEVICELIST_TIMESTAMP_ATTRIBUTE_NAME, + error_types, USERS_TABLE, USERS_TABLE_DEVICELIST_TIMESTAMP_ATTRIBUTE_NAME, USERS_TABLE_PARTITION_KEY, }, error::{DeviceListError, Error}, grpc_services::protos::{self, unauth::DeviceType}, grpc_utils::DeviceKeysInfo, olm::is_valid_olm_key, }; use super::DatabaseClient; // We omit the content and notif one-time key count attributes from this struct // because they are internal helpers and are not provided by users #[derive(Clone, Debug)] pub struct DeviceRow { pub user_id: String, pub device_id: String, pub device_type: DeviceType, pub device_key_info: IdentityKeyInfo, pub content_prekey: Prekey, pub notif_prekey: Prekey, // migration-related data pub code_version: u64, /// Timestamp of last login (access token generation) pub login_time: DateTime, } #[derive(Clone, Debug)] pub struct DeviceListRow { pub user_id: String, pub timestamp: DateTime, pub device_ids: Vec, } #[derive(Clone, Debug)] pub struct IdentityKeyInfo { pub key_payload: String, pub key_payload_signature: String, } #[derive(Clone, Debug)] pub struct Prekey { pub prekey: String, pub prekey_signature: String, } /// A struct representing device list update request /// payload; issued by the primary device #[derive(derive_more::Constructor)] pub struct DeviceListUpdate { pub devices: Vec, pub timestamp: DateTime, } impl DeviceRow { #[tracing::instrument(skip_all)] pub fn from_device_key_upload( user_id: impl Into, upload: FlattenedDeviceKeyUpload, code_version: u64, login_time: DateTime, ) -> Result { if !is_valid_olm_key(&upload.content_prekey) || !is_valid_olm_key(&upload.notif_prekey) { - error!("Invalid prekey format"); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Invalid prekey format" + ); return Err(Error::InvalidFormat); } let device_row = Self { user_id: user_id.into(), device_id: upload.device_id_key, device_type: DeviceType::from_str_name(upload.device_type.as_str_name()) .expect("DeviceType conversion failed. Identity client and server protos mismatch"), device_key_info: IdentityKeyInfo { key_payload: upload.key_payload, key_payload_signature: upload.key_payload_signature, }, content_prekey: Prekey { prekey: upload.content_prekey, prekey_signature: upload.content_prekey_signature, }, notif_prekey: Prekey { prekey: upload.notif_prekey, prekey_signature: upload.notif_prekey_signature, }, code_version, login_time, }; Ok(device_row) } } impl DeviceListRow { /// Generates new device list row from given devices fn new( user_id: impl Into, device_ids: Vec, timestamp: Option>, ) -> Self { Self { user_id: user_id.into(), device_ids, timestamp: timestamp.unwrap_or_else(Utc::now), } } } // helper structs for converting to/from attribute values for sort key (a.k.a itemID) pub struct DeviceIDAttribute(pub String); struct DeviceListKeyAttribute(DateTime); impl From for AttributeValue { fn from(value: DeviceIDAttribute) -> Self { AttributeValue::S(format!("{DEVICE_ITEM_KEY_PREFIX}{}", value.0)) } } impl From for AttributeValue { fn from(value: DeviceListKeyAttribute) -> Self { AttributeValue::S(format!( "{DEVICE_LIST_KEY_PREFIX}{}", value.0.to_rfc3339() )) } } impl TryFrom> for DeviceIDAttribute { type Error = DBItemError; fn try_from(value: Option) -> Result { let item_id = String::try_from_attr(ATTR_ITEM_ID, value)?; // remove the device- prefix let device_id = item_id .strip_prefix(DEVICE_ITEM_KEY_PREFIX) .ok_or_else(|| DBItemError { attribute_name: ATTR_ITEM_ID.to_string(), attribute_value: item_id.clone().into(), attribute_error: DBItemAttributeError::InvalidValue, })? .to_string(); Ok(Self(device_id)) } } impl TryFrom> for DeviceListKeyAttribute { type Error = DBItemError; fn try_from(value: Option) -> Result { let item_id = String::try_from_attr(ATTR_ITEM_ID, value)?; // remove the device-list- prefix, then parse the timestamp let timestamp: DateTime = item_id .strip_prefix(DEVICE_LIST_KEY_PREFIX) .ok_or_else(|| DBItemError { attribute_name: ATTR_ITEM_ID.to_string(), attribute_value: item_id.clone().into(), attribute_error: DBItemAttributeError::InvalidValue, }) .and_then(|s| { s.parse().map_err(|e| { DBItemError::new( ATTR_ITEM_ID.to_string(), item_id.clone().into(), DBItemAttributeError::InvalidTimestamp(e), ) }) })?; Ok(Self(timestamp)) } } impl TryFrom for DeviceRow { type Error = DBItemError; fn try_from(mut attrs: AttributeMap) -> Result { let user_id = attrs.take_attr(ATTR_USER_ID)?; let DeviceIDAttribute(device_id) = attrs.remove(ATTR_ITEM_ID).try_into()?; let raw_device_type: String = attrs.take_attr(ATTR_DEVICE_TYPE)?; let device_type = DeviceType::from_str_name(&raw_device_type).ok_or_else(|| { DBItemError::new( ATTR_DEVICE_TYPE.to_string(), raw_device_type.into(), DBItemAttributeError::InvalidValue, ) })?; let device_key_info = attrs .take_attr::(ATTR_DEVICE_KEY_INFO) .and_then(IdentityKeyInfo::try_from)?; let content_prekey = attrs .take_attr::(ATTR_CONTENT_PREKEY) .and_then(Prekey::try_from)?; let notif_prekey = attrs .take_attr::(ATTR_NOTIF_PREKEY) .and_then(Prekey::try_from)?; let code_version = attrs .remove(ATTR_CODE_VERSION) .and_then(|attr| attr.as_n().ok().cloned()) .and_then(|val| val.parse::().ok()) .unwrap_or_default(); let login_time: DateTime = attrs.take_attr(ATTR_LOGIN_TIME)?; Ok(Self { user_id, device_id, device_type, device_key_info, content_prekey, notif_prekey, code_version, login_time, }) } } impl From for AttributeMap { fn from(value: DeviceRow) -> Self { HashMap::from([ (ATTR_USER_ID.to_string(), AttributeValue::S(value.user_id)), ( ATTR_ITEM_ID.to_string(), DeviceIDAttribute(value.device_id).into(), ), ( ATTR_DEVICE_TYPE.to_string(), AttributeValue::S(value.device_type.as_str_name().to_string()), ), ( ATTR_DEVICE_KEY_INFO.to_string(), value.device_key_info.into(), ), (ATTR_CONTENT_PREKEY.to_string(), value.content_prekey.into()), (ATTR_NOTIF_PREKEY.to_string(), value.notif_prekey.into()), // migration attributes ( ATTR_CODE_VERSION.to_string(), AttributeValue::N(value.code_version.to_string()), ), ( ATTR_LOGIN_TIME.to_string(), AttributeValue::S(value.login_time.to_rfc3339()), ), ]) } } impl From for protos::unauth::IdentityKeyInfo { fn from(value: IdentityKeyInfo) -> Self { Self { payload: value.key_payload, payload_signature: value.key_payload_signature, } } } impl From for AttributeValue { fn from(value: IdentityKeyInfo) -> Self { let attrs = HashMap::from([ ( ATTR_KEY_PAYLOAD.to_string(), AttributeValue::S(value.key_payload), ), ( ATTR_KEY_PAYLOAD_SIGNATURE.to_string(), AttributeValue::S(value.key_payload_signature), ), ]); AttributeValue::M(attrs) } } impl TryFrom for IdentityKeyInfo { type Error = DBItemError; fn try_from(mut attrs: AttributeMap) -> Result { let key_payload = attrs.take_attr(ATTR_KEY_PAYLOAD)?; let key_payload_signature = attrs.take_attr(ATTR_KEY_PAYLOAD_SIGNATURE)?; Ok(Self { key_payload, key_payload_signature, }) } } impl From for AttributeValue { fn from(value: Prekey) -> Self { let attrs = HashMap::from([ (ATTR_PREKEY.to_string(), AttributeValue::S(value.prekey)), ( ATTR_PREKEY_SIGNATURE.to_string(), AttributeValue::S(value.prekey_signature), ), ]); AttributeValue::M(attrs) } } impl From for protos::unauth::Prekey { fn from(value: Prekey) -> Self { Self { prekey: value.prekey, prekey_signature: value.prekey_signature, } } } impl From for Prekey { fn from(value: protos::unauth::Prekey) -> Self { Self { prekey: value.prekey, prekey_signature: value.prekey_signature, } } } impl TryFrom for Prekey { type Error = DBItemError; fn try_from(mut attrs: AttributeMap) -> Result { let prekey = attrs.take_attr(ATTR_PREKEY)?; let prekey_signature = attrs.take_attr(ATTR_PREKEY_SIGNATURE)?; Ok(Self { prekey, prekey_signature, }) } } impl TryFrom for DeviceListRow { type Error = DBItemError; fn try_from(mut attrs: AttributeMap) -> Result { let user_id = attrs.take_attr(ATTR_USER_ID)?; let DeviceListKeyAttribute(timestamp) = attrs.remove(ATTR_ITEM_ID).try_into()?; // validate timestamps are in sync let timestamps_match = attrs .remove(ATTR_TIMESTAMP) .and_then(|attr| attr.as_n().ok().cloned()) .and_then(|val| val.parse::().ok()) .filter(|val| *val == timestamp.timestamp_millis()) .is_some(); if !timestamps_match { warn!( "DeviceList timestamp mismatch for (userID={}, itemID={})", &user_id, timestamp.to_rfc3339() ); } let device_ids: Vec = attrs.take_attr(ATTR_DEVICE_IDS)?; Ok(Self { user_id, timestamp, device_ids, }) } } impl From for AttributeMap { fn from(device_list: DeviceListRow) -> Self { let mut attrs = HashMap::new(); attrs.insert( ATTR_USER_ID.to_string(), AttributeValue::S(device_list.user_id.clone()), ); attrs.insert( ATTR_ITEM_ID.to_string(), DeviceListKeyAttribute(device_list.timestamp).into(), ); attrs.insert( ATTR_TIMESTAMP.to_string(), AttributeValue::N(device_list.timestamp.timestamp_millis().to_string()), ); attrs.insert( ATTR_DEVICE_IDS.to_string(), AttributeValue::L( device_list .device_ids .into_iter() .map(AttributeValue::S) .collect(), ), ); attrs } } impl DatabaseClient { /// Retrieves user's current devices and their full data #[tracing::instrument(skip_all)] pub async fn get_current_devices( &self, user_id: impl Into, ) -> Result, Error> { let response = query_rows_with_prefix(self, user_id, DEVICE_ITEM_KEY_PREFIX) .send() .await .map_err(|e| { - error!("Failed to get current devices: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to get current devices: {:?}", e + ); Error::AwsSdk(e.into()) })?; let Some(rows) = response.items else { return Ok(Vec::new()); }; rows .into_iter() .map(DeviceRow::try_from) .collect::, DBItemError>>() .map_err(Error::from) } /// Gets user's device list history #[tracing::instrument(skip_all)] pub async fn get_device_list_history( &self, user_id: impl Into, since: Option>, ) -> Result, Error> { let rows = if let Some(since) = since { // When timestamp is provided, it's better to query device lists by timestamp LSI self .client .query() .table_name(devices_table::NAME) .index_name(devices_table::TIMESTAMP_INDEX_NAME) .consistent_read(true) .key_condition_expression("#user_id = :user_id AND #timestamp > :since") .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_names("#timestamp", ATTR_TIMESTAMP) .expression_attribute_values( ":user_id", AttributeValue::S(user_id.into()), ) .expression_attribute_values( ":since", AttributeValue::N(since.timestamp_millis().to_string()), ) .send() .await .map_err(|e| { - error!("Failed to query device list updates by index: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to query device list updates by index: {:?}", e + ); Error::AwsSdk(e.into()) })? .items } else { // Query all device lists for user query_rows_with_prefix(self, user_id, DEVICE_LIST_KEY_PREFIX) .send() .await .map_err(|e| { - error!("Failed to query device list updates (all): {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to query device list updates (all): {:?}", e + ); Error::AwsSdk(e.into()) })? .items }; rows .unwrap_or_default() .into_iter() .map(DeviceListRow::try_from) .collect::, DBItemError>>() .map_err(Error::from) } /// Returns all devices' keys for the given user. Response is in the same format /// as [DatabaseClient::get_keys_for_user] for compatibility reasons. #[tracing::instrument(skip_all)] pub async fn get_keys_for_user_devices( &self, user_id: impl Into, ) -> Result { let user_devices = self.get_current_devices(user_id).await?; let user_devices_keys = user_devices .into_iter() .map(|device| (device.device_id.clone(), DeviceKeysInfo::from(device))) .collect(); Ok(user_devices_keys) } #[tracing::instrument(skip_all)] pub async fn update_device_prekeys( &self, user_id: impl Into, device_id: impl Into, content_prekey: Prekey, notif_prekey: Prekey, ) -> Result<(), Error> { if !is_valid_olm_key(&content_prekey.prekey) || !is_valid_olm_key(¬if_prekey.prekey) { - error!("Invalid prekey format"); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Invalid prekey format" + ); return Err(Error::InvalidFormat); } self .client .update_item() .table_name(devices_table::NAME) .key(ATTR_USER_ID, AttributeValue::S(user_id.into())) .key(ATTR_ITEM_ID, DeviceIDAttribute(device_id.into()).into()) .condition_expression( "attribute_exists(#user_id) AND attribute_exists(#item_id)", ) .update_expression( "SET #content_prekey = :content_prekey, #notif_prekey = :notif_prekey", ) .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_names("#item_id", ATTR_ITEM_ID) .expression_attribute_names("#content_prekey", ATTR_CONTENT_PREKEY) .expression_attribute_names("#notif_prekey", ATTR_NOTIF_PREKEY) .expression_attribute_values(":content_prekey", content_prekey.into()) .expression_attribute_values(":notif_prekey", notif_prekey.into()) .send() .await .map_err(|e| { - error!("Failed to update device prekeys: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to update device prekeys: {:?}", e + ); Error::AwsSdk(e.into()) })?; Ok(()) } /// Checks if given device exists on user's current device list #[tracing::instrument(skip_all)] pub async fn device_exists( &self, user_id: impl Into, device_id: impl Into, ) -> Result { let GetItemOutput { item, .. } = self .client .get_item() .table_name(devices_table::NAME) .key(ATTR_USER_ID, AttributeValue::S(user_id.into())) .key(ATTR_ITEM_ID, DeviceIDAttribute(device_id.into()).into()) // only fetch the primary key, we don't need the rest .projection_expression(format!("{ATTR_USER_ID}, {ATTR_ITEM_ID}")) .send() .await .map_err(|e| { - error!("Failed to check if device exists: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to check if device exists: {:?}", e + ); Error::AwsSdk(e.into()) })?; Ok(item.is_some()) } #[tracing::instrument(skip_all)] pub async fn get_device_data( &self, user_id: impl Into, device_id: impl Into, ) -> Result, Error> { let GetItemOutput { item, .. } = self .client .get_item() .table_name(devices_table::NAME) .key(ATTR_USER_ID, AttributeValue::S(user_id.into())) .key(ATTR_ITEM_ID, DeviceIDAttribute(device_id.into()).into()) .send() .await .map_err(|e| { - error!("Failed to fetch device data: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to fetch device data: {:?}", e + ); Error::AwsSdk(e.into()) })?; let Some(attrs) = item else { return Ok(None); }; let device_data = DeviceRow::try_from(attrs)?; Ok(Some(device_data)) } /// Fails if the device list is empty #[tracing::instrument(skip_all)] pub async fn get_primary_device_data( &self, user_id: &str, ) -> Result { let device_list = self.get_current_device_list(user_id).await?; let Some(primary_device_id) = device_list .as_ref() .and_then(|list| list.device_ids.first()) else { - error!(user_id, "Device list is empty. Cannot fetch primary device"); + error!( + user_id, + errorType = error_types::DEVICE_LIST_DB_LOG, + "Device list is empty. Cannot fetch primary device" + ); return Err(Error::DeviceList(DeviceListError::DeviceNotFound)); }; self .get_device_data(user_id, primary_device_id) .await? .ok_or_else(|| { error!( - "Corrupt database. Missing primary device data for user {}", - user_id + errorType = error_types::DEVICE_LIST_DB_LOG, + "Corrupt database. Missing primary device data for user {}", user_id ); Error::MissingItem }) } /// Required only for migration purposes (determining primary device) #[tracing::instrument(skip_all)] pub async fn update_device_login_time( &self, user_id: impl Into, device_id: impl Into, login_time: DateTime, ) -> Result<(), Error> { self .client .update_item() .table_name(devices_table::NAME) .key(ATTR_USER_ID, AttributeValue::S(user_id.into())) .key(ATTR_ITEM_ID, DeviceIDAttribute(device_id.into()).into()) .condition_expression( "attribute_exists(#user_id) AND attribute_exists(#item_id)", ) .update_expression("SET #login_time = :login_time") .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_names("#item_id", ATTR_ITEM_ID) .expression_attribute_names("#login_time", ATTR_LOGIN_TIME) .expression_attribute_values( ":login_time", AttributeValue::S(login_time.to_rfc3339()), ) .send() .await .map_err(|e| { - error!("Failed to update device login time: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to update device login time: {:?}", e + ); Error::AwsSdk(e.into()) })?; Ok(()) } #[tracing::instrument(skip_all)] pub async fn get_current_device_list( &self, user_id: impl Into, ) -> Result, Error> { self .client .query() .table_name(devices_table::NAME) .index_name(devices_table::TIMESTAMP_INDEX_NAME) .consistent_read(true) .key_condition_expression("#user_id = :user_id") // sort descending .scan_index_forward(false) .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_values( ":user_id", AttributeValue::S(user_id.into()), ) .limit(1) .send() .await .map_err(|e| { - error!("Failed to query device list updates by index: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to query device list updates by index: {:?}", e + ); Error::AwsSdk(e.into()) })? .items .and_then(|mut items| items.pop()) .map(DeviceListRow::try_from) .transpose() .map_err(Error::from) } /// Adds device data to devices table. If the device already exists, its /// data is overwritten. This does not update the device list; the device ID /// should already be present in the device list. #[tracing::instrument(skip_all)] pub async fn put_device_data( &self, user_id: impl Into, device_key_upload: FlattenedDeviceKeyUpload, code_version: u64, login_time: DateTime, ) -> Result<(), Error> { let content_one_time_keys = device_key_upload.content_one_time_keys.clone(); let notif_one_time_keys = device_key_upload.notif_one_time_keys.clone(); let user_id_string = user_id.into(); let new_device = DeviceRow::from_device_key_upload( user_id_string.clone(), device_key_upload, code_version, login_time, )?; let device_id = new_device.device_id.clone(); self .client .put_item() .table_name(devices_table::NAME) .set_item(Some(new_device.into())) .send() .await .map_err(|e| { - error!("Failed to put device data: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to put device data: {:?}", e + ); Error::AwsSdk(e.into()) })?; self .append_one_time_prekeys( &user_id_string, &device_id, &content_one_time_keys, ¬if_one_time_keys, ) .await?; Ok(()) } /// Adds new device to user's device list. If the device already exists, the /// operation fails. Transactionally generates new device list version. pub async fn add_device( &self, user_id: impl Into, device_key_upload: FlattenedDeviceKeyUpload, code_version: u64, login_time: DateTime, ) -> Result<(), Error> { let user_id: String = user_id.into(); self .transact_update_devicelist(&user_id, |device_ids, mut devices_data| { let new_device = DeviceRow::from_device_key_upload( &user_id, device_key_upload, code_version, login_time, )?; if device_ids.iter().any(|id| &new_device.device_id == id) { warn!( "Device already exists in user's device list \ (userID={}, deviceID={})", &user_id, &new_device.device_id ); return Err(Error::DeviceList(DeviceListError::DeviceAlreadyExists)); } device_ids.push(new_device.device_id.clone()); // Reorder devices (determine primary device again) devices_data.push(new_device.clone()); migration::reorder_device_list(&user_id, device_ids, &devices_data); // Put new device let put_device = Put::builder() .table_name(devices_table::NAME) .set_item(Some(new_device.into())) .condition_expression( "attribute_not_exists(#user_id) AND attribute_not_exists(#item_id)", ) .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_names("#item_id", ATTR_ITEM_ID) .build(); let put_device_operation = TransactWriteItem::builder().put(put_device).build(); Ok((Some(put_device_operation), None)) }) .await?; Ok(()) } /// Removes device from user's device list. If the device doesn't exist, the /// operation fails. Transactionally generates new device list version. pub async fn remove_device( &self, user_id: impl Into, device_id: impl AsRef, ) -> Result<(), Error> { let user_id: String = user_id.into(); let device_id = device_id.as_ref(); self .transact_update_devicelist(&user_id, |device_ids, mut devices_data| { let device_exists = device_ids.iter().any(|id| id == device_id); if !device_exists { warn!( "Device doesn't exist in user's device list \ (userID={}, deviceID={})", &user_id, device_id ); return Err(Error::DeviceList(DeviceListError::DeviceNotFound)); } device_ids.retain(|id| id != device_id); // Reorder devices (determine primary device again) devices_data.retain(|d| d.device_id != device_id); migration::reorder_device_list(&user_id, device_ids, &devices_data); // Delete device DDB operation let delete_device = Delete::builder() .table_name(devices_table::NAME) .key(ATTR_USER_ID, AttributeValue::S(user_id.clone())) .key( ATTR_ITEM_ID, DeviceIDAttribute(device_id.to_string()).into(), ) .condition_expression( "attribute_exists(#user_id) AND attribute_exists(#item_id)", ) .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_names("#item_id", ATTR_ITEM_ID) .build(); let operation = TransactWriteItem::builder().delete(delete_device).build(); Ok((Some(operation), None)) }) .await?; Ok(()) } /// applies updated device list received from primary device pub async fn apply_devicelist_update( &self, user_id: &str, update: DeviceListUpdate, // A function that receives previous and new device IDs and // returns boolean determining if the new device list is valid. validator_fn: impl Fn(&[&str], &[&str]) -> bool, ) -> Result { let DeviceListUpdate { devices: new_list, timestamp, } = update; self .transact_update_devicelist(user_id, |current_list, _| { let previous_device_ids: Vec<&str> = current_list.iter().map(AsRef::as_ref).collect(); let new_device_ids: Vec<&str> = new_list.iter().map(AsRef::as_ref).collect(); if !validator_fn(&previous_device_ids, &new_device_ids) { warn!("Received invalid device list update"); return Err(Error::DeviceList( DeviceListError::InvalidDeviceListUpdate, )); } debug!("Applying device list update"); *current_list = new_list; Ok((None, Some(timestamp))) }) .await } /// Performs a transactional update of the device list for the user. Afterwards /// generates a new device list and updates the timestamp in the users table. /// This is done in a transaction. Operation fails if the device list has been /// updated concurrently (timestamp mismatch). /// Returns the new device list row that has been saved to database. #[tracing::instrument(skip_all)] async fn transact_update_devicelist( &self, user_id: &str, // The closure performing a transactional update of the device list. // It receives two arguments: // 1. A mutable reference to the current device list (ordered device IDs). // 2. Details (full data) of the current devices (unordered). // The closure should return a two-element tuple: // - (optional) transactional DDB operation to be performed // when updating the device list. // - (optional) new device list timestamp. Defaults to `Utc::now()`. action: impl FnOnce( &mut Vec, Vec, ) -> Result< (Option, Option>), Error, >, ) -> Result { let previous_timestamp = get_current_devicelist_timestamp(self, user_id).await?; let current_devices_data = self.get_current_devices(user_id).await?; let mut device_ids = self .get_current_device_list(user_id) .await? .map(|device_list| device_list.device_ids) .unwrap_or_default(); // Perform the update action, then generate new device list let (operation, timestamp) = action(&mut device_ids, current_devices_data)?; crate::device_list::verify_device_list_timestamp( previous_timestamp.as_ref(), timestamp.as_ref(), )?; let new_device_list = DeviceListRow::new(user_id, device_ids, timestamp); // Update timestamp in users table let timestamp_update_operation = device_list_timestamp_update_operation( user_id, previous_timestamp, new_device_list.timestamp, ); // Put updated device list (a new version) let put_device_list = Put::builder() .table_name(devices_table::NAME) .set_item(Some(new_device_list.clone().into())) .condition_expression( "attribute_not_exists(#user_id) AND attribute_not_exists(#item_id)", ) .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_names("#item_id", ATTR_ITEM_ID) .build(); let put_device_list_operation = TransactWriteItem::builder().put(put_device_list).build(); let operations = if let Some(operation) = operation { vec![ operation, put_device_list_operation, timestamp_update_operation, ] } else { vec![put_device_list_operation, timestamp_update_operation] }; self .client .transact_write_items() .set_transact_items(Some(operations)) .send() .await .map_err(|e| match DynamoDBError::from(e) { DynamoDBError::TransactionCanceledException( TransactionCanceledException { cancellation_reasons: Some(reasons), .. }, ) if reasons .iter() .any(|reason| reason.code() == Some("ConditionalCheckFailed")) => { Error::DeviceList(DeviceListError::ConcurrentUpdateError) } other => { - error!("Device list update transaction failed: {:?}", other); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Device list update transaction failed: {:?}", other + ); Error::AwsSdk(other) } })?; Ok(new_device_list) } /// Deletes all user data from devices table #[tracing::instrument(skip_all)] pub async fn delete_devices_table_rows_for_user( &self, user_id: impl Into, ) -> Result<(), Error> { // 1. get all rows // 2. batch write delete all // we project only the primary keys so we can pass these directly to delete requests let primary_keys = self .client .query() .table_name(devices_table::NAME) .projection_expression("#user_id, #item_id") .key_condition_expression("#user_id = :user_id") .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_names("#item_id", ATTR_ITEM_ID) .expression_attribute_values( ":user_id", AttributeValue::S(user_id.into()), ) .consistent_read(true) .send() .await .map_err(|e| { - error!("Failed to list user's items in devices table: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to list user's items in devices table: {:?}", e + ); Error::AwsSdk(e.into()) })? .items .unwrap_or_default(); let delete_requests = primary_keys .into_iter() .map(|item| { let request = DeleteRequest::builder().set_key(Some(item)).build(); WriteRequest::builder().delete_request(request).build() }) .collect::>(); // TODO: We can use the batch write helper from comm-services-lib when integrated for batch in delete_requests.chunks(25) { self .client .batch_write_item() .request_items(devices_table::NAME, batch.to_vec()) .send() .await .map_err(|e| { - error!("Failed to batch delete items from devices table: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to batch delete items from devices table: {:?}", e + ); Error::AwsSdk(e.into()) })?; } Ok(()) } } /// Gets timestamp of user's current device list. Returns None if the user /// doesn't have a device list yet. Storing the timestamp in the users table is /// required for consistency. It's used as a condition when updating the device /// list. #[tracing::instrument(skip_all)] async fn get_current_devicelist_timestamp( db: &crate::database::DatabaseClient, user_id: impl Into, ) -> Result>, Error> { let response = db .client .get_item() .table_name(USERS_TABLE) .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id.into())) .projection_expression(USERS_TABLE_DEVICELIST_TIMESTAMP_ATTRIBUTE_NAME) .send() .await .map_err(|e| { - error!("Failed to get user's device list timestamp: {:?}", e); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Failed to get user's device list timestamp: {:?}", e + ); Error::AwsSdk(e.into()) })?; let mut user_item = response.item.unwrap_or_default(); let raw_datetime = user_item.remove(USERS_TABLE_DEVICELIST_TIMESTAMP_ATTRIBUTE_NAME); // existing records will not have this field when // updating device list for the first time if raw_datetime.is_none() { return Ok(None); } let timestamp = DateTime::::try_from_attr( USERS_TABLE_DEVICELIST_TIMESTAMP_ATTRIBUTE_NAME, raw_datetime, )?; Ok(Some(timestamp)) } /// Generates update expression for current device list timestamp in users table. /// The previous timestamp is used as a condition to ensure that the value hasn't changed /// since we got it. This avoids race conditions when updating the device list. fn device_list_timestamp_update_operation( user_id: impl Into, previous_timestamp: Option>, new_timestamp: DateTime, ) -> TransactWriteItem { let update_builder = match previous_timestamp { Some(previous_timestamp) => Update::builder() .condition_expression("#device_list_timestamp = :previous_timestamp") .expression_attribute_values( ":previous_timestamp", AttributeValue::S(previous_timestamp.to_rfc3339()), ), // If there's no previous timestamp, the attribute shouldn't exist yet None => Update::builder() .condition_expression("attribute_not_exists(#device_list_timestamp)"), }; let update = update_builder .table_name(USERS_TABLE) .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id.into())) .update_expression("SET #device_list_timestamp = :new_timestamp") .expression_attribute_names( "#device_list_timestamp", USERS_TABLE_DEVICELIST_TIMESTAMP_ATTRIBUTE_NAME, ) .expression_attribute_values( ":new_timestamp", AttributeValue::S(new_timestamp.to_rfc3339()), ) .build(); TransactWriteItem::builder().update(update).build() } /// Helper function to query rows by given sort key prefix fn query_rows_with_prefix( db: &crate::database::DatabaseClient, user_id: impl Into, prefix: &'static str, ) -> QueryFluentBuilder { db.client .query() .table_name(devices_table::NAME) .key_condition_expression( "#user_id = :user_id AND begins_with(#item_id, :device_prefix)", ) .expression_attribute_names("#user_id", ATTR_USER_ID) .expression_attribute_names("#item_id", ATTR_ITEM_ID) .expression_attribute_values(":user_id", AttributeValue::S(user_id.into())) .expression_attribute_values( ":device_prefix", AttributeValue::S(prefix.to_string()), ) .consistent_read(true) } // Helper module for "migration" code into new device list schema. // We can get rid of this when primary device takes over the responsibility // of managing the device list. mod migration { use std::{cmp::Ordering, collections::HashSet}; use tracing::{debug, error, info}; use super::*; #[tracing::instrument(skip_all)] pub(super) fn reorder_device_list( user_id: &str, list: &mut [String], devices_data: &[DeviceRow], ) { if !verify_device_list_match(list, devices_data) { - error!("Found corrupt device list for user (userID={})!", user_id); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Found corrupt device list for user (userID={})!", user_id + ); return; } let Some(first_device) = list.first() else { debug!("Skipping device list rotation. Nothing to reorder."); return; }; let Some(primary_device) = determine_primary_device(devices_data) else { info!( "No valid primary device found for user (userID={}).\ Skipping device list reorder.", user_id ); return; }; if first_device == &primary_device.device_id { debug!("Skipping device list reorder. Primary device is already first"); return; } // swap primary device with the first one let Some(primary_device_idx) = list.iter().position(|id| id == &primary_device.device_id) else { error!( - "Primary device not found in device list (userID={})", - user_id + errorType = error_types::DEVICE_LIST_DB_LOG, + "Primary device not found in device list (userID={})", user_id ); return; }; list.swap(0, primary_device_idx); info!("Reordered device list for user (userID={})", user_id); } // checks if device list matches given devices data #[tracing::instrument(skip_all)] fn verify_device_list_match( list: &[String], devices_data: &[DeviceRow], ) -> bool { if list.len() != devices_data.len() { - error!("Device list length mismatch!"); + error!( + errorType = error_types::DEVICE_LIST_DB_LOG, + "Device list length mismatch!" + ); return false; } let actual_device_ids = devices_data .iter() .map(|device| &device.device_id) .collect::>(); let device_list_set = list.iter().collect::>(); if let Some(corrupt_device_id) = device_list_set .symmetric_difference(&actual_device_ids) .next() { error!( - "Device list is corrupt (unknown deviceID={})", - corrupt_device_id + errorType = error_types::DEVICE_LIST_DB_LOG, + "Device list is corrupt (unknown deviceID={})", corrupt_device_id ); return false; } true } /// Returns reference to primary device (if any) from given list of devices /// or None if there's no valid primary device. fn determine_primary_device(devices: &[DeviceRow]) -> Option<&DeviceRow> { // 1. Find mobile devices with valid token // 2. Prioritize these with latest code version // 3. If there's a tie, select the one with latest login time let mut mobile_devices = devices .iter() .filter(|device| { device.device_type == DeviceType::Ios || device.device_type == DeviceType::Android }) .collect::>(); mobile_devices.sort_by(|a, b| { let code_version_cmp = b.code_version.cmp(&a.code_version); if code_version_cmp == Ordering::Equal { b.login_time.cmp(&a.login_time) } else { code_version_cmp } }); mobile_devices.first().cloned() } #[cfg(test)] mod tests { use super::*; use chrono::Duration; #[test] fn reorder_skips_no_devices() { let mut list = vec![]; reorder_device_list("", &mut list, &[]); assert_eq!(list, Vec::::new()); } #[test] fn reorder_skips_single_device() { let mut list = vec!["test".into()]; let devices_data = vec![create_test_device("test", DeviceType::Web, 0, Utc::now())]; reorder_device_list("", &mut list, &devices_data); assert_eq!(list, vec!["test"]); } #[test] fn reorder_skips_for_valid_list() { let mut list = vec!["mobile".into(), "web".into()]; let devices_data = vec![ create_test_device("mobile", DeviceType::Android, 1, Utc::now()), create_test_device("web", DeviceType::Web, 0, Utc::now()), ]; reorder_device_list("", &mut list, &devices_data); assert_eq!(list, vec!["mobile", "web"]); } #[test] fn reorder_swaps_primary_device_when_possible() { let mut list = vec!["web".into(), "mobile".into()]; let devices_data = vec![ create_test_device("web", DeviceType::Web, 0, Utc::now()), create_test_device("mobile", DeviceType::Android, 1, Utc::now()), ]; reorder_device_list("", &mut list, &devices_data); assert_eq!(list, vec!["mobile", "web"]); } #[test] fn determine_primary_device_returns_none_for_empty_list() { let devices = vec![]; assert!(determine_primary_device(&devices).is_none()); } #[test] fn determine_primary_device_returns_none_for_web_only() { let devices = vec![create_test_device("web", DeviceType::Web, 0, Utc::now())]; assert!( determine_primary_device(&devices).is_none(), "Primary device should be None for web-only devices" ); } #[test] fn determine_primary_device_prioritizes_mobile() { let devices = vec![ create_test_device("mobile", DeviceType::Android, 0, Utc::now()), create_test_device("web", DeviceType::Web, 0, Utc::now()), ]; let primary_device = determine_primary_device(&devices) .expect("Primary device should be present"); assert_eq!( primary_device.device_id, "mobile", "Primary device should be mobile" ); } #[test] fn determine_primary_device_prioritizes_latest_code_version() { let devices_with_latest_code_version = vec![ create_test_device("mobile1", DeviceType::Android, 1, Utc::now()), create_test_device("mobile2", DeviceType::Android, 2, Utc::now()), create_test_device("web", DeviceType::Web, 0, Utc::now()), ]; let primary_device = determine_primary_device(&devices_with_latest_code_version) .expect("Primary device should be present"); assert_eq!( primary_device.device_id, "mobile2", "Primary device should be mobile with latest code version" ); } #[test] fn determine_primary_device_prioritizes_latest_login_time() { let devices = vec![ create_test_device("mobile1_today", DeviceType::Ios, 1, Utc::now()), create_test_device( "mobile2_yesterday", DeviceType::Android, 1, Utc::now() - Duration::days(1), ), create_test_device("web", DeviceType::Web, 0, Utc::now()), ]; let primary_device = determine_primary_device(&devices) .expect("Primary device should be present"); assert_eq!( primary_device.device_id, "mobile1_today", "Primary device should be mobile with latest login time" ); } #[test] fn determine_primary_device_keeps_deterministic_order() { // Given two identical devices, the first one should be selected as primary let today = Utc::now(); let devices_with_latest_code_version = vec![ create_test_device("mobile1", DeviceType::Android, 1, today), create_test_device("mobile2", DeviceType::Android, 1, today), ]; let primary_device = determine_primary_device(&devices_with_latest_code_version) .expect("Primary device should be present"); assert_eq!( primary_device.device_id, "mobile1", "Primary device selection should be deterministic" ); } #[test] fn determine_primary_device_all_rules_together() { use DeviceType::{Android, Ios, Web}; let today = Utc::now(); let yesterday = today - Duration::days(1); let devices = vec![ create_test_device("mobile1_today", Android, 1, today), create_test_device("mobile2_today", Android, 2, today), create_test_device("mobile3_yesterday", Ios, 1, yesterday), create_test_device("mobile4_yesterday", Ios, 2, yesterday), create_test_device("web", Web, 5, today), ]; let primary_device = determine_primary_device(&devices) .expect("Primary device should be present"); assert_eq!( primary_device.device_id, "mobile2_today", "Primary device should be mobile with latest code version and login time" ); } fn create_test_device( id: &str, platform: DeviceType, code_version: u64, login_time: DateTime, ) -> DeviceRow { DeviceRow { user_id: "test".into(), device_id: id.into(), device_type: platform, device_key_info: IdentityKeyInfo { key_payload: "".into(), key_payload_signature: "".into(), }, content_prekey: Prekey { prekey: "".into(), prekey_signature: "".into(), }, notif_prekey: Prekey { prekey: "".into(), prekey_signature: "".into(), }, code_version, login_time, } } } } diff --git a/services/identity/src/database/farcaster.rs b/services/identity/src/database/farcaster.rs index 1eab46a53..f4de0dcdb 100644 --- a/services/identity/src/database/farcaster.rs +++ b/services/identity/src/database/farcaster.rs @@ -1,149 +1,153 @@ use comm_lib::aws::ddb::types::AttributeValue; use comm_lib::aws::ddb::types::ReturnValue; use comm_lib::database::AttributeExtractor; use comm_lib::database::AttributeMap; use comm_lib::database::DBItemAttributeError; use comm_lib::database::DBItemError; use comm_lib::database::Value; use tracing::error; +use crate::constants::error_types; use crate::constants::USERS_TABLE; use crate::constants::USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME; use crate::constants::USERS_TABLE_FARCASTER_ID_INDEX; use crate::constants::USERS_TABLE_PARTITION_KEY; use crate::constants::USERS_TABLE_USERNAME_ATTRIBUTE; use crate::constants::USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE; use crate::grpc_services::protos::unauth::FarcasterUser; use super::DatabaseClient; use super::Error; pub struct FarcasterUserData(pub FarcasterUser); impl DatabaseClient { #[tracing::instrument(skip_all)] pub async fn get_farcaster_users( &self, farcaster_ids: Vec, ) -> Result, Error> { let mut users: Vec = Vec::new(); for id in farcaster_ids { let query_response = self .client .query() .table_name(USERS_TABLE) .index_name(USERS_TABLE_FARCASTER_ID_INDEX) .key_condition_expression(format!( "{} = :val", USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME )) .expression_attribute_values(":val", AttributeValue::S(id)) .send() .await .map_err(|e| { - error!("Failed to query users by farcasterID: {:?}", e); + error!( + errorType = error_types::FARCASTER_DB_LOG, + "Failed to query users by farcasterID: {:?}", e + ); Error::AwsSdk(e.into()) })? .items .and_then(|mut items| items.pop()) .map(FarcasterUserData::try_from) .transpose() .map_err(Error::from)?; if let Some(data) = query_response { users.push(data); } } Ok(users) } pub async fn add_farcaster_id( &self, user_id: String, farcaster_id: String, ) -> Result<(), Error> { let update_expression = format!( "SET {0} = if_not_exists({0}, :val)", USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME, ); let response = self .client .update_item() .table_name(USERS_TABLE) .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id)) .update_expression(update_expression) .expression_attribute_values( ":val", AttributeValue::S(farcaster_id.clone()), ) .return_values(ReturnValue::UpdatedNew) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; match response.attributes { None => return Err(Error::MissingItem), Some(mut attrs) => { let farcaster_id_from_table: String = attrs.take_attr(USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME)?; if farcaster_id_from_table != farcaster_id { return Err(Error::CannotOverwrite); } } } Ok(()) } pub async fn remove_farcaster_id( &self, user_id: String, ) -> Result<(), Error> { let update_expression = format!("REMOVE {}", USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME); self .client .update_item() .table_name(USERS_TABLE) .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id)) .update_expression(update_expression) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; Ok(()) } } impl TryFrom for FarcasterUserData { type Error = DBItemError; fn try_from(mut attrs: AttributeMap) -> Result { let user_id = attrs.take_attr(USERS_TABLE_PARTITION_KEY)?; let maybe_username = attrs.take_attr(USERS_TABLE_USERNAME_ATTRIBUTE)?; let maybe_wallet_address = attrs.take_attr(USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE)?; let username = match (maybe_username, maybe_wallet_address) { (Some(u), _) => u, (_, Some(w)) => w, (_, _) => { return Err(DBItemError { attribute_name: USERS_TABLE_USERNAME_ATTRIBUTE.to_string(), attribute_value: Value::AttributeValue(None), attribute_error: DBItemAttributeError::Missing, }); } }; let farcaster_id = attrs.take_attr(USERS_TABLE_FARCASTER_ID_ATTRIBUTE_NAME)?; Ok(Self(FarcasterUser { user_id, username, farcaster_id, })) } } diff --git a/services/identity/src/database/one_time_keys.rs b/services/identity/src/database/one_time_keys.rs index 4c0601830..2ab4341b7 100644 --- a/services/identity/src/database/one_time_keys.rs +++ b/services/identity/src/database/one_time_keys.rs @@ -1,490 +1,499 @@ use std::collections::HashSet; use comm_lib::{ aws::{ ddb::types::{ AttributeValue, Delete, DeleteRequest, TransactWriteItem, Update, WriteRequest, }, DynamoDBError, }, database::{ batch_operations::{batch_write, ExponentialBackoffConfig}, parse_int_attribute, AttributeExtractor, AttributeMap, DBItemAttributeError, DBItemError, }, }; use tracing::{debug, error, info}; use crate::{ - constants::{MAX_ONE_TIME_KEYS, ONE_TIME_KEY_UPLOAD_LIMIT_PER_ACCOUNT}, + constants::{ + error_types, MAX_ONE_TIME_KEYS, ONE_TIME_KEY_UPLOAD_LIMIT_PER_ACCOUNT, + }, database::DeviceIDAttribute, ddb_utils::{ create_one_time_key_partition_key, into_one_time_put_requests, into_one_time_update_and_delete_requests, is_transaction_retryable, OlmAccountType, }, error::{consume_error, Error}, olm::is_valid_olm_key, }; use super::DatabaseClient; impl DatabaseClient { /// Gets the next one-time key for the account and then, in a transaction, /// deletes the key and updates the key count /// /// Returns the retrieved one-time key if it exists and a boolean indicating /// whether the `spawn_refresh_keys_task`` was called #[tracing::instrument(skip_all)] pub(super) async fn get_one_time_key( &self, user_id: &str, device_id: &str, account_type: OlmAccountType, can_request_more_keys: bool, ) -> Result<(Option, bool), Error> { use crate::constants::devices_table; use crate::constants::retry; use crate::constants::ONE_TIME_KEY_MINIMUM_THRESHOLD; let attr_otk_count = match account_type { OlmAccountType::Content => devices_table::ATTR_CONTENT_OTK_COUNT, OlmAccountType::Notification => devices_table::ATTR_NOTIF_OTK_COUNT, }; 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); let result = crate::tunnelbroker::send_refresh_keys_request(&device_id).await; consume_error(result); }); } // TODO: Introduce `transact_write_helper` similar to `batch_write_helper` // in `comm-lib` to handle transactions with retries let mut attempt = 0; // TODO: Introduce nanny task that handles calling `spawn_refresh_keys_task` let mut requested_more_keys = false; loop { attempt += 1; if attempt > retry::MAX_ATTEMPTS { return Err(Error::MaxRetriesExceeded); } let otk_count = self.get_otk_count(user_id, device_id, account_type).await?; if otk_count < ONE_TIME_KEY_MINIMUM_THRESHOLD && can_request_more_keys { spawn_refresh_keys_task(device_id); requested_more_keys = true; } if otk_count < 1 { return Ok((None, requested_more_keys)); } let Some(otk_row) = self .get_one_time_keys(user_id, device_id, account_type, Some(1)) .await? .pop() else { return Err(Error::NotEnoughOneTimeKeys); }; let delete_otk_operation = otk_row.as_delete_request(); let update_otk_count = Update::builder() .table_name(devices_table::NAME) .key( devices_table::ATTR_USER_ID, AttributeValue::S(user_id.to_string()), ) .key( devices_table::ATTR_ITEM_ID, DeviceIDAttribute(device_id.into()).into(), ) .update_expression(format!("ADD {} :decrement_val", attr_otk_count)) .expression_attribute_values( ":decrement_val", AttributeValue::N("-1".to_string()), ) .condition_expression(format!("{} = :old_val", attr_otk_count)) .expression_attribute_values( ":old_val", AttributeValue::N(otk_count.to_string()), ) .build(); let update_otk_count_operation = TransactWriteItem::builder() .update(update_otk_count) .build(); let transaction = self .client .transact_write_items() .set_transact_items(Some(vec![ delete_otk_operation, update_otk_count_operation, ])) .send() .await; match transaction { Ok(_) => return Ok((Some(otk_row.otk), requested_more_keys)), Err(e) => { let dynamo_db_error = DynamoDBError::from(e); let retryable_codes = HashSet::from([ retry::CONDITIONAL_CHECK_FAILED, retry::TRANSACTION_CONFLICT, ]); if is_transaction_retryable(&dynamo_db_error, &retryable_codes) { info!("Encountered transaction conflict while retrieving one-time key - retrying"); } else { error!( + errorType = error_types::OTK_DB_LOG, "One-time key retrieval transaction failed: {:?}", dynamo_db_error ); return Err(Error::AwsSdk(dynamo_db_error)); } } } } } #[tracing::instrument(skip_all)] async fn get_one_time_keys( &self, user_id: &str, device_id: &str, account_type: OlmAccountType, num_keys: Option, ) -> Result, Error> { use crate::constants::one_time_keys_table::*; let partition_key = create_one_time_key_partition_key(user_id, device_id, account_type); let mut query = self .client .query() .table_name(NAME) .key_condition_expression("#pk = :pk") .expression_attribute_names("#pk", PARTITION_KEY) .expression_attribute_values(":pk", AttributeValue::S(partition_key)); if let Some(limit) = num_keys { // DynamoDB will reject the `query` request if `limit < 1` if limit < 1 { return Ok(Vec::new()); } query = query.limit(limit as i32); } let otk_rows = query .send() .await .map_err(|e| Error::AwsSdk(e.into()))? .items .unwrap_or_default() .into_iter() .map(OTKRow::try_from) .collect::, _>>() .map_err(Error::from)?; if let Some(limit) = num_keys { if otk_rows.len() != limit { - error!("There are fewer one-time keys than the number requested"); + error!( + errorType = error_types::OTK_DB_LOG, + "There are fewer one-time keys than the number requested" + ); return Err(Error::NotEnoughOneTimeKeys); } } Ok(otk_rows) } #[tracing::instrument(skip_all)] pub async fn append_one_time_prekeys( &self, user_id: &str, device_id: &str, content_one_time_keys: &Vec, notif_one_time_keys: &Vec, ) -> Result<(), Error> { use crate::constants::retry; let num_content_keys_to_append = content_one_time_keys.len(); let num_notif_keys_to_append = notif_one_time_keys.len(); if num_content_keys_to_append > ONE_TIME_KEY_UPLOAD_LIMIT_PER_ACCOUNT || num_notif_keys_to_append > ONE_TIME_KEY_UPLOAD_LIMIT_PER_ACCOUNT { return Err(Error::OneTimeKeyUploadLimitExceeded); } if content_one_time_keys .iter() .any(|otk| !is_valid_olm_key(otk)) || notif_one_time_keys.iter().any(|otk| !is_valid_olm_key(otk)) { debug!("Invalid one-time key format"); return Err(Error::InvalidFormat); } let current_time = chrono::Utc::now(); let content_otk_requests = into_one_time_put_requests( user_id, device_id, content_one_time_keys, OlmAccountType::Content, current_time, ); let notif_otk_requests = into_one_time_put_requests( user_id, device_id, notif_one_time_keys, OlmAccountType::Notification, current_time, ); let current_content_otk_count = self .get_otk_count(user_id, device_id, OlmAccountType::Content) .await?; let current_notif_otk_count = self .get_otk_count(user_id, device_id, OlmAccountType::Notification) .await?; let num_content_keys_to_delete = (num_content_keys_to_append + current_content_otk_count) .saturating_sub(MAX_ONE_TIME_KEYS); let num_notif_keys_to_delete = (num_notif_keys_to_append + current_notif_otk_count) .saturating_sub(MAX_ONE_TIME_KEYS); let content_keys_to_delete = self .get_one_time_keys( user_id, device_id, OlmAccountType::Content, Some(num_content_keys_to_delete), ) .await?; let notif_keys_to_delete = self .get_one_time_keys( user_id, device_id, OlmAccountType::Notification, Some(num_notif_keys_to_delete), ) .await?; let update_and_delete_otk_count_operation = into_one_time_update_and_delete_requests( user_id, device_id, num_content_keys_to_append, num_notif_keys_to_append, content_keys_to_delete, notif_keys_to_delete, ); let mut operations = Vec::new(); operations.extend_from_slice(&content_otk_requests); operations.extend_from_slice(¬if_otk_requests); operations.extend_from_slice(&update_and_delete_otk_count_operation); // TODO: Introduce `transact_write_helper` similar to `batch_write_helper` // in `comm-lib` to handle transactions with retries let mut attempt = 0; loop { attempt += 1; if attempt > retry::MAX_ATTEMPTS { return Err(Error::MaxRetriesExceeded); } let transaction = self .client .transact_write_items() .set_transact_items(Some(operations.clone())) .send() .await; match transaction { Ok(_) => break, Err(e) => { let dynamo_db_error = DynamoDBError::from(e); let retryable_codes = HashSet::from([retry::TRANSACTION_CONFLICT]); if is_transaction_retryable(&dynamo_db_error, &retryable_codes) { info!("Encountered transaction conflict while uploading one-time keys - retrying"); } else { error!( - "One-time key upload transaction failed: {:?}", - dynamo_db_error + errorType = error_types::OTK_DB_LOG, + "One-time key upload transaction failed: {:?}", dynamo_db_error ); return Err(Error::AwsSdk(dynamo_db_error)); } } } } Ok(()) } #[tracing::instrument(skip_all)] async fn get_otk_count( &self, user_id: &str, device_id: &str, account_type: OlmAccountType, ) -> Result { use crate::constants::devices_table; let attr_name = match account_type { OlmAccountType::Content => devices_table::ATTR_CONTENT_OTK_COUNT, OlmAccountType::Notification => devices_table::ATTR_NOTIF_OTK_COUNT, }; let response = self .client .get_item() .table_name(devices_table::NAME) .projection_expression(attr_name) .key( devices_table::ATTR_USER_ID, AttributeValue::S(user_id.to_string()), ) .key( devices_table::ATTR_ITEM_ID, DeviceIDAttribute(device_id.into()).into(), ) .send() .await .map_err(|e| { - error!("Failed to get user's OTK count: {:?}", e); + error!( + errorType = error_types::OTK_DB_LOG, + "Failed to get user's OTK count: {:?}", e + ); Error::AwsSdk(e.into()) })?; let mut user_item = response.item.unwrap_or_default(); match parse_int_attribute(attr_name, user_item.remove(attr_name)) { Ok(num) => Ok(num), Err(DBItemError { attribute_error: DBItemAttributeError::Missing, .. }) => Ok(0), Err(e) => Err(Error::Attribute(e)), } } /// Deletes all data for a user's device from one-time keys table pub async fn delete_otks_table_rows_for_user_device( &self, user_id: &str, device_id: &str, ) -> Result<(), Error> { use crate::constants::one_time_keys_table::*; let content_otk_primary_keys = self .get_one_time_keys(user_id, device_id, OlmAccountType::Content, None) .await?; let notif_otk_primary_keys = self .get_one_time_keys(user_id, device_id, OlmAccountType::Notification, None) .await?; let delete_requests = content_otk_primary_keys .into_iter() .chain(notif_otk_primary_keys) .map(|otk_row| { let request = DeleteRequest::builder() .key(PARTITION_KEY, AttributeValue::S(otk_row.partition_key)) .key(SORT_KEY, AttributeValue::S(otk_row.sort_key)) .build(); WriteRequest::builder().delete_request(request).build() }) .collect::>(); batch_write( &self.client, NAME, delete_requests, ExponentialBackoffConfig::default(), ) .await .map_err(Error::from)?; Ok(()) } /// Deletes all data for a user from one-time keys table pub async fn delete_otks_table_rows_for_user( &self, user_id: &str, ) -> Result<(), Error> { let maybe_device_list_row = self.get_current_device_list(user_id).await?; let Some(device_list_row) = maybe_device_list_row else { info!("No devices associated with user. Skipping one-time key removal."); return Ok(()); }; for device_id in device_list_row.device_ids { self .delete_otks_table_rows_for_user_device(user_id, &device_id) .await?; } Ok(()) } } pub struct OTKRow { pub partition_key: String, pub sort_key: String, pub otk: String, } impl OTKRow { pub fn as_delete_request(&self) -> TransactWriteItem { use crate::constants::one_time_keys_table as otk_table; let delete_otk = Delete::builder() .table_name(otk_table::NAME) .key( otk_table::PARTITION_KEY, AttributeValue::S(self.partition_key.to_string()), ) .key( otk_table::SORT_KEY, AttributeValue::S(self.sort_key.to_string()), ) .condition_expression("attribute_exists(#otk)") .expression_attribute_names("#otk", otk_table::ATTR_ONE_TIME_KEY) .build(); TransactWriteItem::builder().delete(delete_otk).build() } } impl TryFrom for OTKRow { type Error = DBItemError; fn try_from(mut attrs: AttributeMap) -> Result { use crate::constants::one_time_keys_table as otk_table; let partition_key = attrs.take_attr(otk_table::PARTITION_KEY)?; let sort_key = attrs.take_attr(otk_table::SORT_KEY)?; let otk: String = attrs.take_attr(otk_table::ATTR_ONE_TIME_KEY)?; Ok(Self { partition_key, sort_key, otk, }) } } diff --git a/services/identity/src/database/token.rs b/services/identity/src/database/token.rs index 97e9a7eda..af320d4e9 100644 --- a/services/identity/src/database/token.rs +++ b/services/identity/src/database/token.rs @@ -1,279 +1,283 @@ use std::collections::HashMap; use chrono::{DateTime, Utc}; use comm_lib::{ aws::ddb::{ operation::{get_item::GetItemOutput, put_item::PutItemOutput}, types::{AttributeValue, DeleteRequest, WriteRequest}, }, database::{ batch_operations::{batch_write, ExponentialBackoffConfig}, DBItemAttributeError, DBItemError, TryFromAttribute, }, }; use constant_time_eq::constant_time_eq; use tracing::{error, info}; use crate::{ + constants::error_types, error::Error, token::{AccessTokenData, AuthType}, }; use super::{create_composite_primary_key, DatabaseClient}; impl DatabaseClient { #[tracing::instrument(skip_all)] pub async fn get_access_token_data( &self, user_id: String, signing_public_key: String, ) -> Result, Error> { use crate::constants::token_table::*; let primary_key = create_composite_primary_key( (PARTITION_KEY.to_string(), user_id.clone()), (SORT_KEY.to_string(), signing_public_key.clone()), ); let get_item_result = self .client .get_item() .table_name(NAME) .set_key(Some(primary_key)) .consistent_read(true) .send() .await; match get_item_result { Ok(GetItemOutput { item: Some(mut item), .. }) => { let created = DateTime::::try_from_attr( ATTR_CREATED, item.remove(ATTR_CREATED), )?; let auth_type = parse_auth_type_attribute(item.remove(ATTR_AUTH_TYPE))?; let valid = parse_valid_attribute(item.remove(ATTR_VALID))?; let access_token = parse_token_attribute(item.remove(ATTR_TOKEN))?; Ok(Some(AccessTokenData { user_id, signing_public_key, access_token, created, auth_type, valid, })) } Ok(_) => { info!( "No item found for user {} and signing public key {} in token table", user_id, signing_public_key ); Ok(None) } Err(e) => { error!( - "DynamoDB client failed to get token for user {} with signing public key {}: {}", + errorType = error_types::TOKEN_DB_LOG, "DynamoDB client failed to get token for user {} with signing public key {}: {}", user_id, signing_public_key, e ); Err(Error::AwsSdk(e.into())) } } } pub async fn verify_access_token( &self, user_id: String, signing_public_key: String, access_token_to_verify: String, ) -> Result { let is_valid = self .get_access_token_data(user_id, signing_public_key) .await? .map(|access_token_data| { constant_time_eq( access_token_data.access_token.as_bytes(), access_token_to_verify.as_bytes(), ) && access_token_data.is_valid() }) .unwrap_or(false); Ok(is_valid) } pub async fn put_access_token_data( &self, access_token_data: AccessTokenData, ) -> Result { use crate::constants::token_table::*; let item = HashMap::from([ ( PARTITION_KEY.to_string(), AttributeValue::S(access_token_data.user_id), ), ( SORT_KEY.to_string(), AttributeValue::S(access_token_data.signing_public_key), ), ( ATTR_TOKEN.to_string(), AttributeValue::S(access_token_data.access_token), ), ( ATTR_CREATED.to_string(), AttributeValue::S(access_token_data.created.to_rfc3339()), ), ( ATTR_AUTH_TYPE.to_string(), AttributeValue::S(match access_token_data.auth_type { AuthType::Password => "password".to_string(), AuthType::Wallet => "wallet".to_string(), }), ), ( ATTR_VALID.to_string(), AttributeValue::Bool(access_token_data.valid), ), ]); self .client .put_item() .table_name(NAME) .set_item(Some(item)) .send() .await .map_err(|e| Error::AwsSdk(e.into())) } pub async fn delete_access_token_data( &self, user_id: String, device_id_key: String, ) -> Result<(), Error> { use crate::constants::token_table::*; self .client .delete_item() .table_name(NAME) .key(PARTITION_KEY.to_string(), AttributeValue::S(user_id)) .key(SORT_KEY.to_string(), AttributeValue::S(device_id_key)) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; Ok(()) } #[tracing::instrument(skip_all)] pub async fn delete_all_tokens_for_user( &self, user_id: &str, ) -> Result<(), Error> { use crate::constants::token_table::*; let primary_keys = self .client .query() .table_name(NAME) .projection_expression("#pk, #sk") .key_condition_expression("#pk = :pk") .expression_attribute_names("#pk", PARTITION_KEY) .expression_attribute_names("#sk", SORT_KEY) .expression_attribute_values( ":pk", AttributeValue::S(user_id.to_string()), ) .send() .await .map_err(|e| { - error!("Failed to list user's items in tokens table: {:?}", e); + error!( + errorType = error_types::TOKEN_DB_LOG, + "Failed to list user's items in tokens table: {:?}", e + ); Error::AwsSdk(e.into()) })? .items .unwrap_or_default(); let delete_requests = primary_keys .into_iter() .map(|item| { let request = DeleteRequest::builder().set_key(Some(item)).build(); WriteRequest::builder().delete_request(request).build() }) .collect::>(); batch_write( &self.client, NAME, delete_requests, ExponentialBackoffConfig::default(), ) .await .map_err(Error::from)?; Ok(()) } } fn parse_auth_type_attribute( attribute: Option, ) -> Result { use crate::constants::token_table::ATTR_AUTH_TYPE; if let Some(AttributeValue::S(auth_type)) = &attribute { match auth_type.as_str() { "password" => Ok(AuthType::Password), "wallet" => Ok(AuthType::Wallet), _ => Err(DBItemError::new( ATTR_AUTH_TYPE.to_string(), attribute.into(), DBItemAttributeError::IncorrectType, )), } } else { Err(DBItemError::new( ATTR_AUTH_TYPE.to_string(), attribute.into(), DBItemAttributeError::Missing, )) } } fn parse_valid_attribute( attribute: Option, ) -> Result { use crate::constants::token_table::ATTR_VALID; match attribute { Some(AttributeValue::Bool(valid)) => Ok(valid), Some(_) => Err(DBItemError::new( ATTR_VALID.to_string(), attribute.into(), DBItemAttributeError::IncorrectType, )), None => Err(DBItemError::new( ATTR_VALID.to_string(), attribute.into(), DBItemAttributeError::Missing, )), } } fn parse_token_attribute( attribute: Option, ) -> Result { use crate::constants::token_table::ATTR_TOKEN; match attribute { Some(AttributeValue::S(token)) => Ok(token), Some(_) => Err(DBItemError::new( ATTR_TOKEN.to_string(), attribute.into(), DBItemAttributeError::IncorrectType, )), None => Err(DBItemError::new( ATTR_TOKEN.to_string(), attribute.into(), DBItemAttributeError::Missing, )), } } diff --git a/services/identity/src/grpc_services/authenticated.rs b/services/identity/src/grpc_services/authenticated.rs index 8e84993db..148318cd0 100644 --- a/services/identity/src/grpc_services/authenticated.rs +++ b/services/identity/src/grpc_services/authenticated.rs @@ -1,844 +1,863 @@ use std::collections::HashMap; use crate::config::CONFIG; use crate::database::{DeviceListRow, DeviceListUpdate}; use crate::{ client_service::{handle_db_error, UpdateState, WorkflowInProgress}, - constants::request_metadata, + constants::{error_types, request_metadata}, database::DatabaseClient, ddb_utils::DateTimeExt, grpc_services::shared::get_value, }; use chrono::{DateTime, Utc}; use comm_opaque2::grpc::protocol_error_to_grpc_status; use tonic::{Request, Response, Status}; use tracing::{debug, error, warn}; use super::protos::auth::{ identity_client_service_server::IdentityClientService, DeletePasswordUserFinishRequest, DeletePasswordUserStartRequest, DeletePasswordUserStartResponse, GetDeviceListRequest, GetDeviceListResponse, InboundKeyInfo, InboundKeysForUserRequest, InboundKeysForUserResponse, KeyserverKeysResponse, LinkFarcasterAccountRequest, OutboundKeyInfo, OutboundKeysForUserRequest, OutboundKeysForUserResponse, PeersDeviceListsRequest, PeersDeviceListsResponse, RefreshUserPrekeysRequest, UpdateDeviceListRequest, UpdateUserPasswordFinishRequest, UpdateUserPasswordStartRequest, UpdateUserPasswordStartResponse, UploadOneTimeKeysRequest, UserIdentityRequest, UserIdentityResponse, }; use super::protos::unauth::Empty; #[derive(derive_more::Constructor)] pub struct AuthenticatedService { db_client: DatabaseClient, } fn get_auth_info(req: &Request<()>) -> Option<(String, String, String)> { debug!("Retrieving auth info for request: {:?}", req); let user_id = get_value(req, request_metadata::USER_ID)?; let device_id = get_value(req, request_metadata::DEVICE_ID)?; let access_token = get_value(req, request_metadata::ACCESS_TOKEN)?; Some((user_id, device_id, access_token)) } pub fn auth_interceptor( req: Request<()>, db_client: &DatabaseClient, ) -> Result, Status> { debug!("Intercepting request to check auth info: {:?}", req); let (user_id, device_id, access_token) = get_auth_info(&req) .ok_or_else(|| Status::unauthenticated("Missing credentials"))?; let handle = tokio::runtime::Handle::current(); let new_db_client = db_client.clone(); // This function cannot be `async`, yet must call the async db call // Force tokio to resolve future in current thread without an explicit .await let valid_token = tokio::task::block_in_place(move || { handle .block_on(new_db_client.verify_access_token( user_id, device_id, access_token, )) .map_err(handle_db_error) })?; if !valid_token { return Err(Status::aborted("Bad Credentials")); } Ok(req) } pub fn get_user_and_device_id( request: &Request, ) -> Result<(String, String), Status> { let user_id = get_value(request, request_metadata::USER_ID) .ok_or_else(|| Status::unauthenticated("Missing user_id field"))?; let device_id = get_value(request, request_metadata::DEVICE_ID) .ok_or_else(|| Status::unauthenticated("Missing device_id field"))?; Ok((user_id, device_id)) } #[tonic::async_trait] impl IdentityClientService for AuthenticatedService { #[tracing::instrument(skip_all)] async fn refresh_user_prekeys( &self, request: Request, ) -> Result, Status> { let (user_id, device_id) = get_user_and_device_id(&request)?; let message = request.into_inner(); debug!("Refreshing prekeys for user: {}", user_id); let content_keys = message .new_content_prekeys .ok_or_else(|| Status::invalid_argument("Missing content keys"))?; let notif_keys = message .new_notif_prekeys .ok_or_else(|| Status::invalid_argument("Missing notification keys"))?; self .db_client .update_device_prekeys( user_id, device_id, content_keys.into(), notif_keys.into(), ) .await .map_err(handle_db_error)?; let response = Response::new(Empty {}); Ok(response) } #[tracing::instrument(skip_all)] async fn get_outbound_keys_for_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let user_id = &message.user_id; let devices_map = self .db_client .get_keys_for_user(user_id, true) .await .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; let transformed_devices = devices_map .into_iter() .map(|(key, device_info)| (key, OutboundKeyInfo::from(device_info))) .collect::>(); Ok(tonic::Response::new(OutboundKeysForUserResponse { devices: transformed_devices, })) } #[tracing::instrument(skip_all)] async fn get_inbound_keys_for_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let user_id = &message.user_id; let devices_map = self .db_client .get_keys_for_user(user_id, false) .await .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; let transformed_devices = devices_map .into_iter() .map(|(key, device_info)| (key, InboundKeyInfo::from(device_info))) .collect::>(); let identifier = self .db_client .get_user_identifier(user_id) .await .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; Ok(tonic::Response::new(InboundKeysForUserResponse { devices: transformed_devices, identity: Some(identifier.into()), })) } #[tracing::instrument(skip_all)] async fn get_keyserver_keys( &self, request: Request, ) -> Result, Status> { let message = request.into_inner(); let identifier = self .db_client .get_user_identifier(&message.user_id) .await .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; let Some(keyserver_info) = self .db_client .get_keyserver_keys_for_user(&message.user_id) .await .map_err(handle_db_error)? else { return Err(Status::not_found("keyserver not found")); }; let primary_device_data = self .db_client .get_primary_device_data(&message.user_id) .await .map_err(handle_db_error)?; let primary_device_keys = primary_device_data.device_key_info; let response = Response::new(KeyserverKeysResponse { keyserver_info: Some(keyserver_info.into()), identity: Some(identifier.into()), primary_device_identity_info: Some(primary_device_keys.into()), }); return Ok(response); } #[tracing::instrument(skip_all)] async fn upload_one_time_keys( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, device_id) = get_user_and_device_id(&request)?; let message = request.into_inner(); debug!("Attempting to update one time keys for user: {}", user_id); self .db_client .append_one_time_prekeys( &user_id, &device_id, &message.content_one_time_prekeys, &message.notif_one_time_prekeys, ) .await .map_err(handle_db_error)?; Ok(tonic::Response::new(Empty {})) } #[tracing::instrument(skip_all)] async fn update_user_password_start( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, _) = get_user_and_device_id(&request)?; let message = request.into_inner(); let server_registration = comm_opaque2::server::Registration::new(); let server_message = server_registration .start( &CONFIG.server_setup, &message.opaque_registration_request, user_id.as_bytes(), ) .map_err(protocol_error_to_grpc_status)?; let update_state = UpdateState { user_id }; let session_id = self .db_client .insert_workflow(WorkflowInProgress::Update(update_state)) .await .map_err(handle_db_error)?; let response = UpdateUserPasswordStartResponse { session_id, opaque_registration_response: server_message, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn update_user_password_finish( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let Some(WorkflowInProgress::Update(state)) = self .db_client .get_workflow(message.session_id) .await .map_err(handle_db_error)? else { return Err(tonic::Status::not_found("session not found")); }; let server_registration = comm_opaque2::server::Registration::new(); let password_file = server_registration .finish(&message.opaque_registration_upload) .map_err(protocol_error_to_grpc_status)?; self .db_client .update_user_password(state.user_id, password_file) .await .map_err(handle_db_error)?; let response = Empty {}; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn log_out_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, device_id) = get_user_and_device_id(&request)?; self .db_client .remove_device(&user_id, &device_id) .await .map_err(handle_db_error)?; self .db_client .delete_otks_table_rows_for_user_device(&user_id, &device_id) .await .map_err(handle_db_error)?; self .db_client .delete_access_token_data(user_id, device_id) .await .map_err(handle_db_error)?; let response = Empty {}; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn delete_wallet_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, _) = get_user_and_device_id(&request)?; debug!("Attempting to delete wallet user: {}", user_id); self .db_client .delete_user(user_id) .await .map_err(handle_db_error)?; let response = Empty {}; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn delete_password_user_start( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, _) = get_user_and_device_id(&request)?; let message = request.into_inner(); debug!("Attempting to start deleting password user: {}", user_id); let maybe_username_and_password_file = self .db_client .get_username_and_password_file(&user_id) .await .map_err(handle_db_error)?; let Some((username, password_file_bytes)) = maybe_username_and_password_file else { return Err(tonic::Status::not_found("user not found")); }; let mut server_login = comm_opaque2::server::Login::new(); let server_response = server_login .start( &CONFIG.server_setup, &password_file_bytes, &message.opaque_login_request, username.as_bytes(), ) .map_err(protocol_error_to_grpc_status)?; let delete_state = construct_delete_password_user_info(server_login); let session_id = self .db_client .insert_workflow(WorkflowInProgress::PasswordUserDeletion(Box::new( delete_state, ))) .await .map_err(handle_db_error)?; let response = Response::new(DeletePasswordUserStartResponse { session_id, opaque_login_response: server_response, }); Ok(response) } #[tracing::instrument(skip_all)] async fn delete_password_user_finish( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, _) = get_user_and_device_id(&request)?; let message = request.into_inner(); debug!("Attempting to finish deleting password user: {}", user_id); let Some(WorkflowInProgress::PasswordUserDeletion(state)) = self .db_client .get_workflow(message.session_id) .await .map_err(handle_db_error)? else { return Err(tonic::Status::not_found("session not found")); }; let mut server_login = state.opaque_server_login; server_login .finish(&message.opaque_login_upload) .map_err(protocol_error_to_grpc_status)?; self .db_client .delete_user(user_id) .await .map_err(handle_db_error)?; let response = Empty {}; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn get_device_list_for_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let GetDeviceListRequest { user_id, since_timestamp, } = request.into_inner(); let since = since_timestamp .map(|timestamp| { DateTime::::from_utc_timestamp_millis(timestamp) .ok_or_else(|| tonic::Status::invalid_argument("Invalid timestamp")) }) .transpose()?; let mut db_result = self .db_client .get_device_list_history(user_id, since) .await .map_err(handle_db_error)?; // these should be sorted already, but just in case db_result.sort_by_key(|list| list.timestamp); let device_list_updates: Vec = db_result .into_iter() .map(RawDeviceList::from) .map(SignedDeviceList::try_from_raw) .collect::, _>>()?; let stringified_updates = device_list_updates .iter() .map(SignedDeviceList::as_json_string) .collect::, _>>()?; Ok(Response::new(GetDeviceListResponse { device_list_updates: stringified_updates, })) } #[tracing::instrument(skip_all)] async fn get_device_lists_for_users( &self, request: tonic::Request, ) -> Result, tonic::Status> { let PeersDeviceListsRequest { user_ids } = request.into_inner(); // do all fetches concurrently let mut fetch_tasks = tokio::task::JoinSet::new(); let mut device_lists = HashMap::with_capacity(user_ids.len()); for user_id in user_ids { let db_client = self.db_client.clone(); fetch_tasks.spawn(async move { let result = db_client.get_current_device_list(&user_id).await; (user_id, result) }); } while let Some(task_result) = fetch_tasks.join_next().await { match task_result { Ok((user_id, Ok(Some(device_list_row)))) => { let raw_list = RawDeviceList::from(device_list_row); let signed_list = SignedDeviceList::try_from_raw(raw_list)?; let serialized_list = signed_list.as_json_string()?; device_lists.insert(user_id, serialized_list); } Ok((user_id, Ok(None))) => { warn!(user_id, "User has no device list, skipping!"); } Ok((user_id, Err(err))) => { - error!(user_id, "Failed fetching device list: {err}"); + error!( + user_id, + errorType = error_types::GRPC_SERVICES_LOG, + "Failed fetching device list: {err}" + ); // abort fetching other users fetch_tasks.abort_all(); return Err(handle_db_error(err)); } Err(join_error) => { - error!("Failed to join device list task: {join_error}"); + error!( + errorType = error_types::GRPC_SERVICES_LOG, + "Failed to join device list task: {join_error}" + ); fetch_tasks.abort_all(); return Err(Status::aborted("unexpected error")); } } } let response = PeersDeviceListsResponse { users_device_lists: device_lists, }; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn update_device_list( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, _device_id) = get_user_and_device_id(&request)?; // TODO: when we stop doing "primary device rotation" (migration procedure) // we should verify if this RPC is called by primary device only let new_list = SignedDeviceList::try_from(request.into_inner())?; let update = DeviceListUpdate::try_from(new_list)?; self .db_client .apply_devicelist_update( &user_id, update, crate::device_list::validation::update_device_list_rpc_validator, ) .await .map_err(handle_db_error)?; Ok(Response::new(Empty {})) } #[tracing::instrument(skip_all)] async fn link_farcaster_account( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, _) = get_user_and_device_id(&request)?; let message = request.into_inner(); let mut get_farcaster_users_response = self .db_client .get_farcaster_users(vec![message.farcaster_id.clone()]) .await .map_err(handle_db_error)?; if get_farcaster_users_response.len() > 1 { - error!("multiple users associated with the same Farcaster ID"); + error!( + errorType = error_types::GRPC_SERVICES_LOG, + "multiple users associated with the same Farcaster ID" + ); return Err(Status::failed_precondition("cannot link Farcaster ID")); } if let Some(u) = get_farcaster_users_response.pop() { if u.0.user_id == user_id { return Ok(Response::new(Empty {})); } else { return Err(Status::already_exists( "farcaster ID already associated with different user", )); } } self .db_client .add_farcaster_id(user_id, message.farcaster_id) .await .map_err(handle_db_error)?; let response = Empty {}; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn unlink_farcaster_account( &self, request: tonic::Request, ) -> Result, tonic::Status> { let (user_id, _) = get_user_and_device_id(&request)?; self .db_client .remove_farcaster_id(user_id) .await .map_err(handle_db_error)?; let response = Empty {}; Ok(Response::new(response)) } #[tracing::instrument(skip_all)] async fn find_user_identity( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let identifier = self .db_client .get_user_identifier(&message.user_id) .await .map_err(handle_db_error)? .ok_or_else(|| tonic::Status::not_found("user not found"))?; let response = UserIdentityResponse { identity: Some(identifier.into()), }; return Ok(Response::new(response)); } } // raw device list that can be serialized to JSON (and then signed in the future) #[derive(serde::Serialize, serde::Deserialize)] struct RawDeviceList { devices: Vec, timestamp: i64, } impl From for RawDeviceList { fn from(row: DeviceListRow) -> Self { Self { devices: row.device_ids, timestamp: row.timestamp.timestamp_millis(), } } } #[derive(serde::Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] struct SignedDeviceList { /// JSON-stringified [`RawDeviceList`] raw_device_list: String, /// Current primary device signature. /// NOTE: Present only when the payload is received from primary device. /// It's `None` for Identity-generated device-lists #[serde(default)] #[serde(skip_serializing_if = "Option::is_none")] cur_primary_signature: Option, /// Previous primary device signature. Present only /// if primary device has changed since last update. #[serde(default)] #[serde(skip_serializing_if = "Option::is_none")] last_primary_signature: Option, } impl SignedDeviceList { /// Serialize (and sign in the future) a [`RawDeviceList`] fn try_from_raw(raw: RawDeviceList) -> Result { let stringified_list = serde_json::to_string(&raw).map_err(|err| { - error!("Failed to serialize raw device list: {}", err); + error!( + errorType = error_types::GRPC_SERVICES_LOG, + "Failed to serialize raw device list: {}", err + ); tonic::Status::failed_precondition("unexpected error") })?; Ok(Self { raw_device_list: stringified_list, cur_primary_signature: None, last_primary_signature: None, }) } fn as_raw(&self) -> Result { // The device list payload is sent as an escaped JSON payload. // Escaped double quotes need to be trimmed before attempting to deserialize serde_json::from_str(&self.raw_device_list.replace(r#"\""#, r#"""#)) .map_err(|err| { warn!("Failed to deserialize raw device list: {}", err); tonic::Status::invalid_argument("invalid device list payload") }) } /// Serializes the signed device list to a JSON string fn as_json_string(&self) -> Result { serde_json::to_string(self).map_err(|err| { - error!("Failed to serialize device list updates: {}", err); + error!( + errorType = error_types::GRPC_SERVICES_LOG, + "Failed to serialize device list updates: {}", err + ); tonic::Status::failed_precondition("unexpected error") }) } } impl TryFrom for SignedDeviceList { type Error = tonic::Status; fn try_from(request: UpdateDeviceListRequest) -> Result { serde_json::from_str(&request.new_device_list).map_err(|err| { warn!("Failed to deserialize device list update: {}", err); tonic::Status::invalid_argument("invalid device list payload") }) } } impl TryFrom for DeviceListUpdate { type Error = tonic::Status; fn try_from(signed_list: SignedDeviceList) -> Result { let RawDeviceList { devices, timestamp: raw_timestamp, } = signed_list.as_raw()?; let timestamp = DateTime::::from_utc_timestamp_millis(raw_timestamp) .ok_or_else(|| { - error!("Failed to parse RawDeviceList timestamp!"); + error!( + errorType = error_types::GRPC_SERVICES_LOG, + "Failed to parse RawDeviceList timestamp!" + ); tonic::Status::invalid_argument("invalid timestamp") })?; Ok(DeviceListUpdate::new(devices, timestamp)) } } #[derive(Clone, serde::Serialize, serde::Deserialize)] pub struct DeletePasswordUserInfo { pub opaque_server_login: comm_opaque2::server::Login, } fn construct_delete_password_user_info( opaque_server_login: comm_opaque2::server::Login, ) -> DeletePasswordUserInfo { DeletePasswordUserInfo { opaque_server_login, } } #[cfg(test)] mod tests { use super::*; #[test] fn deserialize_device_list_signature() { let payload_with_signature = r#"{"rawDeviceList":"{\"devices\":[\"device1\"],\"timestamp\":111111111}","curPrimarySignature":"foo"}"#; let payload_without_signatures = r#"{"rawDeviceList":"{\"devices\":[\"device1\",\"device2\"],\"timestamp\":222222222}"}"#; let list_with_signature: SignedDeviceList = serde_json::from_str(payload_with_signature).unwrap(); let list_without_signatures: SignedDeviceList = serde_json::from_str(payload_without_signatures).unwrap(); assert_eq!( list_with_signature.cur_primary_signature, Some("foo".to_string()) ); assert!(list_with_signature.last_primary_signature.is_none()); assert!(list_without_signatures.cur_primary_signature.is_none()); assert!(list_without_signatures.last_primary_signature.is_none()); } #[test] fn serialize_device_list_signatures() { let raw_list = r#"{"devices":["device1"],"timestamp":111111111}"#; let expected_payload_without_signatures = r#"{"rawDeviceList":"{\"devices\":[\"device1\"],\"timestamp\":111111111}"}"#; let device_list_without_signature = SignedDeviceList { raw_device_list: raw_list.to_string(), cur_primary_signature: None, last_primary_signature: None, }; assert_eq!( device_list_without_signature.as_json_string().unwrap(), expected_payload_without_signatures ); let expected_payload_with_signature = r#"{"rawDeviceList":"{\"devices\":[\"device1\"],\"timestamp\":111111111}","curPrimarySignature":"foo"}"#; let device_list_with_cur_signature = SignedDeviceList { raw_device_list: raw_list.to_string(), cur_primary_signature: Some("foo".to_string()), last_primary_signature: None, }; assert_eq!( device_list_with_cur_signature.as_json_string().unwrap(), expected_payload_with_signature ); } #[test] fn serialize_device_list_updates() { let raw_updates = vec![ RawDeviceList { devices: vec!["device1".into()], timestamp: 111111111, }, RawDeviceList { devices: vec!["device1".into(), "device2".into()], timestamp: 222222222, }, ]; let expected_raw_list1 = r#"{"devices":["device1"],"timestamp":111111111}"#; let expected_raw_list2 = r#"{"devices":["device1","device2"],"timestamp":222222222}"#; let signed_updates = raw_updates .into_iter() .map(SignedDeviceList::try_from_raw) .collect::, _>>() .expect("signing device list updates failed"); assert_eq!(signed_updates[0].raw_device_list, expected_raw_list1); assert_eq!(signed_updates[1].raw_device_list, expected_raw_list2); let stringified_updates = signed_updates .iter() .map(serde_json::to_string) .collect::, _>>() .expect("serialize signed device lists failed"); let expected_stringified_list1 = r#"{"rawDeviceList":"{\"devices\":[\"device1\"],\"timestamp\":111111111}"}"#; let expected_stringified_list2 = r#"{"rawDeviceList":"{\"devices\":[\"device1\",\"device2\"],\"timestamp\":222222222}"}"#; assert_eq!(stringified_updates[0], expected_stringified_list1); assert_eq!(stringified_updates[1], expected_stringified_list2); } #[test] fn deserialize_device_list_update() { let raw_payload = r#"{"rawDeviceList":"{\"devices\":[\"device1\",\"device2\"],\"timestamp\":123456789}"}"#; let request = UpdateDeviceListRequest { new_device_list: raw_payload.to_string(), }; let signed_list = SignedDeviceList::try_from(request) .expect("Failed to parse SignedDeviceList"); let update = DeviceListUpdate::try_from(signed_list) .expect("Failed to parse DeviceListUpdate from signed list"); let expected_timestamp = DateTime::::from_utc_timestamp_millis(123456789).unwrap(); assert_eq!(update.timestamp, expected_timestamp); assert_eq!( update.devices, vec!["device1".to_string(), "device2".to_string()] ); } } diff --git a/services/identity/src/siwe.rs b/services/identity/src/siwe.rs index e8afa6344..45680da02 100644 --- a/services/identity/src/siwe.rs +++ b/services/identity/src/siwe.rs @@ -1,157 +1,169 @@ use std::collections::HashMap; use chrono::Utc; use comm_lib::{ aws::ddb::types::AttributeValue, database::{AttributeExtractor, AttributeMap, TryFromAttribute}, }; use regex::Regex; use siwe::Message; use tonic::Status; use tracing::error; use crate::constants::{ - SOCIAL_PROOF_MESSAGE_ATTRIBUTE, SOCIAL_PROOF_SIGNATURE_ATTRIBUTE, + error_types, SOCIAL_PROOF_MESSAGE_ATTRIBUTE, SOCIAL_PROOF_SIGNATURE_ATTRIBUTE, }; pub fn parse_and_verify_siwe_message( siwe_message: &str, siwe_signature: &str, ) -> Result { let siwe_message: Message = siwe_message.parse().map_err(|e| { - error!("Failed to parse SIWE message: {}", e); + error!( + errorType = error_types::SIWE_LOG, + "Failed to parse SIWE message: {}", e + ); Status::invalid_argument("invalid message") })?; let decoded_signature = hex::decode(siwe_signature.trim_start_matches("0x")) .map_err(|e| { - error!("Failed to decode SIWE signature: {}", e); + error!( + errorType = error_types::SIWE_LOG, + "Failed to decode SIWE signature: {}", e + ); Status::invalid_argument("invalid signature") })?; let signature = decoded_signature.try_into().map_err(|e| { - error!("Conversion to SIWE signature failed: {:?}", e); + error!( + errorType = error_types::SIWE_LOG, + "Conversion to SIWE signature failed: {:?}", e + ); Status::invalid_argument("invalid message") })?; siwe_message .verify(signature, None, None, Some(&Utc::now())) .map_err(|e| { - error!("Signature verification failed: {}", e); + error!( + errorType = error_types::SIWE_LOG, + "Signature verification failed: {}", e + ); Status::unauthenticated("message not authenticated") })?; Ok(siwe_message) } pub fn is_valid_ethereum_address(candidate: &str) -> bool { let ethereum_address_regex = Regex::new(r"^0x[a-fA-F0-9]{40}$").unwrap(); ethereum_address_regex.is_match(candidate) } #[derive(derive_more::Constructor)] pub struct SocialProof { pub message: String, pub signature: String, } impl From for AttributeValue { fn from(value: SocialProof) -> Self { AttributeValue::M(HashMap::from([ ( SOCIAL_PROOF_MESSAGE_ATTRIBUTE.to_string(), AttributeValue::S(value.message), ), ( SOCIAL_PROOF_SIGNATURE_ATTRIBUTE.to_string(), AttributeValue::S(value.signature), ), ])) } } impl TryFrom for SocialProof { type Error = comm_lib::database::DBItemError; fn try_from(mut attrs: AttributeMap) -> Result { let message = attrs.take_attr(SOCIAL_PROOF_MESSAGE_ATTRIBUTE)?; let signature = attrs.take_attr(SOCIAL_PROOF_SIGNATURE_ATTRIBUTE)?; Ok(Self { message, signature }) } } impl TryFromAttribute for SocialProof { fn try_from_attr( attribute_name: impl Into, attribute: Option, ) -> Result { AttributeMap::try_from_attr(attribute_name, attribute) .and_then(SocialProof::try_from) } } #[cfg(test)] mod tests { use crate::constants::USERS_TABLE_SOCIAL_PROOF_ATTRIBUTE_NAME; use super::*; #[test] fn test_valid_ethereum_address() { assert!(is_valid_ethereum_address( "0x1234567890123456789012345678901234567890" ),); assert!(is_valid_ethereum_address( "0xABCDEF123456789012345678901234567890ABCD" )); assert!(is_valid_ethereum_address( "0xabcdef123456789012345678901234567890abcd" )); } #[test] fn test_invalid_ethereum_address() { // Shorter than 42 characters assert_eq!( is_valid_ethereum_address("0x12345678901234567890123456789012345678"), false ); // Longer than 42 characters assert_eq!( is_valid_ethereum_address("0x123456789012345678901234567890123456789012"), false ); // Missing 0x prefix assert_eq!( is_valid_ethereum_address("1234567890123456789012345678901234567890"), false ); // Contains invalid characters assert_eq!( is_valid_ethereum_address("0x1234567890GHIJKL9012345678901234567890"), false ); // Empty string assert_eq!(is_valid_ethereum_address(""), false); } #[test] fn test_social_proof_ddb_format() { let message = "foo"; let signature = "bar"; let social_proof = SocialProof::new(message.to_string(), signature.to_string()); let mut user_item = AttributeMap::from([( USERS_TABLE_SOCIAL_PROOF_ATTRIBUTE_NAME.to_string(), social_proof.into(), )]); let social_proof_from_attr: SocialProof = user_item .take_attr(USERS_TABLE_SOCIAL_PROOF_ATTRIBUTE_NAME) .expect("social proof fetch failed"); assert_eq!(social_proof_from_attr.message, message); assert_eq!(social_proof_from_attr.signature, signature); } } diff --git a/services/identity/src/sync_identity_search.rs b/services/identity/src/sync_identity_search.rs index 5e74fda88..21d7068ce 100644 --- a/services/identity/src/sync_identity_search.rs +++ b/services/identity/src/sync_identity_search.rs @@ -1,105 +1,105 @@ use crate::config::CONFIG; -use crate::constants::IDENTITY_SEARCH_INDEX; +use crate::constants::{error_types, IDENTITY_SEARCH_INDEX}; use crate::database::DatabaseClient; use crate::error; use identity_search_messages::IdentitySearchUser; use serde_json::json; use tracing::error; pub async fn sync_index( database_client: &DatabaseClient, ) -> Result<(), error::Error> { let identity_reserved_users = database_client.get_all_reserved_user_details().await?; let identity_users = database_client.get_all_user_details().await?; let mut identity_search_users: Vec = identity_reserved_users .into_iter() .map(|user| IdentitySearchUser { username: user.username, user_id: user.user_id, }) .collect(); for user in identity_users { identity_search_users.push(IdentitySearchUser { username: user.username, user_id: user.user_id, }); } let client = reqwest::Client::new(); clear_index(&client).await?; restore_index(&client, &identity_search_users).await?; Ok(()) } pub async fn clear_index( reqwest_client: &reqwest::Client, ) -> Result<(), error::Error> { let url = format!( "https://{}/{}/_delete_by_query", &CONFIG.opensearch_endpoint, IDENTITY_SEARCH_INDEX ); let query = serde_json::json!({ "query": { "match_all": {} } }); let response = reqwest_client .post(&url) .header(reqwest::header::CONTENT_TYPE, "application/json") .json(&query) .send() .await .expect("Failed to send clear index request"); if !response.status().is_success() { - error!("Sync Error: Failed to clear index"); + error!(errorType = error_types::SYNC_LOG, "Failed to clear index"); } Ok(()) } pub async fn restore_index( reqwest_client: &reqwest::Client, identity_search_users: &Vec, ) -> Result<(), error::Error> { let mut bulk_data = String::new(); for user in identity_search_users { let action = json!({ "index": { "_index": IDENTITY_SEARCH_INDEX, "_id": user.user_id } }); bulk_data.push_str(&action.to_string()); bulk_data.push('\n'); bulk_data.push_str( &serde_json::to_string(&user) .expect("Failed to serialize identity search index user"), ); bulk_data.push('\n'); } let url = format!( "https://{}/{}/_bulk/", &CONFIG.opensearch_endpoint, IDENTITY_SEARCH_INDEX ); let response = reqwest_client .post(&url) .header(reqwest::header::CONTENT_TYPE, "application/x-ndjson") .body(bulk_data) .send() .await .expect("Failed to send restore index request"); if !response.status().is_success() { - error!("Sync Error: Failed to restore index"); + error!(errorType = error_types::SYNC_LOG, "Failed to restore index"); } Ok(()) } diff --git a/services/identity/src/tunnelbroker.rs b/services/identity/src/tunnelbroker.rs index c234dc23f..8167c6b02 100644 --- a/services/identity/src/tunnelbroker.rs +++ b/services/identity/src/tunnelbroker.rs @@ -1,48 +1,52 @@ use crate::config::CONFIG; use grpc_clients::tunnelbroker::create_tunnelbroker_client as shared_tb_client; use grpc_clients::tunnelbroker::protos; use protos::tunnelbroker_service_client::TunnelbrokerServiceClient; use protos::{Empty, MessageToDevice}; use tonic::transport::Channel; use tonic::Response; use tonic::Status; use tracing::error; use tunnelbroker_messages as messages; +use crate::constants::error_types; use crate::error::Error; pub async fn create_tunnelbroker_client( ) -> Result, Error> { shared_tb_client(&CONFIG.tunnelbroker_endpoint) .await .map_err(|e| { - error!("Unable able to connect to tunnelbroker: {:?}", e); + error!( + errorType = error_types::TUNNELBROKER_LOG, + "Unable able to connect to tunnelbroker: {:?}", e + ); Error::Status(Status::invalid_argument(format!("{}", e))) }) } pub async fn send_refresh_keys_request( device_id: &str, ) -> Result, Error> { use crate::constants::ONE_TIME_KEY_REFRESH_NUMBER; let mut tunnelbroker_client = create_tunnelbroker_client().await?; let refresh_request = messages::RefreshKeyRequest { device_id: device_id.to_string(), number_of_keys: ONE_TIME_KEY_REFRESH_NUMBER, }; let payload = serde_json::to_string(&refresh_request).unwrap(); let request = MessageToDevice { device_id: device_id.to_string(), payload, }; let grpc_message = tonic::Request::new(request); Ok( tunnelbroker_client .send_message_to_device(grpc_message) .await?, ) } diff --git a/services/identity/src/websockets/auth.rs b/services/identity/src/websockets/auth.rs index 13fc2e071..6f820210f 100644 --- a/services/identity/src/websockets/auth.rs +++ b/services/identity/src/websockets/auth.rs @@ -1,81 +1,90 @@ use client_proto::VerifyUserAccessTokenRequest; use grpc_clients::identity; use grpc_clients::tonic::Request; use identity::get_unauthenticated_client; use identity::protos::unauthenticated as client_proto; use identity_search_messages::IdentitySearchAuthMessage; use tracing::{debug, error}; -use crate::constants::DEFAULT_IDENTITY_ENDPOINT; +use crate::constants::{error_types, DEFAULT_IDENTITY_ENDPOINT}; use crate::websockets::errors::WebsocketError; const PLACEHOLDER_CODE_VERSION: u64 = 0; const DEVICE_TYPE: &str = "service"; #[tracing::instrument(skip_all)] async fn verify_user_access_token( user_id: &str, device_id: &str, access_token: &str, ) -> Result { let grpc_client = get_unauthenticated_client( DEFAULT_IDENTITY_ENDPOINT, PLACEHOLDER_CODE_VERSION, DEVICE_TYPE.to_string(), ) .await; let mut grpc_client = match grpc_client { Ok(grpc_client) => grpc_client, Err(e) => { - error!("Failed to get unauthenticated client: {}", e); + error!( + errorType = error_types::SEARCH_LOG, + "Failed to get unauthenticated client: {}", e + ); return Err(WebsocketError::AuthError); } }; let message = VerifyUserAccessTokenRequest { user_id: user_id.to_string(), device_id: device_id.to_string(), access_token: access_token.to_string(), }; let request = Request::new(message); let response = match grpc_client.verify_user_access_token(request).await { Ok(response) => response, Err(_) => { - error!("Failed to verify user access token"); + error!( + errorType = error_types::SEARCH_LOG, + "Failed to verify user access token" + ); return Err(WebsocketError::AuthError); } }; Ok(response.into_inner().token_valid) } #[tracing::instrument(skip_all)] pub async fn handle_auth_message(message: &str) -> Result<(), WebsocketError> { let auth_message = serde_json::from_str(message.trim()); let auth_message: IdentitySearchAuthMessage = match auth_message { Ok(auth_message) => auth_message, Err(_) => { - error!("Failed to parse auth message"); + error!( + errorType = error_types::SEARCH_LOG, + "Failed to parse auth message" + ); return Err(WebsocketError::InvalidMessage); } }; let user_id = auth_message.user_id; let device_id = auth_message.device_id; let access_token = auth_message.access_token; let is_valid_token = verify_user_access_token(&user_id, &device_id, &access_token).await?; if is_valid_token { debug!("User {} authenticated", user_id); } else { debug!("User {} not authenticated", user_id); return Err(WebsocketError::UnauthorizedDevice); } Ok(()) } diff --git a/services/identity/src/websockets/errors.rs b/services/identity/src/websockets/errors.rs index 5bf3805d4..27e15c262 100644 --- a/services/identity/src/websockets/errors.rs +++ b/services/identity/src/websockets/errors.rs @@ -1,26 +1,36 @@ +use crate::constants::error_types; + pub type BoxedError = Box; #[derive( Debug, derive_more::Display, derive_more::From, derive_more::Error, )] pub enum WebsocketError { InvalidMessage, UnauthorizedDevice, SearchError, AuthError, SerializationError, } impl From for WebsocketError { fn from(err: serde_json::Error) -> Self { - tracing::error!("Error serializing: {}", err); + tracing::error!( + errorType = error_types::SEARCH_LOG, + "Error serializing: {}", + err + ); WebsocketError::SerializationError } } impl From for WebsocketError { fn from(err: reqwest::Error) -> Self { - tracing::error!("Error with search request: {}", err); + tracing::error!( + errorType = error_types::SEARCH_LOG, + "Error with search request: {}", + err + ); WebsocketError::SearchError } } diff --git a/services/identity/src/websockets/mod.rs b/services/identity/src/websockets/mod.rs index 6537fc08c..30ca2cb6e 100644 --- a/services/identity/src/websockets/mod.rs +++ b/services/identity/src/websockets/mod.rs @@ -1,342 +1,357 @@ use std::future::Future; use std::net::SocketAddr; use std::pin::Pin; use std::sync::Arc; use futures::lock::Mutex; use futures_util::{SinkExt, StreamExt}; use hyper::{Body, Request, Response, StatusCode}; use hyper_tungstenite::tungstenite::Message; use hyper_tungstenite::HyperWebsocket; use identity_search_messages::{ ConnectionInitializationResponse, ConnectionInitializationStatus, Heartbeat, IdentitySearchFailure, IdentitySearchMethod, IdentitySearchResponse, IdentitySearchResult, IdentitySearchUser, MessagesToServer, }; use serde::{Deserialize, Serialize}; use tokio::net::TcpListener; use tracing::{debug, error, info}; mod auth; mod opensearch; mod send; use crate::config::CONFIG; use crate::constants::{ - IDENTITY_SEARCH_INDEX, IDENTITY_SEARCH_RESULT_SIZE, + error_types, IDENTITY_SEARCH_INDEX, IDENTITY_SEARCH_RESULT_SIZE, IDENTITY_SERVICE_WEBSOCKET_ADDR, SOCKET_HEARTBEAT_TIMEOUT, }; use opensearch::OpenSearchResponse; use send::{send_message, WebsocketSink}; pub mod errors; #[derive(Serialize, Deserialize)] struct Query { size: u32, query: Prefix, } #[derive(Serialize, Deserialize)] struct Prefix { prefix: Username, } #[derive(Serialize, Deserialize)] struct Username { username: String, } struct WebsocketService { addr: SocketAddr, } impl hyper::service::Service> for WebsocketService { type Response = Response; type Error = errors::BoxedError; type Future = Pin> + Send>>; fn poll_ready( &mut self, _: &mut std::task::Context<'_>, ) -> std::task::Poll> { std::task::Poll::Ready(Ok(())) } fn call(&mut self, mut req: Request) -> Self::Future { let addr = self.addr; let future = async move { debug!( "Incoming HTTP request on WebSocket port: {} {}", req.method(), req.uri().path() ); if hyper_tungstenite::is_upgrade_request(&req) { let (response, websocket) = hyper_tungstenite::upgrade(&mut req, None)?; tokio::spawn(async move { accept_connection(websocket, addr).await; }); return Ok(response); } let response = match req.uri().path() { "/health" => Response::new(Body::from("OK")), _ => Response::builder() .status(StatusCode::NOT_FOUND) .body(Body::from("Not found"))?, }; Ok(response) }; Box::pin(future) } } #[tracing::instrument(skip_all)] pub async fn run_server() -> Result<(), errors::BoxedError> { let addr: SocketAddr = IDENTITY_SERVICE_WEBSOCKET_ADDR.parse()?; let listener = TcpListener::bind(&addr).await.expect("Failed to bind"); info!("Listening to WebSocket traffic on {}", addr); let mut http = hyper::server::conn::Http::new(); http.http1_only(true); http.http1_keep_alive(true); while let Ok((stream, addr)) = listener.accept().await { let connection = http .serve_connection(stream, WebsocketService { addr }) .with_upgrades(); tokio::spawn(async move { if let Err(err) = connection.await { - error!("Error serving HTTP/WebSocket connection: {:?}", err); + error!( + errorType = error_types::SEARCH_LOG, + "Error serving HTTP/WebSocket connection: {:?}", err + ); } }); } Ok(()) } #[tracing::instrument(skip_all)] async fn send_search_request( url: &str, json_body: T, ) -> Result { let client = reqwest::Client::new(); client .post(url) .header(reqwest::header::CONTENT_TYPE, "application/json") .json(&json_body) .send() .await } #[tracing::instrument(skip_all)] async fn close_connection(outgoing: WebsocketSink) { if let Err(e) = outgoing.lock().await.close().await { - error!("Error closing connection: {}", e); + error!( + errorType = error_types::SEARCH_LOG, + "Error closing connection: {}", e + ); } } #[tracing::instrument(skip_all)] async fn handle_prefix_search( request_id: &str, prefix_request: identity_search_messages::IdentitySearchPrefix, ) -> Result { let prefix_query = Query { size: IDENTITY_SEARCH_RESULT_SIZE, query: Prefix { prefix: Username { username: prefix_request.prefix.trim().to_string(), }, }, }; let opensearch_url = format!( "https://{}/{}/_search/", &CONFIG.opensearch_endpoint, IDENTITY_SEARCH_INDEX ); let search_response = send_search_request(&opensearch_url, prefix_query) .await? .json::>() .await?; let usernames: Vec = search_response .hits .inner .into_iter() .filter_map(|hit| hit.source) .collect(); let search_result = IdentitySearchResult { id: request_id.to_string(), hits: usernames, }; Ok(search_result) } #[tracing::instrument(skip_all)] async fn handle_websocket_frame( text: String, outgoing: WebsocketSink, ) -> Result<(), errors::WebsocketError> { let Ok(serialized_message) = serde_json::from_str::(&text) else { return Err(errors::WebsocketError::SerializationError); }; match serialized_message { MessagesToServer::Heartbeat(Heartbeat {}) => { debug!("Received heartbeat"); Ok(()) } MessagesToServer::IdentitySearchQuery(search_query) => { let handler_result = match search_query.search_method { IdentitySearchMethod::IdentitySearchPrefix(prefix_query) => { handle_prefix_search(&search_query.id, prefix_query).await } }; let search_response = match handler_result { Ok(search_result) => IdentitySearchResponse::Success(search_result), Err(e) => IdentitySearchResponse::Error(IdentitySearchFailure { id: search_query.id, error: e.to_string(), }), }; let serialized_message = serde_json::to_string(&search_response)?; send_message(Message::Text(serialized_message), outgoing.clone()).await; Ok(()) } _ => Err(errors::WebsocketError::InvalidMessage), } } #[tracing::instrument(skip_all)] async fn accept_connection(hyper_ws: HyperWebsocket, addr: SocketAddr) { debug!("Incoming WebSocket connection from {}", addr); let ws_stream = match hyper_ws.await { Ok(stream) => stream, Err(e) => { - error!("WebSocket handshake error: {}", e); + error!( + errorType = error_types::SEARCH_LOG, + "WebSocket handshake error: {}", e + ); return; } }; let (outgoing, mut incoming) = ws_stream.split(); let outgoing = Arc::new(Mutex::new(outgoing)); if let Some(Ok(auth_message)) = incoming.next().await { match auth_message { Message::Text(text) => { if let Err(auth_error) = auth::handle_auth_message(&text).await { let error_response = ConnectionInitializationResponse { status: ConnectionInitializationStatus::Error( auth_error.to_string(), ), }; let serialized_response = serde_json::to_string(&error_response) .expect("Error serializing auth error response"); send_message(Message::Text(serialized_response), outgoing.clone()) .await; close_connection(outgoing).await; return; } else { let success_response = ConnectionInitializationResponse { status: ConnectionInitializationStatus::Success, }; let serialized_response = serde_json::to_string(&success_response) .expect("Error serializing auth success response"); send_message(Message::Text(serialized_response), outgoing.clone()) .await; } } _ => { - error!("Invalid authentication message from {}", addr); + error!( + errorType = error_types::SEARCH_LOG, + "Invalid authentication message from {}", addr + ); close_connection(outgoing).await; return; } } } else { - error!("No authentication message from {}", addr); + error!( + errorType = error_types::SEARCH_LOG, + "No authentication message from {}", addr + ); close_connection(outgoing).await; return; } let mut ping_timeout = Box::pin(tokio::time::sleep(SOCKET_HEARTBEAT_TIMEOUT)); let mut got_heartbeat_response = true; loop { tokio::select! { client_message = incoming.next() => { let message: Message = match client_message { Some(Ok(msg)) => msg, _ => { debug!("Connection to {} closed remotely.", addr); break; } }; match message { Message::Close(_) => { debug!("Connection to {} closed.", addr); break; } Message::Pong(_) => { debug!("Received Pong message from {}", addr); } Message::Ping(msg) => { debug!("Received Ping message from {}", addr); send_message(Message::Pong(msg), outgoing.clone()).await; } Message::Text(text) => { got_heartbeat_response = true; ping_timeout = Box::pin(tokio::time::sleep(SOCKET_HEARTBEAT_TIMEOUT)); if let Err(e) = handle_websocket_frame(text, outgoing.clone()).await { - error!("Error handling WebSocket frame: {}", e); + error!(errorType = error_types::SEARCH_LOG, "Error handling WebSocket frame: {}", e); continue; }; } _ => { - error!("Client sent invalid message type"); + error!(errorType = error_types::SEARCH_LOG, "Client sent invalid message type"); break; } } } _ = &mut ping_timeout => { if !got_heartbeat_response { - error!("Connection to {} died.", addr); + error!(errorType = error_types::SEARCH_LOG, "Connection to {} died.", addr); break; } let serialized = serde_json::to_string(&Heartbeat {}).unwrap(); send_message(Message::text(serialized), outgoing.clone()).await; got_heartbeat_response = false; ping_timeout = Box::pin(tokio::time::sleep(SOCKET_HEARTBEAT_TIMEOUT)); } else => { debug!("Unhealthy connection for: {}", addr); break; } } } info!("unregistering connection to: {}", addr); close_connection(outgoing).await; } diff --git a/services/identity/src/websockets/send.rs b/services/identity/src/websockets/send.rs index 50660e0fd..8ad122f3e 100644 --- a/services/identity/src/websockets/send.rs +++ b/services/identity/src/websockets/send.rs @@ -1,19 +1,24 @@ use std::sync::Arc; use futures::lock::Mutex; use futures_util::stream::SplitSink; use futures_util::SinkExt; use hyper::upgrade::Upgraded; use hyper_tungstenite::tungstenite::Message; use hyper_tungstenite::WebSocketStream; use tracing::error; +use crate::constants::error_types; + pub type WebsocketSink = Arc, Message>>>; #[tracing::instrument(skip_all)] pub async fn send_message(message: Message, outgoing: WebsocketSink) { if let Err(e) = outgoing.lock().await.send(message).await { - error!("Failed to send message to device: {}", e); + error!( + errorType = error_types::SEARCH_LOG, + "Failed to send message to device: {}", e + ); } }