Skip to content

Commit

Permalink
Harden dataclass utils (#1171)
Browse files Browse the repository at this point in the history
  • Loading branch information
wwwillchen authored Jan 3, 2025
1 parent cbab8bf commit 748e20d
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 0 deletions.
4 changes: 4 additions & 0 deletions mesop/dataclass_utils/dataclass_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
19 changes: 19 additions & 0 deletions mesop/dataclass_utils/dataclass_utils_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
serialize_dataclass,
update_dataclass_from_json,
)
from mesop.exceptions import MesopDeveloperException


@dataclass
Expand Down Expand Up @@ -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__]))

0 comments on commit 748e20d

Please sign in to comment.