diff --git a/services/commtest/src/identity/device.rs b/services/commtest/src/identity/device.rs --- a/services/commtest/src/identity/device.rs +++ b/services/commtest/src/identity/device.rs @@ -2,6 +2,7 @@ use grpc_clients::identity::{get_auth_client, get_unauthenticated_client}; use rand::{distributions::Alphanumeric, Rng}; +use crate::identity::olm_account_infos::generate_random_olm_key; use crate::identity::olm_account_infos::{ ClientPublicKeys, DEFAULT_CLIENT_KEYS, }; @@ -70,11 +71,11 @@ social_proof: None, }), content_upload: Some(Prekey { - prekey: "content_prekey".to_string(), + prekey: generate_random_olm_key(), prekey_signature: "content_prekey_sig".to_string(), }), notif_upload: Some(Prekey { - prekey: "notif_prekey".to_string(), + prekey: generate_random_olm_key(), prekey_signature: "notif_prekey_sig".to_string(), }), one_time_content_prekeys: Vec::new(), @@ -153,11 +154,11 @@ social_proof: None, }), content_upload: Some(Prekey { - prekey: "content_prekey".to_string(), + prekey: generate_random_olm_key(), prekey_signature: "content_prekey_sig".to_string(), }), notif_upload: Some(Prekey { - prekey: "notif_prekey".to_string(), + prekey: generate_random_olm_key(), prekey_signature: "notif_prekey_sig".to_string(), }), one_time_content_prekeys: Vec::new(), diff --git a/services/commtest/src/identity/olm_account_infos.rs b/services/commtest/src/identity/olm_account_infos.rs --- a/services/commtest/src/identity/olm_account_infos.rs +++ b/services/commtest/src/identity/olm_account_infos.rs @@ -48,7 +48,7 @@ }; } -pub fn get_random_otk() -> String { +pub fn generate_random_olm_key() -> String { rand::thread_rng() .sample_iter(&Alphanumeric) .take(43) diff --git a/services/commtest/tests/identity_keyserver_tests.rs b/services/commtest/tests/identity_keyserver_tests.rs --- a/services/commtest/tests/identity_keyserver_tests.rs +++ b/services/commtest/tests/identity_keyserver_tests.rs @@ -1,7 +1,7 @@ use commtest::identity::device::{ register_user_device, DEVICE_TYPE, PLACEHOLDER_CODE_VERSION, }; -use commtest::identity::olm_account_infos::get_random_otk; +use commtest::identity::olm_account_infos::generate_random_olm_key; use commtest::service_addr; use grpc_clients::identity::{ get_auth_client, @@ -26,8 +26,8 @@ .await .expect("Couldn't connect to identity service"); - let content_one_time_prekey = get_random_otk(); - let notif_one_time_prekey = get_random_otk(); + let content_one_time_prekey = generate_random_olm_key(); + let notif_one_time_prekey = generate_random_olm_key(); let upload_request = UploadOneTimeKeysRequest { content_one_time_prekeys: vec![content_one_time_prekey.clone()], diff --git a/services/commtest/tests/identity_one_time_key_tests.rs b/services/commtest/tests/identity_one_time_key_tests.rs --- a/services/commtest/tests/identity_one_time_key_tests.rs +++ b/services/commtest/tests/identity_one_time_key_tests.rs @@ -1,7 +1,7 @@ use commtest::identity::device::{ register_user_device, DEVICE_TYPE, PLACEHOLDER_CODE_VERSION, }; -use commtest::identity::olm_account_infos::get_random_otk; +use commtest::identity::olm_account_infos::generate_random_olm_key; use commtest::service_addr; use grpc_clients::identity::{ get_auth_client, protos::authenticated::UploadOneTimeKeysRequest, @@ -23,8 +23,14 @@ .expect("Couldn't connect to identity service"); let upload_request = UploadOneTimeKeysRequest { - content_one_time_prekeys: vec![get_random_otk(), get_random_otk()], - notif_one_time_prekeys: vec![get_random_otk(), get_random_otk()], + content_one_time_prekeys: vec![ + generate_random_olm_key(), + generate_random_olm_key(), + ], + notif_one_time_prekeys: vec![ + generate_random_olm_key(), + generate_random_olm_key(), + ], }; identity_client diff --git a/services/commtest/tests/identity_tunnelbroker_tests.rs b/services/commtest/tests/identity_tunnelbroker_tests.rs --- a/services/commtest/tests/identity_tunnelbroker_tests.rs +++ b/services/commtest/tests/identity_tunnelbroker_tests.rs @@ -1,7 +1,7 @@ use commtest::identity::device::{ register_user_device, DEVICE_TYPE, PLACEHOLDER_CODE_VERSION, }; -use commtest::identity::olm_account_infos::get_random_otk; +use commtest::identity::olm_account_infos::generate_random_olm_key; use commtest::service_addr; use commtest::tunnelbroker::socket::{create_socket, receive_message}; use futures_util::StreamExt; @@ -48,8 +48,8 @@ .await .expect("Couldn't connect to identity service"); - let content_one_time_prekeys = vec![get_random_otk()]; - let notif_one_time_prekeys = vec![get_random_otk()]; + let content_one_time_prekeys = vec![generate_random_olm_key()]; + let notif_one_time_prekeys = vec![generate_random_olm_key()]; let upload_request = UploadOneTimeKeysRequest { content_one_time_prekeys, notif_one_time_prekeys, diff --git a/services/identity/src/database/device_list.rs b/services/identity/src/database/device_list.rs --- a/services/identity/src/database/device_list.rs +++ b/services/identity/src/database/device_list.rs @@ -26,6 +26,7 @@ error::{DeviceListError, Error}, grpc_services::protos::{self, unauth::DeviceType}, grpc_utils::DeviceKeysInfo, + olm::is_valid_olm_key, }; use super::DatabaseClient; @@ -230,9 +231,17 @@ } } -impl From for AttributeMap { - fn from(value: DeviceRow) -> Self { - HashMap::from([ +impl TryFrom for AttributeMap { + type Error = Error; + + fn try_from(value: DeviceRow) -> Result { + if !is_valid_olm_key(&value.content_prekey.prekey) + || !is_valid_olm_key(&value.notif_prekey.prekey) + { + error!("Invalid prekey format"); + return Err(Error::InvalidFormat); + } + let attribute_map = HashMap::from([ (ATTR_USER_ID.to_string(), AttributeValue::S(value.user_id)), ( ATTR_ITEM_ID.to_string(), @@ -257,7 +266,8 @@ ATTR_LOGIN_TIME.to_string(), AttributeValue::S(value.login_time.to_rfc3339()), ), - ]) + ]); + Ok(attribute_map) } } @@ -506,6 +516,12 @@ 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"); + return Err(Error::InvalidFormat); + } self .client .update_item() @@ -701,7 +717,7 @@ .client .put_item() .table_name(devices_table::NAME) - .set_item(Some(new_device.into())) + .set_item(Some(new_device.try_into()?)) .send() .await .map_err(|e| { @@ -757,7 +773,7 @@ // Put new device let put_device = Put::builder() .table_name(devices_table::NAME) - .set_item(Some(new_device.into())) + .set_item(Some(new_device.try_into()?)) .condition_expression( "attribute_not_exists(#user_id) AND attribute_not_exists(#item_id)", )