From 9b2f9451f1b887a1eaacc24dfc38cfc76c8d7296 Mon Sep 17 00:00:00 2001
From: Lambda <csgl@xiretza.xyz>
Date: Thu, 10 Apr 2025 18:38:07 +0000
Subject: [PATCH] Integrity-check account data events

---
 src/integrity.rs | 137 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 136 insertions(+), 1 deletion(-)

diff --git a/src/integrity.rs b/src/integrity.rs
index 62c905d8..76ae7c37 100644
--- a/src/integrity.rs
+++ b/src/integrity.rs
@@ -1,3 +1,8 @@
+use std::string::FromUtf8Error;
+
+use bstr::BString;
+use serde::Deserialize;
+use serde_json::value::RawValue;
 use thiserror::Error;
 
 use crate::{
@@ -23,6 +28,9 @@ pub(crate) enum IntegrityError {
 
         y_name: &'static str,
     },
+
+    #[error(transparent)]
+    InvalidAccountDataEvent(#[from] InvalidAccountDataEvent),
 }
 
 #[derive(Debug, Error)]
@@ -142,10 +150,137 @@ fn short_ids(
     eventid_symmetry.chain(eventid_foreign).chain(statekey_symmetry)
 }
 
+#[derive(Debug, Error)]
+pub(crate) enum InvalidAccountDataEvent {
+    #[error("missing event type in key {key:?}")]
+    MissingType {
+        key: BString,
+    },
+    #[error("invalid event type string in key {key:?}")]
+    InvalidType {
+        key: BString,
+        #[source]
+        err: FromUtf8Error,
+    },
+
+    #[error("missing count in key {key:?}")]
+    MissingCount {
+        key: BString,
+    },
+
+    #[error("missing user ID in key {key:?}")]
+    MissingUserId {
+        key: BString,
+    },
+    #[error("invalid user ID string in key {key:?}")]
+    InvalidUserId {
+        key: BString,
+        #[source]
+        err: FromUtf8Error,
+    },
+
+    #[error("invalid event data for {key:?}: {value:?}")]
+    InvalidEventData {
+        key: BString,
+        value: BString,
+        #[source]
+        err: serde_json::Error,
+    },
+
+    #[error(
+        "mismatch between event type in column key ({key_type}) and \
+         serialised field ({serialised_type}) for user {user_id}"
+    )]
+    MismatchedType {
+        user_id: String,
+        key_type: String,
+        serialised_type: String,
+    },
+}
+
+fn account_data(
+    database: &'static KeyValueDatabase,
+) -> impl Iterator<Item = Result<IntegrityError>> {
+    let tree = tree!(database, roomuserdataid_accountdata);
+    tree.tree
+        .iter()
+        .map(|(key, value)| {
+            #[allow(dead_code)]
+            #[derive(Deserialize)]
+            struct ExtractEventFields<'a> {
+                #[serde(rename = "type")]
+                event_type: &'a str,
+                content: &'a RawValue,
+            }
+
+            let key = BString::from(key);
+            let mut key_parts = key.rsplit(|&b| b == 0xFF);
+
+            let event_type = String::from_utf8(
+                key_parts
+                    .next()
+                    .ok_or_else(|| InvalidAccountDataEvent::MissingType {
+                        key: key.clone(),
+                    })?
+                    .to_vec(),
+            )
+            .map_err(|err| {
+                InvalidAccountDataEvent::InvalidType {
+                    key: key.clone(),
+                    err,
+                }
+            })?;
+
+            let Some(_) = key_parts.next() else {
+                return Err(InvalidAccountDataEvent::MissingCount {
+                    key: key.clone(),
+                });
+            };
+
+            let user_id = String::from_utf8(
+                key_parts
+                    .next()
+                    .ok_or_else(|| InvalidAccountDataEvent::MissingUserId {
+                        key: key.clone(),
+                    })?
+                    .to_vec(),
+            )
+            .map_err(|err| {
+                InvalidAccountDataEvent::InvalidUserId {
+                    key: key.clone(),
+                    err,
+                }
+            })?;
+
+            let extract: ExtractEventFields<'_> =
+                serde_json::from_slice(&value).map_err(|err| {
+                    InvalidAccountDataEvent::InvalidEventData {
+                        key: key.clone(),
+                        value: value.clone().into(),
+                        err,
+                    }
+                })?;
+
+            if extract.event_type != event_type {
+                return Err(InvalidAccountDataEvent::MismatchedType {
+                    user_id,
+                    key_type: event_type,
+                    serialised_type: extract.event_type.to_owned(),
+                });
+            }
+
+            Ok(())
+        })
+        .filter_map(|r| match r {
+            Ok(()) => None,
+            Err(e) => Some(Ok(IntegrityError::InvalidAccountDataEvent(e))),
+        })
+}
+
 impl CheckIntegrity for KeyValueDatabase {
     fn check_integrity(
         &'static self,
     ) -> Box<dyn Iterator<Item = Result<IntegrityError>>> {
-        Box::new(short_ids(self))
+        Box::new(short_ids(self).chain(account_data(self)))
     }
 }