implement notifications

src/database/key_value/rooms/user.rs

@@ -1,9 +1,20 @@
-use ruma::{OwnedRoomId, OwnedUserId, RoomId, UserId};
+use std::mem::size_of;
+
+use ruma::{
+    api::client::{
+        error::ErrorKind, push::get_notifications::v3::Notification,
+    },
+    push, MilliSecondsSinceUnixEpoch, OwnedRoomId, OwnedUserId, RoomId, UInt,
+    UserId,
+};
 
 use crate::{
     database::KeyValueDatabase,
-    service::{self, rooms::short::ShortStateHash},
-    services, utils, Error, Result,
+    service::{
+        self,
+        rooms::{short::ShortStateHash, timeline::PduCount},
+    },
+    services, utils, Error, PduEvent, Result,
 };
 
 impl service::rooms::user::Data for KeyValueDatabase {
@@ -24,10 +35,40 @@ impl service::rooms::user::Data for KeyValueDatabase {
         self.userroomid_highlightcount
             .insert(&userroom_id, &0_u64.to_be_bytes())?;
 
-        self.roomuserid_lastnotificationread.insert(
-            &roomuser_id,
-            &services().globals.next_count()?.to_be_bytes(),
-        )?;
+        let next_count = services().globals.next_count()?;
+
+        self.roomuserid_lastnotificationread
+            .insert(&roomuser_id, &next_count.to_be_bytes())?;
+
+        let mut userpducount_id = user_id.localpart().as_bytes().to_vec();
+
+        userpducount_id.push(0xFF);
+        userpducount_id.extend_from_slice(&next_count.to_be_bytes());
+
+        let shortroomid =
+            services().rooms.short.get_shortroomid(room_id)?.ok_or_else(
+                || {
+                    Error::bad_database(
+                        "Looked for bad shortroomid for notifications",
+                    )
+                },
+            )?;
+
+        let it =
+            self.userpducountid_notifications.iter_from(&userpducount_id, true);
+
+        for (k, mut v) in it.filter(|(_, v)| {
+            v[2] == 1
+                && u64::from_be_bytes(
+                    v[3..3 + size_of::<u64>()].try_into().unwrap(),
+                ) == shortroomid.get()
+        }) {
+            self.userpducountid_notifications.remove(&k)?;
+
+            v[2] = 0;
+
+            self.userpducountid_notifications.insert(&k, &v)?;
+        }
 
         Ok(())
     }
@@ -93,6 +134,111 @@ impl service::rooms::user::Data for KeyValueDatabase {
             .unwrap_or(0))
     }
 
+    fn store_push_action(
+        &self,
+        pdu: &PduEvent,
+        user_id: &UserId,
+        notify: bool,
+        highlight: bool,
+        actions: &[push::Action],
+    ) -> Result<()> {
+        let Some(PduCount::Normal(pdu_count)) =
+            services().rooms.timeline.get_pdu_count(&pdu.event_id)?
+        else {
+            return Err(Error::BadRequest(
+                ErrorKind::InvalidParam,
+                "Event does not exist.",
+            ));
+        };
+
+        let mut key = user_id.localpart().as_bytes().to_vec();
+
+        key.push(0xFF);
+        key.extend_from_slice(&pdu_count.to_be_bytes());
+
+        let (notify, highlight, unread) =
+            (u8::from(notify), u8::from(highlight), u8::from(true));
+
+        let notification = serde_json::to_vec(&Notification {
+            actions: actions.to_owned(),
+            event: pdu.to_sync_room_event(),
+            profile_tag: None,
+            read: false,
+            room_id: pdu.room_id.clone(),
+            // TODO
+            ts: MilliSecondsSinceUnixEpoch::now(),
+        })
+        .expect("Notification should serialize");
+
+        let shortroomid =
+            services().rooms.short.get_shortroomid(&pdu.room_id)?.ok_or_else(
+                || {
+                    Error::bad_database(
+                        "Looked for bad shortroomid for notifications",
+                    )
+                },
+            )?;
+
+        let mut value = vec![notify, highlight, unread];
+
+        value.extend_from_slice(&shortroomid.get().to_be_bytes());
+
+        value.extend_from_slice(&notification);
+
+        self.userpducountid_notifications.insert(&key, &value)?;
+
+        Ok(())
+    }
+
+    fn get_notifications(
+        &self,
+        user_id: &UserId,
+        from: Option<u64>,
+        limit: Option<UInt>,
+        highlight: bool,
+    ) -> Result<(Vec<Notification>, Option<String>)> {
+        let mut key = user_id.localpart().as_bytes().to_vec();
+
+        key.push(0xFF);
+        key.extend_from_slice(&from.unwrap_or(u64::MAX).to_be_bytes());
+
+        let limit = limit.and_then(|n| n.try_into().ok()).unwrap_or(50);
+
+        let it = self
+            .userpducountid_notifications
+            .iter_from(&key, true)
+            .take_while(move |(_, v)| {
+                v[0] == 1 // notify
+                    && (!highlight || v[1] == 1) // highlight
+                    && v[2] == 1 // unread
+            })
+            .take(limit + 1);
+
+        let mut notifications: Vec<_> = it
+            .map(|(k, v)| {
+                (
+                    k,
+                    Notification {
+                        read: v[2] != 1,
+                        ..serde_json::from_slice(&v[3 + size_of::<u64>()..])
+                            .unwrap()
+                    },
+                )
+            })
+            .collect();
+
+        let next_token = notifications
+            .pop()
+            .and_then(|(k, _)| {
+                k.split(|b| *b == 0xFF).nth(1).map(<[u8]>::to_vec)
+            })
+            .map(|pdu_count| {
+                format!("{}", u64::from_be_bytes(pdu_count.try_into().unwrap()))
+            });
+
+        Ok((notifications.into_iter().map(|(_, n)| n).collect(), next_token))
+    }
+
     fn associate_token_shortstatehash(
         &self,
         room_id: &RoomId,