From 0faee1b50b0da91bde6a44eeee05a723833e9664 Mon Sep 17 00:00:00 2001
From: Ruben De Smet <ruben.de.smet@rubdos.be>
Date: Sun, 14 Jan 2024 20:19:53 +0100
Subject: [PATCH] Deduplicate SignedPreKeyEntity and SignedPreKey

---
 libsignal-service/src/account_manager.rs | 12 +++++++-----
 libsignal-service/src/pre_keys.rs        | 20 +++++---------------
 libsignal-service/src/push_service.rs    |  7 +++----
 3 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/libsignal-service/src/account_manager.rs b/libsignal-service/src/account_manager.rs
index aa1fdd4e0..47934db6b 100644
--- a/libsignal-service/src/account_manager.rs
+++ b/libsignal-service/src/account_manager.rs
@@ -97,7 +97,7 @@ impl<Service: PushService> AccountManager<Service> {
     ) -> Result<
         (
             Vec<PreKeyEntity>,
-            SignedPreKeyRecord,
+            SignedPreKeyEntity,
             Vec<KyberPreKeyEntity>,
             Option<KyberPreKeyEntity>,
         ),
@@ -184,6 +184,8 @@ impl<Service: PushService> AccountManager<Service> {
                     &signed_prekey_record,
                 )
                     .instrument(tracing::trace_span!(parent: &span, "save signed pre key", signed_pre_key_id = ?next_signed_pre_key_id)).await?;
+        let signed_prekey_entity =
+            SignedPreKeyEntity::try_from(signed_prekey_record)?;
 
         let pq_last_resort_key = if use_last_resort_key {
             tracing::warn!("Last resort Kyber key unimplemented");
@@ -204,7 +206,7 @@ impl<Service: PushService> AccountManager<Service> {
 
         Ok((
             pre_key_entities,
-            signed_prekey_record,
+            signed_prekey_entity,
             pq_pre_key_entities,
             pq_last_resort_key,
         ))
@@ -260,8 +262,8 @@ impl<Service: PushService> AccountManager<Service> {
             return Ok(());
         }
 
-        let (pre_keys, signed_pre_key_record, pq_pre_keys, pq_last_resort_key) =
-            self.generate_pre_keys(
+        let (pre_keys, signed_pre_key, pq_pre_keys, pq_last_resort_key) = self
+            .generate_pre_keys(
                 protocol_store,
                 service_id_type,
                 csprng,
@@ -281,7 +283,7 @@ impl<Service: PushService> AccountManager<Service> {
 
         let pre_key_state = PreKeyState {
             pre_keys,
-            signed_pre_key: signed_pre_key_record.try_into()?,
+            signed_pre_key,
             identity_key,
             pq_pre_keys,
             pq_last_resort_key,
diff --git a/libsignal-service/src/pre_keys.rs b/libsignal-service/src/pre_keys.rs
index 0c5c6b4ac..af6c7d93a 100644
--- a/libsignal-service/src/pre_keys.rs
+++ b/libsignal-service/src/pre_keys.rs
@@ -74,24 +74,14 @@ pub struct SignedPreKeyEntity {
     pub signature: Vec<u8>,
 }
 
-#[derive(Debug, Serialize)]
-#[serde(rename_all = "camelCase")]
-pub struct SignedPreKey {
-    key_id: u32,
-    #[serde(with = "serde_public_key")]
-    public_key: PublicKey,
-    #[serde(with = "serde_base64")]
-    signature: Vec<u8>,
-}
-
-impl TryFrom<SignedPreKeyRecord> for SignedPreKey {
+impl TryFrom<SignedPreKeyRecord> for SignedPreKeyEntity {
     type Error = SignalProtocolError;
 
     fn try_from(key: SignedPreKeyRecord) -> Result<Self, Self::Error> {
-        Ok(SignedPreKey {
+        Ok(SignedPreKeyEntity {
             key_id: key.id()?.into(),
-            public_key: key.key_pair()?.public_key,
-            signature: key.signature()?,
+            public_key: key.key_pair()?.public_key.serialize().to_vec(),
+            signature: key.signature()?.to_vec(),
         })
     }
 }
@@ -122,7 +112,7 @@ impl TryFrom<KyberPreKeyRecord> for KyberPreKeyEntity {
 #[serde(rename_all = "camelCase")]
 pub struct PreKeyState {
     pub pre_keys: Vec<PreKeyEntity>,
-    pub signed_pre_key: SignedPreKey,
+    pub signed_pre_key: SignedPreKeyEntity,
     #[serde(with = "serde_public_key")]
     pub identity_key: PublicKey,
     #[serde(skip_serializing_if = "Option::is_none")]
diff --git a/libsignal-service/src/push_service.rs b/libsignal-service/src/push_service.rs
index 5abbcad0b..13bd1c933 100644
--- a/libsignal-service/src/push_service.rs
+++ b/libsignal-service/src/push_service.rs
@@ -5,8 +5,7 @@ use crate::{
     envelope::*,
     groups_v2::GroupDecodingError,
     pre_keys::{
-        KyberPreKeyEntity, PreKeyEntity, PreKeyState, SignedPreKey,
-        SignedPreKeyEntity,
+        KyberPreKeyEntity, PreKeyEntity, PreKeyState, SignedPreKeyEntity,
     },
     profile_cipher::ProfileCipherError,
     proto::{attachment_pointer::AttachmentIdentifier, AttachmentPointer},
@@ -407,8 +406,8 @@ pub struct StaleDevices {
 pub struct LinkRequest {
     pub verification_code: String,
     pub account_attributes: LinkAccountAttributes,
-    pub aci_signed_pre_key: SignedPreKey,
-    pub pni_signed_pre_key: SignedPreKey,
+    pub aci_signed_pre_key: SignedPreKeyEntity,
+    pub pni_signed_pre_key: SignedPreKeyEntity,
     pub aci_pq_last_resort_pre_key: KyberPreKeyEntity,
     pub pni_pq_last_resort_pre_key: KyberPreKeyEntity,
 }