Sliding sync improvements and redaction fixes

src/api/client_server/keys.rs

@@ -17,7 +17,11 @@ use ruma::{
     DeviceKeyAlgorithm, OwnedDeviceId, OwnedUserId, UserId,
 };
 use serde_json::json;
-use std::collections::{BTreeMap, HashMap, HashSet};
+use std::{
+    collections::{hash_map, BTreeMap, HashMap, HashSet},
+    time::{Duration, Instant},
+};
+use tracing::debug;
 
 /// # `POST /_matrix/client/r0/keys/upload`
 ///
@@ -335,31 +339,68 @@ pub(crate) async fn get_keys_helper<F: Fn(&UserId) -> bool>(
 
     let mut failures = BTreeMap::new();
 
+    let back_off = |id| match services()
+        .globals
+        .bad_query_ratelimiter
+        .write()
+        .unwrap()
+        .entry(id)
+    {
+        hash_map::Entry::Vacant(e) => {
+            e.insert((Instant::now(), 1));
+        }
+        hash_map::Entry::Occupied(mut e) => *e.get_mut() = (Instant::now(), e.get().1 + 1),
+    };
+
     let mut futures: FuturesUnordered<_> = get_over_federation
         .into_iter()
         .map(|(server, vec)| async move {
+            if let Some((time, tries)) = services()
+                .globals
+                .bad_query_ratelimiter
+                .read()
+                .unwrap()
+                .get(&*server)
+            {
+                // Exponential backoff
+                let mut min_elapsed_duration = Duration::from_secs(30) * (*tries) * (*tries);
+                if min_elapsed_duration > Duration::from_secs(60 * 60 * 24) {
+                    min_elapsed_duration = Duration::from_secs(60 * 60 * 24);
+                }
+
+                if time.elapsed() < min_elapsed_duration {
+                    debug!("Backing off query from {:?}", server);
+                    return (
+                        server,
+                        Err(Error::BadServerResponse("bad query, still backing off")),
+                    );
+                }
+            }
+
             let mut device_keys_input_fed = BTreeMap::new();
             for (user_id, keys) in vec {
                 device_keys_input_fed.insert(user_id.to_owned(), keys.clone());
             }
             (
                 server,
-                services()
-                    .sending
-                    .send_federation_request(
+                tokio::time::timeout(
+                    Duration::from_secs(25),
+                    services().sending.send_federation_request(
                         server,
                         federation::keys::get_keys::v1::Request {
                             device_keys: device_keys_input_fed,
                         },
-                    )
-                    .await,
+                    ),
+                )
+                .await
+                .map_err(|e| Error::BadServerResponse("Query took too long")),
             )
         })
         .collect();
 
     while let Some((server, response)) = futures.next().await {
         match response {
-            Ok(response) => {
+            Ok(Ok(response)) => {
                 for (user, masterkey) in response.master_keys {
                     let (master_key_id, mut master_key) =
                         services().users.parse_master_key(&user, &masterkey)?;
@@ -386,7 +427,8 @@ pub(crate) async fn get_keys_helper<F: Fn(&UserId) -> bool>(
                 self_signing_keys.extend(response.self_signing_keys);
                 device_keys.extend(response.device_keys);
             }
-            Err(_e) => {
+            _ => {
+                back_off(server.to_owned());
                 failures.insert(server.to_string(), json!({}));
             }
         }