check if we already have a more preferable key backup before adding

This commit is contained in:
cy 2025-03-14 23:01:28 -04:00 committed by June Clementine Strawberry 🍓🦴
parent 4518f55408
commit 658c19d55e
2 changed files with 76 additions and 13 deletions

View file

@ -1,3 +1,5 @@
use std::cmp::Ordering;
use axum::extract::State;
use conduwuit::{Err, err};
use ruma::{
@ -232,16 +234,77 @@ pub(crate) async fn add_backup_keys_for_session_route(
)));
}
services
// Check if we already have a better key
let mut ok_to_replace = true;
if let Some(old_key) = &services
.key_backups
.add_key(
body.sender_user(),
&body.version,
&body.room_id,
&body.session_id,
&body.session_data,
)
.await?;
.get_session(body.sender_user(), &body.version, &body.room_id, &body.session_id)
.await
.ok()
{
let old_is_verified = old_key
.get_field::<bool>("is_verified")?
.unwrap_or_default();
let new_is_verified = body
.session_data
.get_field::<bool>("is_verified")?
.ok_or_else(|| err!(Request(BadJson("`is_verified` field should exist"))))?;
// Prefer key that `is_verified`
if old_is_verified != new_is_verified {
if old_is_verified {
ok_to_replace = false;
}
} else {
// If both have same `is_verified`, prefer the one with lower
// `first_message_index`
let old_first_message_index = old_key
.get_field::<UInt>("first_message_index")?
.unwrap_or(UInt::MAX);
let new_first_message_index = body
.session_data
.get_field::<UInt>("first_message_index")?
.ok_or_else(|| {
err!(Request(BadJson("`first_message_index` field should exist")))
})?;
ok_to_replace = match new_first_message_index.cmp(&old_first_message_index) {
| Ordering::Less => true,
| Ordering::Greater => false,
| Ordering::Equal => {
// If both have same `first_message_index`, prefer the one with lower
// `forwarded_count`
let old_forwarded_count = old_key
.get_field::<UInt>("forwarded_count")?
.unwrap_or(UInt::MAX);
let new_forwarded_count = body
.session_data
.get_field::<UInt>("forwarded_count")?
.ok_or_else(|| {
err!(Request(BadJson("`forwarded_count` field should exist")))
})?;
new_forwarded_count < old_forwarded_count
},
};
};
}
if ok_to_replace {
services
.key_backups
.add_key(
body.sender_user(),
&body.version,
&body.room_id,
&body.session_id,
&body.session_data,
)
.await?;
}
Ok(add_backup_keys_for_session::v3::Response {
count: services

View file

@ -134,10 +134,10 @@
{"Action":"pass","Test":"TestDeviceManagement/PUT_/device/{deviceId}_gives_a_404_for_unknown_devices"}
{"Action":"pass","Test":"TestDeviceManagement/PUT_/device/{deviceId}_updates_device_fields"}
{"Action":"pass","Test":"TestDisplayNameUpdate"}
{"Action":"fail","Test":"TestE2EKeyBackupReplaceRoomKeyRules"}
{"Action":"fail","Test":"TestE2EKeyBackupReplaceRoomKeyRules/parallel"}
{"Action":"fail","Test":"TestE2EKeyBackupReplaceRoomKeyRules/parallel/{isVerified:false_firstMessageIndex:10_forwardedCount:5}"}
{"Action":"fail","Test":"TestE2EKeyBackupReplaceRoomKeyRules/parallel/{isVerified:true_firstMessageIndex:10_forwardedCount:5}"}
{"Action":"pass","Test":"TestE2EKeyBackupReplaceRoomKeyRules"}
{"Action":"pass","Test":"TestE2EKeyBackupReplaceRoomKeyRules/parallel"}
{"Action":"pass","Test":"TestE2EKeyBackupReplaceRoomKeyRules/parallel/{isVerified:false_firstMessageIndex:10_forwardedCount:5}"}
{"Action":"pass","Test":"TestE2EKeyBackupReplaceRoomKeyRules/parallel/{isVerified:true_firstMessageIndex:10_forwardedCount:5}"}
{"Action":"pass","Test":"TestEvent"}
{"Action":"pass","Test":"TestEvent/Parallel"}
{"Action":"pass","Test":"TestEvent/Parallel/Large_Event"}