Skip to content

Commit a4f74d8

Browse files
lemon-mintlemon-mint
committed
feat: enhance WebSocket handling and improve rewriter functionality
1 parent d001b16 commit a4f74d8

5 files changed

Lines changed: 240 additions & 47 deletions

File tree

‎cmd/wasm-kernel/main.go‎

Lines changed: 26 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -259,29 +259,9 @@ func (k *Kernel) tabFromValues(tabID, keyB64 string) *zphttp.TabState {
259259

260260
func newJSWebSocketStream(ctx context.Context, cancel context.CancelFunc, conn *wsproto.Conn) js.Value {
261261
handlers := js.Value{}
262+
var start sync.Once
262263
obj := js.Global().Get("Object").New()
263-
obj.Set("setHandlers", js.FuncOf(func(this js.Value, args []js.Value) any {
264-
if len(args) > 0 {
265-
handlers = args[0]
266-
}
267-
return nil
268-
}))
269-
obj.Set("send", js.FuncOf(func(this js.Value, args []js.Value) any {
270-
if len(args) == 0 {
271-
return nil
272-
}
273-
data, binary := jsPayload(args[0])
274-
op := byte(wsproto.OpText)
275-
if binary {
276-
op = wsproto.OpBinary
277-
}
278-
if err := conn.WriteFrame(op, data); err != nil && handlers.Truthy() {
279-
callHandler(handlers, "error", jsError("TARGET_CONNECT_FAILED"))
280-
}
281-
return nil
282-
}))
283-
obj.Set("close", js.FuncOf(func(this js.Value, args []js.Value) any { cancel(); _ = conn.Close(); return nil }))
284-
go func() {
264+
readLoop := func() {
285265
defer cancel()
286266
defer conn.Close()
287267
for {
@@ -292,9 +272,6 @@ func newJSWebSocketStream(ctx context.Context, cancel context.CancelFunc, conn *
292272
}
293273
return
294274
}
295-
if !handlers.Truthy() {
296-
continue
297-
}
298275
if op == wsproto.OpClose {
299276
callHandler(handlers, "close", js.Null())
300277
return
@@ -307,7 +284,29 @@ func newJSWebSocketStream(ctx context.Context, cancel context.CancelFunc, conn *
307284
js.CopyBytesToJS(arr, payload)
308285
callHandler(handlers, "message", arr.Get("buffer"))
309286
}
310-
}()
287+
}
288+
obj.Set("setHandlers", js.FuncOf(func(this js.Value, args []js.Value) any {
289+
if len(args) > 0 {
290+
handlers = args[0]
291+
}
292+
start.Do(func() { go readLoop() })
293+
return nil
294+
}))
295+
obj.Set("send", js.FuncOf(func(this js.Value, args []js.Value) any {
296+
if len(args) == 0 {
297+
return nil
298+
}
299+
data, binary := jsPayload(args[0])
300+
op := byte(wsproto.OpText)
301+
if binary {
302+
op = wsproto.OpBinary
303+
}
304+
if err := conn.WriteFrame(op, data); err != nil && handlers.Truthy() {
305+
callHandler(handlers, "error", jsError("TARGET_CONNECT_FAILED"))
306+
}
307+
return nil
308+
}))
309+
obj.Set("close", js.FuncOf(func(this js.Value, args []js.Value) any { cancel(); _ = conn.Close(); return nil }))
311310
return obj
312311
}
313312

@@ -346,7 +345,7 @@ func jsPayload(v js.Value) ([]byte, bool) {
346345
return b, true
347346
}
348347
if v.Get("buffer").Truthy() {
349-
arr := js.Global().Get("Uint8Array").New(v.Get("buffer"))
348+
arr := js.Global().Get("Uint8Array").New(v.Get("buffer"), v.Get("byteOffset"), v.Get("byteLength"))
350349
b := make([]byte, arr.Get("byteLength").Int())
351350
js.CopyBytesToGo(b, arr)
352351
return b, true

‎test/e2e/proxy.test.js‎

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ function createTargetServer(requests) {
103103
res.end(`<!doctype html><html><head><title>E2E Home</title></head><body>
104104
<main><h1>E2E Home</h1><a id="next" href="/next">Next page</a></main>
105105
<script>window.__ua = navigator.userAgent; window.__platform = navigator.platform; window.__phase2Location = { href: location.href, windowHref: window.location.href }; try { Function('return location.href')(); window.__phase2FunctionBlocked = ''; } catch (err) { window.__phase2FunctionBlocked = err && err.message || String(err); }</script>
106+
<script src="/rewrite-fixture.js"></script>
106107
</body></html>`);
107108
return;
108109
}
@@ -114,6 +115,24 @@ function createTargetServer(requests) {
114115
</body></html>`);
115116
return;
116117
}
118+
if (url.pathname === '/rewrite-fixture.js') {
119+
res.writeHead(200, { 'Content-Type': 'text/javascript; charset=utf-8', 'Cache-Control': 'no-store' });
120+
res.end(`(() => {
121+
const NativeWebSocket = window.WebSocket;
122+
window.__rewriteAdvanced = { initialHref: window.location.href, constructorSource: NativeWebSocket.toString() };
123+
const ws = new NativeWebSocket('/ws', ['zp-rewrite']);
124+
ws.binaryType = 'arraybuffer';
125+
ws.onopen = () => ws.send('rewrite-script');
126+
ws.onmessage = ev => {
127+
window.__rewriteAdvanced.wsURL = ws.url;
128+
window.__rewriteAdvanced.wsProtocol = ws.protocol;
129+
window.__rewriteAdvanced.wsMessage = String(ev.data);
130+
ws.close(1000, 'done');
131+
};
132+
ws.onerror = () => { window.__rewriteAdvanced.wsError = true; };
133+
})();`);
134+
return;
135+
}
117136
if (url.pathname === '/set-cookie') {
118137
res.writeHead(200, {
119138
'Content-Type': 'text/plain; charset=utf-8',
@@ -161,7 +180,7 @@ function createTargetServer(requests) {
161180
}
162181

163182
function handleWebSocketUpgrade(req, socket, requests) {
164-
requests.push({ url: req.url, method: req.method, host: req.headers.host || '', userAgent: req.headers['user-agent'] || '', cookie: req.headers.cookie || '', upgrade: true });
183+
requests.push({ url: req.url, method: req.method, host: req.headers.host || '', userAgent: req.headers['user-agent'] || '', cookie: req.headers.cookie || '', protocol: req.headers['sec-websocket-protocol'] || '', upgrade: true });
165184
if (new URL(req.url, 'http://target.local').pathname !== '/ws') {
166185
socket.end('HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n');
167186
return;
@@ -361,6 +380,14 @@ test('browser traffic uses test SOCKS5 and covers proxied runtime integrations',
361380
assert.deepEqual(home.phase2Location, { href: `http://e2e.test:${targetPort}/`, windowHref: `http://e2e.test:${targetPort}/` });
362381
assert.equal(home.phase2FunctionBlocked, 'Blocked by ZeroProxy policy');
363382
assert.ok(requests.some(r => r.url === '/' && r.userAgent === TARGET_UA), `target requests: ${JSON.stringify(requests)}`);
383+
await page.waitForFunction(() => window.__rewriteAdvanced && window.__rewriteAdvanced.wsMessage === 'echo:rewrite-script', { timeout: 30000 });
384+
const rewriteAdvanced = await page.evaluate(() => window.__rewriteAdvanced);
385+
assert.equal(rewriteAdvanced.initialHref, `http://e2e.test:${targetPort}/`);
386+
assert.equal(rewriteAdvanced.wsURL, `ws://e2e.test:${targetPort}/ws`);
387+
assert.equal(rewriteAdvanced.wsProtocol, 'zp-rewrite');
388+
assert.equal(rewriteAdvanced.wsMessage, 'echo:rewrite-script');
389+
assert.equal(rewriteAdvanced.wsError, undefined);
390+
assert.ok(requests.some(r => r.upgrade && r.url === '/ws' && r.protocol === 'zp-rewrite' && r.userAgent === TARGET_UA), `target requests: ${JSON.stringify(requests)}`);
364391

365392
const iframeIsolation = await page.evaluate(async target => {
366393
const blockedByPolicy = fn => {

‎test/js/rewriter.test.js‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,3 +46,44 @@ test('OXC rewriter blocks constructor escape compound writes', async () => {
4646
assert.match(out.code, /Blocked by ZeroProxy rewrite policy/);
4747
assert.match(out.code, /__zp_call\(__zp_get\(\(\{\}\),\"constructor\"\),\"constructor\"/);
4848
});
49+
50+
test('OXC rewriter routes location assignments and WebSocket construction through helpers', async () => {
51+
const rewriter = await loadRewriter();
52+
const out = rewriter.rewriteScript(`
53+
const assigned = (window.location = "https://google.com/");
54+
location.href = window.location.href + "#frag";
55+
const ws = new WebSocket("ws://example.test/socket", ["chat"]);
56+
const ws2 = new window.WebSocket("wss://example.test/secure");
57+
window.result = { assigned, href: location.href, wsURL: ws.url, ws2URL: ws2.url };
58+
`, { kind: 'classic' });
59+
assert.equal(out.ok, true, JSON.stringify(out.diagnostics));
60+
assert.match(out.code, /__zp_set\(__zp_get\(globalThis,"window"\),"location","https:\/\/google\.com\/"\)/);
61+
assert.match(out.code, /__zp_set\(__zp_get\(globalThis,"location"\),"href",__zp_get\(__zp_get\(globalThis,"window"\),"location"\)\.href \+ "#frag"\)/);
62+
assert.match(out.code, /__zp_construct\(__zp_get\(globalThis,"WebSocket"\),\["ws:\/\/example\.test\/socket",\["chat"\]\]\)/);
63+
assert.match(out.code, /__zp_construct\(__zp_get\(globalThis,"window"\)\.WebSocket,\["wss:\/\/example\.test\/secure"\]\)/);
64+
65+
const loc = { href: 'https://origin.test/start' };
66+
function FakeWebSocket(url, protocols) {
67+
this.url = url;
68+
this.protocols = protocols;
69+
}
70+
const ctx = {
71+
location: loc,
72+
window: { location: loc, WebSocket: FakeWebSocket },
73+
WebSocket: FakeWebSocket,
74+
};
75+
ctx.globalThis = ctx;
76+
ctx.__zp_get = (base, prop) => base[prop];
77+
ctx.__zp_set = (base, prop, value) => {
78+
if ((base === ctx.window && prop === 'location') || (base === loc && prop === 'href')) loc.href = String(value);
79+
else base[prop] = value;
80+
return value;
81+
};
82+
ctx.__zp_construct = (ctor, args) => new ctor(...args);
83+
vm.runInNewContext(out.code, ctx);
84+
assert.equal(loc.href, 'https://google.com/#frag');
85+
assert.equal(ctx.window.result.assigned, 'https://google.com/');
86+
assert.equal(ctx.window.result.href, 'https://google.com/#frag');
87+
assert.equal(ctx.window.result.wsURL, 'ws://example.test/socket');
88+
assert.equal(ctx.window.result.ws2URL, 'wss://example.test/secure');
89+
});

‎web/js-rewriter.js‎

Lines changed: 71 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
const VERSION = 'phase2-oxc-abi-2';
77
const BLOCK_CODE = "throw new DOMException('Blocked by ZeroProxy rewrite policy','NotSupportedError');";
8-
const GLOBALS = new Set(['window', 'self', 'globalThis', 'location', 'document', 'history', 'top', 'parent', 'opener', 'frames', 'eval', 'Function', 'AsyncFunction', 'GeneratorFunction', 'AsyncGeneratorFunction']);
8+
const GLOBALS = new Set(['window', 'self', 'globalThis', 'location', 'document', 'history', 'top', 'parent', 'opener', 'frames', 'WebSocket', 'eval', 'Function', 'AsyncFunction', 'GeneratorFunction', 'AsyncGeneratorFunction']);
99
const MEMBER_HELPER_PROPS = new Set(['location', 'defaultView', 'contentWindow', 'contentDocument', 'top', 'parent', 'opener', 'frames', 'constructor']);
1010
const CALL_HELPER_PROPS = new Set(['assign', 'replace', 'open', 'get', 'getOwnPropertyDescriptor', 'defineProperty']);
1111
let parser = null;
@@ -159,6 +159,60 @@
159159
renderedCache.set(node, out);
160160
return out;
161161
}
162+
function exprCode(node) {
163+
if (!node) return '';
164+
if (node.type === 'Identifier' || node.type === 'MemberExpression' || node.type === 'OptionalMemberExpression' || node.type === 'ChainExpression') return render(node);
165+
if (node.type === 'BinaryExpression' || node.type === 'LogicalExpression') return exprCode(node.left) + source.slice(node.left.end, node.right.start) + exprCode(node.right);
166+
if (node.type === 'ConditionalExpression') return exprCode(node.test) + source.slice(node.test.end, node.consequent.start) + exprCode(node.consequent) + source.slice(node.consequent.end, node.alternate.start) + exprCode(node.alternate);
167+
if (node.type === 'UnaryExpression' || node.type === 'UpdateExpression') {
168+
if (node.prefix) return source.slice(node.start, node.argument.start) + exprCode(node.argument);
169+
return exprCode(node.argument) + source.slice(node.argument.end, node.end);
170+
}
171+
return src(node);
172+
}
173+
174+
function isVirtualWindowExpr(node) {
175+
if (!node) return false;
176+
if (node.type === 'ChainExpression') return isVirtualWindowExpr(node.expression);
177+
if (node.type === 'Identifier') return isGlobalIdentifier(node) && (node.name === 'window' || node.name === 'self' || node.name === 'globalThis' || node.name === 'top' || node.name === 'parent' || node.name === 'opener' || node.name === 'frames');
178+
if (node.type !== 'MemberExpression' && node.type !== 'OptionalMemberExpression') return false;
179+
const name = propName(node.property, node.computed);
180+
if ((name === 'defaultView' || name === 'contentWindow') && MEMBER_HELPER_PROPS.has(name)) return true;
181+
return (name === 'window' || name === 'self' || name === 'globalThis' || name === 'top' || name === 'parent' || name === 'opener' || name === 'frames') && isVirtualWindowExpr(node.object);
182+
}
183+
184+
function isVirtualLocationExpr(node) {
185+
if (!node) return false;
186+
if (node.type === 'ChainExpression') return isVirtualLocationExpr(node.expression);
187+
if (node.type === 'Identifier') return isGlobalIdentifier(node) && node.name === 'location';
188+
if (node.type !== 'MemberExpression' && node.type !== 'OptionalMemberExpression') return false;
189+
const name = propName(node.property, node.computed);
190+
return name === 'location' && isVirtualWindowExpr(node.object) || node.computed && isVirtualWindowExpr(node.object);
191+
}
192+
193+
function assignmentSetTarget(node) {
194+
if (!node) return null;
195+
if (node.type === 'ChainExpression') return assignmentSetTarget(node.expression);
196+
if (node.type === 'Identifier' && isGlobalIdentifier(node) && (node.name === 'location' || node.name === 'window')) return { base: 'globalThis', prop: JSON.stringify(node.name) };
197+
if (node.type !== 'MemberExpression' && node.type !== 'OptionalMemberExpression') return null;
198+
const name = propName(node.property, node.computed);
199+
if (name === 'location' && isVirtualWindowExpr(node.object)) return { base: render(node.object), prop: propCode(node.property, node.computed) };
200+
if (name === 'href' && isVirtualLocationExpr(node.object)) return { base: render(node.object), prop: propCode(node.property, node.computed) };
201+
if (node.computed && (isVirtualWindowExpr(node.object) || isVirtualLocationExpr(node.object))) return { base: render(node.object), prop: propCode(node.property, node.computed) };
202+
return null;
203+
}
204+
function argList(args) {
205+
return (args || []).map(exprCode).join(',');
206+
}
207+
208+
function constructTarget(node) {
209+
if (!node) return '';
210+
if (node.type === 'ChainExpression') return constructTarget(node.expression);
211+
if (node.type === 'Identifier' && isGlobalIdentifier(node)) return render(node);
212+
if ((node.type === 'MemberExpression' || node.type === 'OptionalMemberExpression') && isVirtualWindowExpr(node.object)) return render(node);
213+
return '';
214+
}
215+
162216

163217
function enterNode(node, parent, key) {
164218
if (!node || typeof node.type !== 'string') return false;
@@ -198,9 +252,12 @@
198252
diagnostics.push({ level: 'warning', message: 'blocked compound assignment touching virtualized browser state', start: node.start, end: node.end });
199253
return true;
200254
}
201-
if (node.left && node.left.type === 'Identifier' && isGlobalIdentifier(node.left) && (node.left.name === 'location' || node.left.name === 'window')) {
202-
addReplacement(node, '__zp_set(globalThis,' + JSON.stringify(node.left.name) + ',' + src(node.right) + ')', 100);
203-
return true;
255+
if (node.operator === '=') {
256+
const target = assignmentSetTarget(node.left);
257+
if (target) {
258+
addReplacement(node, '__zp_set(' + target.base + ',' + target.prop + ',' + exprCode(node.right) + ')', 100);
259+
return true;
260+
}
204261
}
205262
break;
206263
}
@@ -215,16 +272,23 @@
215272
case 'CallExpression':
216273
case 'NewExpression': {
217274
const callee = node.callee;
275+
const args = argList(node.arguments);
218276
if (callee && (callee.type === 'MemberExpression' || callee.type === 'OptionalMemberExpression')) {
219277
const name = propName(callee.property, callee.computed);
220278
if (CALL_HELPER_PROPS.has(name) || MEMBER_HELPER_PROPS.has(name)) {
221-
const args = (node.arguments || []).map(a => src(a)).join(',');
222279
const call = '__zp_call(' + render(callee.object) + ',' + propCode(callee.property, callee.computed) + ',[' + args + '])';
223280
addReplacement(node, node.type === 'NewExpression' ? BLOCK_CODE : call, 90);
224281
if (node.type === 'NewExpression') diagnostics.push({ level: 'warning', message: 'blocked construction through virtualized browser state', start: node.start, end: node.end });
225282
return true;
226283
}
227284
}
285+
if (node.type === 'NewExpression') {
286+
const ctor = constructTarget(callee);
287+
if (ctor) {
288+
addReplacement(node, '__zp_construct(' + ctor + ',[' + args + '])', 90);
289+
return true;
290+
}
291+
}
228292
break;
229293
}
230294
case 'MemberExpression':
@@ -249,8 +313,8 @@
249313

250314
function containsDangerousLHS(node) {
251315
if (!node) return false;
252-
if (node.type === 'Identifier') return isGlobalIdentifier(node) && (node.name === 'location' || node.name === 'window');
253-
if (node.type === 'MemberExpression' || node.type === 'OptionalMemberExpression') return MEMBER_HELPER_PROPS.has(propName(node.property, node.computed)) || containsDangerousLHS(node.object);
316+
if (assignmentSetTarget(node)) return true;
317+
if (node.type === 'MemberExpression' || node.type === 'OptionalMemberExpression') return containsDangerousLHS(node.object);
254318
return false;
255319
}
256320

0 commit comments

Comments
 (0)