fix(codec): prevent pre-allocating a large number of structs (#525)

y-octo/src/doc/codec/delete_set.rs

@@ -135,14 +135,16 @@ impl DeleteSet {
 impl<R: CrdtReader> CrdtRead<R> for DeleteSet {
     fn read(decoder: &mut R) -> JwstCodecResult<Self> {
         let num_of_clients = decoder.read_var_u64()? as usize;
-        let mut map = HashMap::with_capacity(num_of_clients);
+        // See: [Update::read]
+        let mut map = HashMap::with_capacity(num_of_clients.min(1 << 10));
 
         for _ in 0..num_of_clients {
             let client = decoder.read_var_u64()?;
             let deletes = OrderRange::read(decoder)?;
             map.insert(client, deletes);
         }
 
+        map.shrink_to_fit();
         Ok(DeleteSet(map))
     }
 }

y-octo/src/doc/codec/update.rs

@@ -26,33 +26,35 @@ pub struct Update {
 
 impl<R: CrdtReader> CrdtRead<R> for Update {
     fn read(decoder: &mut R) -> JwstCodecResult<Self> {
-        let num_of_clients = decoder.read_var_u64()?;
-
-        let mut map = HashMap::with_capacity(if num_of_clients > CLIENTS_SAFE_CAPACITY {
-            CLIENTS_SAFE_CAPACITY
-        } else {
-            num_of_clients
-        } as usize);
+        let num_of_clients = decoder.read_var_u64()? as usize;
+
+        // NOTE:
+        //   We do not use [HashMap::with_capacity(num_of_clients)] directly here because we don't trust the input data.
+        //   For instance, what if the first u64 was somehow set a very big value?
+        //   A pre-allocated HashMap with a big capacity may cause OOM.
+        //   A kinda safer approach is give it a max capacity of 1024 at first allocation,
+        //     and then let std makes the growth as need.
+        let mut map = HashMap::with_capacity(num_of_clients.min(1 << 10));
         for _ in 0..num_of_clients {
-            let num_of_structs = decoder.read_var_u64()?;
+            let num_of_structs = decoder.read_var_u64()? as usize;
             let client = decoder.read_var_u64()?;
             let mut clock = decoder.read_var_u64()?;
 
-            let mut structs = VecDeque::with_capacity(if num_of_structs > STRUCTS_SAFE_CAPACITY {
-                STRUCTS_SAFE_CAPACITY
-            } else {
-                num_of_structs
-            } as usize);
+            // same reason as above
+            let mut structs = VecDeque::with_capacity(num_of_structs.min(1 << 10));
 
             for _ in 0..num_of_structs {
                 let struct_info = Node::read(decoder, Id::new(client, clock))?;
                 clock += struct_info.len();
                 structs.push_back(struct_info);
             }
 
+            structs.shrink_to_fit();
             map.insert(client, structs);
         }
 
+        map.shrink_to_fit();
+
         let delete_set = DeleteSet::read(decoder)?;
 
         if !decoder.is_empty() {

y-octo/src/doc/common/state.rs

@@ -78,13 +78,14 @@ impl<R: CrdtReader> CrdtRead<R> for StateVector {
     fn read(decoder: &mut R) -> JwstCodecResult<Self> {
         let len = decoder.read_var_u64()? as usize;
 
-        let mut map = HashMap::with_capacity(len);
+        let mut map = HashMap::with_capacity(len.min(1 << 10));
         for _ in 0..len {
             let client = decoder.read_var_u64()?;
             let clock = decoder.read_var_u64()?;
             map.insert(client, clock);
         }
 
+        map.shrink_to_fit();
         Ok(Self(map))
     }
 }