Skip to content

Commit da3a484

Browse files
committed
Coordinate cold sessions before warming browser transports
1 parent 548db53 commit da3a484

4 files changed

Lines changed: 63 additions & 18 deletions

File tree

‎packages/reflex-base/src/reflex_base/.templates/web/utils/state.js‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,10 @@ export const getToken = () => {
115115

116116
/** Get the saved client token to warm the transport before the app mounts. */
117117
const peekToken = () =>
118-
getToken() || window.sessionStorage.getItem("token") || "";
118+
token ||
119+
window.sessionStorage.getItem(`${sessionStorageKey()}:token`) ||
120+
window.sessionStorage.getItem("token") ||
121+
"";
119122

120123
/**
121124
* Get the URL for the backend server
@@ -276,6 +279,16 @@ if (typeof window !== "undefined") {
276279
return;
277280
}
278281
try {
282+
// A cold transport would capture cookies before the shared session lock.
283+
if (
284+
!(
285+
Number(
286+
window.localStorage.getItem(`${sessionStorageKey()}:expires_at`),
287+
) > Date.now()
288+
)
289+
) {
290+
return;
291+
}
279292
warmSocket = createSocket(
280293
getBackendURL(EVENTURL),
281294
[env.TRANSPORT],
@@ -832,7 +845,6 @@ export const connect = async (
832845
warmSocket ?? createSocket(endpoint, transports, session_token);
833846
warmSocket = null;
834847
cancelWarmup();
835-
socket.current.auth = bootAuth(true);
836848
socket.current.wait_connect = false;
837849
let firstConnect = true;
838850
let finishBootstrap;

‎tests/units/compiler/state_js.test.mjs‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ const source = await readFile(
1010
),
1111
"utf8",
1212
);
13+
const tokenKey = "reflex:ws://localhost:8000/_event:token";
1314

1415
/** Evaluate the actual frontend module with controlled transport and browser APIs. */
1516
async function setup({
@@ -26,11 +27,14 @@ async function setup({
2627
const updates = [];
2728
const window = new EventTarget();
2829
window.location = new URL("http://localhost:3000/page?query=value");
29-
const storage = new Map(storedToken ? [["token", storedToken]] : []);
30+
const storage = new Map(storedToken ? [[tokenKey, storedToken]] : []);
3031
window.sessionStorage = {
3132
getItem: (key) => storage.get(key),
3233
setItem: (key, value) => storage.set(key, value),
3334
};
35+
window.localStorage = {
36+
getItem: () => (storedToken ? String(Date.now() + 3600000) : undefined),
37+
};
3438
const document = new EventTarget();
3539
document.cookie = disabled ? "backend-enabled=false" : "";
3640
document.visibilityState = hidden ? "hidden" : "visible";
@@ -194,34 +198,31 @@ test("warm the transport without hydrating, then reuse it with every handler att
194198
assert.deepEqual(app.updates, [7]);
195199
assert.deepEqual(app.firstHydrates, [true]);
196200
assert.equal(app.timers.size, 0);
197-
assert.equal(app.window.sessionStorage.getItem("token"), "assigned-token");
201+
assert.equal(app.window.sessionStorage.getItem(tokenKey), "assigned-token");
198202
assert.equal(app.socket.current.auth.event.router_data.pathname, "/page");
199203
});
200204

201-
test("a new session's token is saved only once the mounted app connects", async () => {
205+
test("a cold session stores only the token assigned by the server", async () => {
202206
const app = await setup({ storedToken: null });
203207
app.flush();
204-
const warm = app.sockets[0];
205-
const token = warm.io.opts.query.token;
206-
assert.ok(token);
207-
// Waiting for the token in session storage must mean the app has mounted.
208-
assert.equal(app.window.sessionStorage.getItem("token"), undefined);
208+
assert.equal(app.sockets.length, 0);
209+
assert.equal(app.window.sessionStorage.getItem(tokenKey), undefined);
209210
const saved = [];
210211
const setItem = app.window.sessionStorage.setItem;
211212
app.window.sessionStorage.setItem = (key, value) => {
212213
saved.push(value);
213214
setItem(key, value);
214215
};
215216
await app.connect();
216-
assert.equal(app.socket.current, warm);
217-
assert.deepEqual(saved, [token, "assigned-token"]);
217+
assert.equal(app.socket.current.io.opts.query.token, "");
218+
assert.deepEqual(saved, ["assigned-token"]);
218219
});
219220

220221
test("a warm transport carrying another session's token is replaced", async () => {
221222
const app = await setup();
222223
app.flush();
223224
const warm = app.sockets[0];
224-
app.window.sessionStorage.setItem("token", "replaced-token");
225+
app.window.sessionStorage.setItem(tokenKey, "replaced-token");
225226
await app.connect();
226227
assert.equal(warm.disconnects, 1);
227228
assert.equal(app.sockets.length, 2);

‎tests/units/reflex_base/client_event_queue_runtime.mjs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { SourceTextModule, SyntheticModule } from "node:vm";
33
/** Evaluate the complete frontend module with isolated dependency stubs. */
44
export async function createQueueRuntime(
55
source,
6-
{ uploadFiles, throttle, io, env, browser } = {},
6+
{ uploadFiles, throttle, io, env, browser, initialState = {} } = {},
77
) {
88
const unused = () => {
99
throw new Error("Unexpected frontend dependency in queue test");
@@ -15,7 +15,7 @@ export async function createQueueRuntime(
1515
payload: first ? { hashes: { test_state: "compiled-defaults" } } : {},
1616
},
1717
],
18-
initialState: {},
18+
initialState,
1919
onLoadInternalEvent: unused,
2020
state_name: "test_state",
2121
exception_state_name: "test_exception_state",

‎tests/units/reflex_base/client_session.mjs‎

Lines changed: 35 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,12 +36,14 @@ async function browserSession({
3636
localStorage = storage(),
3737
sessionStorage = storage(),
3838
webLocks,
39+
warm = false,
3940
eventURL = "ws://localhost:8000/prefix/_event",
4041
} = {}) {
4142
const callbacks = new Map();
4243
const timeouts = new Map();
4344
const errors = [];
4445
const connections = [];
46+
const warmups = [];
4547
const requests = [];
4648
const emitted = [];
4749
const uploads = [];
@@ -76,15 +78,31 @@ async function browserSession({
7678
uploadFiles: (...args) => uploads.push(args),
7779
browser: {
7880
window,
79-
document: { addEventListener() {}, removeEventListener() {} },
81+
document: {
82+
cookie: "",
83+
visibilityState: "visible",
84+
addEventListener() {},
85+
removeEventListener() {},
86+
},
8087
localStorage,
8188
sessionStorage,
8289
},
83-
env: { EVENT: eventURL },
90+
env: { EVENT: eventURL, TRANSPORT: "websocket" },
91+
initialState: warm
92+
? { test_state: {}, test_exception_state: {} }
93+
: undefined,
8494
io: (_url, options) => {
8595
const current = {
8696
connected: false,
87-
io: { opts: options, encoder: {}, decoder: {} },
97+
io: {
98+
opts: options,
99+
encoder: {},
100+
decoder: {},
101+
open(callback) {
102+
warmups.push(options.query.token);
103+
callback();
104+
},
105+
},
88106
on: (event, callback) => callbacks.set(event, callback),
89107
emit: (event, payload) => emitted.push([event, payload]),
90108
connect() {
@@ -107,6 +125,7 @@ async function browserSession({
107125
callbacks,
108126
timeouts,
109127
connections,
128+
warmups,
110129
requests,
111130
emitted,
112131
uploads,
@@ -237,6 +256,7 @@ test("fresh hint reconnects directly without HTTP or a lock", async () => {
237256
await first.connect();
238257
await (await first.accept("existing-tab", "credential")).exchange;
239258
const tab = await browserSession({
259+
warm: true,
240260
localStorage: first.localStorage,
241261
sessionStorage: first.sessionStorage,
242262
webLocks: {
@@ -245,6 +265,9 @@ test("fresh hint reconnects directly without HTTP or a lock", async () => {
245265
},
246266
},
247267
});
268+
await tick();
269+
assert.deepEqual(tab.warmups, ["existing-tab"]);
270+
assert.deepEqual(tab.connections, []);
248271
await tab.connect();
249272
assert.deepEqual(tab.connections, ["existing-tab"]);
250273
assert.equal(tab.requests.length, 0);
@@ -256,6 +279,15 @@ test("fresh hint reconnects directly without HTTP or a lock", async () => {
256279
assert.equal(tab.requests.length, 0);
257280
});
258281

282+
test("cold transport cannot capture cookies before the session lock", async () => {
283+
const tab = await browserSession({ warm: true, webLocks: locks() });
284+
await tick();
285+
assert.deepEqual(tab.warmups, []);
286+
await tab.connect();
287+
assert.deepEqual(tab.connections, [""]);
288+
await (await tab.accept("bound-tab", "credential")).exchange;
289+
});
290+
259291
test("cookie-wins exchange adopts its token and reconnects the provisional socket", async () => {
260292
const tab = await browserSession({
261293
fetch: async () => ({

0 commit comments

Comments
 (0)