From 80533bfab27a95c4ce77beca41dc0d33f2543917 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timo=20K=C3=B6sters?= <timo@koesters.xyz>
Date: Tue, 20 Jul 2021 22:06:42 +0200
Subject: [PATCH] fix: improve code when skipping /state_ids

---
 src/server_server.rs | 65 ++++++++++++++++++++------------------------
 1 file changed, 30 insertions(+), 35 deletions(-)

diff --git a/src/server_server.rs b/src/server_server.rs
index e8c19dbc..5ab8646c 100644
--- a/src/server_server.rs
+++ b/src/server_server.rs
@@ -1027,42 +1027,37 @@ pub fn handle_incoming_pdu<'a>(
                 .map_err(|_| "Failed talking to db".to_owned())?
                 .map(|shortstatehash| db.rooms.state_full_ids(shortstatehash).ok())
                 .flatten();
-            if let Some(mut state) = state {
-                if db
-                    .rooms
-                    .get_pdu(prev_event)
-                    .ok()
-                    .flatten()
-                    .ok_or_else(|| "Could not find prev event, but we know the state.".to_owned())?
-                    .state_key
-                    .is_some()
-                {
-                    state.insert(prev_event.clone());
-                }
-                state_at_incoming_event = Some(
-                    fetch_and_handle_events(
-                        db,
-                        origin,
-                        &state.into_iter().collect::<Vec<_>>(),
-                        &room_id,
-                        pub_key_map,
-                    )
-                    .await
-                    .map_err(|_| "Failed to fetch state events locally".to_owned())?
-                    .into_iter()
-                    .map(|pdu| {
+            if let Some(state) = state {
+                let mut state = fetch_and_handle_events(
+                    db,
+                    origin,
+                    &state.into_iter().collect::<Vec<_>>(),
+                    &room_id,
+                    pub_key_map,
+                )
+                .await
+                .map_err(|_| "Failed to fetch state events locally".to_owned())?
+                .into_iter()
+                .map(|pdu| {
+                    (
                         (
-                            (
-                                pdu.kind.clone(),
-                                pdu.state_key
-                                    .clone()
-                                    .expect("events from state_full_ids are state events"),
-                            ),
-                            pdu,
-                        )
-                    })
-                    .collect(),
-                );
+                            pdu.kind.clone(),
+                            pdu.state_key
+                                .clone()
+                                .expect("events from state_full_ids are state events"),
+                        ),
+                        pdu,
+                    )
+                })
+                .collect::<HashMap<_, _>>();
+
+                let prev_pdu = db.rooms.get_pdu(prev_event).ok().flatten().ok_or_else(|| {
+                    "Could not find prev event, but we know the state.".to_owned()
+                })?;
+
+                if let Some(state_key) = &prev_pdu.state_key {
+                    state.insert((prev_pdu.kind.clone(), state_key.clone()), prev_pdu);
+                }
             }
             // TODO: set incoming_auth_events?
         }