feat: improve request timeout handling

FireBlinkLTD · Nov 2, 2023 · 14cc37a · 14cc37a
1 parent 08c8b9b
commit 14cc37a
Show file tree

Hide file tree

Showing 8 changed files with 154 additions and 44 deletions.
diff --git a/src/Prxi.ts b/src/Prxi.ts
@@ -13,6 +13,7 @@ interface Proxy {
 
 export class Prxi {
   private server: Server = null;
+  private connections = new Set<Socket>();
   private logInfo: (message?: any, ...params: any[]) => void;
   private logError: (message?: any, ...params: any[]) => void;
 
@@ -115,6 +116,14 @@ export class Prxi {
       }
     });
 
+    // keep track of all open connections
+    server.on('connection', (connection: Socket) => {
+      this.connections.add(connection);
+      connection.on('close', () => {
+        this.connections.delete(connection);
+      });
+    });
+
     // handle upgrade action
     server.on('upgrade', (req: IncomingMessage, socket: Socket, head: Buffer) => {
       const requestId = id++;
@@ -246,8 +255,21 @@ export class Prxi {
     /* istanbul ignore next */
     if (this.server) {
       await new Promise<void>((res, rej) => {
+        const timer = setTimeout(() => {
+          this.connections.forEach((connection: Socket) => {
+            try {
+              connection.destroy();
+            } catch (e) {
+              this.logError('Failed to destroy connection', e);
+            }
+          });
+        }, this.configuration.proxyRequestTimeout ?? 60 * 1000);
+
         this.logInfo('Stopping Prxi');
+
         this.server.close((err) => {
+          clearTimeout(timer);
+
           if (err) {
             this.logError('Failed to stop Prxi', err);
             return rej(err);

diff --git a/src/handlers/HttpProxyHandler.ts b/src/handlers/HttpProxyHandler.ts
@@ -4,8 +4,7 @@ import {request as httpsRequest} from 'node:https';
 
 import { Configuration, ProxyRequestConfiguration } from "../interfaces";
 import { UpstreamConfiguration } from "../interfaces/UpstreamConfiguration";
-import { RequestUtils } from "../utils";
-
+import { RequestUtils, Timer } from "../utils";
 
 const emptyObj = {};
 
@@ -45,6 +44,7 @@ export class HttpProxyHandler {
 
     this.logInfo(`[${requestId}] [HttpProxyHandler] Processing HTTP/HTTPS proxy request with method ${method} to ${target}${url}`);
 
+    const proxyRequestTimeout = this.configuration.proxyRequestTimeout ?? 60 * 1000;
     const options: RequestOptions = {
       method,
       host,
@@ -58,40 +58,49 @@ export class HttpProxyHandler {
         proxyConfiguration?.proxyRequestHeaders,
       ),
       path: RequestUtils.concatPath(initialPath, url),
-      timeout: this.configuration.proxyRequestTimeout || 60 * 1000,
+      timeout: proxyRequestTimeout,
     };
 
+    // setup timer to force incoming request to be destroyed after 2x of proxyRequestTimeout configuration setting
+    const timer = new Timer(() => {
+      req.destroy();
+    }, proxyRequestTimeout * 2);
+
     const client = request(options);
 
-    await new Promise<void>((resolve, reject) => {
-      req.pipe(client);
+    try {
+      await new Promise<void>((resolve, reject) => {
+        req.pipe(client);
 
-      client.on('error', (err) => {
-        reject(err);
-      });
+        client.on('error', (err) => {
+          reject(err);
+        });
 
-      client.on('response', (response: IncomingMessage) => {
-        const headersToSet = RequestUtils.prepareProxyHeaders(
-          response.headers,
-          this.configuration.responseHeaders,
-          this.upstream.responseHeaders,
-          // istanbul ignore next
-          proxyConfiguration?.proxyResponseHeaders
-        );
-        RequestUtils.updateResponseHeaders(res, headersToSet);
+        client.on('response', (response: IncomingMessage) => {
+          const headersToSet = RequestUtils.prepareProxyHeaders(
+            response.headers,
+            this.configuration.responseHeaders,
+            this.upstream.responseHeaders,
+            // istanbul ignore next
+            proxyConfiguration?.proxyResponseHeaders
+          );
+          RequestUtils.updateResponseHeaders(res, headersToSet);
 
-        // istanbul ignore else
-        if (!res.writableEnded) {
-          response.on('end', () => {
-            this.logInfo(`[${requestId}] [HttpProxyHandler] Proxy request with method ${method} to ${host}${url} completed`);
-            resolve();
-          });
+          // istanbul ignore else
+          if (!res.writableEnded) {
+            response.on('end', () => {
+              this.logInfo(`[${requestId}] [HttpProxyHandler] Proxy request with method ${method} to ${host}${url} completed`);
+              resolve();
+            });
 
-          response.pipe(res);
-        } else {
-          resolve();
-        }
+            response.pipe(res);
+          } else {
+            resolve();
+          }
+        });
       });
-    });
+    } finally {
+      timer.cancel();
+    }
   }
 }
diff --git a/src/handlers/WebSocketProxyHandler.ts b/src/handlers/WebSocketProxyHandler.ts
@@ -5,7 +5,7 @@ import { Socket } from "node:net";
 
 import { Configuration, ProxyRequestConfiguration } from "../interfaces";
 import { UpstreamConfiguration } from "../interfaces/UpstreamConfiguration";
-import { RequestUtils, WebSocketUtils } from "../utils";
+import { RequestUtils, Timer, WebSocketUtils } from "../utils";
 
 const emptyObj = {};
 
@@ -40,6 +40,13 @@ export class WebSocketProxyHandler {
     head: Buffer,
     proxyConfiguration?: ProxyRequestConfiguration,
   ): Promise<void> {
+    const proxyRequestTimeout = this.configuration.proxyRequestTimeout ?? 60 * 1000;
+
+     // setup timer to force incoming request to be destroyed after 2x of proxyRequestTimeout configuration setting
+     const timer = new Timer(() => {
+      req.destroy();
+    }, proxyRequestTimeout * 2);
+
     try {
       WebSocketProxyHandler.debug.incomingSocket = socket;
       proxyConfiguration = proxyConfiguration || emptyObj;
@@ -71,7 +78,7 @@ export class WebSocketProxyHandler {
           proxyConfiguration?.proxyRequestHeaders,
         ),
         path: RequestUtils.concatPath(initialPath, url),
-        timeout: this.configuration.proxyRequestTimeout || 60 * 1000,
+        timeout: proxyRequestTimeout,
       };
 
       const client = request(options);
@@ -83,9 +90,6 @@ export class WebSocketProxyHandler {
         socket.unshift(head);
       }
 
-      // keep socket alive
-      WebSocketUtils.keepAlive(socket);
-
       await new Promise<void>((resolve, reject) => {
         req.pipe(client);
 
@@ -153,7 +157,9 @@ export class WebSocketProxyHandler {
             reject(err);
           });
 
-          // keep socket alive
+          // keep sockets alive
+          timer.cancel();
+          WebSocketUtils.keepAlive(socket);
           WebSocketUtils.keepAlive(proxySocket);
 
           const headersToSet = RequestUtils.prepareProxyHeaders(
@@ -167,6 +173,7 @@ export class WebSocketProxyHandler {
         });
       });
     } finally {
+      timer.cancel();
       delete WebSocketProxyHandler.debug.incomingSocket;
       delete WebSocketProxyHandler.debug.upstreamSocket;
       delete WebSocketProxyHandler.debug.upstreamRequest;

diff --git a/src/utils/Timer.ts b/src/utils/Timer.ts
@@ -0,0 +1,21 @@
+export class Timer {
+  private timer;
+
+  constructor(fn: Function, delay: number) {
+    if (delay <= 0) {
+      return fn();
+    }
+
+    this.timer = setTimeout(fn, delay);
+  }
+
+  /**
+   * Cancel the timer
+   */
+  public cancel(): void {
+    if (this.timer) {
+      clearTimeout(this.timer);
+      this.timer = null;
+    }
+  }
+}
diff --git a/src/utils/index.ts b/src/utils/index.ts
@@ -1,2 +1,3 @@
 export * from './RequestUtils';
+export * from './Timer';
 export * from './WebSocketUtils';
diff --git a/test/HttpProxy.errors.suite.ts b/test/HttpProxy.errors.suite.ts
@@ -33,7 +33,7 @@ export class HttpProxyErrorSuite {
 
     @test()
     async addressNotFoundFailErrorHandler(): Promise<void> {
-      this.proxy = new TestProxy('non-existing-host');
+      this.proxy = new TestProxy({}, 'non-existing-host');
       await this.proxy.start();
 
       const result = axios.post(`${this.proxyUrl}/echo`, { test: true });
@@ -44,7 +44,7 @@ export class HttpProxyErrorSuite {
     @test()
     async addressNotFoundPassErrorHandler(): Promise<void> {
       const customError = 'Custom Error';
-      this.proxy = new TestProxy('non-existing-host', async (req: IncomingMessage, res: ServerResponse, err: Error): Promise<void> => {
+      this.proxy = new TestProxy({}, 'non-existing-host', async (req: IncomingMessage, res: ServerResponse, err: Error): Promise<void> => {
         match(err.message, /getaddrinfo .* non-existing-host/gi);
 
         await writeJson(res, JSON.stringify({customError}));
@@ -58,7 +58,7 @@ export class HttpProxyErrorSuite {
     @test()
     async missingHandler(): Promise<void> {
       let msg = null;
-      this.proxy = new TestProxy('localhost', async (req: IncomingMessage, res: ServerResponse, err?: Error) => {
+      this.proxy = new TestProxy({}, 'localhost', async (req: IncomingMessage, res: ServerResponse, err?: Error) => {
         msg = err.message;
         throw err;
       }, false);
@@ -72,7 +72,7 @@ export class HttpProxyErrorSuite {
     @test()
     async noHandlers(): Promise<void> {
       let msg = null;
-      this.proxy = new TestProxy('localhost', async (req: IncomingMessage, res: ServerResponse, err?: Error) => {
+      this.proxy = new TestProxy({}, 'localhost', async (req: IncomingMessage, res: ServerResponse, err?: Error) => {
         msg = err.message;
         throw err;
       }, null);
@@ -85,7 +85,7 @@ export class HttpProxyErrorSuite {
 
     @test()
     async noWebSocketHandler(): Promise<void> {
-      this.proxy = new TestProxy('localhost', false, true, false);
+      this.proxy = new TestProxy({}, 'localhost', false, true, false);
       await this.proxy.start();
       const sio = io(`http://localhost:${TestProxy.PORT}`, {
         transports: ['websocket'],
@@ -108,7 +108,7 @@ export class HttpProxyErrorSuite {
 
     @test()
     async failedWebSocketHandler(): Promise<void> {
-      this.proxy = new TestProxy('localhost', null, true, async () => {
+      this.proxy = new TestProxy({}, 'localhost', null, true, async () => {
         throw new Error('test');
       });
       await this.proxy.start();

diff --git a/test/HttpProxy.success.suite.ts b/test/HttpProxy.success.suite.ts
@@ -3,6 +3,7 @@ import { TestServer, TestProxy } from './helpers';
 import axios from 'axios';
 import {deepEqual, strictEqual} from 'assert';
 import {io} from 'socket.io-client';
+import { Configuration } from '../src';
 
 @suite()
 export class HttpProxySuccessSuite {
@@ -34,8 +35,8 @@ export class HttpProxySuccessSuite {
     /**
      * Init proxy server
      */
-    private async initProxy(): Promise<void> {
-      this.proxy = new TestProxy();
+    private async initProxy(configOverride: Partial<Configuration> = {}): Promise<void> {
+      this.proxy = new TestProxy(configOverride);
       await this.proxy.start();
     }
 
@@ -58,7 +59,7 @@ export class HttpProxySuccessSuite {
     async customPath(): Promise<void> {
       await this.after();
       this.server = new TestServer(true, '/api');
-      this.proxy = new TestProxy('localhost', null, true, null, '/api');
+      this.proxy = new TestProxy({}, 'localhost', null, true, null, '/api');
       await this.proxy.start();
 
       await this.server.start();
@@ -86,6 +87,23 @@ export class HttpProxySuccessSuite {
       });
     }
 
+    @test()
+    async destroyRequestDueToTimeoutSettings(): Promise<void> {
+      await this.initProxy({
+        proxyRequestTimeout: 0,
+      });
+
+      let err;
+      try {
+        const resp = await axios.get(`${this.proxyUrl}/headers`);
+        console.log(resp);
+      } catch (e) {
+        err = e;
+      }
+
+      strictEqual(err.message, 'socket hang up');
+    }
+
     @test()
     async headersRequest(): Promise<void> {
       await this.initProxy();
@@ -131,6 +149,35 @@ export class HttpProxySuccessSuite {
       });
     }
 
+    @test()
+    async destroyWebsocketRequestDueToTimeoutSettings(): Promise<void> {
+      await this.initProxy({
+        proxyRequestTimeout: 0,
+      });
+
+      const sio = io(`http://localhost:${TestProxy.PORT}`, {
+        transports: ['websocket'],
+        reconnection: false,
+      });
+
+      let err: Error;
+      await new Promise<void>((res, rej) => {
+        const timeout = setTimeout(() => {
+          sio.disconnect();
+          rej(new Error('Unable to connect to WS'));
+        }, 2000);
+
+        sio.on('connect_error', (e) => {
+          err = e;
+          sio.disconnect();
+          clearTimeout(timeout);
+          res();
+        });
+      });
+
+      strictEqual(err.message, 'websocket error');
+    }
+
     @test()
     async websocket(): Promise<void> {
       await this.initProxy();

diff --git a/test/helpers/TestProxy.ts b/test/helpers/TestProxy.ts
@@ -1,13 +1,14 @@
 import { IncomingMessage, ServerResponse } from 'http';
 import { Duplex } from 'stream';
-import { ErrorHandler, Prxi, ProxyRequest, WebSocketHandlerFunction, WebSocketHandlerConfig } from '../../src';
+import { ErrorHandler, Prxi, ProxyRequest, WebSocketHandlerFunction, WebSocketHandlerConfig, Configuration } from '../../src';
 import { TestServer } from './TestServer';
 
 export class TestProxy {
   public static readonly PORT = 8888;
   private proxy: Prxi;
 
   constructor(
+    private configOverride: Partial<Configuration> = {},
     private host = 'localhost',
     private customErrorHandler: ErrorHandler | false = null,
     private isMatching: boolean | null = true,
@@ -51,6 +52,8 @@ export class TestProxy {
         ResConfigLevelOverwrite: 'CONFIG-RESPONSE-OVERWRITE',
         ResConfigLevelClear: null,
       },
+
+      ...this.configOverride
     });
 
     // start it