Handle arbitrary sized integers in unsigned. (#19769)

Handle arbitrary sized integers in `unsigned` (and other Rust objects
that use `serde_json::Value`)

rust/src/acl/mod.rs

@@ -47,7 +47,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
 }
 
 #[derive(Debug, Clone)]
-#[pyclass(frozen)]
+#[pyclass(frozen, skip_from_py_object)]
 pub struct ServerAclEvaluator {
     allow_ip_literals: bool,
     allow: Vec<Regex>,

rust/src/canonical_json.rs

@@ -29,7 +29,7 @@ use std::{
     io::{self, Write},
 };
 
-use serde::ser::SerializeMap;
+use serde::{ser::SerializeMap, Serializer as _};
 use serde::{
     ser::{Error as _, SerializeStruct},
     Serialize,
@@ -37,7 +37,7 @@ use serde::{
 use serde_json::{
     ser::{Formatter, Serializer},
     value::RawValue,
-    Value,
+    Number, Value,
 };
 
 /// The minimum integer that can be used in canonical JSON.
@@ -46,6 +46,12 @@ pub const MIN_VALID_INTEGER: i64 = -(2i64.pow(53)) + 1;
 /// The maximum integer that can be used in canonical JSON.
 pub const MAX_VALID_INTEGER: i64 = (2i64.pow(53)) - 1;
 
+/// A token used by `serde_json` to identify its internal `Number` type when the
+/// `arbitrary_precision` feature is enabled. This is a copy from serde_json's
+/// internal `TOKEN` for `Number`, which unfortunately isn't exported by the
+/// crate.
+const SERDE_JSON_NUMBER_TOKEN: &str = "$serde_json::private::Number";
+
 /// Options to control how strict JSON canonicalization is.
 #[derive(Clone, Debug)]
 pub struct CanonicalizationOptions {
@@ -236,7 +242,7 @@ where
 
     type SerializeMap = CanonicalSerializeMap<'a, W>;
 
-    type SerializeStruct = CanonicalSerializeMap<'a, W>;
+    type SerializeStruct = CanonicalSerializeStruct<'a, W>;
 
     type SerializeStructVariant =
         <&'a mut Serializer<W, CanonicalFormatter> as serde::Serializer>::SerializeStructVariant;
@@ -426,7 +432,7 @@ where
     fn serialize_struct(
         self,
         name: &'static str,
-        _len: usize,
+        len: usize,
     ) -> Result<Self::SerializeStruct, Self::Error> {
         // We want to disallow `RawValue` as we don't know if its contents is
         // canonical JSON.
@@ -436,10 +442,7 @@ where
         if name == "$serde_json::private::RawValue" {
             return Err(Self::Error::custom("`RawValue` is not supported"));
         }
-        Ok(CanonicalSerializeMap::new(
-            &mut self.inner,
-            self.options.clone(),
-        ))
+        CanonicalSerializeStruct::new(name, len, &mut self.inner, self.options.clone())
     }
 
     fn serialize_struct_variant(
@@ -554,7 +557,42 @@ where
     }
 }
 
-impl<'a, W> SerializeStruct for CanonicalSerializeMap<'a, W>
+/// A helper type for [`CanonicalSerializer`] that serializes structs in
+/// lexicographic order.
+#[doc(hidden)]
+pub struct CanonicalSerializeStruct<'a, W: Write> {
+    name: &'static str,
+    // We buffer up the key and serialized value for each field we see.
+    // The BTreeMap will then serialize in lexicographic order.
+    map: BTreeMap<&'static str, Box<RawValue>>,
+    options: CanonicalizationOptions,
+    // The serializer to use to write the sorted map too.
+    struct_serializer:
+        <&'a mut Serializer<W, CanonicalFormatter> as serde::Serializer>::SerializeStruct,
+}
+
+impl<'a, W> CanonicalSerializeStruct<'a, W>
+where
+    W: Write,
+{
+    fn new(
+        name: &'static str,
+        len: usize,
+        ser: &'a mut Serializer<W, CanonicalFormatter>,
+        options: CanonicalizationOptions,
+    ) -> Result<Self, serde_json::Error> {
+        let struct_serializer = ser.serialize_struct(name, len)?;
+
+        Ok(Self {
+            name,
+            map: BTreeMap::new(),
+            options,
+            struct_serializer,
+        })
+    }
+}
+
+impl<'a, W> SerializeStruct for CanonicalSerializeStruct<'a, W>
 where
     W: Write,
 {
@@ -566,20 +604,69 @@ where
     where
         T: Serialize + ?Sized,
     {
-        let key_string = key.to_string();
+        // Check if this is the special case of `SERDE_JSON_NUMBER_TOKEN`,
+        // which is used when serializing numbers with the `arbitrary_precision`
+        // feature. If so, we can just serialize it directly without
+        // canonicalizing it first, as `serde_json` will have already serialized
+        // it in a canonical way.
+        if key == SERDE_JSON_NUMBER_TOKEN && self.name == SERDE_JSON_NUMBER_TOKEN {
+            if self.options.enforce_int_range {
+                // We need to check that the number is in the valid range, as
+                // `serde_json` won't have done this for us as we're using the
+                // `arbitrary_precision` feature.
+
+                // The value here will be something that serializes to a JSON
+                // string containing the number, so we first serialize it to a
+                // Value and pull the string out, then parse it as a `Number`.
+
+                let serde_val = serde_json::to_value(value)?;
+                let serde_json::Value::String(number_str) = serde_val else {
+                    return Err(serde_json::Error::custom("invalid number"));
+                };
+
+                let number: Number = number_str
+                    .parse()
+                    .map_err(|_| serde_json::Error::custom("invalid number"))?;
+
+                // Now check that the number is an integer in the valid range.
+                if let Some(int) = number.as_i64() {
+                    assert_integer_in_range(int)?;
+                } else {
+                    // Can't be cast to an i64, so it must be out of range.
+                    return Err(serde_json::Error::custom("integer out of range"));
+                }
+            }
+
+            self.struct_serializer.serialize_field(key, value)?;
+            return Ok(());
+        }
 
         // We serialize the value canonically, then store it as a `RawValue` in
         // the buffer map.
         let value_string = to_string_canonical(value, self.options.clone())?;
 
-        self.map
-            .insert(key_string, RawValue::from_string(value_string)?);
+        self.map.insert(key, RawValue::from_string(value_string)?);
 
         Ok(())
     }
 
-    fn end(self) -> Result<Self::Ok, Self::Error> {
-        self.map.serialize(self.ser)?;
+    fn end(mut self) -> Result<Self::Ok, Self::Error> {
+        if self.name == SERDE_JSON_NUMBER_TOKEN {
+            // Map must be empty in this case, as `SERDE_JSON_NUMBER_TOKEN`
+            // only has one field and we've handled it in `serialize_field`.
+            if !self.map.is_empty() {
+                return Err(Self::Error::custom(format!(
+                    "unexpected fields in `{}`",
+                    SERDE_JSON_NUMBER_TOKEN
+                )));
+            }
+        }
+
+        for (key, value) in self.map {
+            self.struct_serializer.serialize_field(key, &value)?;
+        }
+
+        SerializeStruct::end(self.struct_serializer)?;
 
         Ok(())
     }
@@ -737,6 +824,23 @@ mod tests {
         assert!(to_string_canonical(&-(2i128.pow(60)), CanonicalizationOptions::strict()).is_err());
     }
 
+    #[test]
+    fn bigints() {
+        // Create a `serde_json::Number` that is too big to be represented as an
+        // i64, but can be represented as a string.
+        let bigint_string = "10000000000000000000000000000000000000";
+        let value: serde_json::Number = bigint_string.parse().unwrap();
+
+        // This should work with relaxed option.
+        assert_eq!(
+            to_string_canonical(&value, CanonicalizationOptions::relaxed()).unwrap(),
+            bigint_string
+        );
+
+        // But should fail with strict option, as it's out of range.
+        assert!(to_string_canonical(&value, CanonicalizationOptions::strict()).is_err());
+    }
+
     #[test]
     fn backwards_compatibility() {
         assert_eq!(

rust/src/events/internal_metadata.rs

@@ -476,7 +476,7 @@ impl EventInternalMetadataInner {
     }
 }
 
-#[pyclass(frozen)]
+#[pyclass(frozen, skip_from_py_object)]
 #[derive(Clone)]
 pub struct EventInternalMetadata {
     inner: Arc<RwLock<EventInternalMetadataInner>>,

rust/src/events/unsigned.rs

@@ -23,6 +23,7 @@ use pyo3::{
 };
 use pythonize::{depythonize, pythonize};
 use serde::{Deserialize, Serialize};
+use serde_json::Number;
 
 #[pyclass(frozen, skip_from_py_object)]
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
@@ -36,7 +37,7 @@ pub struct Unsigned {
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 struct PersistedUnsignedFields {
     #[serde(skip_serializing_if = "Option::is_none")]
-    age_ts: Option<i64>,
+    age_ts: Option<Number>,
     #[serde(skip_serializing_if = "Option::is_none")]
     replaces_state: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
@@ -129,11 +130,14 @@ impl Unsigned {
         let unsigned = self.py_read()?;
 
         match field {
-            UnsignedField::AgeTs => Ok(unsigned
-                .persisted_fields
-                .age_ts
-                .ok_or_else(|| PyKeyError::new_err("age_ts"))?
-                .into_bound_py_any(py)?),
+            UnsignedField::AgeTs => {
+                let age_ts = &unsigned
+                    .persisted_fields
+                    .age_ts
+                    .as_ref()
+                    .ok_or_else(|| PyKeyError::new_err("age_ts"))?;
+                Ok(pythonize(py, age_ts)?)
+            }
             UnsignedField::ReplacesState => Ok((unsigned.persisted_fields.replaces_state)
                 .as_ref()
                 .ok_or_else(|| PyKeyError::new_err("replaces_state"))?
@@ -203,7 +207,7 @@ impl Unsigned {
         let mut unsigned = self.py_write()?;
 
         match field {
-            UnsignedField::AgeTs => unsigned.persisted_fields.age_ts = Some(value.extract()?),
+            UnsignedField::AgeTs => unsigned.persisted_fields.age_ts = Some(depythonize(&value)?),
             UnsignedField::ReplacesState => {
                 unsigned.persisted_fields.replaces_state = Some(value.extract()?)
             }
@@ -339,7 +343,7 @@ mod tests {
     #[test]
     fn test_persisted_fields_serialize_populated() {
         let fields = PersistedUnsignedFields {
-            age_ts: Some(1234),
+            age_ts: Some(1234.into()),
             replaces_state: Some("$prev:example.com".to_string()),
             invite_room_state: Some(vec![json!({"type": "m.room.name"})]),
             knock_room_state: Some(vec![json!({"type": "m.room.topic"})]),
@@ -360,7 +364,7 @@ mod tests {
     fn test_unsigned_inner_flattens_persisted_fields() {
         let inner = UnsignedInner {
             persisted_fields: PersistedUnsignedFields {
-                age_ts: Some(99),
+                age_ts: Some(99.into()),
                 ..Default::default()
             },
             prev_content: Some(Box::new(json!({"body": "hi"}))),
@@ -382,7 +386,7 @@ mod tests {
     fn test_unsigned_inner_roundtrip() {
         let original = UnsignedInner {
             persisted_fields: PersistedUnsignedFields {
-                age_ts: Some(10),
+                age_ts: Some(10.into()),
                 replaces_state: Some("$state:example.com".to_string()),
                 invite_room_state: None,
                 knock_room_state: None,
@@ -394,7 +398,7 @@ mod tests {
         let json = serde_json::to_string(&original).unwrap();
         let roundtripped: UnsignedInner = serde_json::from_str(&json).unwrap();
 
-        assert_eq!(roundtripped.persisted_fields.age_ts, Some(10));
+        assert_eq!(roundtripped.persisted_fields.age_ts, Some(10.into()));
         assert_eq!(
             roundtripped.persisted_fields.replaces_state.as_deref(),
             Some("$state:example.com")
@@ -423,7 +427,7 @@ mod tests {
         });
         let unsigned: Unsigned = serde_json::from_value(json).unwrap();
         let inner = unsigned.inner.read().unwrap();
-        assert_eq!(inner.persisted_fields.age_ts, Some(5));
+        assert_eq!(inner.persisted_fields.age_ts, Some(5.into()));
         assert_eq!(inner.prev_sender.as_deref(), Some("@bob:example.com"));
     }
 }

rust/src/push/mod.rs

@@ -104,7 +104,7 @@ fn get_base_rule_ids() -> HashSet<&'static str> {
 
 /// A single push rule for a user.
 #[derive(Debug, Clone)]
-#[pyclass(frozen)]
+#[pyclass(frozen, from_py_object)]
 pub struct PushRule {
     /// A unique ID for this rule
     pub rule_id: Cow<'static, str>,
@@ -462,7 +462,7 @@ pub struct RelatedEventMatchTypeCondition {
 
 /// The collection of push rules for a user.
 #[derive(Debug, Clone, Default)]
-#[pyclass(frozen)]
+#[pyclass(frozen, from_py_object)]
 pub struct PushRules {
     /// Custom push rules that override a base rule.
     overridden_base_rules: HashMap<Cow<'static, str>, PushRule>,
@@ -549,7 +549,7 @@ impl PushRules {
 /// A wrapper around `PushRules` that checks the enabled state of rules and
 /// filters out disabled experimental rules.
 #[derive(Debug, Clone, Default)]
-#[pyclass(frozen)]
+#[pyclass(frozen, skip_from_py_object)]
 pub struct FilteredPushRules {
     push_rules: PushRules,
     enabled_map: BTreeMap<String, bool>,

rust/src/room_versions.rs

@@ -93,7 +93,7 @@ impl PushRuleRoomFlag {
 }
 
 /// An object which describes the unique attributes of a room version.
-#[pyclass(frozen, eq, hash, get_all)]
+#[pyclass(frozen, eq, hash, get_all, skip_from_py_object)]
 #[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct RoomVersion {
     /// The identifier for this version.
@@ -622,7 +622,7 @@ const ROOM_VERSION_MSC4242V12: RoomVersion = RoomVersion {
 /// Note: room versions can be added to this mapping at startup (allowing
 /// support for experimental room versions to be behind experimental feature
 /// flags).
-#[pyclass(frozen, mapping)]
+#[pyclass(frozen, mapping, skip_from_py_object)]
 #[derive(Clone)]
 pub struct KnownRoomVersionsMapping {
     // Note we use a Vec here to ensure that the order of keys is
@@ -637,19 +637,22 @@ pub struct KnownRoomVersionsMapping {
 impl KnownRoomVersionsMapping {
     /// Add a new room version to the mapping, indicating that this instance
     /// supports it.
-    fn add_room_version(&self, version: RoomVersion) -> PyResult<()> {
+    fn add_room_version(&self, version: Bound<'_, RoomVersion>) -> PyResult<()> {
         let mut versions = self
             .versions
             .write()
             .map_err(|_| PyRuntimeError::new_err("KnownRoomVersionsMapping lock poisoned"))?;
 
-        if versions.iter().any(|v| v.identifier == version.identifier) {
+        if versions
+            .iter()
+            .any(|v| v.identifier == version.get().identifier)
+        {
             // We already have this room version, so we don't add it again (as
             // otherwise we'd end up with duplicates).
             return Ok(());
         }
 
-        versions.push(version);
+        versions.push(*version.get());
         Ok(())
     }