From 044a1fd31ced38292261c0245d2d4af3547891ee Mon Sep 17 00:00:00 2001
From: Zixuan Chen <remch183@outlook.com>
Date: Sat, 9 Nov 2024 00:51:21 +0800
Subject: [PATCH] fix: avoid creating non-root containers that doesn't exist by
 get_container api (#541)

---
 crates/loro-internal/src/loro.rs    | 18 ++++++++++++++++++
 crates/loro-internal/src/state.rs   |  5 +++++
 crates/loro/tests/loro_rust_test.rs | 12 ++++++++++--
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/crates/loro-internal/src/loro.rs b/crates/loro-internal/src/loro.rs
index a954de732..f1384322b 100644
--- a/crates/loro-internal/src/loro.rs
+++ b/crates/loro-internal/src/loro.rs
@@ -618,6 +618,7 @@ impl LoroDoc {
 
     #[inline]
     pub fn get_handler(&self, id: ContainerID) -> Handler {
+        self.assert_container_exists(&id);
         Handler::new_attached(
             id,
             self.arena.clone(),
@@ -631,6 +632,7 @@ impl LoroDoc {
     #[inline]
     pub fn get_text<I: IntoContainerId>(&self, id: I) -> TextHandler {
         let id = id.into_container_id(&self.arena, ContainerType::Text);
+        self.assert_container_exists(&id);
         Handler::new_attached(
             id,
             self.arena.clone(),
@@ -646,6 +648,7 @@ impl LoroDoc {
     #[inline]
     pub fn get_list<I: IntoContainerId>(&self, id: I) -> ListHandler {
         let id = id.into_container_id(&self.arena, ContainerType::List);
+        self.assert_container_exists(&id);
         Handler::new_attached(
             id,
             self.arena.clone(),
@@ -661,6 +664,7 @@ impl LoroDoc {
     #[inline]
     pub fn get_movable_list<I: IntoContainerId>(&self, id: I) -> MovableListHandler {
         let id = id.into_container_id(&self.arena, ContainerType::MovableList);
+        self.assert_container_exists(&id);
         Handler::new_attached(
             id,
             self.arena.clone(),
@@ -676,6 +680,7 @@ impl LoroDoc {
     #[inline]
     pub fn get_map<I: IntoContainerId>(&self, id: I) -> MapHandler {
         let id = id.into_container_id(&self.arena, ContainerType::Map);
+        self.assert_container_exists(&id);
         Handler::new_attached(
             id,
             self.arena.clone(),
@@ -691,6 +696,7 @@ impl LoroDoc {
     #[inline]
     pub fn get_tree<I: IntoContainerId>(&self, id: I) -> TreeHandler {
         let id = id.into_container_id(&self.arena, ContainerType::Tree);
+        self.assert_container_exists(&id);
         Handler::new_attached(
             id,
             self.arena.clone(),
@@ -707,6 +713,7 @@ impl LoroDoc {
         id: I,
     ) -> crate::handler::counter::CounterHandler {
         let id = id.into_container_id(&self.arena, ContainerType::Counter);
+        self.assert_container_exists(&id);
         Handler::new_attached(
             id,
             self.arena.clone(),
@@ -717,6 +724,17 @@ impl LoroDoc {
         .unwrap()
     }
 
+    fn assert_container_exists(&self, id: &ContainerID) {
+        if id.is_root() {
+            return;
+        }
+
+        let exist = self.state.try_lock().unwrap().does_container_exist(id);
+        if !exist {
+            panic!("Target container does not exist: {:?}", id);
+        }
+    }
+
     /// Undo the operations between the given id_span. It can be used even in a collaborative environment.
     ///
     /// This is an internal API. You should NOT use it directly.
diff --git a/crates/loro-internal/src/state.rs b/crates/loro-internal/src/state.rs
index 9a38ea1be..feb4cdf96 100644
--- a/crates/loro-internal/src/state.rs
+++ b/crates/loro-internal/src/state.rs
@@ -764,6 +764,11 @@ impl DocState {
         self.store.iter_all_containers()
     }
 
+    pub fn does_container_exist(&self, id: &ContainerID) -> bool {
+        // TODO: we may need a better way to handle this in the future when we need to enable fully lazy loading on state
+        self.arena.id_to_idx(id).is_some()
+    }
+
     pub(crate) fn init_container(
         &mut self,
         cid: ContainerID,
diff --git a/crates/loro/tests/loro_rust_test.rs b/crates/loro/tests/loro_rust_test.rs
index 8973fe6dc..f11e55892 100644
--- a/crates/loro/tests/loro_rust_test.rs
+++ b/crates/loro/tests/loro_rust_test.rs
@@ -9,8 +9,8 @@ use std::{
 };
 
 use loro::{
-    awareness::Awareness, loro_value, Frontiers, FrontiersNotIncluded, LoroDoc, LoroError,
-    LoroList, LoroMap, LoroText, ToJson,
+    awareness::Awareness, loro_value, ContainerID, ContainerType, Frontiers, FrontiersNotIncluded,
+    LoroDoc, LoroError, LoroList, LoroMap, LoroText, ToJson,
 };
 use loro_internal::{handler::TextDelta, id::ID, vv, LoroResult};
 use rand::{Rng, SeedableRng};
@@ -1993,3 +1993,11 @@ fn test_fork_when_detached() {
     doc.checkout_to_latest();
     assert_eq!(doc.get_text("text").to_string(), "Hello, world! Alice!");
 }
+
+#[test]
+#[should_panic]
+fn should_avoid_initialize_new_container_accidentally() {
+    let doc = LoroDoc::new();
+    let id = ContainerID::new_normal(ID::new(0, 0), ContainerType::Text);
+    let _text = doc.get_text(id);
+}