Skip to content

Commit fb85436

Browse files
shadowusrshadowusr
committed
fix: correctly handle __Host and __Secure cookies when restoring state
1 parent 0d06d9f commit fb85436

2 files changed

Lines changed: 245 additions & 6 deletions

File tree

‎src/browser/commands/restoreState/index.ts‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,30 @@ import type { Browser } from "../../types";
66
import { DEVTOOLS_PROTOCOL, WEBDRIVER_PROTOCOL } from "../../../constants/config";
77
import { getOverridesProtocol, getWebdriverFrames, SaveStateData } from "../saveState";
88
import { getActivePuppeteerPage } from "../../existing-browser";
9-
import { Cookie } from "@testplane/wdio-protocols";
109
import { StateOpts } from "../../../config/types";
10+
import type { Cookie as WebdriverCookie } from "@testplane/wdio-protocols";
11+
import type { Cookie } from "../../../types";
1112

1213
export type RestoreStateOptions = Omit<StateOpts, "keepFile"> & {
1314
data?: SaveStateData;
1415
refresh?: boolean;
1516
};
1617

18+
const normalizeCookiePrefixConstraints = (cookie: Cookie): Cookie => {
19+
if (cookie.name.startsWith("__Host-")) {
20+
const cookieWithoutDomain = { ...cookie };
21+
delete cookieWithoutDomain.domain;
22+
23+
return { ...cookieWithoutDomain, path: "/", secure: true };
24+
}
25+
26+
if (cookie.name.startsWith("__Secure-")) {
27+
return { ...cookie, secure: true };
28+
}
29+
30+
return cookie;
31+
};
32+
1733
export default (browser: Browser): void => {
1834
const { publicAPI: session } = browser;
1935

@@ -40,13 +56,15 @@ export default (browser: Browser): void => {
4056
restoreState.cookies = restoreState?.cookies.filter(options.cookieFilter);
4157
}
4258

59+
const cookies = restoreState.cookies?.map(normalizeCookiePrefixConstraints);
60+
4361
switch (getOverridesProtocol(browser)) {
4462
case WEBDRIVER_PROTOCOL: {
4563
await session.switchToParentFrame();
4664

47-
if (restoreState.cookies && options.cookies) {
65+
if (cookies && options.cookies) {
4866
await session.setCookies(
49-
restoreState.cookies.map(
67+
cookies.map(
5068
cookie =>
5169
({
5270
...cookie,
@@ -55,7 +73,7 @@ export default (browser: Browser): void => {
5573
cookie.sameSite && session.isBidi
5674
? cookie.sameSite.toLowerCase()
5775
: cookie.sameSite,
58-
} as Cookie),
76+
} as WebdriverCookie),
5977
),
6078
);
6179
}
@@ -114,8 +132,8 @@ export default (browser: Browser): void => {
114132

115133
const frames = page.frames();
116134

117-
if (restoreState.cookies && options.cookies) {
118-
await page.setCookie(...restoreState.cookies);
135+
if (cookies && options.cookies) {
136+
await page.setCookie(...cookies);
119137
}
120138

121139
for (const frame of frames) {

‎test/src/browser/commands/restoreState.js‎

Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
1+
"use strict";
2+
3+
const restoreStateCommand = require("src/browser/commands/restoreState").default;
4+
const { DEVTOOLS_PROTOCOL, WEBDRIVER_PROTOCOL } = require("src/constants/config");
5+
const { mkSessionStub_ } = require("../utils");
6+
7+
describe('"restoreState" command', () => {
8+
const sandbox = sinon.createSandbox();
9+
10+
const mkBrowser_ = ({ session, automationProtocol = WEBDRIVER_PROTOCOL } = {}) => ({
11+
publicAPI: session,
12+
config: {
13+
automationProtocol,
14+
isolation: false,
15+
stateOpts: {
16+
cookies: true,
17+
localStorage: true,
18+
sessionStorage: true,
19+
},
20+
},
21+
});
22+
23+
const initWebdriverSession_ = ({ isBidi = false } = {}) => {
24+
const session = mkSessionStub_();
25+
session.getUrl.resolves("https://example.com/page");
26+
session.isBidi = isBidi;
27+
session.setCookies = sandbox.stub().named("setCookies").resolves();
28+
session.refresh = sandbox.stub().named("refresh").resolves();
29+
30+
restoreStateCommand(mkBrowser_({ session }));
31+
32+
return session;
33+
};
34+
35+
const initDevtoolsSession_ = () => {
36+
const session = mkSessionStub_();
37+
const page = {
38+
setCookie: sandbox.stub().named("setCookie").resolves(),
39+
frames: sandbox.stub().named("frames").returns([]),
40+
reload: sandbox.stub().named("reload").resolves(),
41+
target: sandbox.stub().named("target").returns({ _targetId: "active-target" }),
42+
};
43+
44+
session.getUrl.resolves("https://example.com/page");
45+
session.getWindowHandle = sandbox.stub().named("getWindowHandle").resolves("active-target");
46+
session.getPuppeteer.resolves({
47+
pages: sandbox.stub().named("pages").resolves([page]),
48+
});
49+
50+
restoreStateCommand(mkBrowser_({ session, automationProtocol: DEVTOOLS_PROTOCOL }));
51+
52+
return { page, session };
53+
};
54+
55+
afterEach(() => sandbox.restore());
56+
57+
it("should normalize cookie prefix constraints before webdriver restore", async () => {
58+
const session = initWebdriverSession_();
59+
const hostCookie = {
60+
name: "__Host-csrf-token",
61+
value: "host-value",
62+
domain: "www.chromatic.com",
63+
path: "/custom",
64+
secure: false,
65+
httpOnly: true,
66+
sameSite: "Lax",
67+
expires: 12345,
68+
};
69+
const secureCookie = {
70+
name: "__Secure-session",
71+
value: "secure-value",
72+
domain: "www.chromatic.com",
73+
path: "/auth",
74+
secure: false,
75+
sameSite: "Strict",
76+
};
77+
const ordinaryCookie = {
78+
name: "regular",
79+
value: "regular-value",
80+
domain: "www.chromatic.com",
81+
path: "/custom",
82+
secure: false,
83+
httpOnly: true,
84+
sameSite: "Lax",
85+
};
86+
87+
await session.restoreState({
88+
data: { cookies: [hostCookie, secureCookie, ordinaryCookie] },
89+
cookies: true,
90+
refresh: false,
91+
});
92+
93+
assert.calledOnce(session.setCookies);
94+
assert.deepEqual(session.setCookies.firstCall.args[0], [
95+
{
96+
name: "__Host-csrf-token",
97+
value: "host-value",
98+
path: "/",
99+
secure: true,
100+
httpOnly: true,
101+
sameSite: "Lax",
102+
expires: 12345,
103+
},
104+
{
105+
name: "__Secure-session",
106+
value: "secure-value",
107+
domain: "www.chromatic.com",
108+
path: "/auth",
109+
secure: true,
110+
sameSite: "Strict",
111+
},
112+
ordinaryCookie,
113+
]);
114+
assert.property(hostCookie, "domain");
115+
});
116+
117+
it("should keep webdriver sameSite/BiDi normalization", async () => {
118+
const session = initWebdriverSession_({ isBidi: true });
119+
120+
await session.restoreState({
121+
data: {
122+
cookies: [
123+
{
124+
name: "same-site-none",
125+
value: "value",
126+
secure: false,
127+
sameSite: "None",
128+
},
129+
],
130+
},
131+
cookies: true,
132+
refresh: false,
133+
});
134+
135+
assert.calledOnceWith(session.setCookies, [
136+
{
137+
name: "same-site-none",
138+
value: "value",
139+
secure: true,
140+
sameSite: "none",
141+
},
142+
]);
143+
});
144+
145+
it("should normalize cookie prefix constraints before devtools restore", async () => {
146+
const { page, session } = initDevtoolsSession_();
147+
const ordinaryCookie = {
148+
name: "regular",
149+
value: "regular-value",
150+
domain: "www.chromatic.com",
151+
path: "/custom",
152+
secure: false,
153+
sameSite: "Lax",
154+
};
155+
156+
await session.restoreState({
157+
data: {
158+
cookies: [
159+
{
160+
name: "__Host-csrf-token",
161+
value: "host-value",
162+
domain: "www.chromatic.com",
163+
path: "/custom",
164+
secure: false,
165+
},
166+
{
167+
name: "__Secure-session",
168+
value: "secure-value",
169+
domain: "www.chromatic.com",
170+
path: "/auth",
171+
secure: false,
172+
},
173+
ordinaryCookie,
174+
],
175+
},
176+
cookies: true,
177+
refresh: false,
178+
});
179+
180+
assert.calledOnce(page.setCookie);
181+
assert.deepEqual(page.setCookie.firstCall.args, [
182+
{
183+
name: "__Host-csrf-token",
184+
value: "host-value",
185+
path: "/",
186+
secure: true,
187+
},
188+
{
189+
name: "__Secure-session",
190+
value: "secure-value",
191+
domain: "www.chromatic.com",
192+
path: "/auth",
193+
secure: true,
194+
},
195+
ordinaryCookie,
196+
]);
197+
});
198+
199+
it("should not ignore unrelated cookie restore errors", async () => {
200+
const session = initWebdriverSession_();
201+
const error = new Error("unable to set cookie");
202+
203+
session.setCookies.rejects(error);
204+
205+
await assert.isRejected(
206+
session.restoreState({
207+
data: {
208+
cookies: [
209+
{
210+
name: "invalid-cookie",
211+
value: "value",
212+
},
213+
],
214+
},
215+
cookies: true,
216+
refresh: false,
217+
}),
218+
/unable to set cookie/,
219+
);
220+
});
221+
});

0 commit comments

Comments
 (0)