diff --git a/src/server/chromium.ts b/src/server/chromium.ts index e7690a9713..a11a6d8cc9 100644 --- a/src/server/chromium.ts +++ b/src/server/chromium.ts @@ -208,16 +208,42 @@ export class Chromium extends AbstractBrowserType { } } +type SessionData = { + socket: ws, + children: Set, + isBrowserSession: boolean, + parent?: string, +}; + function wrapTransportWithWebSocket(transport: ConnectionTransport, logger: InnerLogger, port: number): WebSocketWrapper { const server = new ws.Server({ port }); const guid = helper.guid(); const awaitingBrowserTarget = new Map(); - const sessionToSocket = new Map(); + const sessionToData = new Map(); const socketToBrowserSession = new Map(); - const browserSessions = new Set(); let lastSequenceNumber = 1; + function addSession(sessionId: string, socket: ws, parentSessionId?: string) { + sessionToData.set(sessionId, { + socket, + children: new Set(), + isBrowserSession: !parentSessionId, + parent: parentSessionId + }); + if (parentSessionId) + sessionToData.get(parentSessionId)!.children.add(sessionId); + } + + function removeSession(sessionId: string) { + const data = sessionToData.get(sessionId)!; + for (const child of data.children) + removeSession(child); + if (data.parent) + sessionToData.get(data.parent)!.children.delete(sessionId); + sessionToData.delete(sessionId); + } + transport.onmessage = message => { if (typeof message.id === 'number' && awaitingBrowserTarget.has(message.id)) { const freshSocket = awaitingBrowserTarget.get(message.id)!; @@ -225,14 +251,13 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, logger: Inne const sessionId = message.result.sessionId; if (freshSocket.readyState !== ws.CLOSED && freshSocket.readyState !== ws.CLOSING) { - sessionToSocket.set(sessionId, freshSocket); const { queue } = socketToBrowserSession.get(freshSocket)!; for (const item of queue!) { item.sessionId = sessionId; transport.send(item); } socketToBrowserSession.set(freshSocket, { sessionId }); - browserSessions.add(sessionId); + addSession(sessionId, freshSocket); } else { transport.send({ id: ++lastSequenceNumber, @@ -248,16 +273,16 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, logger: Inne if (!message.sessionId) return; - const socket = sessionToSocket.get(message.sessionId); - if (socket && socket.readyState !== ws.CLOSING) { + const data = sessionToData.get(message.sessionId); + if (data && data.socket.readyState !== ws.CLOSING) { if (message.method === 'Target.attachedToTarget') - sessionToSocket.set(message.params.sessionId, socket); + addSession(message.params.sessionId, data.socket, message.sessionId); if (message.method === 'Target.detachedFromTarget') - sessionToSocket.delete(message.params.sessionId); + removeSession(message.params.sessionId); // Strip session ids from the browser sessions. - if (browserSessions.has(message.sessionId)) + if (data.isBrowserSession) delete message.sessionId; - socket.send(JSON.stringify(message)); + data.socket.send(JSON.stringify(message)); } }; @@ -311,8 +336,7 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, logger: Inne const session = socketToBrowserSession.get(socket); if (!session || !session.sessionId) return; - sessionToSocket.delete(session.sessionId); - browserSessions.delete(session.sessionId); + removeSession(session.sessionId); socketToBrowserSession.delete(socket); transport.send({ id: ++lastSequenceNumber, @@ -324,7 +348,7 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, logger: Inne const address = server.address(); const wsEndpoint = typeof address === 'string' ? `${address}/${guid}` : `ws://127.0.0.1:${address.port}/${guid}`; - return new WebSocketWrapper(wsEndpoint, [awaitingBrowserTarget, sessionToSocket, socketToBrowserSession, browserSessions]); + return new WebSocketWrapper(wsEndpoint, [awaitingBrowserTarget, sessionToData, socketToBrowserSession]); }