0 errors left!

src/database/key_value/sending.rs

@@ -0,0 +1,203 @@
+use ruma::{ServerName, UserId};
+
+use crate::{
+    database::KeyValueDatabase,
+    service::{
+        self,
+        sending::{OutgoingKind, SendingEventType},
+    },
+    utils, Error, Result,
+};
+
+impl service::sending::Data for KeyValueDatabase {
+    fn active_requests<'a>(
+        &'a self,
+    ) -> Box<dyn Iterator<Item = Result<(Vec<u8>, OutgoingKind, SendingEventType)>> + 'a> {
+        Box::new(
+            self.servercurrentevent_data
+                .iter()
+                .map(|(key, v)| parse_servercurrentevent(&key, v).map(|(k, e)| (key, k, e))),
+        )
+    }
+
+    fn active_requests_for<'a>(
+        &'a self,
+        outgoing_kind: &OutgoingKind,
+    ) -> Box<dyn Iterator<Item = Result<(Vec<u8>, SendingEventType)>> + 'a> {
+        let prefix = outgoing_kind.get_prefix();
+        Box::new(
+            self.servercurrentevent_data
+                .scan_prefix(prefix)
+                .map(|(key, v)| parse_servercurrentevent(&key, v).map(|(_, e)| (key, e))),
+        )
+    }
+
+    fn delete_active_request(&self, key: Vec<u8>) -> Result<()> {
+        self.servercurrentevent_data.remove(&key)
+    }
+
+    fn delete_all_active_requests_for(&self, outgoing_kind: &OutgoingKind) -> Result<()> {
+        let prefix = outgoing_kind.get_prefix();
+        for (key, _) in self.servercurrentevent_data.scan_prefix(prefix.clone()) {
+            self.servercurrentevent_data.remove(&key)?;
+        }
+
+        Ok(())
+    }
+
+    fn delete_all_requests_for(&self, outgoing_kind: &OutgoingKind) -> Result<()> {
+        let prefix = outgoing_kind.get_prefix();
+        for (key, _) in self.servercurrentevent_data.scan_prefix(prefix.clone()) {
+            self.servercurrentevent_data.remove(&key).unwrap();
+        }
+
+        for (key, _) in self.servernameevent_data.scan_prefix(prefix.clone()) {
+            self.servernameevent_data.remove(&key).unwrap();
+        }
+
+        Ok(())
+    }
+
+    fn queue_requests(
+        &self,
+        requests: &[(&OutgoingKind, SendingEventType)],
+    ) -> Result<Vec<Vec<u8>>> {
+        let mut batch = Vec::new();
+        let mut keys = Vec::new();
+        for (outgoing_kind, event) in requests {
+            let mut key = outgoing_kind.get_prefix();
+            key.push(0xff);
+            key.extend_from_slice(if let SendingEventType::Pdu(value) = &event {
+                &**value
+            } else {
+                &[]
+            });
+            let value = if let SendingEventType::Edu(value) = &event {
+                &**value
+            } else {
+                &[]
+            };
+            batch.push((key.clone(), value.to_owned()));
+            keys.push(key);
+        }
+        self.servernameevent_data
+            .insert_batch(&mut batch.into_iter())?;
+        Ok(keys)
+    }
+
+    fn queued_requests<'a>(
+        &'a self,
+        outgoing_kind: &OutgoingKind,
+    ) -> Box<dyn Iterator<Item = Result<(SendingEventType, Vec<u8>)>> + 'a> {
+        let prefix = outgoing_kind.get_prefix();
+        return Box::new(
+            self.servernameevent_data
+                .scan_prefix(prefix.clone())
+                .map(|(k, v)| parse_servercurrentevent(&k, v).map(|(_, ev)| (ev, k))),
+        );
+    }
+
+    fn mark_as_active(&self, events: &[(SendingEventType, Vec<u8>)]) -> Result<()> {
+        for (e, key) in events {
+            let value = if let SendingEventType::Edu(value) = &e {
+                &**value
+            } else {
+                &[]
+            };
+            self.servercurrentevent_data.insert(key, value)?;
+            self.servernameevent_data.remove(key)?;
+        }
+
+        Ok(())
+    }
+
+    fn set_latest_educount(&self, server_name: &ServerName, last_count: u64) -> Result<()> {
+        self.servername_educount
+            .insert(server_name.as_bytes(), &last_count.to_be_bytes())
+    }
+
+    fn get_latest_educount(&self, server_name: &ServerName) -> Result<u64> {
+        self.servername_educount
+            .get(server_name.as_bytes())?
+            .map_or(Ok(0), |bytes| {
+                utils::u64_from_bytes(&bytes)
+                    .map_err(|_| Error::bad_database("Invalid u64 in servername_educount."))
+            })
+    }
+}
+
+#[tracing::instrument(skip(key))]
+fn parse_servercurrentevent(
+    key: &[u8],
+    value: Vec<u8>,
+) -> Result<(OutgoingKind, SendingEventType)> {
+    // Appservices start with a plus
+    Ok::<_, Error>(if key.starts_with(b"+") {
+        let mut parts = key[1..].splitn(2, |&b| b == 0xff);
+
+        let server = parts.next().expect("splitn always returns one element");
+        let event = parts
+            .next()
+            .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?;
+        let server = utils::string_from_bytes(server).map_err(|_| {
+            Error::bad_database("Invalid server bytes in server_currenttransaction")
+        })?;
+
+        (
+            OutgoingKind::Appservice(server),
+            if value.is_empty() {
+                SendingEventType::Pdu(event.to_vec())
+            } else {
+                SendingEventType::Edu(value)
+            },
+        )
+    } else if key.starts_with(b"$") {
+        let mut parts = key[1..].splitn(3, |&b| b == 0xff);
+
+        let user = parts.next().expect("splitn always returns one element");
+        let user_string = utils::string_from_bytes(&user)
+            .map_err(|_| Error::bad_database("Invalid user string in servercurrentevent"))?;
+        let user_id = UserId::parse(user_string)
+            .map_err(|_| Error::bad_database("Invalid user id in servercurrentevent"))?;
+
+        let pushkey = parts
+            .next()
+            .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?;
+        let pushkey_string = utils::string_from_bytes(pushkey)
+            .map_err(|_| Error::bad_database("Invalid pushkey in servercurrentevent"))?;
+
+        let event = parts
+            .next()
+            .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?;
+        (
+            OutgoingKind::Push(user_id, pushkey_string),
+            if value.is_empty() {
+                SendingEventType::Pdu(event.to_vec())
+            } else {
+                // I'm pretty sure this should never be called
+                SendingEventType::Edu(value)
+            },
+        )
+    } else {
+        let mut parts = key.splitn(2, |&b| b == 0xff);
+
+        let server = parts.next().expect("splitn always returns one element");
+        let event = parts
+            .next()
+            .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?;
+        let server = utils::string_from_bytes(server).map_err(|_| {
+            Error::bad_database("Invalid server bytes in server_currenttransaction")
+        })?;
+
+        (
+            OutgoingKind::Normal(ServerName::parse(server).map_err(|_| {
+                Error::bad_database("Invalid server string in server_currenttransaction")
+            })?),
+            if value.is_empty() {
+                SendingEventType::Pdu(event.to_vec())
+            } else {
+                SendingEventType::Edu(value)
+            },
+        )
+    })
+}