diff --git a/services/identity/src/client_service.rs b/services/identity/src/client_service.rs index 9d729a986..5c9fb0edd 100644 --- a/services/identity/src/client_service.rs +++ b/services/identity/src/client_service.rs @@ -1,892 +1,885 @@ pub mod client_proto { tonic::include_proto!("identity.client"); } use std::str::FromStr; use crate::{ client_service::client_proto::{ AddReservedUsernamesRequest, DeleteUserRequest, Empty, GenerateNonceResponse, InboundKeysForUserRequest, InboundKeysForUserResponse, OpaqueLoginFinishRequest, OpaqueLoginFinishResponse, OpaqueLoginStartRequest, OpaqueLoginStartResponse, OutboundKeysForUserRequest, OutboundKeysForUserResponse, OutboundKeyserverResponse, RefreshUserPreKeysRequest, RegistrationFinishRequest, RegistrationFinishResponse, RegistrationStartRequest, RegistrationStartResponse, RemoveReservedUsernameRequest, ReservedRegistrationStartRequest, UpdateUserPasswordFinishRequest, UpdateUserPasswordStartRequest, UpdateUserPasswordStartResponse, UploadOneTimeKeysRequest, VerifyUserAccessTokenRequest, VerifyUserAccessTokenResponse, WalletLoginRequest, WalletLoginResponse, }, config::CONFIG, database::{DatabaseClient, Error as DBError, KeyPayload}, id::generate_uuid, nonce::generate_nonce_data, reserved_users::{ validate_add_reserved_usernames_message, validate_remove_reserved_username_message, validate_signed_account_ownership_message, }, siwe::parse_and_verify_siwe_message, token::{AccessTokenData, AuthType}, }; use aws_sdk_dynamodb::Error as DynamoDBError; pub use client_proto::identity_client_service_server::{ IdentityClientService, IdentityClientServiceServer, }; use comm_opaque2::grpc::protocol_error_to_grpc_status; use constant_time_eq::constant_time_eq; use moka::future::Cache; use rand::rngs::OsRng; use tonic::Response; use tracing::{debug, error}; #[derive(Clone)] pub enum WorkflowInProgress { Registration(UserRegistrationInfo), Login(UserLoginInfo), Update(UpdateState), } #[derive(Clone)] pub struct UserRegistrationInfo { pub username: String, pub flattened_device_key_upload: FlattenedDeviceKeyUpload, } #[derive(Clone)] pub struct UserLoginInfo { pub user_id: String, pub flattened_device_key_upload: FlattenedDeviceKeyUpload, pub opaque_server_login: comm_opaque2::server::Login, } #[derive(Clone)] pub struct UpdateState { pub user_id: String, } #[derive(Clone)] 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_onetime_keys: Vec, pub notif_prekey: String, pub notif_prekey_signature: String, pub notif_onetime_keys: Vec, } #[derive(derive_more::Constructor)] pub struct ClientService { client: DatabaseClient, cache: Cache, } #[tonic::async_trait] impl IdentityClientService for ClientService { 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); let username_taken = self .client .username_taken(message.username.clone()) .await .map_err(handle_db_error)?; let username_in_reserved_usernames_table = self .client .username_in_reserved_usernames_table(&message.username) .await .map_err(handle_db_error)?; if username_taken || username_in_reserved_usernames_table { return Err(tonic::Status::already_exists("username already exists")); } if CONFIG.reserved_usernames.contains(&message.username) { return Err(tonic::Status::invalid_argument("username reserved")); } if let client_proto::RegistrationStartRequest { opaque_registration_request: register_message, username, device_key_upload: Some(client_proto::DeviceKeyUpload { device_key_info: Some(client_proto::IdentityKeyInfo { payload, payload_signature, social_proof: _social_proof, }), content_upload: Some(client_proto::PreKey { pre_key: content_prekey, pre_key_signature: content_prekey_signature, }), notif_upload: Some(client_proto::PreKey { pre_key: notif_prekey, pre_key_signature: notif_prekey_signature, }), onetime_content_prekeys, onetime_notif_prekeys, }), } = message { let server_registration = comm_opaque2::server::Registration::new(); let server_message = server_registration .start(&CONFIG.server_setup, ®ister_message, username.as_bytes()) .map_err(protocol_error_to_grpc_status)?; let key_info = KeyPayload::from_str(&payload) .map_err(|_| tonic::Status::invalid_argument("malformed payload"))?; let registration_state = UserRegistrationInfo { username, flattened_device_key_upload: FlattenedDeviceKeyUpload { device_id_key: key_info.primary_identity_public_keys.ed25519, key_payload: payload, key_payload_signature: payload_signature, content_prekey, content_prekey_signature, content_onetime_keys: onetime_content_prekeys, notif_prekey, notif_prekey_signature, notif_onetime_keys: onetime_notif_prekeys, }, }; let session_id = generate_uuid(); self .cache .insert( session_id.clone(), WorkflowInProgress::Registration(registration_state), ) .await; let response = RegistrationStartResponse { session_id, opaque_registration_response: server_message, }; Ok(Response::new(response)) } else { Err(tonic::Status::invalid_argument("unexpected message data")) } } async fn register_reserved_password_user_start( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let username_taken = self .client .username_taken(message.username.clone()) .await .map_err(handle_db_error)?; if username_taken { return Err(tonic::Status::already_exists("username already exists")); } if CONFIG.reserved_usernames.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 { validate_signed_account_ownership_message( &message.username, &message.keyserver_message, &message.keyserver_signature, )?; } else { return Err(tonic::Status::permission_denied("username not reserved")); } if let client_proto::ReservedRegistrationStartRequest { opaque_registration_request: register_message, username, device_key_upload: Some(client_proto::DeviceKeyUpload { device_key_info: Some(client_proto::IdentityKeyInfo { payload, payload_signature, social_proof: _social_proof, }), content_upload: Some(client_proto::PreKey { pre_key: content_prekey, pre_key_signature: content_prekey_signature, }), notif_upload: Some(client_proto::PreKey { pre_key: notif_prekey, pre_key_signature: notif_prekey_signature, }), onetime_content_prekeys, onetime_notif_prekeys, }), .. } = message { let server_registration = comm_opaque2::server::Registration::new(); let server_message = server_registration .start(&CONFIG.server_setup, ®ister_message, username.as_bytes()) .map_err(protocol_error_to_grpc_status)?; let key_info = KeyPayload::from_str(&payload) .map_err(|_| tonic::Status::invalid_argument("malformed payload"))?; let registration_state = UserRegistrationInfo { username, flattened_device_key_upload: FlattenedDeviceKeyUpload { device_id_key: key_info.primary_identity_public_keys.ed25519, key_payload: payload, key_payload_signature: payload_signature, content_prekey, content_prekey_signature, content_onetime_keys: onetime_content_prekeys, notif_prekey, notif_prekey_signature, notif_onetime_keys: onetime_notif_prekeys, }, }; let session_id = generate_uuid(); self .cache .insert( session_id.clone(), WorkflowInProgress::Registration(registration_state), ) .await; let response = RegistrationStartResponse { session_id, opaque_registration_response: server_message, }; Ok(Response::new(response)) } else { Err(tonic::Status::invalid_argument("unexpected message data")) } } async fn register_password_user_finish( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); if let Some(WorkflowInProgress::Registration(state)) = self.cache.get(&message.session_id) { self.cache.invalidate(&message.session_id).await; 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 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) .await .map_err(handle_db_error)?; // Create access token let token = AccessTokenData::new( user_id.clone(), device_id, 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 = RegistrationFinishResponse { user_id, access_token, }; Ok(Response::new(response)) } else { Err(tonic::Status::not_found("session not found")) } } async fn update_user_password_start( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let access_token = self .client .get_access_token_data(message.user_id.clone(), message.device_id_key) .await .map_err(handle_db_error)?; if let Some(token) = access_token { if !token.is_valid() || !constant_time_eq( token.access_token.as_bytes(), message.access_token.as_bytes(), ) { return Err(tonic::Status::permission_denied("bad token")); } let server_registration = comm_opaque2::server::Registration::new(); let server_message = server_registration .start( &CONFIG.server_setup, &message.opaque_registration_request, message.user_id.as_bytes(), ) .map_err(protocol_error_to_grpc_status)?; let update_state = UpdateState { user_id: message.user_id, }; let session_id = generate_uuid(); self .cache .insert(session_id.clone(), WorkflowInProgress::Update(update_state)) .await; let response = UpdateUserPasswordStartResponse { session_id, opaque_registration_response: server_message, }; Ok(Response::new(response)) } else { Err(tonic::Status::permission_denied("bad token")) } } async fn update_user_password_finish( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); if let Some(WorkflowInProgress::Update(state)) = self.cache.get(&message.session_id) { self.cache.invalidate(&message.session_id).await; 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 .client .update_user_password(state.user_id, password_file) .await .map_err(handle_db_error)?; let response = Empty {}; Ok(Response::new(response)) } else { Err(tonic::Status::not_found("session not found")) } } async fn login_password_user_start( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); debug!("Attempting to login 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::failed_precondition( "need keyserver message to claim username", )); } return Err(tonic::Status::not_found("user not found")); }; if let client_proto::OpaqueLoginStartRequest { opaque_login_request: login_message, username, device_key_upload: Some(client_proto::DeviceKeyUpload { device_key_info: Some(client_proto::IdentityKeyInfo { payload, payload_signature, social_proof: _social_proof, }), content_upload: Some(client_proto::PreKey { pre_key: content_prekey, pre_key_signature: content_prekey_signature, }), notif_upload: Some(client_proto::PreKey { pre_key: notif_prekey, pre_key_signature: notif_prekey_signature, }), onetime_content_prekeys, onetime_notif_prekeys, }), } = message { let mut server_login = comm_opaque2::server::Login::new(); let server_response = server_login .start( &CONFIG.server_setup, &password_file_bytes, &login_message, username.as_bytes(), ) .map_err(protocol_error_to_grpc_status)?; let key_info = KeyPayload::from_str(&payload) .map_err(|_| tonic::Status::invalid_argument("malformed payload"))?; let login_state = UserLoginInfo { user_id, opaque_server_login: server_login, flattened_device_key_upload: FlattenedDeviceKeyUpload { device_id_key: key_info.primary_identity_public_keys.ed25519, key_payload: payload, key_payload_signature: payload_signature, content_prekey, content_prekey_signature, content_onetime_keys: onetime_content_prekeys, notif_prekey, notif_prekey_signature, notif_onetime_keys: onetime_notif_prekeys, }, }; let session_id = generate_uuid(); self .cache .insert(session_id.clone(), WorkflowInProgress::Login(login_state)) .await; let response = Response::new(OpaqueLoginStartResponse { session_id, opaque_login_response: server_response, }); Ok(response) } else { Err(tonic::Status::invalid_argument("unexpected message data")) } } async fn login_password_user_finish( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); if let Some(WorkflowInProgress::Login(state)) = self.cache.get(&message.session_id) { self.cache.invalidate(&message.session_id).await; let mut server_login = state.opaque_server_login.clone(); server_login .finish(&message.opaque_login_upload) .map_err(protocol_error_to_grpc_status)?; self .client .add_password_user_device_to_users_table( state.user_id.clone(), state.flattened_device_key_upload.clone(), ) .await .map_err(handle_db_error)?; // Create access token let token = AccessTokenData::new( state.user_id.clone(), state.flattened_device_key_upload.device_id_key, 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 = OpaqueLoginFinishResponse { user_id: state.user_id, access_token, }; Ok(Response::new(response)) } else { Err(tonic::Status::not_found("session not found")) } } async fn login_wallet_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let wallet_address = parse_and_verify_siwe_message( &message.siwe_message, &message.siwe_signature, )?; let (flattened_device_key_upload, social_proof) = if let client_proto::WalletLoginRequest { siwe_message: _, siwe_signature: _, device_key_upload: Some(client_proto::DeviceKeyUpload { device_key_info: Some(client_proto::IdentityKeyInfo { payload, payload_signature, social_proof: Some(social_proof), }), content_upload: Some(client_proto::PreKey { pre_key: content_prekey, pre_key_signature: content_prekey_signature, }), notif_upload: Some(client_proto::PreKey { pre_key: notif_prekey, pre_key_signature: notif_prekey_signature, }), onetime_content_prekeys, onetime_notif_prekeys, }), } = message { let key_info = KeyPayload::from_str(&payload) .map_err(|_| tonic::Status::invalid_argument("malformed payload"))?; ( FlattenedDeviceKeyUpload { device_id_key: key_info.primary_identity_public_keys.ed25519, key_payload: payload, key_payload_signature: payload_signature, content_prekey, content_prekey_signature, content_onetime_keys: onetime_content_prekeys, notif_prekey, notif_prekey_signature, notif_onetime_keys: onetime_notif_prekeys, }, social_proof, ) } else { return Err(tonic::Status::invalid_argument("unexpected message data")); }; let user_id = match self .client .get_user_id_from_user_info(wallet_address.clone(), AuthType::Wallet) .await .map_err(handle_db_error)? { Some(id) => { // User already exists, so we should update the DDB item self .client .add_wallet_user_device_to_users_table( id.clone(), flattened_device_key_upload.clone(), social_proof, ) .await .map_err(handle_db_error)?; id } None => { // User doesn't exist yet, so we should add a new user in DDB self .client .add_wallet_user_to_users_table( flattened_device_key_upload.clone(), wallet_address, social_proof, ) .await .map_err(handle_db_error)? } }; // Create access token let token = AccessTokenData::new( user_id.clone(), flattened_device_key_upload.device_id_key, 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 = WalletLoginResponse { user_id, access_token, }; Ok(Response::new(response)) } async fn delete_user( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let access_token = self .client .get_access_token_data(message.user_id.clone(), message.device_id_key) .await .map_err(handle_db_error)?; if let Some(token) = access_token { if !token.is_valid() || !constant_time_eq( token.access_token.as_bytes(), message.access_token.as_bytes(), ) { return Err(tonic::Status::permission_denied("bad token")); } self .client .delete_user(message.user_id) .await .map_err(handle_db_error)?; let response = Empty {}; Ok(Response::new(response)) } else { Err(tonic::Status::permission_denied("bad token")) } } 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)), } } async fn get_outbound_keys_for_user( &self, _request: tonic::Request, ) -> Result, tonic::Status> { unimplemented!(); } async fn get_inbound_keys_for_user( &self, _request: tonic::Request, ) -> Result, tonic::Status> { unimplemented!(); } async fn get_keyserver_keys( &self, _request: tonic::Request, ) -> Result, tonic::Status> { unimplemented!(); } async fn upload_one_time_keys( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); debug!("Validating token: {:?}", message); let token_valid = self .client .verify_access_token( message.user_id.clone(), message.device_id.clone(), message.access_token, ) .await .map_err(handle_db_error)?; if !token_valid { return Err(tonic::Status::unauthenticated("Invalid token")); } debug!( "Attempting to update one time keys for user: {}", message.user_id ); self .client .append_one_time_prekeys( message.user_id, message.device_id, message.content_one_time_pre_keys, message.notif_one_time_pre_keys, ) .await .map_err(handle_db_error)?; Ok(tonic::Response::new(Empty {})) } async fn refresh_user_pre_keys( &self, _request: tonic::Request, ) -> Result, tonic::Status> { unimplemented!(); } async fn verify_user_access_token( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let token_valid = self .client .verify_access_token( message.user_id, message.signing_public_key, message.access_token, ) .await .map_err(handle_db_error)?; let response = Response::new(VerifyUserAccessTokenResponse { token_valid }); Ok(response) } async fn add_reserved_usernames( &self, request: tonic::Request, ) -> Result, tonic::Status> { let message = request.into_inner(); let usernames = validate_add_reserved_usernames_message( &message.message, &message.signature, )?; - let mut filtered_usernames = Vec::new(); - - for username in usernames { - if !self - .client - .username_taken(username.clone()) - .await - .map_err(handle_db_error)? - { - filtered_usernames.push(username); - } - } + let filtered_usernames = self + .client + .filter_out_taken_usernames(usernames) + .await + .map_err(handle_db_error)?; self .client .add_usernames_to_reserved_usernames_table(filtered_usernames) .await .map_err(handle_db_error)?; let response = Response::new(Empty {}); Ok(response) } 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) } } 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") } e => { error!("Encountered an unexpected error: {}", e); tonic::Status::failed_precondition("unexpected error") } } } diff --git a/services/identity/src/database.rs b/services/identity/src/database.rs index 57dac41cd..0b4a1d392 100644 --- a/services/identity/src/database.rs +++ b/services/identity/src/database.rs @@ -1,1111 +1,1108 @@ use constant_time_eq::constant_time_eq; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::fmt::{Display, Formatter, Result as FmtResult}; use std::str::FromStr; use std::sync::Arc; use aws_config::SdkConfig; use aws_sdk_dynamodb::model::{AttributeValue, PutRequest, WriteRequest}; use aws_sdk_dynamodb::output::{ DeleteItemOutput, GetItemOutput, PutItemOutput, QueryOutput, }; use aws_sdk_dynamodb::types::Blob; use aws_sdk_dynamodb::{Client, Error as DynamoDBError}; 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::{ ACCESS_TOKEN_SORT_KEY, ACCESS_TOKEN_TABLE, ACCESS_TOKEN_TABLE_AUTH_TYPE_ATTRIBUTE, ACCESS_TOKEN_TABLE_CREATED_ATTRIBUTE, ACCESS_TOKEN_TABLE_PARTITION_KEY, ACCESS_TOKEN_TABLE_TOKEN_ATTRIBUTE, ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE, NONCE_TABLE, NONCE_TABLE_CREATED_ATTRIBUTE, NONCE_TABLE_PARTITION_KEY, RESERVED_USERNAMES_TABLE, RESERVED_USERNAMES_TABLE_PARTITION_KEY, USERS_TABLE, USERS_TABLE_DEVICES_ATTRIBUTE, USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_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::{AccessTokenData, AuthType}; #[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 enum Device { Client, Keyserver, } impl Display for Device { fn fmt(&self, f: &mut Formatter) -> FmtResult { match self { Device::Client => write!(f, "client"), Device::Keyserver => write!(f, "keyserver"), } } } #[derive(Clone)] pub struct DatabaseClient { client: Arc, } impl DatabaseClient { pub fn new(aws_config: &SdkConfig) -> Self { let client = match &CONFIG.localstack_endpoint { Some(endpoint) => { info!( "Configuring DynamoDB client to use LocalStack endpoint: {}", endpoint ); let ddb_config_builder = aws_sdk_dynamodb::config::Builder::from(aws_config) .endpoint_url(endpoint); Client::from_conf(ddb_config_builder.build()) } None => Client::new(aws_config), }; DatabaseClient { client: Arc::new(client), } } pub async fn add_password_user_to_users_table( &self, registration_state: UserRegistrationInfo, password_file: Vec, ) -> Result { self .add_user_to_users_table( registration_state.flattened_device_key_upload, Some((registration_state.username, Blob::new(password_file))), None, None, ) .await } pub async fn add_wallet_user_to_users_table( &self, flattened_device_key_upload: FlattenedDeviceKeyUpload, wallet_address: String, social_proof: String, ) -> Result { self .add_user_to_users_table( flattened_device_key_upload, None, Some(wallet_address), Some(social_proof), ) .await } async fn add_user_to_users_table( &self, flattened_device_key_upload: FlattenedDeviceKeyUpload, username_and_password_file: Option<(String, Blob)>, wallet_address: Option, social_proof: Option, ) -> Result { let user_id = generate_uuid(); let device_info = create_device_info(flattened_device_key_upload.clone(), social_proof); let devices = HashMap::from([( flattened_device_key_upload.device_id_key, AttributeValue::M(device_info), )]); let mut user = HashMap::from([ ( USERS_TABLE_PARTITION_KEY.to_string(), AttributeValue::S(user_id.clone()), ), ( USERS_TABLE_DEVICES_ATTRIBUTE.to_string(), AttributeValue::M(devices), ), ]); 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(address) = wallet_address { user.insert( USERS_TABLE_WALLET_ADDRESS_ATTRIBUTE.to_string(), AttributeValue::S(address), ); } self .client .put_item() .table_name(USERS_TABLE) .set_item(Some(user)) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; Ok(user_id) } pub async fn add_password_user_device_to_users_table( &self, user_id: String, flattened_device_key_upload: FlattenedDeviceKeyUpload, ) -> Result<(), Error> { self .add_device_to_users_table(user_id, flattened_device_key_upload, None) .await } pub async fn add_wallet_user_device_to_users_table( &self, user_id: String, flattened_device_key_upload: FlattenedDeviceKeyUpload, social_proof: String, ) -> Result<(), Error> { self .add_device_to_users_table( user_id, flattened_device_key_upload, Some(social_proof), ) .await } pub async fn append_one_time_prekeys( &self, user_id: String, device_id: String, content_one_time_keys: Vec, notif_one_time_keys: Vec, ) -> Result<(), Error> { let notif_keys_av: Vec = notif_one_time_keys .into_iter() .map(AttributeValue::S) .collect(); let content_keys_av: Vec = content_one_time_keys .into_iter() .map(AttributeValue::S) .collect(); let update_expression = format!("SET {0}.#{1}.{2} = list_append({0}.#{1}.{2}, :n), {0}.#{1}.{3} = list_append({0}.#{1}.{3}, :i)", USERS_TABLE_DEVICES_ATTRIBUTE, "deviceID", USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME, USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME ); let expression_attribute_names = HashMap::from([(format!("#{}", "deviceID"), device_id)]); let expression_attribute_values = HashMap::from([ (":n".to_string(), AttributeValue::L(notif_keys_av)), (":i".to_string(), AttributeValue::L(content_keys_av)), ]); self .client .update_item() .table_name(USERS_TABLE) .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id)) .update_expression(update_expression) .set_expression_attribute_names(Some(expression_attribute_names)) .set_expression_attribute_values(Some(expression_attribute_values)) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; Ok(()) } async fn add_device_to_users_table( &self, user_id: String, flattened_device_key_upload: FlattenedDeviceKeyUpload, social_proof: Option, ) -> Result<(), Error> { let device_info = create_device_info(flattened_device_key_upload.clone(), social_proof); let update_expression = format!("SET {}.#{} = :v", USERS_TABLE_DEVICES_ATTRIBUTE, "deviceID",); let expression_attribute_names = HashMap::from([( format!("#{}", "deviceID"), flattened_device_key_upload.device_id_key, )]); let expression_attribute_values = HashMap::from([(":v".to_string(), AttributeValue::M(device_info))]); self .client .update_item() .table_name(USERS_TABLE) .key(USERS_TABLE_PARTITION_KEY, AttributeValue::S(user_id)) .update_expression(update_expression) .set_expression_attribute_names(Some(expression_attribute_names)) .set_expression_attribute_values(Some(expression_attribute_values)) .send() .await .map_err(|e| Error::AwsSdk(e.into()))?; Ok(()) } 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(()) } pub async fn delete_user( &self, user_id: String, ) -> Result { debug!("Attempting to delete user: {}", user_id); 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); Err(Error::AwsSdk(e.into())) } } } pub async fn get_access_token_data( &self, user_id: String, signing_public_key: String, ) -> Result, Error> { let primary_key = create_composite_primary_key( ( ACCESS_TOKEN_TABLE_PARTITION_KEY.to_string(), user_id.clone(), ), ( ACCESS_TOKEN_SORT_KEY.to_string(), signing_public_key.clone(), ), ); let get_item_result = self .client .get_item() .table_name(ACCESS_TOKEN_TABLE) .set_key(Some(primary_key)) .consistent_read(true) .send() .await; match get_item_result { Ok(GetItemOutput { item: Some(mut item), .. }) => { let created = parse_created_attribute( item.remove(ACCESS_TOKEN_TABLE_CREATED_ATTRIBUTE), )?; let auth_type = parse_auth_type_attribute( item.remove(ACCESS_TOKEN_TABLE_AUTH_TYPE_ATTRIBUTE), )?; let valid = parse_valid_attribute( item.remove(ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE), )?; let access_token = parse_token_attribute( item.remove(ACCESS_TOKEN_TABLE_TOKEN_ATTRIBUTE), )?; 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 {}: {}", 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 { let item = HashMap::from([ ( ACCESS_TOKEN_TABLE_PARTITION_KEY.to_string(), AttributeValue::S(access_token_data.user_id), ), ( ACCESS_TOKEN_SORT_KEY.to_string(), AttributeValue::S(access_token_data.signing_public_key), ), ( ACCESS_TOKEN_TABLE_TOKEN_ATTRIBUTE.to_string(), AttributeValue::S(access_token_data.access_token), ), ( ACCESS_TOKEN_TABLE_CREATED_ATTRIBUTE.to_string(), AttributeValue::S(access_token_data.created.to_rfc3339()), ), ( ACCESS_TOKEN_TABLE_AUTH_TYPE_ATTRIBUTE.to_string(), AttributeValue::S(match access_token_data.auth_type { AuthType::Password => "password".to_string(), AuthType::Wallet => "wallet".to_string(), }), ), ( ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE.to_string(), AttributeValue::Bool(access_token_data.valid), ), ]); self .client .put_item() .table_name(ACCESS_TOKEN_TABLE) .set_item(Some(item)) .send() .await .map_err(|e| Error::AwsSdk(e.into())) } 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, + usernames: Vec, + ) -> Result, Error> { + let db_usernames = self.get_all_usernames().await?; + + let db_usernames_set: HashSet = db_usernames.into_iter().collect(); + let usernames_set: HashSet = usernames.into_iter().collect(); + + let available_usernames: Vec = usernames_set + .difference(&db_usernames_set) + .cloned() + .collect(); + + Ok(available_usernames) + } + 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, attribute_value: None, attribute_error: DBItemAttributeError::Missing, })? .as_s() .map_err(|_| DBItemError { attribute_name: USERS_TABLE_PARTITION_KEY, attribute_value: first_item.get(USERS_TABLE_PARTITION_KEY).cloned(), 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!( "DynamoDB client failed to get user from {} {}: {}", attribute_name, user_info, e ); Err(Error::AwsSdk(e.into())) } } } 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)) => parse_string_attribute( USERS_TABLE_PARTITION_KEY, user.remove(USERS_TABLE_PARTITION_KEY), ) .map(Some) .map_err(Error::Attribute), Ok(_) => Ok(None), Err(e) => Err(e), } } 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 = parse_string_attribute( USERS_TABLE_PARTITION_KEY, user.remove(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!( "DynamoDB client failed to get registration data for user {}: {}", username, e ); Err(e) } } } pub 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())) } - pub async fn get_users(&self) -> Result, Error> { + async fn get_all_usernames(&self) -> Result, Error> { let scan_output = self .client .scan() .table_name(USERS_TABLE) - .projection_expression(USERS_TABLE_PARTITION_KEY) + .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 { - let id = parse_string_attribute( - USERS_TABLE_PARTITION_KEY, - attribute.remove(USERS_TABLE_PARTITION_KEY), - ) - .map_err(Error::Attribute)?; - result.push(id); + if let Ok(username) = parse_string_attribute( + USERS_TABLE_USERNAME_ATTRIBUTE, + attribute.remove(USERS_TABLE_USERNAME_ATTRIBUTE), + ) { + result.push(username); + } } } 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()), ), ]); self .client .put_item() .table_name(NONCE_TABLE) .set_item(Some(item)) .send() .await .map_err(|e| Error::AwsSdk(e.into())) } pub async fn add_usernames_to_reserved_usernames_table( &self, usernames: Vec, ) -> Result<(), Error> { - let mut write_requests = vec![]; - - for username in usernames { - let item: HashMap = vec![( - RESERVED_USERNAMES_TABLE_PARTITION_KEY.to_string(), - AttributeValue::S(username), - )] - .into_iter() - .collect(); - - let write_request = WriteRequest::builder() - .put_request(PutRequest::builder().set_item(Some(item)).build()) - .build(); - - write_requests.push(write_request); - } - - loop { - let output = self + // A single call to BatchWriteItem can consist of up to 25 operations + for usernames_chunk in usernames.chunks(25) { + let write_requests = usernames_chunk + .iter() + .map(|username| { + let put_request = PutRequest::builder() + .item( + RESERVED_USERNAMES_TABLE_PARTITION_KEY, + AttributeValue::S(username.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()))?; - - let unprocessed_items_map = match output.unprocessed_items() { - Some(map) => map, - None => break, - }; - - let unprocessed_requests = - match unprocessed_items_map.get(RESERVED_USERNAMES_TABLE) { - Some(requests) => requests, - None => break, - }; - - info!( - "{} unprocessed items, retrying...", - unprocessed_requests.len() - ); - - write_requests = unprocessed_requests.to_vec(); } info!("Batch write item to reserved usernames table succeeded"); Ok(()) } 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); 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())), } } } #[derive( Debug, derive_more::Display, derive_more::From, derive_more::Error, )] pub enum Error { #[display(...)] AwsSdk(DynamoDBError), #[display(...)] Attribute(DBItemError), } #[derive(Debug, derive_more::Error, derive_more::Constructor)] pub struct DBItemError { attribute_name: &'static str, attribute_value: Option, attribute_error: DBItemAttributeError, } impl Display for DBItemError { fn fmt(&self, f: &mut Formatter) -> FmtResult { match &self.attribute_error { DBItemAttributeError::Missing => { write!(f, "Attribute {} is missing", self.attribute_name) } DBItemAttributeError::IncorrectType => write!( f, "Value for attribute {} has incorrect type: {:?}", self.attribute_name, self.attribute_value ), error => write!( f, "Error regarding attribute {} with value {:?}: {}", self.attribute_name, self.attribute_value, error ), } } } #[derive(Debug, derive_more::Display, derive_more::Error)] pub enum DBItemAttributeError { #[display(...)] Missing, #[display(...)] IncorrectType, #[display(...)] InvalidTimestamp(chrono::ParseError), } type AttributeName = String; 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_created_attribute( attribute: Option, ) -> Result, DBItemError> { if let Some(AttributeValue::S(created)) = &attribute { created.parse().map_err(|e| { DBItemError::new( ACCESS_TOKEN_TABLE_CREATED_ATTRIBUTE, attribute, DBItemAttributeError::InvalidTimestamp(e), ) }) } else { Err(DBItemError::new( ACCESS_TOKEN_TABLE_CREATED_ATTRIBUTE, attribute, DBItemAttributeError::Missing, )) } } fn parse_auth_type_attribute( attribute: Option, ) -> Result { if let Some(AttributeValue::S(auth_type)) = &attribute { match auth_type.as_str() { "password" => Ok(AuthType::Password), "wallet" => Ok(AuthType::Wallet), _ => Err(DBItemError::new( ACCESS_TOKEN_TABLE_AUTH_TYPE_ATTRIBUTE, attribute, DBItemAttributeError::IncorrectType, )), } } else { Err(DBItemError::new( ACCESS_TOKEN_TABLE_AUTH_TYPE_ATTRIBUTE, attribute, DBItemAttributeError::Missing, )) } } fn parse_valid_attribute( attribute: Option, ) -> Result { match attribute { Some(AttributeValue::Bool(valid)) => Ok(valid), Some(_) => Err(DBItemError::new( ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE, attribute, DBItemAttributeError::IncorrectType, )), None => Err(DBItemError::new( ACCESS_TOKEN_TABLE_VALID_ATTRIBUTE, attribute, DBItemAttributeError::Missing, )), } } fn parse_token_attribute( attribute: Option, ) -> Result { match attribute { Some(AttributeValue::S(token)) => Ok(token), Some(_) => Err(DBItemError::new( ACCESS_TOKEN_TABLE_TOKEN_ATTRIBUTE, attribute, DBItemAttributeError::IncorrectType, )), None => Err(DBItemError::new( ACCESS_TOKEN_TABLE_TOKEN_ATTRIBUTE, attribute, DBItemAttributeError::Missing, )), } } 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, attribute, DBItemAttributeError::IncorrectType, )), None => Err(DBItemError::new( USERS_TABLE_REGISTRATION_ATTRIBUTE, attribute, DBItemAttributeError::Missing, )), } } fn parse_map_attribute( attribute_name: &'static str, attribute_value: Option, ) -> Result, DBItemError> { match attribute_value { Some(AttributeValue::M(map)) => Ok(map), Some(_) => Err(DBItemError::new( attribute_name, attribute_value, DBItemAttributeError::IncorrectType, )), None => Err(DBItemError::new( attribute_name, attribute_value, DBItemAttributeError::Missing, )), } } fn parse_string_attribute( attribute_name: &'static str, attribute_value: Option, ) -> Result { match attribute_value { Some(AttributeValue::S(value)) => Ok(value), Some(_) => Err(DBItemError::new( attribute_name, attribute_value, DBItemAttributeError::IncorrectType, )), None => Err(DBItemError::new( attribute_name, attribute_value, DBItemAttributeError::Missing, )), } } fn create_device_info( flattened_device_key_upload: FlattenedDeviceKeyUpload, social_proof: Option, ) -> HashMap { let mut device_info = HashMap::from([ ( USERS_TABLE_DEVICES_MAP_DEVICE_TYPE_ATTRIBUTE_NAME.to_string(), AttributeValue::S(Device::Client.to_string()), ), ( USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_ATTRIBUTE_NAME.to_string(), AttributeValue::S(flattened_device_key_upload.key_payload), ), ( USERS_TABLE_DEVICES_MAP_KEY_PAYLOAD_SIGNATURE_ATTRIBUTE_NAME.to_string(), AttributeValue::S(flattened_device_key_upload.key_payload_signature), ), ( USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_ATTRIBUTE_NAME.to_string(), AttributeValue::S(flattened_device_key_upload.content_prekey), ), ( USERS_TABLE_DEVICES_MAP_CONTENT_PREKEY_SIGNATURE_ATTRIBUTE_NAME .to_string(), AttributeValue::S(flattened_device_key_upload.content_prekey_signature), ), ( USERS_TABLE_DEVICES_MAP_CONTENT_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(), AttributeValue::L( flattened_device_key_upload .content_onetime_keys .into_iter() .map(AttributeValue::S) .collect(), ), ), ( USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_ATTRIBUTE_NAME.to_string(), AttributeValue::S(flattened_device_key_upload.notif_prekey), ), ( USERS_TABLE_DEVICES_MAP_NOTIF_PREKEY_SIGNATURE_ATTRIBUTE_NAME.to_string(), AttributeValue::S(flattened_device_key_upload.notif_prekey_signature), ), ( USERS_TABLE_DEVICES_MAP_NOTIF_ONETIME_KEYS_ATTRIBUTE_NAME.to_string(), AttributeValue::L( flattened_device_key_upload .notif_onetime_keys .into_iter() .map(AttributeValue::S) .collect(), ), ), ]); if let Some(social_proof) = social_proof { device_info.insert( USERS_TABLE_DEVICES_MAP_SOCIAL_PROOF_ATTRIBUTE_NAME.to_string(), AttributeValue::S(social_proof), ); } device_info } #[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" ); } }