refactor vless for node.js (#93)

refactor vless code

Co-authored-by: zizifn <515868058@qq.com>

libs/vless-js/src/lib/vless-js.ts

@@ -19,7 +19,7 @@ export async function processWebSocket({
   libs: { uuid: any; lodash: any };
 }) {
   let address = '';
-  let port = 0;
+  let portWithRandomLog = '';
   let remoteConnection: {
     readable: any;
     writable: any;
@@ -29,7 +29,7 @@ export async function processWebSocket({
   let remoteConnectionReadyResolve: Function;
   try {
     const log = (info: string, event?: any) => {
-      console.log(`[${address}:${port}] ${info}`, event || '');
+      console.log(`[${address}:${portWithRandomLog}] ${info}`, event || '');
     };
     const readableWebSocketStream = makeReadableWebSocketStream(webSocket, log);
     let vlessResponseHeader: Uint8Array | null = null;
@@ -46,130 +46,50 @@ export async function processWebSocket({
               );
               return;
             }
-            if (vlessBuffer.byteLength < 24) {
-              console.log('invalid data');
-              controller.error('invalid data');
-              return;
-            }
-            const version = new Uint8Array(vlessBuffer.slice(0, 1));
-            let isValidUser = false;
-            if (
-              uuid.stringify(new Uint8Array(vlessBuffer.slice(1, 17))) ===
-              userID
-            ) {
-              isValidUser = true;
-            }
-            if (!isValidUser) {
-              console.log('in valid user');
-              controller.error('in valid user');
-              return;
-            }
-
-            const optLength = new Uint8Array(vlessBuffer.slice(17, 18))[0];
-            //skip opt for now
-
-            const command = new Uint8Array(
-              vlessBuffer.slice(18 + optLength, 18 + optLength + 1)
-            )[0];
-            // 0x01 TCP
-            // 0x02 UDP
-            // 0x03 MUX
-            if (command === 1) {
-            } else {
-              controller.error(
-                `command ${command} is not support, command 01-tcp,02-udp,03-mux`
-              );
-              return;
-            }
-            const portIndex = 18 + optLength + 1;
-            const portBuffer = vlessBuffer.slice(portIndex, portIndex + 2);
-            // port is big-Endian in raw data etc 80 == 0x005d
-            const portRemote = new DataView(portBuffer).getInt16(0);
-            port = portRemote;
-            let addressIndex = portIndex + 2;
-            const addressBuffer = new Uint8Array(
-              vlessBuffer.slice(addressIndex, addressIndex + 1)
-            );
-
-            // 1--> ipv4  addressLength =4
-            // 2--> domain name addressLength=addressBuffer[1]
-            // 3--> ipv6  addressLength =16
-            const addressType = addressBuffer[0];
-            let addressLength = 0;
-            let addressValueIndex = addressIndex + 1;
-            let addressValue = '';
-            switch (addressType) {
-              case 1:
-                addressLength = 4;
-                addressValue = new Uint8Array(
-                  vlessBuffer.slice(
-                    addressValueIndex,
-                    addressValueIndex + addressLength
-                  )
-                ).join('.');
-                break;
-              case 2:
-                addressLength = new Uint8Array(
-                  vlessBuffer.slice(addressValueIndex, addressValueIndex + 1)
-                )[0];
-                addressValueIndex += 1;
-                addressValue = new TextDecoder().decode(
-                  vlessBuffer.slice(
-                    addressValueIndex,
-                    addressValueIndex + addressLength
-                  )
-                );
-                break;
-              case 3:
-                addressLength = 16;
-                const addressChunkBy2: number[][] = lodash.chunk(
-                  new Uint8Array(
-                    vlessBuffer.slice(
-                      addressValueIndex,
-                      addressValueIndex + addressLength
-                    )
-                  ),
-                  2,
-                  null
-                );
-                // 2001:0db8:85a3:0000:0000:8a2e:0370:7334
-                addressValue = addressChunkBy2
-                  .map((items) =>
-                    items
-                      .map((item) => item.toString(16).padStart(2, '0'))
-                      .join('')
-                  )
-                  .join(':');
-                break;
-              default:
-                console.log(`[${address}:${port}] invild address`);
-            }
-            address = addressValue;
-            if (!addressValue) {
-              // console.log(`[${address}:${port}] addressValue is empty`);
-              controller.error(`[${address}:${port}] addressValue is empty`);
-              return;
+            const {
+              hasError,
+              message,
+              portRemote,
+              addressRemote,
+              rawDataIndex,
+              vlessVersion,
+            } = processVlessHeader(vlessBuffer, userID, uuid, lodash);
+            address = addressRemote || '';
+            portWithRandomLog = `${portRemote}--${Math.random()}`;
+            if (hasError) {
+              controller.error(`[${address}:${portWithRandomLog}] ${message} `);
             }
             // const addressType = requestAddr >> 4;
             // const addressLength = requestAddr & 0x0f;
-            console.log(`[${addressValue}:${port}] connecting`);
-            remoteConnection = await rawTCPFactory(port, addressValue);
-            vlessResponseHeader = new Uint8Array([version[0], 0]);
-            const rawDataIndex = addressValueIndex + addressLength;
-            const rawClientData = vlessBuffer.slice(rawDataIndex);
+            console.log(`[${address}:${portWithRandomLog}] connecting`);
+            remoteConnection = await rawTCPFactory(portRemote!, address!);
+            vlessResponseHeader = new Uint8Array([vlessVersion![0], 0]);
+            const rawClientData = vlessBuffer.slice(rawDataIndex!);
             await remoteConnection!.write(new Uint8Array(rawClientData));
             remoteConnectionReadyResolve(remoteConnection);
           },
+          close() {
+            console.log(
+              `[${address}:${portWithRandomLog}] readableWebSocketStream is close`
+            );
+          },
+          abort(reason) {
+            console.log(
+              `[${address}:${portWithRandomLog}] readableWebSocketStream is abort`,
+              JSON.stringify(reason)
+            );
+          },
         })
       )
       .catch((error) => {
-        console.log(
-          `[${address}:${port}] readableWebSocketStream pipeto has exception`,
+        console.error(
+          `[${address}:${portWithRandomLog}] readableWebSocketStream pipeto has exception`,
           error.stack || error
         );
-        closeWebSocket(webSocket);
+        // error is cancel readable stream anyway, no need close websocket in here
+        // closeWebSocket(webSocket);
         // close remote conn
-        remoteConnection?.close();
+        // remoteConnection?.close();
       });
     await new Promise((resolve) => (remoteConnectionReadyResolve = resolve));
     let remoteChunkCount = 0;
@@ -178,13 +98,15 @@ export async function processWebSocket({
     await remoteConnection!.readable.pipeTo(
       new WritableStream({
         start() {
-          webSocket.send(vlessResponseHeader!);
+          if (webSocket.readyState === webSocket.OPEN) {
+            webSocket.send(vlessResponseHeader!);
+          }
         },
         async write(chunk: Uint8Array, controller) {
           function send2WebSocket() {
             if (webSocket.readyState !== webSocket.OPEN) {
               controller.error(
-                `[${address}:${port}] abort when webSocket is close can't accept data from remoteConnection!.readable`
+                `can't accept data from remoteConnection!.readable when client webSocket is close early`
               );
               return;
             }
@@ -219,13 +141,13 @@ export async function processWebSocket({
         },
         close() {
           console.log(
-            `[${address}:${port}] remoteConnection!.readable is close`
+            `[${address}:${portWithRandomLog}] remoteConnection!.readable is close`
           );
         },
         abort(reason) {
           closeWebSocket(webSocket);
           console.error(
-            `[${address}:${port}] remoteConnection!.readable abort`,
+            `[${address}:${portWithRandomLog}] remoteConnection!.readable abort`,
             reason
           );
         },
@@ -233,7 +155,7 @@ export async function processWebSocket({
     );
   } catch (error: any) {
     console.error(
-      `[${address}:${port}] processWebSocket has esception `,
+      `[${address}:${portWithRandomLog}] processWebSocket has exception `,
       error.stack || error
     );
     closeWebSocket(webSocket);
@@ -242,7 +164,8 @@ export async function processWebSocket({
 }
 
 function makeReadableWebSocketStream(ws: WebSocket, log: Function) {
-  return new ReadableStream({
+  let readableStreamCancel = false;
+  return new ReadableStream<ArrayBuffer>({
     start(controller) {
       ws.addEventListener('message', async (e) => {
         const vlessBuffer: ArrayBuffer = e.data;
@@ -250,22 +173,31 @@ function makeReadableWebSocketStream(ws: WebSocket, log: Function) {
         controller.enqueue(vlessBuffer);
       });
       ws.addEventListener('error', (e) => {
-        log('socket has error', e);
+        log('socket has error');
+        readableStreamCancel = true;
         controller.error(e);
       });
       ws.addEventListener('close', () => {
         try {
-          log('socket is close');
+          log('webSocket is close');
+          // is stream is cancel, skill controller.close
+          if (readableStreamCancel) {
+            return;
+          }
           controller.close();
         } catch (error) {
-          log(`websocketStream can't close`);
+          log(`websocketStream can't close DUE to `, error);
         }
       });
     },
     pull(controller) {},
     cancel(reason) {
-      log(`websocketStream is cancel`, reason);
-      ws.close();
+      log(`websocketStream is cancel DUE to `, reason);
+      if (readableStreamCancel) {
+        return;
+      }
+      readableStreamCancel = true;
+      closeWebSocket(ws);
     },
   });
 }
@@ -275,3 +207,127 @@ function closeWebSocket(socket: WebSocket) {
     socket.close();
   }
 }
+
+function processVlessHeader(
+  vlessBuffer: ArrayBuffer,
+  userID: string,
+  uuidLib: any,
+  lodash: any
+) {
+  if (vlessBuffer.byteLength < 24) {
+    // console.log('invalid data');
+    // controller.error('invalid data');
+    return {
+      hasError: true,
+      message: 'invalid data',
+    };
+  }
+  const version = new Uint8Array(vlessBuffer.slice(0, 1));
+  let isValidUser = false;
+  if (uuidLib.stringify(new Uint8Array(vlessBuffer.slice(1, 17))) === userID) {
+    isValidUser = true;
+  }
+  if (!isValidUser) {
+    // console.log('in valid user');
+    // controller.error('in valid user');
+    return {
+      hasError: true,
+      message: 'in valid user',
+    };
+  }
+
+  const optLength = new Uint8Array(vlessBuffer.slice(17, 18))[0];
+  //skip opt for now
+
+  const command = new Uint8Array(
+    vlessBuffer.slice(18 + optLength, 18 + optLength + 1)
+  )[0];
+  // 0x01 TCP
+  // 0x02 UDP
+  // 0x03 MUX
+  if (command === 1) {
+  } else {
+    // controller.error(
+    //   `command ${command} is not support, command 01-tcp,02-udp,03-mux`
+    // );
+    return {
+      hasError: true,
+      message: `command ${command} is not support, command 01-tcp,02-udp,03-mux`,
+    };
+  }
+  const portIndex = 18 + optLength + 1;
+  const portBuffer = vlessBuffer.slice(portIndex, portIndex + 2);
+  // port is big-Endian in raw data etc 80 == 0x005d
+  const portRemote = new DataView(portBuffer).getInt16(0);
+
+  let addressIndex = portIndex + 2;
+  const addressBuffer = new Uint8Array(
+    vlessBuffer.slice(addressIndex, addressIndex + 1)
+  );
+
+  // 1--> ipv4  addressLength =4
+  // 2--> domain name addressLength=addressBuffer[1]
+  // 3--> ipv6  addressLength =16
+  const addressType = addressBuffer[0];
+  let addressLength = 0;
+  let addressValueIndex = addressIndex + 1;
+  let addressValue = '';
+  switch (addressType) {
+    case 1:
+      addressLength = 4;
+      addressValue = new Uint8Array(
+        vlessBuffer.slice(addressValueIndex, addressValueIndex + addressLength)
+      ).join('.');
+      break;
+    case 2:
+      addressLength = new Uint8Array(
+        vlessBuffer.slice(addressValueIndex, addressValueIndex + 1)
+      )[0];
+      addressValueIndex += 1;
+      addressValue = new TextDecoder().decode(
+        vlessBuffer.slice(addressValueIndex, addressValueIndex + addressLength)
+      );
+      break;
+    case 3:
+      addressLength = 16;
+      const addressChunkBy2: number[][] = lodash.chunk(
+        new Uint8Array(
+          vlessBuffer.slice(
+            addressValueIndex,
+            addressValueIndex + addressLength
+          )
+        ),
+        2,
+        null
+      );
+      // 2001:0db8:85a3:0000:0000:8a2e:0370:7334
+      addressValue = addressChunkBy2
+        .map((items) =>
+          items.map((item) => item.toString(16).padStart(2, '0')).join('')
+        )
+        .join(':');
+      if (addressValue) {
+        addressValue = `[${addressValue}]`;
+      }
+
+      break;
+    default:
+      console.log(`invild  addressType is ${addressType}`);
+  }
+  if (!addressValue) {
+    // console.log(`[${address}:${port}] addressValue is empty`);
+    // controller.error(`[${address}:${portWithRandomLog}] addressValue is empty`);
+    return {
+      hasError: true,
+      message: `addressValue is empty, addressType is ${addressType}`,
+    };
+  }
+
+  return {
+    hasError: false,
+    addressRemote: addressValue,
+    portRemote,
+    rawDataIndex: addressValueIndex + addressLength,
+    vlessVersion: version,
+  };
+}