From c9f4ff5cf8e20dba0e6dfc24de6acb83458e3b2d Mon Sep 17 00:00:00 2001
From: Devin Ragotzy <devin.ragotzy@gmail.com>
Date: Mon, 1 Mar 2021 08:23:28 -0500
Subject: [PATCH] Ask multiple servers for keys when not known or sending
 server failed

---
 src/database/rooms.rs |  2 +-
 src/server_server.rs  | 61 ++++++++++++++++++++++++++++++++++++-------
 2 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/src/database/rooms.rs b/src/database/rooms.rs
index 4ad499c4..992c97cf 100644
--- a/src/database/rooms.rs
+++ b/src/database/rooms.rs
@@ -1605,7 +1605,7 @@ impl Rooms {
             })
     }
 
-    /// Returns an iterator over all joined members of a room.
+    /// Returns an iterator of all servers participating in this room.
     pub fn room_servers(&self, room_id: &RoomId) -> impl Iterator<Item = Result<Box<ServerName>>> {
         let mut prefix = room_id.as_bytes().to_vec();
         prefix.push(0xff);
diff --git a/src/server_server.rs b/src/server_server.rs
index 1e81d5eb..58c4b33c 100644
--- a/src/server_server.rs
+++ b/src/server_server.rs
@@ -8,8 +8,8 @@ use ruma::{
         federation::{
             directory::{get_public_rooms, get_public_rooms_filtered},
             discovery::{
-                get_server_keys, get_server_version::v1 as get_server_version, ServerSigningKeys,
-                VerifyKey,
+                get_remote_server_keys, get_server_keys,
+                get_server_version::v1 as get_server_version, ServerSigningKeys, VerifyKey,
             },
             event::{get_event, get_missing_events, get_room_state_ids},
             query::get_profile_information,
@@ -575,7 +575,7 @@ pub async fn send_transaction_message_route<'a>(
                 return None;
             }
 
-            Some((event_id, value))
+            Some((event_id, room_id, value))
         })
         .collect::<Vec<_>>();
 
@@ -586,7 +586,7 @@ pub async fn send_transaction_message_route<'a>(
     // events over federation. For example, the Federation API's /send endpoint would
     // discard the event whereas the Client Server API's /send/{eventType} endpoint
     // would return a M_BAD_JSON error.
-    'main_pdu_loop: for (event_id, value) in pdus_to_resolve {
+    'main_pdu_loop: for (event_id, room_id, value) in pdus_to_resolve {
         let server_name = &body.body.origin;
         let mut pub_key_map = BTreeMap::new();
 
@@ -595,7 +595,7 @@ pub async fn send_transaction_message_route<'a>(
                 UserId::try_from(sender.as_str()).expect("All PDUs have a valid sender field");
             let origin = sender.server_name();
 
-            let keys = match fetch_signing_keys(&db, origin).await {
+            let keys = match fetch_signing_keys(&db, &room_id, origin).await {
                 Ok(keys) => keys,
                 Err(_) => {
                     resolved_map.insert(
@@ -1122,18 +1122,61 @@ pub(crate) async fn fetch_events(
 /// fetch them from the server and save to our DB.
 pub(crate) async fn fetch_signing_keys(
     db: &Database,
+    room_id: &RoomId,
     origin: &ServerName,
 ) -> Result<BTreeMap<ServerSigningKeyId, VerifyKey>> {
     match db.globals.signing_keys_for(origin)? {
         keys if !keys.is_empty() => Ok(keys),
         _ => {
-            let keys = db
+            match db
                 .sending
                 .send_federation_request(&db.globals, origin, get_server_keys::v2::Request::new())
                 .await
-                .map_err(|_| Error::BadServerResponse("Failed to request server keys"))?;
-            db.globals.add_signing_key(origin, &keys.server_key)?;
-            Ok(keys.server_key.verify_keys)
+            {
+                Ok(keys) => {
+                    db.globals.add_signing_key(origin, &keys.server_key)?;
+                    Ok(keys.server_key.verify_keys)
+                }
+                _ => {
+                    for server in db.rooms.room_servers(room_id) {
+                        let server = server?;
+                        if let Ok(keys) = db
+                            .sending
+                            .send_federation_request(
+                                &db.globals,
+                                &server,
+                                get_remote_server_keys::v2::Request::new(
+                                    &server,
+                                    SystemTime::now()
+                                        .checked_add(Duration::from_secs(3600))
+                                        .expect("SystemTime to large"),
+                                ),
+                            )
+                            .await
+                        {
+                            let keys: Vec<ServerSigningKeys> = keys.server_keys;
+                            let key = keys.into_iter().fold(None, |mut key, next| {
+                                if let Some(verified) = &key {
+                                    // rustc cannot elide this type for some reason
+                                    let v: &ServerSigningKeys = verified;
+                                    if v.verify_keys
+                                        .iter()
+                                        .zip(next.verify_keys.iter())
+                                        .all(|(a, b)| a.1.key == b.1.key)
+                                    {
+                                    }
+                                } else {
+                                    key = Some(next)
+                                }
+                                key
+                            });
+                        }
+                    }
+                    Err(Error::BadServerResponse(
+                        "Failed to find public key for server",
+                    ))
+                }
+            }
         }
     }
 }