From 748e20d4a363d89b841d62213f5b0c6b4bed788f Mon Sep 17 00:00:00 2001
From: Will Chen <willchen90@gmail.com>
Date: Fri, 3 Jan 2025 09:15:57 -0800
Subject: [PATCH] Harden dataclass utils (#1171)

---
 mesop/dataclass_utils/dataclass_utils.py      |  4 ++++
 mesop/dataclass_utils/dataclass_utils_test.py | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/mesop/dataclass_utils/dataclass_utils.py b/mesop/dataclass_utils/dataclass_utils.py
index 6ad757b16..580ee20fb 100644
--- a/mesop/dataclass_utils/dataclass_utils.py
+++ b/mesop/dataclass_utils/dataclass_utils.py
@@ -129,6 +129,10 @@ def update_dataclass_from_json(instance: Any, json_string: str):
 
 def _recursive_update_dataclass_from_json_obj(instance: Any, json_dict: Any):
   for key, value in json_dict.items():
+    if key.startswith("__") and key.endswith("__"):
+      raise MesopDeveloperException(
+        f"Cannot use dunder property: {key} in stateclass"
+      )
     if hasattr(instance, key):
       attr = getattr(instance, key)
       if isinstance(value, dict):
diff --git a/mesop/dataclass_utils/dataclass_utils_test.py b/mesop/dataclass_utils/dataclass_utils_test.py
index 03b1c5194..f04a7a1c4 100644
--- a/mesop/dataclass_utils/dataclass_utils_test.py
+++ b/mesop/dataclass_utils/dataclass_utils_test.py
@@ -14,6 +14,7 @@
   serialize_dataclass,
   update_dataclass_from_json,
 )
+from mesop.exceptions import MesopDeveloperException
 
 
 @dataclass
@@ -593,5 +594,23 @@ class ChildClass(ParentClass):
   assert has_parent(ParentClass) is False
 
 
+def test_globals_pollution():
+  @dataclass
+  class A:
+    val: str
+
+  initial_name = __name__
+  obj = A(val="default")
+  with pytest.raises(MesopDeveloperException) as exc_info:
+    update_dataclass_from_json(
+      obj, '{"__init__": {"__globals__": {"__name__": "polluted"}}}'
+    )
+  assert "Cannot use dunder property: __init__ in stateclass" in str(
+    exc_info.value
+  )
+  # Make sure __name__ has not been modified via the __globals__ pollution attempt
+  assert __name__ == initial_name
+
+
 if __name__ == "__main__":
   raise SystemExit(pytest.main(["-vv", __file__]))