fix: websocket error

Author: GrinZero

packages/network-debugger/src/core/request.ts

@@ -4,7 +4,7 @@ import { RequestDetail } from '../common'
 import { getTimestamp } from '../utils'
 import { MainProcess } from './fork'
 import { BINARY_TYPES } from './ws/constants'
-import { Receiver } from './ws/reveiver'
+import { Receiver } from './ws/receiver'
 
 export interface RequestFn {
   (options: RequestOptions | string | URL, callback?: (res: IncomingMessage) => void): ClientRequest
@@ -38,14 +38,14 @@ function proxyClientRequestFactory(
   })
 
   if (requestDetail.isWebSocket()) {
-    actualRequest.on('upgrade', (res: IncomingMessage, socket: Socket, head: Buffer) => {
+    actualRequest.on('upgrade', async (res: IncomingMessage, socket: Socket, head: Buffer) => {
       const originalWrite = socket.write
 
       if (requestDetail.isHiden()) {
         return
       }
 
-      mainProcess.send({
+      await mainProcess.send({
         type: 'Network.webSocketCreated',
         data: {
           requestId: requestDetail.id,

packages/network-debugger/src/core/ws/reveiver.ts renamed to packages/network-debugger/src/core/ws/receiver.ts

@@ -17,11 +17,14 @@ const resetPluginContext = () => (currentPluginContext = null)
 export interface CoreCotext {
   devtool: DevtoolServer
   core: RequestCenter
+  plugins: PluginInstance<any>[]
 }
 
-export type EffectCleaner = () => void
-export type PluginHandler = (props: CoreCotext) => EffectCleaner | void
-export type PluginInstance = (props: CoreCotext) => EffectCleaner | void
+export type PluginHandler<T> = (props: CoreCotext) => T
+export type PluginInstance<T> = {
+  (props: CoreCotext): T
+  id: string
+}
 
 /**
  * @description create a plugin for devtool
@@ -48,13 +51,15 @@ export type PluginInstance = (props: CoreCotext) => EffectCleaner | void
  * ```
  * @returns PluginInstance
  */
-export const createPlugin = (fn: PluginHandler): PluginInstance => {
-  return (props: CoreCotext) => {
+export const createPlugin = <T>(id: string, fn: PluginHandler<T>) => {
+  const plugin: Omit<PluginInstance<T>, 'id'> = (props: CoreCotext) => {
     initPluginContext(props.core, props.devtool)
-    const effectDestroy = fn(props)
+    const output = fn(props)
     resetPluginContext()
-    return effectDestroy
+    return output
   }
+  const instance = Object.assign(plugin, { id })
+  return instance as PluginInstance<T>
 }
 
 /**

packages/network-debugger/src/fork/module/common.ts

@@ -11,7 +11,7 @@ export interface ScriptSourceData {
   scriptId: string
 }
 
-export const debuggerPlugin = createPlugin(({ devtool, core }) => {
+export const debuggerPlugin = createPlugin('debugger', ({ devtool, core }) => {
   useHandler<ScriptSourceData>('Debugger.getScriptSource', ({ id, data }) => {
     if (!id) {
       return

packages/network-debugger/src/fork/module/debugger/index.ts

@@ -8,7 +8,7 @@ import { toMimeType } from '../../devtool'
 const frameId = '517.528'
 const loaderId = '517.529'
 
-export const networkPlugin = createPlugin(({ devtool, core }) => {
+export const networkPlugin = createPlugin('network', ({ devtool, core }) => {
   const requests: Record<string, RequestDetail> = {}
 
   const getRequest = (id: string) => requests[id]
@@ -214,4 +214,10 @@ export const networkPlugin = createPlugin(({ devtool, core }) => {
       })
     }
   })
+
+  return {
+    getRequest
+  }
 })
+
+export type NetworkPluginCore = ReturnType<typeof networkPlugin>

packages/network-debugger/src/fork/module/network/index.ts

@@ -7,6 +7,7 @@ import {
   stringifyNestedObj
 } from '../../../utils'
 import { createPlugin, useHandler } from '../common'
+import { NetworkPluginCore } from '../network'
 
 export interface WebSocketFrameSent {
   requestId: string
@@ -29,10 +30,16 @@ export interface WebSocketCreated {
   response: IncomingMessage
 }
 
-export const websocketPlugin = createPlugin(({ devtool }) => {
+export const websocketPlugin = createPlugin('websocket', ({ devtool, core }) => {
+  const networkPlugin = core.usePlugin<NetworkPluginCore>('network')
+
   useHandler<WebSocketCreated>(
     'Network.webSocketCreated',
-    async ({ request, data: { response } }) => {
+    async ({ data: { response, requestId } }) => {
+      if (!requestId) {
+        return
+      }
+      const request = networkPlugin.getRequest(requestId)
       if (!request) {
         return
       }
@@ -99,8 +106,8 @@ export const websocketPlugin = createPlugin(({ devtool }) => {
     })
   })
 
-  useHandler<WebSocketFrameSent>('Network.webSocketFrameReceived', async ({ request, data }) => {
-    if (!request) {
+  useHandler<WebSocketFrameSent>('Network.webSocketFrameReceived', async ({ data }) => {
+    if (!data.requestId) {
       return
     }
     await devtool.send({
@@ -113,14 +120,14 @@ export const websocketPlugin = createPlugin(({ devtool }) => {
     })
   })
 
-  useHandler<WebSocketFrameSent>('Network.webSocketClosed', async ({ request }) => {
-    if (!request) {
+  useHandler<WebSocketFrameSent>('Network.webSocketClosed', async ({ data }) => {
+    if (!data.requestId) {
       return
     }
     await devtool.send({
       method: 'Network.webSocketClosed',
       params: {
-        requestId: request.id,
+        requestId: data.requestId,
         timestamp: getTimestamp()
       }
     })

packages/network-debugger/src/fork/module/websocket/index.ts

@@ -1,10 +1,9 @@
 import { DevtoolServer } from './devtool'
 import { PORT, READY_MESSAGE, RequestDetail } from '../common'
-import zlib from 'node:zlib'
 import { Server } from 'ws'
 import { log } from '../utils'
 import { ResourceService } from './resource-service'
-import { EffectCleaner, PluginInstance } from './module/common'
+import { PluginInstance } from './module/common'
 
 export interface RequestCenterInitOptions {
   port: number
@@ -26,7 +25,6 @@ export class RequestCenter {
   public resourceService: ResourceService
   private devtool: DevtoolServer
   private server: Server
-  private effects: Array<EffectCleaner> = []
   private listeners: Record<string, Set<DevtoolMessageListener> | undefined> = {}
   private options: RequestCenterInitOptions
   constructor(options: RequestCenterInitOptions) {
@@ -67,18 +65,29 @@ export class RequestCenter {
     this.server = this.initServer()
   }
 
-  #plugins: PluginInstance[] = []
-  public loadPlugins(plugins: PluginInstance[]) {
+  #plugins: PluginInstance<any>[] = []
+  #pluginOutputs?: Map<string, any>
+  public loadPlugins(plugins: PluginInstance<any>[]) {
     this.#plugins = plugins
-    const effects = plugins
-      .map((plugin) =>
-        plugin({
-          devtool: this.devtool,
-          core: this
-        })
-      )
-      .filter(Boolean) as Array<EffectCleaner>
-    this.effects.push(...effects)
+    this.#pluginOutputs = new Map()
+
+    plugins.forEach((plugin) => {
+      const output = plugin({
+        devtool: this.devtool,
+        core: this,
+        plugins: this.#plugins
+      })
+      this.#pluginOutputs!.set(plugin.id, output)
+    })
+  }
+
+  public usePlugin<T = null>(id: string) {
+    if (!this.#pluginOutputs) {
+      return null as T
+    }
+
+    const output = this.#pluginOutputs.get(id)
+    return output as T
   }
 
   public on(method: string, listener: DevtoolMessageListener) {
@@ -94,7 +103,6 @@ export class RequestCenter {
   public close() {
     this.server.close()
     this.devtool.close()
-    this.effects.forEach((effect) => effect())
   }
 
   private initServer() {