Use IdentityKey where it makes sense

whisperfish · Jan 30, 2024 · d6a6fc2 · d6a6fc2
1 parent c2ceb80
commit d6a6fc2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 14 deletions.
diff --git a/libsignal-service/src/account_manager.rs b/libsignal-service/src/account_manager.rs
@@ -6,7 +6,8 @@ use aes::cipher::{KeyIvInit, StreamCipher as _};
 use hmac::digest::Output;
 use hmac::{Hmac, Mac};
 use libsignal_protocol::{
-    IdentityKeyStore, KeyPair, PrivateKey, PublicKey, SignalProtocolError,
+    IdentityKey, IdentityKeyStore, KeyPair, PrivateKey, PublicKey,
+    SignalProtocolError,
 };
 use prost::Message;
 use serde::{Deserialize, Serialize};
@@ -408,7 +409,7 @@ impl<Service: PushService> AccountManager<Service> {
     pub async fn update_device_name(
         &mut self,
         device_name: &str,
-        public_key: &PublicKey,
+        public_key: &IdentityKey,
     ) -> Result<(), ServiceError> {
         let encrypted_device_name = encrypt_device_name(
             &mut rand::thread_rng(),
@@ -585,14 +586,14 @@ fn calculate_hmac256(
 pub fn encrypt_device_name<R: rand::Rng + rand::CryptoRng>(
     csprng: &mut R,
     device_name: &str,
-    identity_public: &PublicKey,
+    identity_public: &IdentityKey,
 ) -> Result<DeviceName, ServiceError> {
     let plaintext = device_name.as_bytes().to_vec();
     let ephemeral_key_pair = KeyPair::generate(csprng);
 
     let master_secret = ephemeral_key_pair
         .private_key
-        .calculate_agreement(identity_public)?;
+        .calculate_agreement(identity_public.public_key())?;
 
     let key1 = calculate_hmac256(&master_secret, b"auth")?;
     let synthetic_iv = calculate_hmac256(&key1, &plaintext)?;

diff --git a/libsignal-service/src/provisioning/mod.rs b/libsignal-service/src/provisioning/mod.rs
@@ -99,10 +99,10 @@ pub struct NewDeviceRegistration {
     pub service_ids: ServiceIds,
     #[derivative(Debug = "ignore")]
     pub aci_private_key: PrivateKey,
-    pub aci_public_key: PublicKey,
+    pub aci_public_key: IdentityKey,
     #[derivative(Debug = "ignore")]
     pub pni_private_key: PrivateKey,
-    pub pni_public_key: PublicKey,
+    pub pni_public_key: IdentityKey,
     #[derivative(Debug = "ignore")]
     pub profile_key: ProfileKey,
 }
@@ -165,6 +165,7 @@ pub async fn link_device<
                     reason: "missing public key".into(),
                 },
             )?)?;
+        let aci_public_key = IdentityKey::new(aci_public_key);
 
         let aci_private_key =
             PrivateKey::deserialize(&message.aci_identity_key_private.ok_or(
@@ -179,6 +180,7 @@ pub async fn link_device<
                     reason: "missing public key".into(),
                 },
             )?)?;
+        let pni_public_key = IdentityKey::new(pni_public_key);
 
         let pni_private_key =
             PrivateKey::deserialize(&message.pni_identity_key_private.ok_or(
@@ -210,14 +212,10 @@ pub async fn link_device<
             },
         )?;
 
-        let aci_key_pair = IdentityKeyPair::new(
-            IdentityKey::new(aci_public_key),
-            aci_private_key,
-        );
-        let pni_key_pair = IdentityKeyPair::new(
-            IdentityKey::new(pni_public_key),
-            pni_private_key,
-        );
+        let aci_key_pair =
+            IdentityKeyPair::new(aci_public_key, aci_private_key);
+        let pni_key_pair =
+            IdentityKeyPair::new(pni_public_key, pni_private_key);
 
         let (
             _aci_pre_keys,