Refactoring ACI/PNI storage types

libsignal-service-actix/examples/registering.rs

@@ -58,7 +58,8 @@ async fn main() -> Result<(), Error> {
     let mut push_service = AwcPushService::new(
         servers,
         Some(ServiceCredentials {
-            uuid: None,
+            aci: None,
+            pni: None,
             phonenumber: phonenumber.clone(),
             password,
             signaling_key: None,

libsignal-service-hyper/examples/registering.rs

@@ -148,7 +148,8 @@ fn create_push_service(
     HyperPushService::new(
         SignalServers::Staging, // You might want to switch to Production servers
         Some(ServiceCredentials {
-            uuid: None,
+            aci: None,
+            pni: None,
             phonenumber,
             password: Some(password),
             signaling_key: None,

libsignal-service/src/account_manager.rs

@@ -7,7 +7,7 @@ use hmac::digest::Output;
 use hmac::{Hmac, Mac};
 use libsignal_protocol::{
     kem, GenericSignedPreKey, IdentityKeyStore, KeyPair, KyberPreKeyRecord,
-    PreKeyRecord, PrivateKey, ProtocolStore, PublicKey, SignalProtocolError,
+    PreKeyRecord, PrivateKey, PublicKey, SignalProtocolError,
     SignedPreKeyRecord,
 };
 use prost::Message;
@@ -16,7 +16,7 @@ use sha2::Sha256;
 use tracing_futures::Instrument;
 use zkgroup::profiles::ProfileKey;
 
-use crate::pre_keys::KyberPreKeyEntity;
+use crate::pre_keys::{KyberPreKeyEntity, PreKeysStore};
 use crate::proto::DeviceName;
 use crate::push_service::{AvatarWrite, RecaptchaAttributes, ServiceIdType};
 use crate::ServiceAddress;
@@ -93,19 +93,22 @@ impl<Service: PushService> AccountManager<Service> {
     #[tracing::instrument(skip(self, protocol_store, csprng))]
     pub async fn update_pre_key_bundle<
         R: rand::Rng + rand::CryptoRng,
-        P: ProtocolStore,
+        P: PreKeysStore,
     >(
         &mut self,
         protocol_store: &mut P,
+        service_id_type: ServiceIdType,
         csprng: &mut R,
-        pre_keys_offset_id: u32,
-        next_signed_pre_key_id: u32,
-        pq_pre_keys_offset_id: u32,
         use_last_resort_key: bool,
     ) -> Result<(u32, u32, u32), ServiceError> {
+        let pre_keys_offset_id = protocol_store.next_pre_key_id().await?;
+        let next_signed_pre_key_id =
+            protocol_store.next_signed_pre_key_id().await?;
+        let pq_pre_keys_offset_id = protocol_store.next_pq_pre_key_id().await?;
+
         let prekey_status = match self
             .service
-            .get_pre_key_status(ServiceIdType::AccountIdentity)
+            .get_pre_key_status(service_id_type)
             .instrument(tracing::span!(
                 tracing::Level::DEBUG,
                 "Fetching pre key status"
@@ -239,7 +242,7 @@ impl<Service: PushService> AccountManager<Service> {
         };
 
         self.service
-            .register_pre_keys(ServiceIdType::AccountIdentity, pre_key_state)
+            .register_pre_keys(service_id_type, pre_key_state)
             .instrument(tracing::span!(
                 tracing::Level::DEBUG,
                 "Uploading pre keys"
@@ -315,7 +318,8 @@ impl<Service: PushService> AccountManager<Service> {
     pub async fn link_device(
         &mut self,
         url: url::Url,
-        identity_store: &dyn IdentityKeyStore,
+        aci_identity_store: &dyn IdentityKeyStore,
+        pni_identity_store: &dyn IdentityKeyStore,
         credentials: ServiceCredentials,
     ) -> Result<(), LinkError> {
         let query: HashMap<_, _> = url.query_pairs().collect();
@@ -327,27 +331,36 @@ impl<Service: PushService> AccountManager<Service> {
         let pub_key = PublicKey::deserialize(&pub_key)
             .map_err(|_e| LinkError::InvalidPublicKey)?;
 
-        let identity_key_pair = identity_store.get_identity_key_pair().await?;
+        let aci_identity_key_pair =
+            aci_identity_store.get_identity_key_pair().await?;
+        let pni_identity_key_pair =
+            pni_identity_store.get_identity_key_pair().await?;
 
-        if credentials.uuid.is_none() {
-            tracing::warn!("No local UUID set");
+        if credentials.aci.is_none() {
+            tracing::warn!("No local ACI set");
+        }
+        if credentials.pni.is_none() {
+            tracing::warn!("No local PNI set");
         }
 
         let provisioning_code = self.new_device_provisioning_code().await?;
 
         let msg = ProvisionMessage {
-            aci: credentials.uuid.as_ref().map(|u| u.to_string()),
+            aci: credentials.aci.as_ref().map(|u| u.to_string()),
             aci_identity_key_public: Some(
-                identity_key_pair.public_key().serialize().into_vec(),
+                aci_identity_key_pair.public_key().serialize().into_vec(),
             ),
             aci_identity_key_private: Some(
-                identity_key_pair.private_key().serialize(),
+                aci_identity_key_pair.private_key().serialize(),
             ),
             number: Some(credentials.e164()),
-            // TODO: implement pni fields
-            pni_identity_key_public: None,
-            pni_identity_key_private: None,
-            pni: None,
+            pni_identity_key_public: Some(
+                pni_identity_key_pair.public_key().serialize().into_vec(),
+            ),
+            pni_identity_key_private: Some(
+                pni_identity_key_pair.private_key().serialize(),
+            ),
+            pni: credentials.pni.as_ref().map(uuid::Uuid::to_string),
             profile_key: self.profile_key.as_ref().map(|x| x.bytes.to_vec()),
             // CURRENT is not exposed by prost :(
             provisioning_version: Some(i32::from(

libsignal-service/src/configuration.rs

@@ -25,7 +25,8 @@ pub type SignalingKey = [u8; CIPHER_KEY_SIZE + MAC_KEY_SIZE];
 
 #[derive(Clone)]
 pub struct ServiceCredentials {
-    pub uuid: Option<uuid::Uuid>,
+    pub aci: Option<uuid::Uuid>,
+    pub pni: Option<uuid::Uuid>,
     pub phonenumber: phonenumber::PhoneNumber,
     pub password: Option<String>,
     pub signaling_key: Option<SignalingKey>,
@@ -49,7 +50,7 @@ impl ServiceCredentials {
 
     pub fn login(&self) -> String {
         let identifier = {
-            if let Some(uuid) = self.uuid.as_ref() {
+            if let Some(uuid) = self.aci.as_ref() {
                 uuid.to_string()
             } else {
                 self.e164()

libsignal-service/src/pre_keys.rs

@@ -1,43 +1,45 @@
 use std::{convert::TryFrom, time::SystemTime};
 
 use crate::utils::{serde_base64, serde_public_key};
+use async_trait::async_trait;
 use libsignal_protocol::{
-    error::SignalProtocolError, kem, GenericSignedPreKey, KeyPair,
-    KyberPreKeyRecord, KyberPreKeyStore, PreKeyRecord, PreKeyStore, PublicKey,
-    SignedPreKeyRecord, SignedPreKeyStore,
+    error::SignalProtocolError, kem, GenericSignedPreKey, IdentityKeyStore,
+    KeyPair, KyberPreKeyRecord, KyberPreKeyStore, PreKeyRecord, PreKeyStore,
+    PublicKey, SignedPreKeyRecord, SignedPreKeyStore,
 };
 
 use serde::{Deserialize, Serialize};
 
 /// Stores the ID of keys published ahead of time
 ///
 /// <https://signal.org/docs/specifications/x3dh/>
+#[async_trait(?Send)]
 pub trait PreKeysStore:
-    PreKeyStore + SignedPreKeyStore + KyberPreKeyStore
+    PreKeyStore + IdentityKeyStore + SignedPreKeyStore + KyberPreKeyStore
 {
     /// ID of the next pre key
-    fn pre_keys_offset_id(&self) -> Result<u32, SignalProtocolError>;
+    async fn next_pre_key_id(&self) -> Result<u32, SignalProtocolError>;
 
     /// ID of the next signed pre key
-    fn next_signed_pre_key_id(&self) -> Result<u32, SignalProtocolError>;
+    async fn next_signed_pre_key_id(&self) -> Result<u32, SignalProtocolError>;
 
     /// ID of the next PQ pre key
-    fn next_pq_pre_key_id(&self) -> Result<u32, SignalProtocolError>;
+    async fn next_pq_pre_key_id(&self) -> Result<u32, SignalProtocolError>;
 
     /// set the ID of the next pre key
-    fn set_pre_keys_offset_id(
+    async fn set_next_pre_key_id(
         &mut self,
         id: u32,
     ) -> Result<(), SignalProtocolError>;
 
     /// set the ID of the next signed pre key
-    fn set_next_signed_pre_key_id(
+    async fn set_next_signed_pre_key_id(
         &mut self,
         id: u32,
     ) -> Result<(), SignalProtocolError>;
 
     /// set the ID of the next PQ pre key
-    fn set_next_pq_pre_key_id(
+    async fn set_next_pq_pre_key_id(
         &mut self,
         id: u32,
     ) -> Result<(), SignalProtocolError>;
@@ -132,7 +134,7 @@ pub(crate) async fn generate_last_resort_kyber_key<S: PreKeysStore>(
     store: &mut S,
     identity_key: &KeyPair,
 ) -> Result<KyberPreKeyRecord, SignalProtocolError> {
-    let id = store.next_pq_pre_key_id()?;
+    let id = store.next_pq_pre_key_id().await?;
     let id = id.max(1); // TODO: Hack, keys start with 1
 
     let record = KyberPreKeyRecord::generate(
@@ -142,7 +144,7 @@ pub(crate) async fn generate_last_resort_kyber_key<S: PreKeysStore>(
     )?;
 
     store.save_kyber_pre_key(id.into(), &record).await?;
-    store.set_next_pq_pre_key_id(id + 1)?;
+    store.set_next_pq_pre_key_id(id + 1).await?;
 
     Ok(record)
 }
@@ -155,7 +157,7 @@ pub(crate) async fn generate_signed_pre_key<
     csprng: &mut R,
     identity_key: &KeyPair,
 ) -> Result<SignedPreKeyRecord, SignalProtocolError> {
-    let id = store.next_signed_pre_key_id()?;
+    let id = store.next_signed_pre_key_id().await?;
     let id = id.max(1); // TODO: Hack, keys start with 1
 
     let key_pair = KeyPair::generate(csprng);
@@ -171,7 +173,7 @@ pub(crate) async fn generate_signed_pre_key<
         SignedPreKeyRecord::new(id.into(), unix_time, &key_pair, &signature);
 
     store.save_signed_pre_key(id.into(), &record).await?;
-    store.set_next_signed_pre_key_id(id + 1)?;
+    store.set_next_signed_pre_key_id(id + 1).await?;
 
     Ok(record)
 }

libsignal-service/src/provisioning/mod.rs

@@ -108,11 +108,12 @@ pub struct NewDeviceRegistration {
 
 pub async fn link_device<
     R: rand::Rng + rand::CryptoRng,
-    S: PreKeysStore,
+    Aci: PreKeysStore,
+    Pni: PreKeysStore,
     P: PushService,
 >(
-    aci_store: &mut S,
-    pni_store: &mut S,
+    aci_store: &mut Aci,
+    pni_store: &mut Pni,
     csprng: &mut R,
     mut push_service: P,
     password: &str,

libsignal-service/src/push_service.rs

@@ -73,7 +73,7 @@ pub const STICKER_PATH: &str = "stickers/%s/full/%d";
 pub const KEEPALIVE_TIMEOUT_SECONDS: Duration = Duration::from_secs(55);
 pub const DEFAULT_DEVICE_ID: u32 = 1;
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
 pub enum ServiceIdType {
     /// Account Identity (ACI)
     ///

libsignal-service/src/service_address.rs

@@ -29,6 +29,10 @@ impl ServiceAddress {
     pub fn aci(&self) -> libsignal_protocol::Aci {
         libsignal_protocol::Aci::from_uuid_bytes(self.uuid.into_bytes())
     }
+
+    pub fn pni(&self) -> libsignal_protocol::Pni {
+        libsignal_protocol::Pni::from_uuid_bytes(self.uuid.into_bytes())
+    }
 }
 
 impl From<Uuid> for ServiceAddress {