From ddbcd3b640dbafa8dab2494ea22d95dc66ee5cce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gabriel=20F=C3=A9ron?= <g@leirbag.net>
Date: Mon, 21 Oct 2024 19:19:47 +0200
Subject: [PATCH] Fix some occurences of not using PNI correctly

---
 src/push_service/keys.rs | 12 ++++++++----
 src/sender.rs            |  6 ++++--
 src/utils.rs             | 25 +++++++++++++++++++++++++
 src/websocket/sender.rs  | 10 ++++++++--
 4 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/src/push_service/keys.rs b/src/push_service/keys.rs
index 89f324aec..82c9198f3 100644
--- a/src/push_service/keys.rs
+++ b/src/push_service/keys.rs
@@ -70,8 +70,8 @@ impl PushService {
         device_id: u32,
     ) -> Result<PreKeyBundle, ServiceError> {
         let path = format!(
-            "/v2/keys/{}/{}?pq=true",
-            destination.raw_uuid(),
+            "/v2/keys/{}/{}",
+            destination.service_id_string(),
             device_id
         );
 
@@ -101,9 +101,13 @@ impl PushService {
         device_id: u32,
     ) -> Result<Vec<PreKeyBundle>, ServiceError> {
         let path = if device_id == 1 {
-            format!("/v2/keys/{}/*?pq=true", destination.raw_uuid())
+            format!("/v2/keys/{}/*", destination.service_id_string())
         } else {
-            format!("/v2/keys/{}/{}?pq=true", destination.raw_uuid(), device_id)
+            format!(
+                "/v2/keys/{}/{}",
+                destination.service_id_string(),
+                device_id
+            )
         };
         let pre_key_response: PreKeyResponse = self
             .request(
diff --git a/src/sender.rs b/src/sender.rs
index c18f562d7..0e0a352f8 100644
--- a/src/sender.rs
+++ b/src/sender.rs
@@ -28,6 +28,7 @@ use crate::{
     service_address::ServiceIdExt,
     session_store::SessionStoreExt,
     unidentified_access::UnidentifiedAccess,
+    utils::serde_service_id,
     websocket::SignalWebSocket,
 };
 
@@ -44,7 +45,8 @@ pub struct OutgoingPushMessage {
 
 #[derive(serde::Serialize, Debug)]
 pub struct OutgoingPushMessages {
-    pub destination: uuid::Uuid,
+    #[serde(with = "serde_service_id")]
+    pub destination: ServiceId,
     pub timestamp: u64,
     pub messages: Vec<OutgoingPushMessage>,
     pub online: bool,
@@ -551,7 +553,7 @@ where
                 .await?;
 
             let messages = OutgoingPushMessages {
-                destination: recipient.raw_uuid(),
+                destination: recipient,
                 timestamp,
                 messages,
                 online,
diff --git a/src/utils.rs b/src/utils.rs
index 236b6968b..4d42ef4f7 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -290,3 +290,28 @@ pub mod serde_phone_number {
             .map_err(serde::de::Error::custom)
     }
 }
+
+pub mod serde_service_id {
+    use libsignal_protocol::ServiceId;
+    use serde::{Deserialize, Deserializer, Serializer};
+
+    pub fn serialize<S>(
+        service_id: &ServiceId,
+        serializer: S,
+    ) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        serializer.serialize_str(&service_id.service_id_string())
+    }
+
+    pub fn deserialize<'de, D>(deserializer: D) -> Result<ServiceId, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        ServiceId::parse_from_service_id_string(&String::deserialize(
+            deserializer,
+        )?)
+        .ok_or_else(|| serde::de::Error::custom("invalid service ID string"))
+    }
+}
diff --git a/src/websocket/sender.rs b/src/websocket/sender.rs
index 55cd4fbdd..36a1a120b 100644
--- a/src/websocket/sender.rs
+++ b/src/websocket/sender.rs
@@ -13,7 +13,10 @@ impl SignalWebSocket {
         messages: OutgoingPushMessages,
     ) -> Result<SendMessageResponse, ServiceError> {
         let request = WebSocketRequestMessage::new(Method::PUT)
-            .path(format!("/v1/messages/{}", messages.destination))
+            .path(format!(
+                "/v1/messages/{}",
+                messages.destination.service_id_string()
+            ))
             .json(&messages)?;
         self.request_json(request).await
     }
@@ -24,7 +27,10 @@ impl SignalWebSocket {
         access: &UnidentifiedAccess,
     ) -> Result<SendMessageResponse, ServiceError> {
         let request = WebSocketRequestMessage::new(Method::PUT)
-            .path(format!("/v1/messages/{}", messages.destination))
+            .path(format!(
+                "/v1/messages/{}",
+                messages.destination.service_id_string()
+            ))
             .header(
                 "Unidentified-Access-Key",
                 BASE64_RELAXED.encode(&access.key),