Prune one time keys based on reasonable configured limit.

Signed-off-by: Jason Volk <jason@zemos.net>

src/api/client/keys.rs

@@ -40,7 +40,11 @@ pub(crate) async fn upload_keys_route(
 ) -> Result<upload_keys::v3::Response> {
 	let (sender_user, sender_device) = body.sender();
 
-	for (key_id, one_time_key) in &body.one_time_keys {
+	for (key_id, one_time_key) in body
+		.one_time_keys
+		.iter()
+		.take(services.config.one_time_key_limit)
+	{
 		if one_time_key
 			.deserialize()
 			.inspect_err(|e| {

src/core/config/mod.rs

@@ -1999,6 +1999,16 @@ pub struct Config {
 	#[serde(default)]
 	pub delete_rooms_after_leave: bool,
 
+	/// Limits the number of One Time Keys per device (not per-algorithm). The
+	/// reference implementation maintains 50 OTK's at any given time, therefor
+	/// our default is at least ten times that. There is no known reason for an
+	/// administrator to adjust this value; it is provided here rather than
+	/// hardcoding it.
+	///
+	/// default: 512
+	#[serde(default = "default_one_time_key_limit")]
+	pub one_time_key_limit: usize,
+
 	// external structure; separate section
 	#[serde(default)]
 	pub blurhashing: BlurhashConfig,
@@ -2822,3 +2832,5 @@ fn default_access_token_ttl() -> u64 { 604_800 }
 fn default_deprioritize_joins_through_servers() -> RegexSet {
 	RegexSet::new([r"matrix\.org"]).unwrap()
 }
+
+fn default_one_time_key_limit() -> usize { 512 }

src/service/users/keys.rs

@@ -152,9 +152,36 @@ pub async fn count_one_time_keys(
 		})
 		.await;
 
+	let total = algorithm_counts
+		.values()
+		.copied()
+		.map(TryInto::try_into)
+		.filter_map(Result::ok)
+		.fold(0_usize, usize::saturating_add);
+
+	if total > self.services.config.one_time_key_limit {
+		self.prune_one_time_keys(user_id, device_id).await;
+	}
+
 	algorithm_counts
 }
 
+#[implement(super::Service)]
+pub async fn prune_one_time_keys(&self, user_id: &UserId, device_id: &DeviceId) {
+	use tuwunel_database::keyval::Key;
+
+	let query = (user_id, device_id);
+	self.db
+		.onetimekeyid_onetimekeys
+		.keys_prefix(&query)
+		.ignore_err()
+		.skip(self.services.config.one_time_key_limit)
+		.ready_for_each(|key: Key<'_>| {
+			self.db.onetimekeyid_onetimekeys.remove(key);
+		})
+		.await;
+}
+
 #[implement(super::Service)]
 pub async fn add_device_keys(
 	&self,

tuwunel-example.toml

@@ -1718,6 +1718,14 @@
 #
 #delete_rooms_after_leave = false
 
+# Limits the number of One Time Keys per device (not per-algorithm). The
+# reference implementation maintains 50 OTK's at any given time, therefor
+# our default is at least ten times that. There is no known reason for an
+# administrator to adjust this value; it is provided here rather than
+# hardcoding it.
+#
+#one_time_key_limit = 512
+
 #[global.tls]
 
 # Path to a valid TLS certificate file.