From 79c9de98cd6699df8647b70ef24d9dd0889a497a Mon Sep 17 00:00:00 2001
From: Devin Ragotzy <devin.ragotzy@gmail.com>
Date: Mon, 1 Mar 2021 09:17:53 -0500
Subject: [PATCH] Add trusted_servers, filter servers to query keys by
 trusted_servers

---
 src/database.rs         |  2 ++
 src/database/globals.rs |  4 ++++
 src/server_server.rs    | 16 +++++++++++++---
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/database.rs b/src/database.rs
index 35b7bcd..20cc7e1 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -39,6 +39,8 @@ pub struct Config {
     #[serde(default = "false_fn")]
     allow_federation: bool,
     jwt_secret: Option<String>,
+    #[serde(default = "Vec::new")]
+    trusted_servers: Vec<Box<ServerName>>,
 }
 
 fn false_fn() -> bool {
diff --git a/src/database/globals.rs b/src/database/globals.rs
index 00b4568..3c65e74 100644
--- a/src/database/globals.rs
+++ b/src/database/globals.rs
@@ -139,6 +139,10 @@ impl Globals {
         self.config.allow_federation
     }
 
+    pub fn trusted_servers(&self) -> &[Box<ServerName>] {
+        &self.config.trusted_servers
+    }
+
     pub fn dns_resolver(&self) -> &TokioAsyncResolver {
         &self.dns_resolver
     }
diff --git a/src/server_server.rs b/src/server_server.rs
index 58c4b33..dcd72f7 100644
--- a/src/server_server.rs
+++ b/src/server_server.rs
@@ -1138,7 +1138,9 @@ pub(crate) async fn fetch_signing_keys(
                     Ok(keys.server_key.verify_keys)
                 }
                 _ => {
-                    for server in db.rooms.room_servers(room_id) {
+                    for server in db.rooms.room_servers(room_id).filter(
+                        |ser| matches!(ser, Ok(s) if db.globals.trusted_servers().contains(s)),
+                    ) {
                         let server = server?;
                         if let Ok(keys) = db
                             .sending
@@ -1154,8 +1156,9 @@ pub(crate) async fn fetch_signing_keys(
                             )
                             .await
                         {
+                            let mut trust = 0;
                             let keys: Vec<ServerSigningKeys> = keys.server_keys;
-                            let key = keys.into_iter().fold(None, |mut key, next| {
+                            let key = keys.iter().fold(None, |mut key, next| {
                                 if let Some(verified) = &key {
                                     // rustc cannot elide this type for some reason
                                     let v: &ServerSigningKeys = verified;
@@ -1164,12 +1167,19 @@ pub(crate) async fn fetch_signing_keys(
                                         .zip(next.verify_keys.iter())
                                         .all(|(a, b)| a.1.key == b.1.key)
                                     {
+                                        trust += 1;
                                     }
                                 } else {
-                                    key = Some(next)
+                                    key = Some(next.clone())
                                 }
                                 key
                             });
+
+                            if trust == (keys.len() - 1) && key.is_some() {
+                                let k = key.unwrap();
+                                db.globals.add_signing_key(origin, &k)?;
+                                return Ok(k.verify_keys);
+                            }
                         }
                     }
                     Err(Error::BadServerResponse(