fix: state size was not checked for dill (#4431)

This commit is contained in:
benedikt-bartscher 2024-11-25 19:38:52 +01:00 committed by GitHub
parent 697e26c25b
commit d7d46e431b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -2175,11 +2175,9 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
Returns: Returns:
The serialized state. The serialized state.
""" """
payload = b""
try: try:
pickle_state = pickle.dumps((self._to_schema(), self)) payload = pickle.dumps((self._to_schema(), self))
if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
self._check_state_size(len(pickle_state))
return pickle_state
except HANDLED_PICKLE_ERRORS as og_pickle_error: except HANDLED_PICKLE_ERRORS as og_pickle_error:
error = ( error = (
f"Failed to serialize state {self.get_full_name()} due to unpicklable object. " f"Failed to serialize state {self.get_full_name()} due to unpicklable object. "
@ -2188,7 +2186,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
try: try:
import dill import dill
return dill.dumps((self._to_schema(), self)) payload = dill.dumps((self._to_schema(), self))
except ImportError: except ImportError:
error += ( error += (
f"Pickle error: {og_pickle_error}. " f"Pickle error: {og_pickle_error}. "
@ -2196,8 +2194,10 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
) )
except HANDLED_PICKLE_ERRORS as ex: except HANDLED_PICKLE_ERRORS as ex:
error += f"Dill was also unable to pickle the state: {ex}" error += f"Dill was also unable to pickle the state: {ex}"
console.warn(error) console.warn(error)
return b"" if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
self._check_state_size(len(payload))
return payload
@classmethod @classmethod
def _deserialize( def _deserialize(