Merge pull request #61 from VKCOM/pavelnikitin/fix/tunnel-hot-reload/…

…MA-20342

Hot reload support

CHANGELOG.md

@@ -1,3 +1,11 @@
+## v0.2.3
+
+- Добавили поддержку `hot-reload`.
+- Добавили параметр запуска `ws-origin`:  
+  Используется для указания, что `WebSocket`-соединение должно использовать тот же `Origin`, что и проксируемый сервер.
+  - Значение `1` — использовать `Origin` проксируемого сервера (значение по умолчанию).
+  - Значение `0` — отключить это поведение.
+
 ## v0.2.2
 
 - Убрали node-fetch из зависимостей.

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@vkontakte/vk-tunnel",
   "license": "MIT",
-  "version": "0.2.2",
+  "version": "0.2.3",
   "bin": {
     "vk-tunnel": "./bin/vk-tunnel.js"
   },

src/constants.ts

@@ -11,6 +11,7 @@ export const DEFAULT_USER_PROXY_APP_SETTINGS = {
   'port': Number(process.env.PROXY_PORT ?? 10888),
   'host': process.env.PROXY_HOST ?? 'localhost',
   'insecure': 0,
+  'ws-origin': 1,
   'app_id': undefined,
   'staging': undefined,
   'endpoints': undefined,

src/entities/WsProxy.ts

@@ -12,64 +12,86 @@ export class WsProxy {
   private static DISABLE_COMPRESS = 'gzip;q=0,deflate;q=0';
 
   private connections: Record<string, WebSocket> = {};
-  private userSettings: UserProxyAppSettings;
-  private logger: Logger;
 
-  public constructor(userSettings: UserProxyAppSettings, logger: Logger) {
-    this.userSettings = userSettings;
-    this.logger = logger;
-  }
+  public constructor(
+    private readonly userSettings: UserProxyAppSettings,
+    private readonly logger: Logger,
+  ) {}
 
-  private filterHeaders(payload: string, proxyHost: string) {
+  private transformPayload(payload: string, proxyHost: string) {
     return payload
       .replace(/Accept-Encoding:.*/, WsProxy.ACCEPT_ENCODING + ': ' + WsProxy.DISABLE_COMPRESS)
       .replace(/Host: .*/, 'Host: ' + proxyHost);
   }
 
+  private filterWebSocketHeaders(headers: Record<string, string>) {
+    const allowedHeaders = [
+      'Sec-WebSocket-Protocol',
+      'Sec-WebSocket-Extensions',
+      'Sec-WebSocket-Key',
+      'Sec-WebSocket-Version',
+    ];
+    return Object.fromEntries(
+      Object.entries(headers).filter(([key]) => allowedHeaders.includes(key)),
+    );
+  }
+
   private closeConnection(seq: string) {
     this.connections[seq].close();
   }
 
   private createConnection(
     seq: string,
     proxiedServerUrl: string,
+    headers: Record<string, string>,
     sendResponseToVkProxyServer: SendResponseToProxyServer,
   ) {
-    this.connections[seq] = new WebSocket(proxiedServerUrl, [], {});
-    this.connections[seq].on('error', (msg) => {
-      this.logger.error('Connection error for ' + seq, msg);
-    });
+    const subprotocol = headers['Sec-Websocket-Protocol'];
+    const host = this.userSettings.wsOrigin ? this.userSettings.host : undefined;
+    const origin = this.userSettings.wsOrigin
+      ? `${this.userSettings.wsProtocol}://${this.userSettings.host}:${this.userSettings.port}`
+      : undefined;
 
-    this.connections[seq].on('upgrade', (msg) => {
-      let response = ['HTTP/1.1 101 Switching Protocols'];
-      let keys = Object.keys(msg.headers);
-      for (let i = 0; i < keys.length; i++) {
-        response.push(`${keys[i]}:${msg.headers[keys[i]]}`);
-      }
-      response.push('\n');
-      sendResponseToVkProxyServer(seq + MessageType.HTTP + response.join('\n'), () => {
-        this.logger.debug('send reply upgrade', seq, response.toString());
-      });
+    const websocket = new WebSocket(proxiedServerUrl, subprotocol, {
+      host,
+      origin,
+      headers: this.filterWebSocketHeaders(headers),
     });
 
-    this.connections[seq].on('open', () => {
+    websocket.on('error', (msg) => this.logger.error('Connection error for ' + seq, msg));
+
+    websocket.on('open', () => {
       this.connections[seq].on('message', (data) => {
         this.logger.debug('incoming ws message from service', seq, data);
         sendResponseToVkProxyServer(`${seq}${MessageType.WEBSOCKET}${data}`, () => {
           this.logger.debug('send reply', seq, data);
         });
       });
     });
+
+    websocket.on('upgrade', (msg) => {
+      const responseHeaders = [
+        'HTTP/1.1 101 Switching Protocols',
+        ...Object.entries(msg.headers).map(([key, value]) => `${key}: ${value}`),
+      ];
+      const response = responseHeaders.join('\n') + '\n\n';
+
+      sendResponseToVkProxyServer(seq + MessageType.HTTP + response, () => {
+        this.logger.debug('send reply upgrade', seq, response.toString());
+      });
+    });
+
+    this.connections[seq] = websocket;
   }
 
   public async proxy(
     request: ProxiedNetworkPacket,
     sendResponseToVkProxyServer: SendResponseToProxyServer,
   ) {
-    const { messageType, payload, isWebsocketUpgrade, seq, endpoint } = request;
+    const { messageType, payload, isWebsocketUpgrade, seq, endpoint, parsedRequest } = request;
 
     if (messageType !== MessageType.HTTP) {
-      const filteredPayload = this.filterHeaders(payload, this.userSettings.host);
+      const filteredPayload = this.transformPayload(payload, this.userSettings.host);
 
       if (messageType === MessageType.WEBSOCKET_CLOSE) {
         return this.closeConnection(seq);
@@ -82,7 +104,13 @@ export class WsProxy {
 
     if (isWebsocketUpgrade) {
       const proxiedServerUrl = `${this.userSettings.wsProtocol}://${this.userSettings.host}:${this.userSettings.port}${endpoint}`;
-      this.createConnection(seq, proxiedServerUrl, sendResponseToVkProxyServer);
+
+      this.createConnection(
+        seq,
+        proxiedServerUrl,
+        parsedRequest.headers,
+        sendResponseToVkProxyServer,
+      );
     }
   }
 }

src/helpers/getUserProxyAppSettings.ts

@@ -16,8 +16,9 @@ export function getUserProxyAppSettings(): UserProxyAppSettings {
   const {
     'http-protocol': httpProtocol,
     'ws-protocol': wsProtocol,
+    'ws-origin': wsOrigin,
     ...userSettings
   } = Object.assign({ ...DEFAULT_USER_PROXY_APP_SETTINGS, ...userConfigFile }, cliArgs);
 
-  return { httpProtocol, wsProtocol, ...userSettings };
+  return { wsOrigin, httpProtocol, wsProtocol, ...userSettings };
 }

src/types.ts

@@ -23,15 +23,17 @@ export interface UserProxyAppSettings {
   host: string;
   port: number;
   insecure: 0 | 1;
+  wsOrigin: 0 | 1;
   app_id?: number;
   staging?: boolean;
   endpoints?: string[];
 }
 
 export interface UserProxyAppSettingsArgs
-  extends Omit<UserProxyAppSettings, 'httpProtocol' | 'wsProtocol'> {
+  extends Omit<UserProxyAppSettings, 'httpProtocol' | 'wsProtocol' | 'wsOrigin'> {
   'http-protocol': HttpProtocol;
   'ws-protocol': WsProtocol;
+  'ws-origin': 0 | 1;
 }
 
 export interface TunnelConnectionData {