diff --git a/services/commtest/src/identity/olm_account.rs b/services/commtest/src/identity/olm_account.rs --- a/services/commtest/src/identity/olm_account.rs +++ b/services/commtest/src/identity/olm_account.rs @@ -3,6 +3,8 @@ use rand::{distributions::Alphanumeric, rngs::OsRng, Rng}; use serde::{Deserialize, Serialize}; +use grpc_clients::identity::protos::unauth::Prekey as GrpcPrekey; + #[derive(Serialize, Deserialize, Debug, Clone)] pub struct IdentityPublicKeys { pub ed25519: String, @@ -20,14 +22,18 @@ impl ClientPublicKeys { /// Generates random keys with given `ed25519` primary account /// signing public key. Use [`ClientPublicKeys::default`] for random key. - pub fn new(primary_signing_public_key: impl Into<String>) -> Self { + pub fn new( + primary_signing_public_key: impl Into<String>, + notif_signing_public_key: Option<String>, + ) -> Self { Self { primary_identity_public_keys: IdentityPublicKeys { ed25519: primary_signing_public_key.into(), curve25519: generate_random_olm_key(), }, notification_identity_public_keys: IdentityPublicKeys { - ed25519: generate_random_olm_key(), + ed25519: notif_signing_public_key + .unwrap_or_else(generate_random_olm_key), curve25519: generate_random_olm_key(), }, } @@ -40,38 +46,78 @@ impl Default for ClientPublicKeys { fn default() -> Self { - Self::new(generate_random_olm_key()) + Self::new(generate_random_olm_key(), None) } } -/// Struct that simulates Olm account +#[derive(Default)] +pub enum AccountType { + #[default] + Content, + Notif, +} + +/// Struct that simulates client Olm account. +/// Specifically it represents a pair of Content and Notif accounts. pub struct MockOlmAccount { // primary account ed25519 keypair - signing_key: Keypair, + content_keypair: Keypair, + notif_keypair: Keypair, } impl MockOlmAccount { pub fn new() -> Self { let mut rng = OsRng {}; - let signing_key = Keypair::generate(&mut rng); - Self { signing_key } + let content_keypair = Keypair::generate(&mut rng); + let notif_keypair = Keypair::generate(&mut rng); + Self { + content_keypair, + notif_keypair, + } } /// returns device public keys, required for device key upload pub fn public_keys(&self) -> ClientPublicKeys { - let signing_public_key = self.signing_key.public.to_bytes(); - let ed25519 = base64::engine::general_purpose::STANDARD_NO_PAD - .encode(signing_public_key); + let base64_engine = &base64::engine::general_purpose::STANDARD_NO_PAD; + + let signing_public_key = self.content_keypair.public.to_bytes(); + let content_ed25519 = base64_engine.encode(signing_public_key); + + let notif_public_key = self.notif_keypair.public.to_bytes(); + let notif_ed25519 = base64_engine.encode(notif_public_key); - ClientPublicKeys::new(ed25519) + ClientPublicKeys::new(content_ed25519, Some(notif_ed25519)) } - /// signs message, returns signature + /// signs message with Content Olm account, returns signature pub fn sign_message(&self, message: &str) -> String { - let signature: Signature = self.signing_key.sign(message.as_bytes()); + self.sign_message_with_account(message, AccountType::Content) + } + + /// signs message, returns signature + pub fn sign_message_with_account( + &self, + message: &str, + account: AccountType, + ) -> String { + let keypair: &Keypair = match account { + AccountType::Content => &self.content_keypair, + AccountType::Notif => &self.notif_keypair, + }; + let signature: Signature = keypair.sign(message.as_bytes()); base64::engine::general_purpose::STANDARD_NO_PAD .encode(signature.to_bytes()) } + + /// generates random prekey with valid signature + pub fn generate_prekey(&self, account: AccountType) -> GrpcPrekey { + let prekey = generate_random_olm_key(); + let prekey_signature = self.sign_message_with_account(&prekey, account); + GrpcPrekey { + prekey, + prekey_signature, + } + } } impl Default for MockOlmAccount { diff --git a/services/commtest/tests/identity_prekey_tests.rs b/services/commtest/tests/identity_prekey_tests.rs --- a/services/commtest/tests/identity_prekey_tests.rs +++ b/services/commtest/tests/identity_prekey_tests.rs @@ -1,16 +1,18 @@ use commtest::identity::device::{ register_user_device, DEVICE_TYPE, PLACEHOLDER_CODE_VERSION, }; +use commtest::identity::olm_account::{AccountType, MockOlmAccount}; use commtest::service_addr; use grpc_clients::identity::PlatformMetadata; use grpc_clients::identity::{ - get_auth_client, - protos::{authenticated::RefreshUserPrekeysRequest, unauth::Prekey}, + get_auth_client, protos::authenticated::RefreshUserPrekeysRequest, }; #[tokio::test] async fn set_prekey() { - let device_info = register_user_device(None, None).await; + let olm_account = MockOlmAccount::new(); + let device_keys = olm_account.public_keys(); + let device_info = register_user_device(Some(&device_keys), None).await; let mut client = get_auth_client( &service_addr::IDENTITY_GRPC.to_string(), @@ -22,18 +24,15 @@ .await .expect("Couldn't connect to identity service"); + let new_content_prekey = olm_account.generate_prekey(AccountType::Content); + let new_notif_prekey = olm_account.generate_prekey(AccountType::Notif); + let upload_request = RefreshUserPrekeysRequest { - new_content_prekey: Some(Prekey { - prekey: "content_prekey".to_string(), - prekey_signature: "content_prekey_signature".to_string(), - }), - new_notif_prekey: Some(Prekey { - prekey: "content_prekey".to_string(), - prekey_signature: "content_prekey_signature".to_string(), - }), + new_content_prekey: Some(new_content_prekey), + new_notif_prekey: Some(new_notif_prekey), }; - // This send will fail if the one-time keys weren't successfully added + // This send will fail if prekeys weren't successfully updated println!( "Error: {:?}", client.refresh_user_prekeys(upload_request).await