introduce PreKeyId wrapper struct

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

@@ -194,12 +194,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(
@@ -218,12 +218,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(
@@ -237,11 +237,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(

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

@@ -277,27 +277,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())?)
     }
 }
 

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))
     }

rust/bridge/shared/src/protocol.rs

@@ -286,7 +286,7 @@ fn PreKeySignalMessage_New(
     PreKeySignalMessage::new(
         message_version,
         registration_id,
-        pre_key_id,
+        pre_key_id.map(|id| id.into()),
         signed_pre_key_id.into(),
         *base_key,
         IdentityKey::new(*identity_key),
@@ -520,9 +520,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(),
@@ -589,7 +589,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);

rust/protocol/src/lib.rs

@@ -75,7 +75,9 @@ pub use {
     session_cipher::{
         message_decrypt, message_decrypt_prekey, message_decrypt_signal, message_encrypt,
     },
-    state::{PreKeyBundle, PreKeyRecord, SessionRecord, SignedPreKeyId, SignedPreKeyRecord},
+    state::{
+        PreKeyBundle, PreKeyId, PreKeyRecord, SessionRecord, SignedPreKeyId, SignedPreKeyRecord,
+    },
     storage::{
         Context, Direction, IdentityKeyStore, InMemIdentityKeyStore, InMemPreKeyStore,
         InMemSenderKeyStore, InMemSessionStore, InMemSignalProtocolStore, InMemSignedPreKeyStore,

rust/protocol/src/protocol.rs

@@ -4,7 +4,7 @@
 //
 
 use crate::proto;
-use crate::state::SignedPreKeyId;
+use crate::state::{PreKeyId, SignedPreKeyId};
 use crate::{IdentityKey, PrivateKey, PublicKey, Result, SignalProtocolError};
 
 use std::convert::TryFrom;
@@ -238,7 +238,7 @@ impl TryFrom<&[u8]> for SignalMessage {
 pub struct PreKeySignalMessage {
     message_version: u8,
     registration_id: u32,
-    pre_key_id: Option<u32>,
+    pre_key_id: Option<PreKeyId>,
     signed_pre_key_id: SignedPreKeyId,
     base_key: PublicKey,
     identity_key: IdentityKey,
@@ -250,15 +250,15 @@ impl PreKeySignalMessage {
     pub fn new(
         message_version: u8,
         registration_id: u32,
-        pre_key_id: Option<u32>,
+        pre_key_id: Option<PreKeyId>,
         signed_pre_key_id: SignedPreKeyId,
         base_key: PublicKey,
         identity_key: IdentityKey,
         message: SignalMessage,
     ) -> Result<Self> {
         let proto_message = proto::wire::PreKeySignalMessage {
             registration_id: Some(registration_id),
-            pre_key_id,
+            pre_key_id: pre_key_id.map(|id| id.into()),
             signed_pre_key_id: Some(signed_pre_key_id.into()),
             base_key: Some(base_key.serialize().into_vec()),
             identity_key: Some(identity_key.serialize().into_vec()),
@@ -290,7 +290,7 @@ impl PreKeySignalMessage {
     }
 
     #[inline]
-    pub fn pre_key_id(&self) -> Option<u32> {
+    pub fn pre_key_id(&self) -> Option<PreKeyId> {
         self.pre_key_id
     }
 
@@ -366,7 +366,7 @@ impl TryFrom<&[u8]> for PreKeySignalMessage {
         Ok(PreKeySignalMessage {
             message_version,
             registration_id: proto_structure.registration_id.unwrap_or(0),
-            pre_key_id: proto_structure.pre_key_id,
+            pre_key_id: proto_structure.pre_key_id.map(|id| id.into()),
             signed_pre_key_id: signed_pre_key_id.into(),
             base_key,
             identity_key: IdentityKey::try_from(identity_key.as_ref())?,

rust/protocol/src/session.rs

@@ -4,13 +4,13 @@
 //
 
 use crate::{
-    Context, Direction, IdentityKeyStore, KeyPair, PreKeyBundle, PreKeySignalMessage, PreKeyStore,
-    ProtocolAddress, Result, SessionRecord, SessionStore, SignalProtocolError, SignedPreKeyStore,
+    Context, Direction, IdentityKeyStore, KeyPair, PreKeyBundle, PreKeyId, PreKeySignalMessage,
+    PreKeyStore, ProtocolAddress, Result, SessionRecord, SessionStore, SignalProtocolError,
+    SignedPreKeyStore,
 };
 
 use crate::ratchet;
 use crate::ratchet::{AliceSignalProtocolParameters, BobSignalProtocolParameters};
-use crate::state::PreKeyId;
 use rand::{CryptoRng, Rng};
 
 /*
@@ -178,7 +178,7 @@ pub async fn process_prekey_bundle<R: Rng + CryptoRng>(
     log::info!(
         "set_unacknowledged_pre_key_message for: {} with preKeyId: {}",
         remote_address,
-        their_one_time_prekey_id.map_or_else(|| "<none>".to_string(), |id| id.to_string())
+        their_one_time_prekey_id.map_or_else(|| "<none>".to_string(), |id| format!("{:?}", id))
     );
 
     session.set_unacknowledged_pre_key_message(

rust/protocol/src/session_cipher.rs

@@ -53,7 +53,7 @@ pub async fn message_encrypt(
             remote_address,
             items
                 .pre_key_id()?
-                .map_or_else(|| "<none>".to_string(), |id| id.to_string())
+                .map_or_else(|| "<none>".to_string(), |id| format!("{:?}", id))
         );
 
         let message = SignalMessage::new(

rust/protocol/src/state/prekey.rs

@@ -7,7 +7,21 @@ use crate::proto::storage::PreKeyRecordStructure;
 use crate::{KeyPair, PrivateKey, PublicKey, Result};
 use prost::Message;
 
-pub type PreKeyId = u32;
+/// A unique identifier selecting among this client's known pre-keys.
+#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, PartialOrd, Ord)]
+pub struct PreKeyId(pub u32);
+
+impl From<u32> for PreKeyId {
+    fn from(value: u32) -> Self {
+        Self(value)
+    }
+}
+
+impl From<PreKeyId> for u32 {
+    fn from(value: PreKeyId) -> Self {
+        value.0
+    }
+}
 
 #[derive(Debug, Clone)]
 pub struct PreKeyRecord {
@@ -20,7 +34,7 @@ impl PreKeyRecord {
         let private_key = key.private_key.serialize().to_vec();
         Self {
             pre_key: PreKeyRecordStructure {
-                id,
+                id: id.into(),
                 public_key,
                 private_key,
             },
@@ -34,7 +48,7 @@ impl PreKeyRecord {
     }
 
     pub fn id(&self) -> Result<PreKeyId> {
-        Ok(self.pre_key.id)
+        Ok(self.pre_key.id.into())
     }
 
     pub fn key_pair(&self) -> Result<KeyPair> {

rust/protocol/src/state/session.rs

@@ -388,7 +388,7 @@ impl SessionState {
     ) -> Result<()> {
         let signed_pre_key_id: u32 = signed_pre_key_id.into();
         let pending = session_structure::PendingPreKey {
-            pre_key_id: pre_key_id.unwrap_or(0),
+            pre_key_id: pre_key_id.unwrap_or(0.into()).into(),
             signed_pre_key_id: signed_pre_key_id as i32,
             base_key: base_key.serialize().to_vec(),
         };
@@ -403,7 +403,7 @@ impl SessionState {
             Ok(Some(UnacknowledgedPreKeyMessageItems::new(
                 match pending_pre_key.pre_key_id {
                     0 => None,
-                    v => Some(v),
+                    v => Some(v.into()),
                 },
                 (pending_pre_key.signed_pre_key_id as u32).into(),
                 PublicKey::deserialize(&pending_pre_key.base_key)?,

rust/protocol/src/storage/inmem.rs

@@ -4,11 +4,10 @@
 //
 
 use crate::{
-    IdentityKey, IdentityKeyPair, PreKeyRecord, ProtocolAddress, Result, SenderKeyRecord,
-    SessionRecord, SignalProtocolError, SignedPreKeyRecord,
+    IdentityKey, IdentityKeyPair, PreKeyId, PreKeyRecord, ProtocolAddress, Result, SenderKeyRecord,
+    SessionRecord, SignalProtocolError, SignedPreKeyId, SignedPreKeyRecord,
 };
 
-use crate::state::{PreKeyId, SignedPreKeyId};
 use crate::storage::traits;
 use crate::storage::Context;