fix(keys): only use keys valid at the time of PDU or transaction, and actually refresh keys

Previously, we only fetched keys once, only requesting them again if we have any missing, allowing for ancient keys to be used to sign PDUs and transactions
Now we refresh keys that either have or are about to expire, preventing attacks that make use of leaked private keys of a homeserver
We also ensure that when validating PDUs or transactions, that they are valid at the origin_server_ts or time of us receiving the transaction respectfully
As to not break event authorization for old rooms, we need to keep old keys around
We move verify_keys which we no longer see in direct requests to the origin to old_verify_keys
We keep old_verify_keys indefinitely as mentioned above, as to not break event authorization (at least until a future MSC addresses this)

src/api/client_server/membership.rs

@@ -18,9 +18,8 @@ use ruma::{
         },
         StateEventType, TimelineEventType,
     },
-    serde::Base64,
-    state_res, CanonicalJsonObject, CanonicalJsonValue, EventId, OwnedEventId, OwnedRoomId,
-    OwnedServerName, OwnedUserId, RoomId, RoomVersionId, UserId,
+    state_res, CanonicalJsonObject, CanonicalJsonValue, EventId, MilliSecondsSinceUnixEpoch,
+    OwnedEventId, OwnedRoomId, OwnedServerName, OwnedUserId, RoomId, RoomVersionId, UserId,
 };
 use serde_json::value::{to_raw_value, RawValue as RawJsonValue};
 use std::{
@@ -32,7 +31,10 @@ use tokio::sync::RwLock;
 use tracing::{debug, error, info, warn};
 
 use crate::{
-    service::pdu::{gen_event_id_canonical_json, PduBuilder},
+    service::{
+        globals::SigningKeys,
+        pdu::{gen_event_id_canonical_json, PduBuilder},
+    },
     services, utils, Error, PduEvent, Result, Ruma,
 };
 
@@ -1130,7 +1132,7 @@ async fn make_join_request(
 async fn validate_and_add_event_id(
     pdu: &RawJsonValue,
     room_version: &RoomVersionId,
-    pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+    pub_key_map: &RwLock<BTreeMap<String, SigningKeys>>,
 ) -> Result<(OwnedEventId, CanonicalJsonObject)> {
     let mut value: CanonicalJsonObject = serde_json::from_str(pdu.get()).map_err(|e| {
         error!("Invalid PDU in server response: {:?}: {:?}", pdu, e);
@@ -1177,8 +1179,35 @@ async fn validate_and_add_event_id(
         }
     }
 
-    if let Err(e) = ruma::signatures::verify_event(&*pub_key_map.read().await, &value, room_version)
-    {
+    let origin_server_ts = value.get("origin_server_ts").ok_or_else(|| {
+        error!("Invalid PDU, no origin_server_ts field");
+        Error::BadRequest(
+            ErrorKind::MissingParam,
+            "Invalid PDU, no origin_server_ts field",
+        )
+    })?;
+
+    let origin_server_ts: MilliSecondsSinceUnixEpoch = {
+        let ts = origin_server_ts.as_integer().ok_or_else(|| {
+            Error::BadRequest(
+                ErrorKind::InvalidParam,
+                "origin_server_ts must be an integer",
+            )
+        })?;
+
+        MilliSecondsSinceUnixEpoch(i64::from(ts).try_into().map_err(|_| {
+            Error::BadRequest(ErrorKind::InvalidParam, "Time must be after the unix epoch")
+        })?)
+    };
+
+    let unfiltered_keys = (*pub_key_map.read().await).clone();
+
+    let keys =
+        services()
+            .globals
+            .filter_keys_server_map(unfiltered_keys, origin_server_ts, room_version);
+
+    if let Err(e) = ruma::signatures::verify_event(&keys, &value, room_version) {
         warn!("Event {} failed verification {:?} {}", event_id, pdu, e);
         back_off(event_id).await;
         return Err(Error::BadServerResponse("Event failed verification."));

src/api/ruma_wrapper/axum.rs

@@ -14,7 +14,7 @@ use http::{Request, StatusCode};
 use ruma::{
     api::{client::error::ErrorKind, AuthScheme, IncomingRequest, OutgoingResponse},
     server_util::authorization::XMatrix,
-    CanonicalJsonValue, OwnedDeviceId, OwnedUserId, UserId,
+    CanonicalJsonValue, MilliSecondsSinceUnixEpoch, OwnedDeviceId, OwnedUserId, UserId,
 };
 use serde::Deserialize;
 use tracing::{debug, error, warn};
@@ -231,7 +231,7 @@ where
                     let keys_result = services()
                         .rooms
                         .event_handler
-                        .fetch_signing_keys(&x_matrix.origin, vec![x_matrix.key.to_string()])
+                        .fetch_signing_keys(&x_matrix.origin, vec![x_matrix.key.to_string()], false)
                         .await;
 
                     let keys = match keys_result {
@@ -245,8 +245,19 @@ where
                         }
                     };
 
-                    let pub_key_map =
-                        BTreeMap::from_iter([(x_matrix.origin.as_str().to_owned(), keys)]);
+                    // Only verify_keys that are currently valid should be used for validating requests
+                    // as per MSC4029
+                    let pub_key_map = BTreeMap::from_iter([(
+                        x_matrix.origin.as_str().to_owned(),
+                        if keys.valid_until_ts > MilliSecondsSinceUnixEpoch::now() {
+                            keys.verify_keys
+                                .into_iter()
+                                .map(|(id, key)| (id, key.key))
+                                .collect()
+                        } else {
+                            BTreeMap::new()
+                        },
+                    )]);
 
                     match ruma::signatures::verify_json(&pub_key_map, &request_map) {
                         Ok(()) => (None, None, Some(x_matrix.origin), None),

src/database/key_value/globals.rs

@@ -1,15 +1,19 @@
-use std::collections::{BTreeMap, HashMap};
+use std::collections::HashMap;
 
 use async_trait::async_trait;
 use futures_util::{stream::FuturesUnordered, StreamExt};
 use lru_cache::LruCache;
 use ruma::{
-    api::federation::discovery::{ServerSigningKeys, VerifyKey},
+    api::federation::discovery::{OldVerifyKey, ServerSigningKeys},
     signatures::Ed25519KeyPair,
-    DeviceId, MilliSecondsSinceUnixEpoch, OwnedServerSigningKeyId, ServerName, UserId,
+    DeviceId, ServerName, UserId,
 };
 
-use crate::{database::KeyValueDatabase, service, services, utils, Error, Result};
+use crate::{
+    database::KeyValueDatabase,
+    service::{self, globals::SigningKeys},
+    services, utils, Error, Result,
+};
 
 pub const COUNTER: &[u8] = b"c";
 pub const LAST_CHECK_FOR_UPDATES_COUNT: &[u8] = b"u";
@@ -237,64 +241,97 @@ lasttimelinecount_cache: {lasttimelinecount_cache}\n"
         self.global.remove(b"keypair")
     }
 
-    fn add_signing_key(
+    fn add_signing_key_from_trusted_server(
         &self,
         origin: &ServerName,
         new_keys: ServerSigningKeys,
-    ) -> Result<BTreeMap<OwnedServerSigningKeyId, VerifyKey>> {
-        // Not atomic, but this is not critical
-        let signingkeys = self.server_signingkeys.get(origin.as_bytes())?;
-
-        let mut keys = signingkeys
-            .and_then(|keys| serde_json::from_slice(&keys).ok())
-            .unwrap_or_else(|| {
-                // Just insert "now", it doesn't matter
-                ServerSigningKeys::new(origin.to_owned(), MilliSecondsSinceUnixEpoch::now())
-            });
-
-        let ServerSigningKeys {
-            verify_keys,
-            old_verify_keys,
-            ..
-        } = new_keys;
-
-        keys.verify_keys.extend(verify_keys);
-        keys.old_verify_keys.extend(old_verify_keys);
-
-        self.server_signingkeys.insert(
-            origin.as_bytes(),
-            &serde_json::to_vec(&keys).expect("serversigningkeys can be serialized"),
-        )?;
-
-        let mut tree = keys.verify_keys;
-        tree.extend(
-            keys.old_verify_keys
-                .into_iter()
-                .map(|old| (old.0, VerifyKey::new(old.1.key))),
-        );
-
-        Ok(tree)
+    ) -> Result<SigningKeys> {
+        let prev_keys = self.server_signingkeys.get(origin.as_bytes())?;
+
+        Ok(
+            if let Some(mut prev_keys) =
+                prev_keys.and_then(|keys| serde_json::from_slice::<ServerSigningKeys>(&keys).ok())
+            {
+                let ServerSigningKeys {
+                    verify_keys,
+                    old_verify_keys,
+                    ..
+                } = new_keys;
+
+                prev_keys.verify_keys.extend(verify_keys);
+                prev_keys.old_verify_keys.extend(old_verify_keys);
+                prev_keys.valid_until_ts = new_keys.valid_until_ts;
+
+                self.server_signingkeys.insert(
+                    origin.as_bytes(),
+                    &serde_json::to_vec(&prev_keys).expect("serversigningkeys can be serialized"),
+                )?;
+
+                prev_keys.into()
+            } else {
+                self.server_signingkeys.insert(
+                    origin.as_bytes(),
+                    &serde_json::to_vec(&new_keys).expect("serversigningkeys can be serialized"),
+                )?;
+
+                new_keys.into()
+            },
+        )
     }
 
-    /// This returns an empty `Ok(BTreeMap<..>)` when there are no keys found for the server.
-    fn signing_keys_for(
+    fn add_signing_key_from_origin(
         &self,
         origin: &ServerName,
-    ) -> Result<BTreeMap<OwnedServerSigningKeyId, VerifyKey>> {
+        new_keys: ServerSigningKeys,
+    ) -> Result<SigningKeys> {
+        let prev_keys = self.server_signingkeys.get(origin.as_bytes())?;
+
+        Ok(
+            if let Some(mut prev_keys) =
+                prev_keys.and_then(|keys| serde_json::from_slice::<ServerSigningKeys>(&keys).ok())
+            {
+                let ServerSigningKeys {
+                    verify_keys,
+                    old_verify_keys,
+                    ..
+                } = new_keys;
+
+                // Moving `verify_keys` no longer present to `old_verify_keys`
+                for (key_id, key) in prev_keys.verify_keys {
+                    if !verify_keys.contains_key(&key_id) {
+                        prev_keys
+                            .old_verify_keys
+                            .insert(key_id, OldVerifyKey::new(prev_keys.valid_until_ts, key.key));
+                    }
+                }
+
+                prev_keys.verify_keys = verify_keys;
+                prev_keys.old_verify_keys.extend(old_verify_keys);
+                prev_keys.valid_until_ts = new_keys.valid_until_ts;
+
+                self.server_signingkeys.insert(
+                    origin.as_bytes(),
+                    &serde_json::to_vec(&prev_keys).expect("serversigningkeys can be serialized"),
+                )?;
+
+                prev_keys.into()
+            } else {
+                self.server_signingkeys.insert(
+                    origin.as_bytes(),
+                    &serde_json::to_vec(&new_keys).expect("serversigningkeys can be serialized"),
+                )?;
+
+                new_keys.into()
+            },
+        )
+    }
+
+    /// This returns an empty `Ok(BTreeMap<..>)` when there are no keys found for the server.
+    fn signing_keys_for(&self, origin: &ServerName) -> Result<Option<SigningKeys>> {
         let signingkeys = self
             .server_signingkeys
             .get(origin.as_bytes())?
-            .and_then(|bytes| serde_json::from_slice(&bytes).ok())
-            .map(|keys: ServerSigningKeys| {
-                let mut tree = keys.verify_keys;
-                tree.extend(
-                    keys.old_verify_keys
-                        .into_iter()
-                        .map(|old| (old.0, VerifyKey::new(old.1.key))),
-                );
-                tree
-            })
-            .unwrap_or_else(BTreeMap::new);
+            .and_then(|bytes| serde_json::from_slice::<SigningKeys>(&bytes).ok());
 
         Ok(signingkeys)
     }

src/service/admin/mod.rs

@@ -19,7 +19,8 @@ use ruma::{
         },
         TimelineEventType,
     },
-    EventId, OwnedRoomAliasId, OwnedRoomId, RoomAliasId, RoomId, RoomVersionId, ServerName, UserId,
+    EventId, MilliSecondsSinceUnixEpoch, OwnedRoomAliasId, OwnedRoomId, RoomAliasId, RoomId,
+    RoomVersionId, ServerName, UserId,
 };
 use serde_json::value::to_raw_value;
 use tokio::sync::{mpsc, Mutex, RwLock};
@@ -858,15 +859,46 @@ impl Service {
                             services()
                                 .rooms
                                 .event_handler
+                                // Generally we shouldn't be checking against expired keys unless required, so in the admin
+                                // room it might be best to not allow expired keys
                                 .fetch_required_signing_keys(&value, &pub_key_map)
                                 .await?;
 
-                            let pub_key_map = pub_key_map.read().await;
-                            match ruma::signatures::verify_json(&pub_key_map, &value) {
-                                Ok(_) => RoomMessageEventContent::text_plain("Signature correct"),
-                                Err(e) => RoomMessageEventContent::text_plain(format!(
+                            let mut expired_key_map = BTreeMap::new();
+                            let mut valid_key_map = BTreeMap::new();
+
+                            for (server, keys) in pub_key_map.into_inner().into_iter() {
+                                if keys.valid_until_ts > MilliSecondsSinceUnixEpoch::now() {
+                                    valid_key_map.insert(
+                                        server,
+                                        keys.verify_keys
+                                            .into_iter()
+                                            .map(|(id, key)| (id, key.key))
+                                            .collect(),
+                                    );
+                                } else {
+                                    expired_key_map.insert(
+                                        server,
+                                        keys.verify_keys
+                                            .into_iter()
+                                            .map(|(id, key)| (id, key.key))
+                                            .collect(),
+                                    );
+                                }
+                            }
+
+                            if ruma::signatures::verify_json(&valid_key_map, &value).is_ok() {
+                                RoomMessageEventContent::text_plain("Signature correct")
+                            } else if let Err(e) =
+                                ruma::signatures::verify_json(&expired_key_map, &value)
+                            {
+                                RoomMessageEventContent::text_plain(format!(
                                     "Signature verification failed: {e}"
-                                )),
+                                ))
+                            } else {
+                                RoomMessageEventContent::text_plain(
+                                    "Signature correct (with expired keys)",
+                                )
                             }
                         }
                         Err(e) => RoomMessageEventContent::text_plain(format!("Invalid json: {e}")),

src/service/globals/data.rs

@@ -1,13 +1,71 @@
-use std::collections::BTreeMap;
+use std::{
+    collections::BTreeMap,
+    time::{Duration, SystemTime},
+};
 
+use crate::{services, Result};
 use async_trait::async_trait;
 use ruma::{
-    api::federation::discovery::{ServerSigningKeys, VerifyKey},
+    api::federation::discovery::{OldVerifyKey, ServerSigningKeys, VerifyKey},
+    serde::Base64,
     signatures::Ed25519KeyPair,
-    DeviceId, OwnedServerSigningKeyId, ServerName, UserId,
+    DeviceId, MilliSecondsSinceUnixEpoch, ServerName, UserId,
 };
+use serde::Deserialize;
+
+/// Similar to ServerSigningKeys, but drops a few unnecessary fields we don't require post-validation
+#[derive(Deserialize, Debug, Clone)]
+pub struct SigningKeys {
+    pub verify_keys: BTreeMap<String, VerifyKey>,
+    pub old_verify_keys: BTreeMap<String, OldVerifyKey>,
+    pub valid_until_ts: MilliSecondsSinceUnixEpoch,
+}
+
+impl SigningKeys {
+    /// Creates the SigningKeys struct, using the keys of the current server
+    pub fn load_own_keys() -> Self {
+        let mut keys = Self {
+            verify_keys: BTreeMap::new(),
+            old_verify_keys: BTreeMap::new(),
+            valid_until_ts: MilliSecondsSinceUnixEpoch::from_system_time(
+                SystemTime::now() + Duration::from_secs(7 * 86400),
+            )
+            .expect("Should be valid until year 500,000,000"),
+        };
+
+        keys.verify_keys.insert(
+            format!("ed25519:{}", services().globals.keypair().version()),
+            VerifyKey {
+                key: Base64::new(services().globals.keypair.public_key().to_vec()),
+            },
+        );
+
+        keys
+    }
+}
 
-use crate::Result;
+impl From<ServerSigningKeys> for SigningKeys {
+    fn from(value: ServerSigningKeys) -> Self {
+        let ServerSigningKeys {
+            verify_keys,
+            old_verify_keys,
+            valid_until_ts,
+            ..
+        } = value;
+
+        Self {
+            verify_keys: verify_keys
+                .into_iter()
+                .map(|(id, key)| (id.to_string(), key))
+                .collect(),
+            old_verify_keys: old_verify_keys
+                .into_iter()
+                .map(|(id, key)| (id.to_string(), key))
+                .collect(),
+            valid_until_ts,
+        }
+    }
+}
 
 #[async_trait]
 pub trait Data: Send + Sync {
@@ -21,17 +79,23 @@ pub trait Data: Send + Sync {
     fn clear_caches(&self, amount: u32);
     fn load_keypair(&self) -> Result<Ed25519KeyPair>;
     fn remove_keypair(&self) -> Result<()>;
-    fn add_signing_key(
+    /// Only extends the cached keys, not moving any verify_keys to old_verify_keys, as if we suddenly
+    /// recieve requests from the origin server, we want to be able to accept requests from them
+    fn add_signing_key_from_trusted_server(
         &self,
         origin: &ServerName,
         new_keys: ServerSigningKeys,
-    ) -> Result<BTreeMap<OwnedServerSigningKeyId, VerifyKey>>;
-
-    /// This returns an empty `Ok(BTreeMap<..>)` when there are no keys found for the server.
-    fn signing_keys_for(
+    ) -> Result<SigningKeys>;
+    /// Extends cached keys, as well as moving verify_keys that are not present in these new keys to
+    /// old_verify_keys, so that potnetially comprimised keys cannot be used to make requests
+    fn add_signing_key_from_origin(
         &self,
         origin: &ServerName,
-    ) -> Result<BTreeMap<OwnedServerSigningKeyId, VerifyKey>>;
+        new_keys: ServerSigningKeys,
+    ) -> Result<SigningKeys>;
+
+    /// This returns an empty `Ok(BTreeMap<..>)` when there are no keys found for the server.
+    fn signing_keys_for(&self, origin: &ServerName) -> Result<Option<SigningKeys>>;
     fn database_version(&self) -> Result<u64>;
     fn bump_database_version(&self, new_version: u64) -> Result<()>;
 }

src/service/globals/mod.rs

@@ -1,9 +1,8 @@
 mod data;
 pub use data::Data;
-use ruma::{
-    serde::Base64, OwnedDeviceId, OwnedEventId, OwnedRoomId, OwnedServerName,
-    OwnedServerSigningKeyId, OwnedUserId,
-};
+pub use data::SigningKeys;
+use ruma::MilliSecondsSinceUnixEpoch;
+use ruma::{serde::Base64, OwnedDeviceId, OwnedEventId, OwnedRoomId, OwnedServerName, OwnedUserId};
 use ruma::{OwnedRoomAliasId, RoomAliasId};
 
 use crate::api::server_server::FedDest;
@@ -14,10 +13,7 @@ use hickory_resolver::TokioAsyncResolver;
 use hyper_util::client::legacy::connect::dns::{GaiResolver, Name as HyperName};
 use reqwest::dns::{Addrs, Name, Resolve, Resolving};
 use ruma::{
-    api::{
-        client::sync::sync_events,
-        federation::discovery::{ServerSigningKeys, VerifyKey},
-    },
+    api::{client::sync::sync_events, federation::discovery::ServerSigningKeys},
     DeviceId, RoomVersionId, ServerName, UserId,
 };
 use std::str::FromStr;
@@ -393,36 +389,89 @@ impl Service {
         room_versions
     }
 
-    /// TODO: the key valid until timestamp is only honored in room version > 4
-    /// Remove the outdated keys and insert the new ones.
-    ///
     /// This doesn't actually check that the keys provided are newer than the old set.
-    pub fn add_signing_key(
+    pub fn add_signing_key_from_trusted_server(
         &self,
         origin: &ServerName,
         new_keys: ServerSigningKeys,
-    ) -> Result<BTreeMap<OwnedServerSigningKeyId, VerifyKey>> {
-        self.db.add_signing_key(origin, new_keys)
+    ) -> Result<SigningKeys> {
+        self.db
+            .add_signing_key_from_trusted_server(origin, new_keys)
     }
 
-    /// This returns an empty `Ok(BTreeMap<..>)` when there are no keys found for the server.
-    pub fn signing_keys_for(
+    /// Same as from_trusted_server, except it will move active keys not present in `new_keys` to old_signing_keys
+    pub fn add_signing_key_from_origin(
         &self,
         origin: &ServerName,
-    ) -> Result<BTreeMap<OwnedServerSigningKeyId, VerifyKey>> {
-        let mut keys = self.db.signing_keys_for(origin)?;
-        if origin == self.server_name() {
-            keys.insert(
-                format!("ed25519:{}", services().globals.keypair().version())
-                    .try_into()
-                    .expect("found invalid server signing keys in DB"),
-                VerifyKey {
-                    key: Base64::new(self.keypair.public_key().to_vec()),
-                },
-            );
-        }
+        new_keys: ServerSigningKeys,
+    ) -> Result<SigningKeys> {
+        self.db.add_signing_key_from_origin(origin, new_keys)
+    }
+
+    /// This returns Ok(None) when there are no keys found for the server.
+    pub fn signing_keys_for(&self, origin: &ServerName) -> Result<Option<SigningKeys>> {
+        Ok(self.db.signing_keys_for(origin)?.or_else(|| {
+            if origin == self.server_name() {
+                Some(SigningKeys::load_own_keys())
+            } else {
+                None
+            }
+        }))
+    }
+
+    /// Filters the key map of multiple servers down to keys that should be accepted given the expiry time,
+    /// room version, and timestamp of the paramters
+    pub fn filter_keys_server_map(
+        &self,
+        keys: BTreeMap<String, SigningKeys>,
+        timestamp: MilliSecondsSinceUnixEpoch,
+        room_version_id: &RoomVersionId,
+    ) -> BTreeMap<String, BTreeMap<String, Base64>> {
+        keys.into_iter()
+            .filter_map(|(server, keys)| {
+                self.filter_keys_single_server(keys, timestamp, room_version_id)
+                    .map(|keys| (server, keys))
+            })
+            .collect()
+    }
 
-        Ok(keys)
+    /// Filters the keys of a single server down to keys that should be accepted given the expiry time,
+    /// room version, and timestamp of the paramters
+    pub fn filter_keys_single_server(
+        &self,
+        keys: SigningKeys,
+        timestamp: MilliSecondsSinceUnixEpoch,
+        room_version_id: &RoomVersionId,
+    ) -> Option<BTreeMap<String, Base64>> {
+        if keys.valid_until_ts > timestamp
+            // valid_until_ts MUST be ignored in room versions 1, 2, 3, and 4.
+            // https://spec.matrix.org/v1.10/server-server-api/#get_matrixkeyv2server
+                || matches!(room_version_id, RoomVersionId::V1
+                    | RoomVersionId::V2
+                    | RoomVersionId::V4
+                    | RoomVersionId::V3)
+        {
+            // Given that either the room version allows stale keys, or the valid_until_ts is
+            // in the future, all verify_keys are valid
+            let mut map: BTreeMap<_, _> = keys
+                .verify_keys
+                .into_iter()
+                .map(|(id, key)| (id, key.key))
+                .collect();
+
+            map.extend(keys.old_verify_keys.into_iter().filter_map(|(id, key)| {
+                // Even on old room versions, we don't allow old keys if they are expired
+                if key.expired_ts > timestamp {
+                    Some((id, key.key))
+                } else {
+                    None
+                }
+            }));
+
+            Some(map)
+        } else {
+            None
+        }
     }
 
     pub fn database_version(&self) -> Result<u64> {

src/service/rooms/event_handler/mod.rs

@@ -9,6 +9,7 @@ use std::{
 };
 
 use futures_util::{stream::FuturesUnordered, Future, StreamExt};
+use globals::SigningKeys;
 use ruma::{
     api::{
         client::error::ErrorKind,
@@ -30,7 +31,6 @@ use ruma::{
         StateEventType, TimelineEventType,
     },
     int,
-    serde::Base64,
     state_res::{self, RoomVersion, StateMap},
     uint, CanonicalJsonObject, CanonicalJsonValue, EventId, MilliSecondsSinceUnixEpoch,
     OwnedServerName, OwnedServerSigningKeyId, RoomId, RoomVersionId, ServerName,
@@ -78,7 +78,7 @@ impl Service {
         room_id: &'a RoomId,
         value: BTreeMap<String, CanonicalJsonValue>,
         is_timeline_event: bool,
-        pub_key_map: &'a RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &'a RwLock<BTreeMap<String, SigningKeys>>,
     ) -> Result<Option<Vec<u8>>> {
         // 0. Check the server is in the room
         if !services().rooms.metadata.exists(room_id)? {
@@ -304,19 +304,12 @@ impl Service {
         room_id: &'a RoomId,
         mut value: BTreeMap<String, CanonicalJsonValue>,
         auth_events_known: bool,
-        pub_key_map: &'a RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &'a RwLock<BTreeMap<String, SigningKeys>>,
     ) -> AsyncRecursiveType<'a, Result<(Arc<PduEvent>, BTreeMap<String, CanonicalJsonValue>)>> {
         Box::pin(async move {
             // 1.1. Remove unsigned field
             value.remove("unsigned");
 
-            // TODO: For RoomVersion6 we must check that Raw<..> is canonical do we anywhere?: https://matrix.org/docs/spec/rooms/v6#canonical-json
-
-            // We go through all the signatures we see on the value and fetch the corresponding signing
-            // keys
-            self.fetch_required_signing_keys(&value, pub_key_map)
-                .await?;
-
             // 2. Check signatures, otherwise drop
             // 3. check content hash, redact if doesn't match
             let create_event_content: RoomCreateEventContent =
@@ -329,41 +322,80 @@ impl Service {
             let room_version =
                 RoomVersion::new(room_version_id).expect("room version is supported");
 
-            let guard = pub_key_map.read().await;
-            let mut val = match ruma::signatures::verify_event(&guard, &value, room_version_id) {
-                Err(e) => {
-                    // Drop
-                    warn!("Dropping bad event {}: {}", event_id, e,);
-                    return Err(Error::BadRequest(
+            // TODO: For RoomVersion6 we must check that Raw<..> is canonical do we anywhere?: https://matrix.org/docs/spec/rooms/v6#canonical-json
+
+            // We go through all the signatures we see on the value and fetch the corresponding signing
+            // keys
+            self.fetch_required_signing_keys(&value, pub_key_map)
+                .await?;
+
+            let origin_server_ts = value.get("origin_server_ts").ok_or_else(|| {
+                error!("Invalid PDU, no origin_server_ts field");
+                Error::BadRequest(
+                    ErrorKind::MissingParam,
+                    "Invalid PDU, no origin_server_ts field",
+                )
+            })?;
+
+            let origin_server_ts: MilliSecondsSinceUnixEpoch = {
+                let ts = origin_server_ts.as_integer().ok_or_else(|| {
+                    Error::BadRequest(
                         ErrorKind::InvalidParam,
-                        "Signature verification failed",
-                    ));
-                }
-                Ok(ruma::signatures::Verified::Signatures) => {
-                    // Redact
-                    warn!("Calculated hash does not match: {}", event_id);
-                    let obj = match ruma::canonical_json::redact(value, room_version_id, None) {
-                        Ok(obj) => obj,
-                        Err(_) => {
-                            return Err(Error::BadRequest(
-                                ErrorKind::InvalidParam,
-                                "Redaction failed",
-                            ))
-                        }
-                    };
+                        "origin_server_ts must be an integer",
+                    )
+                })?;
+
+                MilliSecondsSinceUnixEpoch(i64::from(ts).try_into().map_err(|_| {
+                    Error::BadRequest(ErrorKind::InvalidParam, "Time must be after the unix epoch")
+                })?)
+            };
 
-                    // Skip the PDU if it is redacted and we already have it as an outlier event
-                    if services().rooms.timeline.get_pdu_json(event_id)?.is_some() {
+            let guard = pub_key_map.read().await;
+
+            let pkey_map = (*guard).clone();
+
+            // Removing all the expired keys, unless the room version allows stale keys
+            let filtered_keys = services().globals.filter_keys_server_map(
+                pkey_map,
+                origin_server_ts,
+                room_version_id,
+            );
+
+            let mut val =
+                match ruma::signatures::verify_event(&filtered_keys, &value, room_version_id) {
+                    Err(e) => {
+                        // Drop
+                        warn!("Dropping bad event {}: {}", event_id, e,);
                         return Err(Error::BadRequest(
                             ErrorKind::InvalidParam,
-                            "Event was redacted and we already knew about it",
+                            "Signature verification failed",
                         ));
                     }
+                    Ok(ruma::signatures::Verified::Signatures) => {
+                        // Redact
+                        warn!("Calculated hash does not match: {}", event_id);
+                        let obj = match ruma::canonical_json::redact(value, room_version_id, None) {
+                            Ok(obj) => obj,
+                            Err(_) => {
+                                return Err(Error::BadRequest(
+                                    ErrorKind::InvalidParam,
+                                    "Redaction failed",
+                                ))
+                            }
+                        };
 
-                    obj
-                }
-                Ok(ruma::signatures::Verified::All) => value,
-            };
+                        // Skip the PDU if it is redacted and we already have it as an outlier event
+                        if services().rooms.timeline.get_pdu_json(event_id)?.is_some() {
+                            return Err(Error::BadRequest(
+                                ErrorKind::InvalidParam,
+                                "Event was redacted and we already knew about it",
+                            ));
+                        }
+
+                        obj
+                    }
+                    Ok(ruma::signatures::Verified::All) => value,
+                };
 
             drop(guard);
 
@@ -487,7 +519,7 @@ impl Service {
         create_event: &PduEvent,
         origin: &ServerName,
         room_id: &RoomId,
-        pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &RwLock<BTreeMap<String, SigningKeys>>,
     ) -> Result<Option<Vec<u8>>> {
         // Skip the PDU if we already have it as a timeline event
         if let Ok(Some(pduid)) = services().rooms.timeline.get_pdu_id(&incoming_pdu.event_id) {
@@ -1097,7 +1129,7 @@ impl Service {
         create_event: &'a PduEvent,
         room_id: &'a RoomId,
         room_version_id: &'a RoomVersionId,
-        pub_key_map: &'a RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &'a RwLock<BTreeMap<String, SigningKeys>>,
     ) -> AsyncRecursiveType<'a, Vec<(Arc<PduEvent>, Option<BTreeMap<String, CanonicalJsonValue>>)>>
     {
         Box::pin(async move {
@@ -1280,7 +1312,7 @@ impl Service {
         create_event: &PduEvent,
         room_id: &RoomId,
         room_version_id: &RoomVersionId,
-        pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &RwLock<BTreeMap<String, SigningKeys>>,
         initial_set: Vec<Arc<EventId>>,
     ) -> Result<(
         Vec<Arc<EventId>>,
@@ -1378,7 +1410,7 @@ impl Service {
     pub(crate) async fn fetch_required_signing_keys(
         &self,
         event: &BTreeMap<String, CanonicalJsonValue>,
-        pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &RwLock<BTreeMap<String, SigningKeys>>,
     ) -> Result<()> {
         let signatures = event
             .get("signatures")
@@ -1407,6 +1439,7 @@ impl Service {
                         )
                     })?,
                     signature_ids,
+                    true,
                 )
                 .await;
 
@@ -1434,7 +1467,7 @@ impl Service {
         pdu: &RawJsonValue,
         servers: &mut BTreeMap<OwnedServerName, BTreeMap<OwnedServerSigningKeyId, QueryCriteria>>,
         room_version: &RoomVersionId,
-        pub_key_map: &mut RwLockWriteGuard<'_, BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &mut RwLockWriteGuard<'_, BTreeMap<String, SigningKeys>>,
     ) -> Result<()> {
         let value: CanonicalJsonObject = serde_json::from_str(pdu.get()).map_err(|e| {
             error!("Invalid PDU in server response: {:?}: {:?}", pdu, e);
@@ -1485,8 +1518,18 @@ impl Service {
 
             let signature_ids = signature_object.keys().cloned().collect::<Vec<_>>();
 
-            let contains_all_ids = |keys: &BTreeMap<String, Base64>| {
-                signature_ids.iter().all(|id| keys.contains_key(id))
+            let contains_all_ids = |keys: &SigningKeys| {
+                signature_ids.iter().all(|id| {
+                    keys.verify_keys
+                        .keys()
+                        .map(ToString::to_string)
+                        .any(|key_id| id == &key_id)
+                        || keys
+                            .old_verify_keys
+                            .keys()
+                            .map(ToString::to_string)
+                            .any(|key_id| id == &key_id)
+                })
             };
 
             let origin = <&ServerName>::try_from(signature_server.as_str()).map_err(|_| {
@@ -1499,19 +1542,14 @@ impl Service {
 
             trace!("Loading signing keys for {}", origin);
 
-            let result: BTreeMap<_, _> = services()
-                .globals
-                .signing_keys_for(origin)?
-                .into_iter()
-                .map(|(k, v)| (k.to_string(), v.key))
-                .collect();
-
-            if !contains_all_ids(&result) {
-                trace!("Signing key not loaded for {}", origin);
-                servers.insert(origin.to_owned(), BTreeMap::new());
-            }
+            if let Some(result) = services().globals.signing_keys_for(origin)? {
+                if !contains_all_ids(&result) {
+                    trace!("Signing key not loaded for {}", origin);
+                    servers.insert(origin.to_owned(), BTreeMap::new());
+                }
 
-            pub_key_map.insert(origin.to_string(), result);
+                pub_key_map.insert(origin.to_string(), result);
+            }
         }
 
         Ok(())
@@ -1521,7 +1559,7 @@ impl Service {
         &self,
         event: &create_join_event::v2::Response,
         room_version: &RoomVersionId,
-        pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &RwLock<BTreeMap<String, SigningKeys>>,
     ) -> Result<()> {
         let mut servers: BTreeMap<
             OwnedServerName,
@@ -1584,10 +1622,7 @@ impl Service {
 
                     let result = services()
                         .globals
-                        .add_signing_key(&k.server_name, k.clone())?
-                        .into_iter()
-                        .map(|(k, v)| (k.to_string(), v.key))
-                        .collect::<BTreeMap<_, _>>();
+                        .add_signing_key_from_trusted_server(&k.server_name, k.clone())?;
 
                     pkm.insert(k.server_name.to_string(), result);
                 }
@@ -1618,12 +1653,9 @@ impl Service {
             if let (Ok(get_keys_response), origin) = result {
                 info!("Result is from {origin}");
                 if let Ok(key) = get_keys_response.server_key.deserialize() {
-                    let result: BTreeMap<_, _> = services()
+                    let result = services()
                         .globals
-                        .add_signing_key(&origin, key)?
-                        .into_iter()
-                        .map(|(k, v)| (k.to_string(), v.key))
-                        .collect();
+                        .add_signing_key_from_origin(&origin, key)?;
                     pub_key_map.write().await.insert(origin.to_string(), result);
                 }
             }
@@ -1681,9 +1713,23 @@ impl Service {
         &self,
         origin: &ServerName,
         signature_ids: Vec<String>,
-    ) -> Result<BTreeMap<String, Base64>> {
-        let contains_all_ids =
-            |keys: &BTreeMap<String, Base64>| signature_ids.iter().all(|id| keys.contains_key(id));
+        // Whether to ask for keys from trusted servers. Should be false when getting
+        // keys for validating requests, as per MSC4029
+        query_via_trusted_servers: bool,
+    ) -> Result<SigningKeys> {
+        let contains_all_ids = |keys: &SigningKeys| {
+            signature_ids.iter().all(|id| {
+                keys.verify_keys
+                    .keys()
+                    .map(ToString::to_string)
+                    .any(|key_id| id == &key_id)
+                    || keys
+                        .old_verify_keys
+                        .keys()
+                        .map(ToString::to_string)
+                        .any(|key_id| id == &key_id)
+            })
+        };
 
         let permit = services()
             .globals
@@ -1744,94 +1790,172 @@ impl Service {
 
         trace!("Loading signing keys for {}", origin);
 
-        let mut result: BTreeMap<_, _> = services()
-            .globals
-            .signing_keys_for(origin)?
-            .into_iter()
-            .map(|(k, v)| (k.to_string(), v.key))
-            .collect();
+        let result = services().globals.signing_keys_for(origin)?;
+
+        let mut expires_soon_or_has_expired = false;
 
-        if contains_all_ids(&result) {
-            return Ok(result);
+        if let Some(result) = result.clone() {
+            let ts_threshold = MilliSecondsSinceUnixEpoch::from_system_time(
+                SystemTime::now() + Duration::from_secs(30 * 60),
+            )
+            .expect("Should be valid until year 500,000,000");
+
+            debug!(
+                "The treshhold is {:?}, found time is {:?} for server {}",
+                ts_threshold, result.valid_until_ts, origin
+            );
+
+            if contains_all_ids(&result) {
+                // We want to ensure that the keys remain valid by the time the other functions that handle signatures reach them
+                if result.valid_until_ts > ts_threshold {
+                    debug!(
+                        "Keys for {} are deemed as valid, as they expire at {:?}",
+                        &origin, &result.valid_until_ts
+                    );
+                    return Ok(result);
+                }
+
+                expires_soon_or_has_expired = true;
+            }
         }
 
+        let mut keys = result.unwrap_or_else(|| SigningKeys {
+            verify_keys: BTreeMap::new(),
+            old_verify_keys: BTreeMap::new(),
+            valid_until_ts: MilliSecondsSinceUnixEpoch::now(),
+        });
+
+        // We want to set this to the max, and then lower it whenever we see older keys
+        keys.valid_until_ts = MilliSecondsSinceUnixEpoch::from_system_time(
+            SystemTime::now() + Duration::from_secs(7 * 86400),
+        )
+        .expect("Should be valid until year 500,000,000");
+
         debug!("Fetching signing keys for {} over federation", origin);
 
-        if let Some(server_key) = services()
+        if let Some(mut server_key) = services()
             .sending
             .send_federation_request(origin, get_server_keys::v2::Request::new())
             .await
             .ok()
             .and_then(|resp| resp.server_key.deserialize().ok())
         {
+            // Keys should only be valid for a maximum of seven days
+            server_key.valid_until_ts = server_key.valid_until_ts.min(
+                MilliSecondsSinceUnixEpoch::from_system_time(
+                    SystemTime::now() + Duration::from_secs(7 * 86400),
+                )
+                .expect("Should be valid until year 500,000,000"),
+            );
+
             services()
                 .globals
-                .add_signing_key(origin, server_key.clone())?;
+                .add_signing_key_from_origin(origin, server_key.clone())?;
 
-            result.extend(
+            if keys.valid_until_ts > server_key.valid_until_ts {
+                keys.valid_until_ts = server_key.valid_until_ts;
+            }
+
+            keys.verify_keys.extend(
                 server_key
                     .verify_keys
                     .into_iter()
-                    .map(|(k, v)| (k.to_string(), v.key)),
+                    .map(|(id, key)| (id.to_string(), key)),
             );
-            result.extend(
+            keys.old_verify_keys.extend(
                 server_key
                     .old_verify_keys
                     .into_iter()
-                    .map(|(k, v)| (k.to_string(), v.key)),
+                    .map(|(id, key)| (id.to_string(), key)),
             );
 
-            if contains_all_ids(&result) {
-                return Ok(result);
+            if contains_all_ids(&keys) {
+                return Ok(keys);
             }
         }
 
-        for server in services().globals.trusted_servers() {
-            debug!("Asking {} for {}'s signing key", server, origin);
-            if let Some(server_keys) = services()
-                .sending
-                .send_federation_request(
-                    server,
-                    get_remote_server_keys::v2::Request::new(
-                        origin.to_owned(),
-                        MilliSecondsSinceUnixEpoch::from_system_time(
-                            SystemTime::now()
-                                .checked_add(Duration::from_secs(3600))
-                                .expect("SystemTime to large"),
-                        )
-                        .expect("time is valid"),
-                    ),
-                )
-                .await
-                .ok()
-                .map(|resp| {
-                    resp.server_keys
-                        .into_iter()
-                        .filter_map(|e| e.deserialize().ok())
-                        .collect::<Vec<_>>()
-                })
-            {
-                trace!("Got signing keys: {:?}", server_keys);
-                for k in server_keys {
-                    services().globals.add_signing_key(origin, k.clone())?;
-                    result.extend(
-                        k.verify_keys
-                            .into_iter()
-                            .map(|(k, v)| (k.to_string(), v.key)),
-                    );
-                    result.extend(
-                        k.old_verify_keys
+        if query_via_trusted_servers {
+            for server in services().globals.trusted_servers() {
+                debug!("Asking {} for {}'s signing key", server, origin);
+                if let Some(server_keys) = services()
+                    .sending
+                    .send_federation_request(
+                        server,
+                        get_remote_server_keys::v2::Request::new(
+                            origin.to_owned(),
+                            MilliSecondsSinceUnixEpoch::from_system_time(
+                                SystemTime::now()
+                                    .checked_add(Duration::from_secs(3600))
+                                    .expect("SystemTime to large"),
+                            )
+                            .expect("time is valid"),
+                        ),
+                    )
+                    .await
+                    .ok()
+                    .map(|resp| {
+                        resp.server_keys
                             .into_iter()
-                            .map(|(k, v)| (k.to_string(), v.key)),
-                    );
-                }
+                            .filter_map(|e| e.deserialize().ok())
+                            .collect::<Vec<_>>()
+                    })
+                {
+                    trace!("Got signing keys: {:?}", server_keys);
+                    for mut k in server_keys {
+                        if k.valid_until_ts
+                        // Half an hour should give plenty of time for the server to respond with keys that are still
+                        // valid, given we requested keys which are valid at least an hour from now
+                            < MilliSecondsSinceUnixEpoch::from_system_time(
+                                SystemTime::now() + Duration::from_secs(30 * 60),
+                            )
+                            .expect("Should be valid until year 500,000,000")
+                        {
+                            // Keys should only be valid for a maximum of seven days
+                            k.valid_until_ts = k.valid_until_ts.min(
+                                MilliSecondsSinceUnixEpoch::from_system_time(
+                                    SystemTime::now() + Duration::from_secs(7 * 86400),
+                                )
+                                .expect("Should be valid until year 500,000,000"),
+                            );
 
-                if contains_all_ids(&result) {
-                    return Ok(result);
+                            if keys.valid_until_ts > k.valid_until_ts {
+                                keys.valid_until_ts = k.valid_until_ts;
+                            }
+
+                            services()
+                                .globals
+                                .add_signing_key_from_trusted_server(origin, k.clone())?;
+                            keys.verify_keys.extend(
+                                k.verify_keys
+                                    .into_iter()
+                                    .map(|(id, key)| (id.to_string(), key)),
+                            );
+                            keys.old_verify_keys.extend(
+                                k.old_verify_keys
+                                    .into_iter()
+                                    .map(|(id, key)| (id.to_string(), key)),
+                            );
+                        } else {
+                            warn!(
+                                "Server {} gave us keys older than we requested, valid until: {:?}",
+                                origin, k.valid_until_ts
+                            );
+                        }
+
+                        if contains_all_ids(&keys) {
+                            return Ok(keys);
+                        }
+                    }
                 }
             }
         }
 
+        // We should return these keys if fresher keys were not found
+        if expires_soon_or_has_expired {
+            info!("Returning stale keys for {}", origin);
+            return Ok(keys);
+        }
+
         drop(permit);
 
         back_off(signature_ids).await;

src/service/rooms/timeline/mod.rs

@@ -21,7 +21,6 @@ use ruma::{
         GlobalAccountDataEventType, StateEventType, TimelineEventType,
     },
     push::{Action, Ruleset, Tweak},
-    serde::Base64,
     state_res::{self, Event, RoomVersion},
     uint, user_id, CanonicalJsonObject, CanonicalJsonValue, EventId, OwnedEventId, OwnedRoomId,
     OwnedServerName, RoomId, RoomVersionId, ServerName, UserId,
@@ -33,7 +32,10 @@ use tracing::{error, info, warn};
 
 use crate::{
     api::server_server,
-    service::pdu::{EventHash, PduBuilder},
+    service::{
+        globals::SigningKeys,
+        pdu::{EventHash, PduBuilder},
+    },
     services, utils, Error, PduEvent, Result,
 };
 
@@ -1214,7 +1216,7 @@ impl Service {
         &self,
         origin: &ServerName,
         pdu: Box<RawJsonValue>,
-        pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
+        pub_key_map: &RwLock<BTreeMap<String, SigningKeys>>,
     ) -> Result<()> {
         let (event_id, value, room_id) = server_server::parse_incoming_pdu(&pdu)?;