introduce wrapper structs for DeviceId, SignedPreKeyId, and PreKeyId

rust/bridge/ffi/src/lib.rs

@@ -182,7 +182,7 @@ pub unsafe extern "C" fn signal_sealed_session_cipher_decrypt(
             timestamp,
             local_e164,
             local_uuid,
-            local_device_id,
+            local_device_id.into(),
             &mut identity_store,
             &mut session_store,
             &mut prekey_store,
@@ -194,7 +194,7 @@ pub unsafe extern "C" fn signal_sealed_session_cipher_decrypt(
 
         write_optional_cstr_to(sender_e164, Ok(decrypted.sender_e164))?;
         write_cstr_to(sender_uuid, Ok(decrypted.sender_uuid))?;
-        write_result_to(sender_device_id, decrypted.device_id)?;
+        write_result_to(sender_device_id, u32::from(decrypted.device_id))?;
         write_bytearray_to(out, out_len, Some(decrypted.message))
     })
 }

rust/bridge/shared/src/ffi/storage.rs

@@ -197,12 +197,12 @@ pub struct FfiPreKeyStoreStruct {
 impl PreKeyStore for &FfiPreKeyStoreStruct {
     async fn get_pre_key(
         &self,
-        prekey_id: u32,
+        prekey_id: PreKeyId,
         ctx: Context,
     ) -> Result<PreKeyRecord, SignalProtocolError> {
         let ctx = ctx.unwrap_or(std::ptr::null_mut());
         let mut record = std::ptr::null_mut();
-        let result = (self.load_pre_key)(self.ctx, &mut record, prekey_id, ctx);
+        let result = (self.load_pre_key)(self.ctx, &mut record, prekey_id.into(), ctx);
 
         if let Some(error) = CallbackError::check(result) {
             return Err(SignalProtocolError::ApplicationCallbackError(
@@ -221,12 +221,12 @@ impl PreKeyStore for &FfiPreKeyStoreStruct {
 
     async fn save_pre_key(
         &mut self,
-        prekey_id: u32,
+        prekey_id: PreKeyId,
         record: &PreKeyRecord,
         ctx: Context,
     ) -> Result<(), SignalProtocolError> {
         let ctx = ctx.unwrap_or(std::ptr::null_mut());
-        let result = (self.store_pre_key)(self.ctx, prekey_id, &*record, ctx);
+        let result = (self.store_pre_key)(self.ctx, prekey_id.into(), &*record, ctx);
 
         if let Some(error) = CallbackError::check(result) {
             return Err(SignalProtocolError::ApplicationCallbackError(
@@ -240,11 +240,11 @@ impl PreKeyStore for &FfiPreKeyStoreStruct {
 
     async fn remove_pre_key(
         &mut self,
-        prekey_id: u32,
+        prekey_id: PreKeyId,
         ctx: Context,
     ) -> Result<(), SignalProtocolError> {
         let ctx = ctx.unwrap_or(std::ptr::null_mut());
-        let result = (self.remove_pre_key)(self.ctx, prekey_id, ctx);
+        let result = (self.remove_pre_key)(self.ctx, prekey_id.into(), ctx);
 
         if let Some(error) = CallbackError::check(result) {
             return Err(SignalProtocolError::ApplicationCallbackError(
@@ -282,12 +282,12 @@ pub struct FfiSignedPreKeyStoreStruct {
 impl SignedPreKeyStore for &FfiSignedPreKeyStoreStruct {
     async fn get_signed_pre_key(
         &self,
-        prekey_id: u32,
+        prekey_id: SignedPreKeyId,
         ctx: Context,
     ) -> Result<SignedPreKeyRecord, SignalProtocolError> {
         let ctx = ctx.unwrap_or(std::ptr::null_mut());
         let mut record = std::ptr::null_mut();
-        let result = (self.load_signed_pre_key)(self.ctx, &mut record, prekey_id, ctx);
+        let result = (self.load_signed_pre_key)(self.ctx, &mut record, prekey_id.into(), ctx);
 
         if let Some(error) = CallbackError::check(result) {
             return Err(SignalProtocolError::ApplicationCallbackError(
@@ -307,12 +307,12 @@ impl SignedPreKeyStore for &FfiSignedPreKeyStoreStruct {
 
     async fn save_signed_pre_key(
         &mut self,
-        prekey_id: u32,
+        prekey_id: SignedPreKeyId,
         record: &SignedPreKeyRecord,
         ctx: Context,
     ) -> Result<(), SignalProtocolError> {
         let ctx = ctx.unwrap_or(std::ptr::null_mut());
-        let result = (self.store_signed_pre_key)(self.ctx, prekey_id, &*record, ctx);
+        let result = (self.store_signed_pre_key)(self.ctx, prekey_id.into(), &*record, ctx);
 
         if let Some(error) = CallbackError::check(result) {
             return Err(SignalProtocolError::ApplicationCallbackError(

rust/bridge/shared/src/jni/storage.rs

@@ -253,27 +253,27 @@ impl<'a> JniPreKeyStore<'a> {
 impl<'a> PreKeyStore for JniPreKeyStore<'a> {
     async fn get_pre_key(
         &self,
-        prekey_id: u32,
+        prekey_id: PreKeyId,
         _ctx: Context,
     ) -> Result<PreKeyRecord, SignalProtocolError> {
-        Ok(self.do_get_pre_key(prekey_id)?)
+        Ok(self.do_get_pre_key(prekey_id.into())?)
     }
 
     async fn save_pre_key(
         &mut self,
-        prekey_id: u32,
+        prekey_id: PreKeyId,
         record: &PreKeyRecord,
         _ctx: Context,
     ) -> Result<(), SignalProtocolError> {
-        Ok(self.do_save_pre_key(prekey_id, record)?)
+        Ok(self.do_save_pre_key(prekey_id.into(), record)?)
     }
 
     async fn remove_pre_key(
         &mut self,
-        prekey_id: u32,
+        prekey_id: PreKeyId,
         _ctx: Context,
     ) -> Result<(), SignalProtocolError> {
-        Ok(self.do_remove_pre_key(prekey_id)?)
+        Ok(self.do_remove_pre_key(prekey_id.into())?)
     }
 }
 
@@ -331,19 +331,19 @@ impl<'a> JniSignedPreKeyStore<'a> {
 impl<'a> SignedPreKeyStore for JniSignedPreKeyStore<'a> {
     async fn get_signed_pre_key(
         &self,
-        prekey_id: u32,
+        prekey_id: SignedPreKeyId,
         _ctx: Context,
     ) -> Result<SignedPreKeyRecord, SignalProtocolError> {
-        Ok(self.do_get_signed_pre_key(prekey_id)?)
+        Ok(self.do_get_signed_pre_key(prekey_id.into())?)
     }
 
     async fn save_signed_pre_key(
         &mut self,
-        prekey_id: u32,
+        prekey_id: SignedPreKeyId,
         record: &SignedPreKeyRecord,
         _ctx: Context,
     ) -> Result<(), SignalProtocolError> {
-        Ok(self.do_save_signed_pre_key(prekey_id, record)?)
+        Ok(self.do_save_signed_pre_key(prekey_id.into(), record)?)
     }
 }
 

rust/bridge/shared/src/node/storage.rs

@@ -105,31 +105,31 @@ impl Finalize for NodePreKeyStore {
 impl PreKeyStore for NodePreKeyStore {
     async fn get_pre_key(
         &self,
-        pre_key_id: u32,
+        pre_key_id: PreKeyId,
         _ctx: libsignal_protocol::Context,
     ) -> Result<PreKeyRecord, SignalProtocolError> {
-        self.do_get_pre_key(pre_key_id)
+        self.do_get_pre_key(pre_key_id.into())
             .await
             .map_err(|s| js_error_to_rust("getPreKey", s))
     }
 
     async fn save_pre_key(
         &mut self,
-        pre_key_id: u32,
+        pre_key_id: PreKeyId,
         record: &PreKeyRecord,
         _ctx: libsignal_protocol::Context,
     ) -> Result<(), SignalProtocolError> {
-        self.do_save_pre_key(pre_key_id, record.clone())
+        self.do_save_pre_key(pre_key_id.into(), record.clone())
             .await
             .map_err(|s| js_error_to_rust("savePreKey", s))
     }
 
     async fn remove_pre_key(
         &mut self,
-        pre_key_id: u32,
+        pre_key_id: PreKeyId,
         _ctx: libsignal_protocol::Context,
     ) -> Result<(), SignalProtocolError> {
-        self.do_remove_pre_key(pre_key_id)
+        self.do_remove_pre_key(pre_key_id.into())
             .await
             .map_err(|s| js_error_to_rust("removePreKey", s))
     }
@@ -210,21 +210,21 @@ impl Finalize for NodeSignedPreKeyStore {
 impl SignedPreKeyStore for NodeSignedPreKeyStore {
     async fn get_signed_pre_key(
         &self,
-        signed_pre_key_id: u32,
+        signed_pre_key_id: SignedPreKeyId,
         _ctx: libsignal_protocol::Context,
     ) -> Result<SignedPreKeyRecord, SignalProtocolError> {
-        self.do_get_signed_pre_key(signed_pre_key_id)
+        self.do_get_signed_pre_key(signed_pre_key_id.into())
             .await
             .map_err(|s| js_error_to_rust("getSignedPreKey", s))
     }
 
     async fn save_signed_pre_key(
         &mut self,
-        signed_pre_key_id: u32,
+        signed_pre_key_id: SignedPreKeyId,
         record: &SignedPreKeyRecord,
         _ctx: libsignal_protocol::Context,
     ) -> Result<(), SignalProtocolError> {
-        self.do_save_signed_pre_key(signed_pre_key_id, record.clone())
+        self.do_save_signed_pre_key(signed_pre_key_id.into(), record.clone())
             .await
             .map_err(|s| js_error_to_rust("saveSignedPreKey", s))
     }

rust/bridge/shared/src/protocol.rs

@@ -1,5 +1,5 @@
 //
-// Copyright 2021 Signal Messenger, LLC.
+// Copyright 2021-2022 Signal Messenger, LLC.
 // SPDX-License-Identifier: AGPL-3.0-only
 //
 
@@ -87,7 +87,7 @@ fn HKDF_Derive(output: &mut [u8], ikm: &[u8], label: &[u8], salt: &[u8]) -> Resu
 
 #[bridge_fn(ffi = "address_new")]
 fn ProtocolAddress_New(name: String, device_id: u32) -> ProtocolAddress {
-    ProtocolAddress::new(name, device_id)
+    ProtocolAddress::new(name, device_id.into())
 }
 
 #[bridge_fn(ffi = "publickey_deserialize", jni = false)]
@@ -316,8 +316,8 @@ fn PreKeySignalMessage_New(
     PreKeySignalMessage::new(
         message_version,
         registration_id,
-        pre_key_id,
-        signed_pre_key_id,
+        pre_key_id.map(|id| id.into()),
+        signed_pre_key_id.into(),
         *base_key,
         IdentityKey::new(*identity_key),
         signal_message.clone(),
@@ -496,9 +496,9 @@ fn PreKeyBundle_New(
 ) -> Result<PreKeyBundle> {
     let identity_key = IdentityKey::new(*identity_key);
 
-    let prekey = match (prekey, prekey_id) {
+    let prekey: Option<(PreKeyId, PublicKey)> = match (prekey, prekey_id) {
         (None, None) => None,
-        (Some(k), Some(id)) => Some((id, *k)),
+        (Some(k), Some(id)) => Some((id.into(), *k)),
         _ => {
             return Err(SignalProtocolError::InvalidArgument(
                 "Must supply both or neither of prekey and prekey_id".to_owned(),
@@ -508,9 +508,9 @@ fn PreKeyBundle_New(
 
     PreKeyBundle::new(
         registration_id,
-        device_id,
+        device_id.into(),
         prekey,
-        signed_prekey_id,
+        signed_prekey_id.into(),
         *signed_prekey,
         signed_prekey_signature.to_vec(),
         identity_key,
@@ -550,7 +550,7 @@ fn SignedPreKeyRecord_New(
     signature: &[u8],
 ) -> SignedPreKeyRecord {
     let keypair = KeyPair::new(*pub_key, *priv_key);
-    SignedPreKeyRecord::new(id, timestamp.as_millis(), &keypair, signature)
+    SignedPreKeyRecord::new(id.into(), timestamp.as_millis(), &keypair, signature)
 }
 
 bridge_deserialize!(PreKeyRecord::deserialize);
@@ -565,7 +565,7 @@ bridge_get!(PreKeyRecord::private_key -> PrivateKey);
 #[bridge_fn]
 fn PreKeyRecord_New(id: u32, pub_key: &PublicKey, priv_key: &PrivateKey) -> PreKeyRecord {
     let keypair = KeyPair::new(*pub_key, *priv_key);
-    PreKeyRecord::new(id, &keypair)
+    PreKeyRecord::new(id.into(), &keypair)
 }
 
 bridge_deserialize!(SenderKeyRecord::deserialize);
@@ -631,7 +631,7 @@ fn SenderCertificate_New(
         sender_uuid,
         sender_e164,
         *sender_key,
-        sender_device_id,
+        sender_device_id.into(),
         expiration.as_millis(),
         signer_cert.clone(),
         signer_key,
@@ -1101,7 +1101,7 @@ async fn SealedSender_DecryptMessage(
         timestamp.as_millis(),
         local_e164,
         local_uuid,
-        local_device_id,
+        local_device_id.into(),
         identity_store,
         session_store,
         prekey_store,

rust/protocol/benches/ratchet.rs

@@ -1,5 +1,5 @@
 //
-// Copyright 2020 Signal Messenger, LLC.
+// Copyright 2020-2022 Signal Messenger, LLC.
 // SPDX-License-Identifier: AGPL-3.0-only
 //
 
@@ -20,7 +20,7 @@ pub fn ratchet_forward_result(c: &mut Criterion) -> Result<(), SignalProtocolErr
 
     let mut csprng = rand::rngs::OsRng;
 
-    let sender_address = ProtocolAddress::new("+14159999111".to_owned(), 1);
+    let sender_address = ProtocolAddress::new("+14159999111".to_owned(), 1.into());
     let distribution_id = Uuid::from_u128(0xd1d1d1d1_7000_11eb_b32a_33b8a8a487a6);
 
     let mut alice_store = support::test_in_memory_protocol_store()?;

rust/protocol/benches/sealed_sender.rs

@@ -16,8 +16,10 @@ mod support;
 pub fn v1(c: &mut Criterion) {
     let mut rng = OsRng;
 
-    let alice_address = ProtocolAddress::new("9d0652a3-dcc3-4d11-975f-74d61598733f".to_owned(), 1);
-    let bob_address = ProtocolAddress::new("796abedb-ca4e-4f18-8803-1fde5b921f9f".to_owned(), 1);
+    let alice_address =
+        ProtocolAddress::new("9d0652a3-dcc3-4d11-975f-74d61598733f".to_owned(), 1.into());
+    let bob_address =
+        ProtocolAddress::new("796abedb-ca4e-4f18-8803-1fde5b921f9f".to_owned(), 1.into());
 
     let mut alice_store = support::test_in_memory_protocol_store().expect("brand new store");
     let mut bob_store = support::test_in_memory_protocol_store().expect("brand new store");
@@ -104,8 +106,10 @@ pub fn v1(c: &mut Criterion) {
 pub fn v2(c: &mut Criterion) {
     let mut rng = OsRng;
 
-    let alice_address = ProtocolAddress::new("9d0652a3-dcc3-4d11-975f-74d61598733f".to_owned(), 1);
-    let bob_address = ProtocolAddress::new("796abedb-ca4e-4f18-8803-1fde5b921f9f".to_owned(), 1);
+    let alice_address =
+        ProtocolAddress::new("9d0652a3-dcc3-4d11-975f-74d61598733f".to_owned(), 1.into());
+    let bob_address =
+        ProtocolAddress::new("796abedb-ca4e-4f18-8803-1fde5b921f9f".to_owned(), 1.into());
 
     let mut alice_store = support::test_in_memory_protocol_store().expect("brand new store");
     let mut bob_store = support::test_in_memory_protocol_store().expect("brand new store");
@@ -201,7 +205,7 @@ pub fn v2(c: &mut Criterion) {
     // Fill out additional recipients.
     let mut recipients = vec![bob_address.clone()];
     while recipients.len() < 10 {
-        let next_address = ProtocolAddress::new(Uuid::from_bytes(rng.gen()).to_string(), 1);
+        let next_address = ProtocolAddress::new(Uuid::from_bytes(rng.gen()).to_string(), 1.into());
 
         let mut next_store = support::test_in_memory_protocol_store().expect("brand new store");
 

rust/protocol/benches/session.rs

@@ -14,8 +14,8 @@ mod support;
 pub fn session_encrypt_result(c: &mut Criterion) -> Result<(), SignalProtocolError> {
     let (alice_session_record, bob_session_record) = support::initialize_sessions_v3()?;
 
-    let alice_address = ProtocolAddress::new("+14159999999".to_owned(), 1);
-    let bob_address = ProtocolAddress::new("+14158888888".to_owned(), 1);
+    let alice_address = ProtocolAddress::new("+14159999999".to_owned(), 1.into());
+    let bob_address = ProtocolAddress::new("+14158888888".to_owned(), 1.into());
 
     let mut alice_store = support::test_in_memory_protocol_store()?;
     let mut bob_store = support::test_in_memory_protocol_store()?;
@@ -98,9 +98,9 @@ pub fn session_encrypt_result(c: &mut Criterion) -> Result<(), SignalProtocolErr
             .get_local_registration_id(None)
             .now_or_never()
             .expect("sync")?,
-        1,                 // device id
-        None,              // pre key
-        signed_pre_key_id, // signed pre key id
+        1.into(),                 // device id
+        None,                     // pre key
+        signed_pre_key_id.into(), // signed pre key id
         bob_signed_pre_key_pair.public_key,
         bob_signed_pre_key_signature.to_vec(),
         *bob_store
@@ -112,9 +112,9 @@ pub fn session_encrypt_result(c: &mut Criterion) -> Result<(), SignalProtocolErr
 
     bob_store
         .save_signed_pre_key(
-            signed_pre_key_id,
+            signed_pre_key_id.into(),
             &SignedPreKeyRecord::new(
-                signed_pre_key_id,
+                signed_pre_key_id.into(),
                 /*timestamp*/ 42,
                 &bob_signed_pre_key_pair,
                 &bob_signed_pre_key_signature,
@@ -183,8 +183,8 @@ pub fn session_encrypt_result(c: &mut Criterion) -> Result<(), SignalProtocolErr
 pub fn session_encrypt_decrypt_result(c: &mut Criterion) -> Result<(), SignalProtocolError> {
     let (alice_session_record, bob_session_record) = support::initialize_sessions_v3()?;
 
-    let alice_address = ProtocolAddress::new("+14159999999".to_owned(), 1);
-    let bob_address = ProtocolAddress::new("+14158888888".to_owned(), 1);
+    let alice_address = ProtocolAddress::new("+14159999999".to_owned(), 1.into());
+    let bob_address = ProtocolAddress::new("+14158888888".to_owned(), 1.into());
 
     let mut alice_store = support::test_in_memory_protocol_store()?;
     let mut bob_store = support::test_in_memory_protocol_store()?;