Add PduId wrapper struct

Death to Vec<u8>

src/service/rooms/event_handler.rs

@@ -37,7 +37,7 @@ use serde_json::value::RawValue as RawJsonValue;
 use tokio::sync::{RwLock, RwLockWriteGuard, Semaphore};
 use tracing::{debug, error, info, trace, warn};
 
-use super::state_compressor::CompressedStateEvent;
+use super::{state_compressor::CompressedStateEvent, timeline::PduId};
 use crate::{
     service::{globals::SigningKeys, pdu},
     services,
@@ -89,7 +89,7 @@ impl Service {
         value: CanonicalJsonObject,
         is_timeline_event: bool,
         pub_key_map: &'a RwLock<BTreeMap<String, SigningKeys>>,
-    ) -> Result<Option<Vec<u8>>> {
+    ) -> Result<Option<PduId>> {
         // 0. Check the server is in the room
         if !services().rooms.metadata.exists(room_id)? {
             return Err(Error::BadRequest(
@@ -565,7 +565,7 @@ impl Service {
         origin: &ServerName,
         room_id: &RoomId,
         pub_key_map: &RwLock<BTreeMap<String, SigningKeys>>,
-    ) -> Result<Option<Vec<u8>>> {
+    ) -> Result<Option<PduId>> {
         // Skip the PDU if we already have it as a timeline event
         if let Ok(Some(pduid)) =
             services().rooms.timeline.get_pdu_id(&incoming_pdu.event_id)

src/service/rooms/search/data.rs

@@ -1,19 +1,19 @@
 use ruma::RoomId;
 
-use crate::Result;
+use crate::{service::rooms::timeline::PduId, Result};
 
 pub(crate) trait Data: Send + Sync {
     fn index_pdu(
         &self,
         shortroomid: u64,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
         message_body: &str,
     ) -> Result<()>;
 
     fn deindex_pdu(
         &self,
         shortroomid: u64,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
         message_body: &str,
     ) -> Result<()>;
 
@@ -22,5 +22,5 @@ pub(crate) trait Data: Send + Sync {
         &'a self,
         room_id: &RoomId,
         search_string: &str,
-    ) -> Result<Option<(Box<dyn Iterator<Item = Vec<u8>> + 'a>, Vec<String>)>>;
+    ) -> Result<Option<(Box<dyn Iterator<Item = PduId> + 'a>, Vec<String>)>>;
 }

src/service/rooms/threads/data.rs

@@ -3,7 +3,7 @@ use ruma::{
     UserId,
 };
 
-use crate::{PduEvent, Result};
+use crate::{service::rooms::timeline::PduId, PduEvent, Result};
 
 pub(crate) trait Data: Send + Sync {
     #[allow(clippy::type_complexity)]
@@ -17,11 +17,11 @@ pub(crate) trait Data: Send + Sync {
 
     fn update_participants(
         &self,
-        root_id: &[u8],
+        root_id: &PduId,
         participants: &[OwnedUserId],
     ) -> Result<()>;
     fn get_participants(
         &self,
-        root_id: &[u8],
+        root_id: &PduId,
     ) -> Result<Option<Vec<OwnedUserId>>>;
 }

src/service/rooms/timeline.rs

@@ -44,6 +44,23 @@ use crate::{
     Error, PduEvent, Result,
 };
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(crate) struct PduId {
+    inner: Vec<u8>,
+}
+
+impl PduId {
+    pub(crate) fn new(inner: Vec<u8>) -> Self {
+        Self {
+            inner,
+        }
+    }
+
+    pub(crate) fn as_bytes(&self) -> &[u8] {
+        &self.inner
+    }
+}
+
 #[derive(Hash, PartialEq, Eq, Clone, Copy, Debug)]
 pub(crate) enum PduCount {
     Backfilled(u64),
@@ -146,7 +163,7 @@ impl Service {
     pub(crate) fn get_pdu_id(
         &self,
         event_id: &EventId,
-    ) -> Result<Option<Vec<u8>>> {
+    ) -> Result<Option<PduId>> {
         self.db.get_pdu_id(event_id)
     }
 
@@ -165,7 +182,7 @@ impl Service {
     /// This does __NOT__ check the outliers `Tree`.
     pub(crate) fn get_pdu_from_id(
         &self,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
     ) -> Result<Option<PduEvent>> {
         self.db.get_pdu_from_id(pdu_id)
     }
@@ -173,7 +190,7 @@ impl Service {
     /// Returns the pdu as a `BTreeMap<String, CanonicalJsonValue>`.
     pub(crate) fn get_pdu_json_from_id(
         &self,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
     ) -> Result<Option<CanonicalJsonObject>> {
         self.db.get_pdu_json_from_id(pdu_id)
     }
@@ -182,7 +199,7 @@ impl Service {
     #[tracing::instrument(skip(self))]
     pub(crate) fn replace_pdu(
         &self,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
         pdu_json: &CanonicalJsonObject,
         pdu: &PduEvent,
     ) -> Result<()> {
@@ -202,7 +219,7 @@ impl Service {
         mut pdu_json: CanonicalJsonObject,
         leaves: Vec<OwnedEventId>,
         room_id: &KeyToken<OwnedRoomId, marker::State>,
-    ) -> Result<Vec<u8>> {
+    ) -> Result<PduId> {
         assert_eq!(*pdu.room_id, **room_id, "Token for incorrect room passed");
 
         let shortroomid = services()
@@ -282,6 +299,7 @@ impl Service {
         let count2 = services().globals.next_count()?;
         let mut pdu_id = shortroomid.to_be_bytes().to_vec();
         pdu_id.extend_from_slice(&count2.to_be_bytes());
+        let pdu_id = PduId::new(pdu_id);
 
         // Insert pdu
         self.db.append_pdu(&pdu_id, pdu, &pdu_json, count2)?;
@@ -1106,7 +1124,7 @@ impl Service {
         state_ids_compressed: Arc<HashSet<CompressedStateEvent>>,
         soft_fail: bool,
         room_id: &KeyToken<OwnedRoomId, marker::State>,
-    ) -> Result<Option<Vec<u8>>> {
+    ) -> Result<Option<PduId>> {
         assert_eq!(*pdu.room_id, **room_id, "Token for incorrect room passed");
 
         // We append to state before appending the pdu, so we don't have a
@@ -1344,6 +1362,7 @@ impl Service {
         let mut pdu_id = shortroomid.to_be_bytes().to_vec();
         pdu_id.extend_from_slice(&0_u64.to_be_bytes());
         pdu_id.extend_from_slice(&(u64::MAX - count).to_be_bytes());
+        let pdu_id = PduId::new(pdu_id);
 
         // Insert pdu
         self.db.prepend_backfill_pdu(&pdu_id, &event_id, &value)?;

src/service/rooms/timeline/data.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use ruma::{CanonicalJsonObject, EventId, OwnedUserId, RoomId, UserId};
 
 use super::PduCount;
-use crate::{PduEvent, Result};
+use crate::{service::rooms::timeline::PduId, PduEvent, Result};
 
 pub(crate) trait Data: Send + Sync {
     fn last_timeline_count(
@@ -28,7 +28,7 @@ pub(crate) trait Data: Send + Sync {
     ) -> Result<Option<CanonicalJsonObject>>;
 
     /// Returns the pdu's id.
-    fn get_pdu_id(&self, event_id: &EventId) -> Result<Option<Vec<u8>>>;
+    fn get_pdu_id(&self, event_id: &EventId) -> Result<Option<PduId>>;
 
     /// Returns the pdu.
     ///
@@ -46,18 +46,18 @@ pub(crate) trait Data: Send + Sync {
     /// Returns the pdu.
     ///
     /// This does __NOT__ check the outliers `Tree`.
-    fn get_pdu_from_id(&self, pdu_id: &[u8]) -> Result<Option<PduEvent>>;
+    fn get_pdu_from_id(&self, pdu_id: &PduId) -> Result<Option<PduEvent>>;
 
     /// Returns the pdu as a `BTreeMap<String, CanonicalJsonValue>`.
     fn get_pdu_json_from_id(
         &self,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
     ) -> Result<Option<CanonicalJsonObject>>;
 
     /// Adds a new pdu to the timeline
     fn append_pdu(
         &self,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
         pdu: &PduEvent,
         json: &CanonicalJsonObject,
         count: u64,
@@ -66,7 +66,7 @@ pub(crate) trait Data: Send + Sync {
     // Adds a new pdu to the backfilled timeline
     fn prepend_backfill_pdu(
         &self,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
         event_id: &EventId,
         json: &CanonicalJsonObject,
     ) -> Result<()>;
@@ -74,7 +74,7 @@ pub(crate) trait Data: Send + Sync {
     /// Removes a pdu and creates a new one with the same id.
     fn replace_pdu(
         &self,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
         pdu_json: &CanonicalJsonObject,
         pdu: &PduEvent,
     ) -> Result<()>;

src/service/sending.rs

@@ -37,6 +37,7 @@ use tokio::{
 };
 use tracing::{debug, error, warn, Span};
 
+use super::rooms::timeline::PduId;
 use crate::{
     api::{appservice_server, server_server},
     services,
@@ -83,7 +84,7 @@ impl Destination {
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub(crate) enum SendingEventType {
     // pduid
-    Pdu(Vec<u8>),
+    Pdu(PduId),
     // pdu json
     Edu(Vec<u8>),
 }
@@ -565,7 +566,7 @@ impl Service {
     #[tracing::instrument(skip(self, pdu_id, user, pushkey))]
     pub(crate) fn send_push_pdu(
         &self,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
         user: &UserId,
         pushkey: String,
     ) -> Result<()> {
@@ -589,7 +590,7 @@ impl Service {
     pub(crate) fn send_pdu<I: Iterator<Item = OwnedServerName>>(
         &self,
         servers: I,
-        pdu_id: &[u8],
+        pdu_id: &PduId,
     ) -> Result<()> {
         let requests = servers
             .into_iter()
@@ -644,7 +645,7 @@ impl Service {
     pub(crate) fn send_pdu_appservice(
         &self,
         appservice_id: String,
-        pdu_id: Vec<u8>,
+        pdu_id: PduId,
     ) -> Result<()> {
         let destination = Destination::Appservice(appservice_id);
         let event_type = SendingEventType::Pdu(pdu_id);
@@ -758,8 +759,8 @@ async fn handle_appservice_event(
                     &events
                         .iter()
                         .map(|e| match e {
-                            SendingEventType::Edu(b)
-                            | SendingEventType::Pdu(b) => &**b,
+                            SendingEventType::Edu(b) => &**b,
+                            SendingEventType::Pdu(b) => b.as_bytes(),
                         })
                         .collect::<Vec<_>>(),
                 ))
@@ -905,8 +906,8 @@ async fn handle_federation_event(
                     &events
                         .iter()
                         .map(|e| match e {
-                            SendingEventType::Edu(b)
-                            | SendingEventType::Pdu(b) => &**b,
+                            SendingEventType::Edu(b) => &**b,
+                            SendingEventType::Pdu(b) => b.as_bytes(),
                         })
                         .collect::<Vec<_>>(),
                 ))