From a77fcd106e5f3f0c7bb3ada59ebd86d548d56213 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timo=20K=C3=B6sters?= <timo@koesters.xyz>
Date: Thu, 18 Mar 2021 00:09:57 +0100
Subject: [PATCH] feat: implement /state_ids and fix federation stuff

---
 src/database.rs       | 12 ++++-----
 src/database/rooms.rs | 16 +++---------
 src/main.rs           |  1 +
 src/server_server.rs  | 57 +++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 68 insertions(+), 18 deletions(-)

diff --git a/src/database.rs b/src/database.rs
index f65d5e0..67f888f 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -110,11 +110,6 @@ impl Database {
         let (admin_sender, admin_receiver) = mpsc::unbounded();
 
         let db = Self {
-            globals: globals::Globals::load(
-                db.open_tree("global")?,
-                db.open_tree("servertimeout_signingkey")?,
-                config,
-            )?,
             users: users::Users {
                 userid_password: db.open_tree("userid_password")?,
                 userid_displayname: db.open_tree("userid_displayname")?,
@@ -191,7 +186,7 @@ impl Database {
             sending: sending::Sending {
                 servernamepduids: db.open_tree("servernamepduids")?,
                 servercurrentpdus: db.open_tree("servercurrentpdus")?,
-                maximum_requests: Arc::new(Semaphore::new(10)),
+                maximum_requests: Arc::new(Semaphore::new(config.max_concurrent_requests as usize)),
             },
             admin: admin::Admin {
                 sender: admin_sender,
@@ -201,6 +196,11 @@ impl Database {
                 id_appserviceregistrations: db.open_tree("id_appserviceregistrations")?,
             },
             pusher: pusher::PushData::new(&db)?,
+            globals: globals::Globals::load(
+                db.open_tree("global")?,
+                db.open_tree("servertimeout_signingkey")?,
+                config,
+            )?,
             _db: db,
         };
 
diff --git a/src/database/rooms.rs b/src/database/rooms.rs
index a342566..2450622 100644
--- a/src/database/rooms.rs
+++ b/src/database/rooms.rs
@@ -85,19 +85,10 @@ impl Rooms {
     /// Builds a StateMap by iterating over all keys that start
     /// with state_hash, this gives the full state for the given state_hash.
     #[tracing::instrument(skip(self))]
-    pub fn state_full_ids(
-        &self,
-        room_id: &RoomId,
-        state_hash: &StateHashId,
-    ) -> Result<Vec<EventId>> {
-        let shortstatehash = self
-            .statehash_shortstatehash
-            .get(state_hash)?
-            .ok_or_else(|| Error::bad_database("Asked for statehash that does not exist."))?;
-
+    pub fn state_full_ids(&self, shortstatehash: u64) -> Result<Vec<EventId>> {
         Ok(self
             .stateid_shorteventid
-            .scan_prefix(&shortstatehash)
+            .scan_prefix(&shortstatehash.to_be_bytes())
             .values()
             .filter_map(|r| r.ok())
             .map(|bytes| self.shorteventid_eventid.get(&bytes).ok().flatten())
@@ -895,7 +886,8 @@ impl Rooms {
             redacts,
         } = pdu_builder;
         // TODO: Make sure this isn't called twice in parallel
-        let prev_events = self.get_pdu_leaves(&room_id)?;
+        let mut prev_events = self.get_pdu_leaves(&room_id)?;
+        prev_events.truncate(20);
 
         let auth_events = self.get_auth_events(
             &room_id,
diff --git a/src/main.rs b/src/main.rs
index 893273f..1aa4d54 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -165,6 +165,7 @@ fn setup_rocket() -> (rocket::Rocket, Config) {
                 server_server::get_public_rooms_filtered_route,
                 server_server::send_transaction_message_route,
                 server_server::get_missing_events_route,
+                server_server::get_room_state_ids_route,
                 server_server::get_profile_information_route,
             ],
         )
diff --git a/src/server_server.rs b/src/server_server.rs
index 2f32b63..da9928c 100644
--- a/src/server_server.rs
+++ b/src/server_server.rs
@@ -6,6 +6,7 @@ use regex::Regex;
 use rocket::{get, post, put, response::content::Json, State};
 use ruma::{
     api::{
+        client::error::ErrorKind,
         federation::{
             directory::{get_public_rooms, get_public_rooms_filtered},
             discovery::{
@@ -1543,6 +1544,62 @@ pub fn get_missing_events_route<'a>(
     Ok(get_missing_events::v1::Response { events }.into())
 }
 
+#[cfg_attr(
+    feature = "conduit_bin",
+    get("/_matrix/federation/v1/state_ids/<_>", data = "<body>")
+)]
+#[tracing::instrument(skip(db, body))]
+pub fn get_room_state_ids_route<'a>(
+    db: State<'a, Database>,
+    body: Ruma<get_room_state_ids::v1::Request<'_>>,
+) -> ConduitResult<get_room_state_ids::v1::Response> {
+    if !db.globals.allow_federation() {
+        return Err(Error::bad_config("Federation is disabled."));
+    }
+
+    let shortstatehash = db
+        .rooms
+        .pdu_shortstatehash(&body.event_id)?
+        .ok_or(Error::BadRequest(
+            ErrorKind::NotFound,
+            "Pdu state not found.",
+        ))?;
+
+    let pdu_ids = db.rooms.state_full_ids(shortstatehash)?;
+
+    let mut auth_chain_ids = BTreeSet::<EventId>::new();
+    let mut todo = BTreeSet::new();
+    todo.insert(body.event_id.clone());
+
+    loop {
+        if let Some(event_id) = todo.iter().next().cloned() {
+            if let Some(pdu) = db.rooms.get_pdu(&event_id)? {
+                todo.extend(
+                    pdu.auth_events
+                        .clone()
+                        .into_iter()
+                        .collect::<BTreeSet<_>>()
+                        .difference(&auth_chain_ids)
+                        .cloned(),
+                );
+                auth_chain_ids.extend(pdu.auth_events.into_iter());
+            } else {
+                warn!("Could not find pdu mentioned in auth events.");
+            }
+
+            todo.remove(&event_id);
+        } else {
+            break;
+        }
+    }
+
+    Ok(get_room_state_ids::v1::Response {
+        auth_chain_ids: auth_chain_ids.into_iter().collect(),
+        pdu_ids,
+    }
+    .into())
+}
+
 #[cfg_attr(
     feature = "conduit_bin",
     get("/_matrix/federation/v1/query/profile", data = "<body>")