fix: correct target stats when sockets are reused

Author: jirimoravcik

package.json

@@ -1,6 +1,6 @@
 {
   "name": "proxy-chain",
-  "version": "2.5.7",
+  "version": "2.5.8",
   "description": "Node.js implementation of a proxy server (think Squid) with support for SSL, authentication, upstream proxy chaining, and protocol tunneling.",
   "main": "dist/index.js",
   "keywords": [

src/chain.ts

@@ -7,6 +7,7 @@ import type { URL } from 'url';
 
 import type { Socket } from './socket';
 import { badGatewayStatusCodes, createCustomStatusHttpResponse, errorCodeToStatusCode } from './statuses';
+import type { SocketWithPreviousStats } from './utils/count_target_bytes';
 import { countTargetBytes } from './utils/count_target_bytes';
 import { getBasicAuthorizationHeader } from './utils/get_basic';
 
@@ -85,9 +86,15 @@ export const chain = (
     const fn = proxy.protocol === 'https:' ? https.request : http.request;
     const client = fn(proxy.origin, options as unknown as http.ClientRequestArgs);
 
-    client.on('connect', (response, targetSocket, clientHead) => {
+    client.once('socket', (targetSocket: SocketWithPreviousStats) => {
+        // Socket can be re-used by multiple requests.
+        // That's why we need to track the previous stats.
+        targetSocket.previousBytesRead = targetSocket.bytesRead;
+        targetSocket.previousBytesWritten = targetSocket.bytesWritten;
         countTargetBytes(sourceSocket, targetSocket);
+    });
 
+    client.on('connect', (response, targetSocket, clientHead) => {
         if (sourceSocket.readyState !== 'open') {
             // Sanity check, should never reach.
             targetSocket.destroy();

src/forward.ts

@@ -6,6 +6,7 @@ import type { URL } from 'url';
 import util from 'util';
 
 import { badGatewayStatusCodes, errorCodeToStatusCode } from './statuses';
+import type { SocketWithPreviousStats } from './utils/count_target_bytes';
 import { countTargetBytes } from './utils/count_target_bytes';
 import { getBasicAuthorizationHeader } from './utils/get_basic';
 import { validHeadersOnly } from './utils/valid_headers_only';
@@ -114,8 +115,12 @@ export const forward = async (
         }
     });
 
-    client.once('socket', (socket) => {
-        countTargetBytes(request.socket, socket);
+    client.once('socket', (socket: SocketWithPreviousStats) => {
+        // Socket can be re-used by multiple requests.
+        // That's why we need to track the previous stats.
+        socket.previousBytesRead = socket.bytesRead;
+        socket.previousBytesWritten = socket.bytesWritten;
+        countTargetBytes(request.socket, socket, (handler) => response.once('close', handler));
     });
 
     // Can't use pipeline here as it automatically destroys the streams

src/utils/count_target_bytes.ts

@@ -7,17 +7,27 @@ const calculateTargetStats = Symbol('calculateTargetStats');
 
 type Stats = { bytesWritten: number | null, bytesRead: number | null };
 
+/**
+ * Socket object extended with previous read and written bytes.
+ * Necessary due to target socket re-use.
+ */
+export type SocketWithPreviousStats = net.Socket & { previousBytesWritten?: number, previousBytesRead?: number };
+
 interface Extras {
     [targetBytesWritten]: number;
     [targetBytesRead]: number;
-    [targets]: Set<net.Socket>;
+    [targets]: Set<SocketWithPreviousStats>;
     [calculateTargetStats]: () => Stats;
 }
 
 // @ts-expect-error TS is not aware that `source` is used in the assertion.
 function typeSocket(source: unknown): asserts source is net.Socket & Extras {}
 
-export const countTargetBytes = (source: net.Socket, target: net.Socket): void => {
+export const countTargetBytes = (
+    source: net.Socket,
+    target: SocketWithPreviousStats,
+    registerCloseHandler?: (handler: () => void) => void,
+): void => {
     typeSocket(source);
 
     source[targetBytesWritten] = source[targetBytesWritten] || 0;
@@ -26,21 +36,24 @@ export const countTargetBytes = (source: net.Socket, target: net.Socket): void =
 
     source[targets].add(target);
 
-    target.once('close', () => {
-        source[targetBytesWritten] += target.bytesWritten;
-        source[targetBytesRead] += target.bytesRead;
-
+    const closeHandler = () => {
+        source[targetBytesWritten] += (target.bytesWritten - (target.previousBytesWritten || 0));
+        source[targetBytesRead] += (target.bytesRead - (target.previousBytesRead || 0));
         source[targets].delete(target);
-    });
+    };
+    if (!registerCloseHandler) {
+        registerCloseHandler = (handler: () => void) => target.once('close', handler);
+    }
+    registerCloseHandler(closeHandler);
 
     if (!source[calculateTargetStats]) {
         source[calculateTargetStats] = () => {
             let bytesWritten = source[targetBytesWritten];
             let bytesRead = source[targetBytesRead];
 
             for (const socket of source[targets]) {
-                bytesWritten += socket.bytesWritten;
-                bytesRead += socket.bytesRead;
+                bytesWritten += (socket.bytesWritten - (socket.previousBytesWritten || 0));
+                bytesRead += (socket.bytesRead - (socket.previousBytesRead || 0));
             }
 
             return {