Skip to content

Commit 5f2ebf1

Browse files
authored
do not broadcast event to all clients if modify state is misused (#5322)
* do not broadcast event to all clients if modify state is misused * fix tests * maybe?
1 parent f9b0d4d commit 5f2ebf1

File tree

2 files changed

+32
-4
lines changed

2 files changed

+32
-4
lines changed

reflex/app.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1996,6 +1996,13 @@ async def emit_update(self, update: StateUpdate, sid: str) -> None:
19961996
update: The state update to send.
19971997
sid: The Socket.IO session id.
19981998
"""
1999+
if not sid:
2000+
# If the sid is None, we are not connected to a client. Prevent sending
2001+
# updates to all clients.
2002+
return
2003+
if sid not in self.sid_to_token:
2004+
console.warn(f"Attempting to send delta to disconnected websocket {sid}")
2005+
return
19992006
# Creating a task prevents the update from being blocked behind other coroutines.
20002007
await asyncio.create_task(
20012008
self.emit(str(constants.SocketEvent.EVENT), update, to=sid)

tests/units/test_state.py

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1948,21 +1948,37 @@ class ModelDC:
19481948

19491949

19501950
@pytest.mark.asyncio
1951-
async def test_state_proxy(grandchild_state: GrandchildState, mock_app: rx.App):
1951+
async def test_state_proxy(
1952+
grandchild_state: GrandchildState, mock_app: rx.App, token: str
1953+
):
19521954
"""Test that the state proxy works.
19531955
19541956
Args:
19551957
grandchild_state: A grandchild state.
19561958
mock_app: An app that will be returned by `get_app()`
1959+
token: A token.
19571960
"""
19581961
child_state = grandchild_state.parent_state
19591962
assert child_state is not None
19601963
parent_state = child_state.parent_state
19611964
assert parent_state is not None
1965+
router_data = RouterData({"query": {}, "token": token, "sid": "test_sid"})
1966+
grandchild_state.router = router_data
1967+
namespace = mock_app.event_namespace
1968+
assert namespace is not None
1969+
namespace.sid_to_token[router_data.session.session_id] = token
19621970
if isinstance(mock_app.state_manager, (StateManagerMemory, StateManagerDisk)):
19631971
mock_app.state_manager.states[parent_state.router.session.client_token] = (
19641972
parent_state
19651973
)
1974+
elif isinstance(mock_app.state_manager, StateManagerRedis):
1975+
pickle_state = parent_state._serialize()
1976+
if pickle_state:
1977+
await mock_app.state_manager.redis.set(
1978+
_substate_key(parent_state.router.session.client_token, parent_state),
1979+
pickle_state,
1980+
ex=mock_app.state_manager.token_expiration,
1981+
)
19661982

19671983
sp = StateProxy(grandchild_state)
19681984
assert sp.__wrapped__ == grandchild_state
@@ -2029,6 +2045,7 @@ async def test_state_proxy(grandchild_state: GrandchildState, mock_app: rx.App):
20292045
assert mcall.args[0] == str(SocketEvent.EVENT)
20302046
assert mcall.args[1] == StateUpdate(
20312047
delta={
2048+
TestState.get_full_name(): {"router": router_data},
20322049
grandchild_state.get_full_name(): {
20332050
"value2": "42",
20342051
},
@@ -2154,7 +2171,11 @@ async def test_background_task_no_block(mock_app: rx.App, token: str):
21542171
mock_app: An app that will be returned by `get_app()`
21552172
token: A token.
21562173
"""
2157-
router_data = {"query": {}}
2174+
router_data = {"query": {}, "token": token}
2175+
sid = "test_sid"
2176+
namespace = mock_app.event_namespace
2177+
assert namespace is not None
2178+
namespace.sid_to_token[sid] = token
21582179
mock_app.state_manager.state = mock_app._state = BackgroundTaskState
21592180
async for update in rx.app.process(
21602181
mock_app,
@@ -2164,7 +2185,7 @@ async def test_background_task_no_block(mock_app: rx.App, token: str):
21642185
router_data=router_data,
21652186
payload={},
21662187
),
2167-
sid="",
2188+
sid=sid,
21682189
headers={},
21692190
client_ip="",
21702191
):
@@ -2184,7 +2205,7 @@ async def test_background_task_no_block(mock_app: rx.App, token: str):
21842205
router_data=router_data,
21852206
payload={},
21862207
),
2187-
sid="",
2208+
sid=sid,
21882209
headers={},
21892210
client_ip="",
21902211
):

0 commit comments

Comments
 (0)