Fix potential panic in rbx_binary when serializing default SharedStri…

…ng (#363)
rojo-rbx · Sep 23, 2023 · 9716af3 · 9716af3
1 parent 8c4328e
commit 9716af3
Show file tree

Hide file tree

Showing 7 changed files with 102 additions and 6 deletions.
diff --git a/rbx_binary/CHANGELOG.md b/rbx_binary/CHANGELOG.md
@@ -2,8 +2,10 @@
 
 ## Unreleased
 * Added support for `SecurityCapabilities` values. ([#361])
+* Fixed potential crash when serializing `SharedString` values ([#363])
 
 [#361]: https://github.com/rojo-rbx/rbx-dom/pull/361
+[#363]: https://github.com/rojo-rbx/rbx-dom/pull/363
 
 ## 0.7.1 (2023-08-09)
 * Added support for `UniqueId` values. ([#271])

diff --git a/rbx_binary/src/serializer/state.rs b/rbx_binary/src/serializer/state.rs
@@ -262,7 +262,10 @@ impl<'dom, W: Write> SerializerState<'dom, W> {
             self.shared_string_ids.insert(shared_string, id as u32);
         }
 
-        log::debug!("Type info discovered: {:#?}", self.type_infos);
+        log::debug!(
+            "Discovered {} unique TypeInfos",
+            self.type_infos.values.len()
+        );
 
         Ok(())
     }
@@ -403,6 +406,16 @@ impl<'dom, W: Write> SerializerState<'dom, W> {
                         }
                     })?;
 
+                // There's no assurance that the default SharedString value
+                // will actually get serialized inside of the SSTR chunk, so we
+                // check here just to make sure.
+                if let Cow::Owned(Variant::SharedString(sstr)) = &default_value {
+                    if !self.shared_string_ids.contains_key(sstr) {
+                        self.shared_string_ids.insert(sstr.clone(), 0);
+                        self.shared_strings.push(sstr.clone());
+                    }
+                }
+
                 let ser_type = Type::from_rbx_type(serialized_ty).ok_or_else(|| {
                     // This is a known value type, but rbx_binary doesn't have a
                     // binary type value for it. rbx_binary might be out of
@@ -1102,8 +1115,14 @@ impl<'dom, W: Write> SerializerState<'dom, W> {
 
                         for (i, rbx_value) in values {
                             if let Variant::SharedString(value) = rbx_value.as_ref() {
-                                let id = &self.shared_string_ids[value];
-                                entries.push(*id);
+                                if let Some(id) = self.shared_string_ids.get(value) {
+                                    entries.push(*id);
+                                } else {
+                                    panic!(
+                                        "SharedString {} was not found during type collection",
+                                        value.hash()
+                                    )
+                                }
                             } else {
                                 return type_mismatch(i, &rbx_value, "SharedString");
                             }

diff --git a/rbx_binary/src/tests/serializer.rs b/rbx_binary/src/tests/serializer.rs
@@ -1,5 +1,5 @@
 use rbx_dom_weak::{
-    types::{BrickColor, Color3, Color3uint8, Enum, Font, Ref, Region3, Vector3},
+    types::{BrickColor, Color3, Color3uint8, Enum, Font, Ref, Region3, SharedString, Vector3},
     InstanceBuilder, WeakDom,
 };
 
@@ -173,3 +173,23 @@ fn part_color() {
     let decoded = DecodedModel::from_reader(buf.as_slice());
     insta::assert_yaml_snapshot!(decoded);
 }
+
+#[test]
+fn default_shared_string() {
+    let mut tree = WeakDom::new(InstanceBuilder::new("Folder"));
+    let ref_1 = tree.insert(
+        tree.root_ref(),
+        InstanceBuilder::new("Model").with_property(
+            // This is the first SharedString property I saw in the database
+            "ModelMeshData",
+            SharedString::new(b"arbitrary string".to_vec()),
+        ),
+    );
+    let ref_2 = tree.insert(tree.root_ref(), InstanceBuilder::new("Model"));
+
+    let mut buf = Vec::new();
+    let _ = to_writer(&mut buf, &tree, &[ref_1, ref_2]);
+
+    let decoded = DecodedModel::from_reader(buf.as_slice());
+    insta::assert_yaml_snapshot!(decoded);
+}
diff --git a/rbx_binary/src/tests/snapshots/rbx_binary__tests__serializer__default_shared_string.snap b/rbx_binary/src/tests/snapshots/rbx_binary__tests__serializer__default_shared_string.snap
@@ -0,0 +1,44 @@
+---
+source: rbx_binary/src/tests/serializer.rs
+expression: decoded
+---
+num_types: 1
+num_instances: 2
+chunks:
+  - Sstr:
+      version: 0
+      entries:
+        - len: 16
+          hash: 5425b48244fac32fd24d4e33c3c97aa7ae75fc79ce0542878fd08db58e82449d
+        - len: 0
+          hash: af1349b9f5f9a1a6a0404dea36dcc9499bcb25c9adc112b7cc9a93cae41f3262
+  - Inst:
+      type_id: 0
+      type_name: Model
+      object_format: 0
+      referents:
+        - 0
+        - 1
+  - Prop:
+      type_id: 0
+      prop_name: ModelMeshData
+      prop_type: SharedString
+      values:
+        - 0
+        - 1
+  - Prop:
+      type_id: 0
+      prop_name: Name
+      prop_type: String
+      values:
+        - Model
+        - Model
+  - Prnt:
+      version: 0
+      links:
+        - - 0
+          - -1
+        - - 1
+          - -1
+  - End
+
diff --git a/rbx_binary/src/tests/snapshots/rbx_binary__tests__util__unions__encoded.snap b/rbx_binary/src/tests/snapshots/rbx_binary__tests__util__unions__encoded.snap
@@ -8,6 +8,8 @@ chunks:
   - Sstr:
       version: 0
       entries:
+        - len: 0
+          hash: af1349b9f5f9a1a6a0404dea36dcc9499bcb25c9adc112b7cc9a93cae41f3262
         - len: 1367
           hash: cc62d7045ec238081b7ebbd97ce5ce7dc15166e688168c307247089135718bdb
         - len: 1803
@@ -348,9 +350,9 @@ chunks:
       prop_name: PhysicalConfigData
       prop_type: SharedString
       values:
+        - 2
+        - 2
         - 1
-        - 1
-        - 0
   - Prop:
       type_id: 0
       prop_name: PhysicsData

diff --git a/rbx_types/CHANGELOG.md b/rbx_types/CHANGELOG.md
@@ -4,10 +4,12 @@
 * Implemented `FromStr` for `TerrainMaterials`. ([#354])
 * `MaterialColorsError` and `UniqueIdError` are no longer publicly exposed. ([#355])
 * Implemented barebones `SecurityCapabilities`. ([#358])
+* Implement `Display` for `SharedStringHash` ([#363])
 
 [#354]: https://github.com/rojo-rbx/rbx-dom/pull/354
 [#355]: https://github.com/rojo-rbx/rbx-dom/pull/355
 [#358]: https://github.com/rojo-rbx/rbx-dom/pull/358
+[#363]: https://github.com/rojo-rbx/rbx-dom/pull/363
 
 ## 1.6.0 (2023-08-09)
 * Added support for `UniqueId` values. ([#271])

diff --git a/rbx_types/src/shared_string.rs b/rbx_types/src/shared_string.rs
@@ -1,6 +1,7 @@
 use std::{
     cmp::Ordering,
     collections::{hash_map::Entry, HashMap},
+    fmt,
     hash::{Hash, Hasher},
     sync::{Arc, Mutex, Weak},
 };
@@ -140,6 +141,12 @@ impl PartialOrd for SharedStringHash {
     }
 }
 
+impl fmt::Display for SharedStringHash {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(self.0.to_hex().as_str())
+    }
+}
+
 #[cfg(feature = "serde")]
 pub(crate) mod variant_serialization {
     use super::*;