From 359eee9df5c3a887914cf4913e26a8bc8d87e571 Mon Sep 17 00:00:00 2001
From: Mikko Ahlroth <mikko@ahlroth.fi>
Date: Thu, 18 Jul 2024 01:24:37 +0300
Subject: [PATCH] Fix bugz

---
 index.js     | 55 ++++++++++++++++++++++++++--------------------------
 package.json |  3 ++-
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/index.js b/index.js
index c6b2dc9..b766b64 100644
--- a/index.js
+++ b/index.js
@@ -113,6 +113,9 @@ class Connection {
   /** @type {NodeJS.Timer | null} */
   idleTimeout = null;
 
+  /** @type {() => void} */
+  idleTimeoutCallback;
+
   /** @type {NodeJS.Timer | null} */
   retryTimeout = null;
 
@@ -124,10 +127,12 @@ class Connection {
   /**
    * @param {Target} source
    * @param {Socket} writeSocket
+   * @param {() => void} idleTimeoutCallback
    */
-  constructor(source, writeSocket) {
+  constructor(source, writeSocket, idleTimeoutCallback) {
     this.source = source;
     this.writeSocket = writeSocket;
+    this.idleTimeoutCallback = idleTimeoutCallback;
   }
 
   initTargetSocket() {
@@ -136,8 +141,8 @@ class Connection {
     this.writeSocket.on("connect", () => {
       console.log(
         "Opened new connection to target from",
-        source.address,
-        source.port
+        this.source.address,
+        this.source.port
       );
     });
     this.writeSocket.on("message", (fromMsg) => {
@@ -146,17 +151,22 @@ class Connection {
         "Got",
         fromMsg.byteLength,
         "bytes from target for",
-        source.address,
-        source.port
+        this.source.address,
+        this.source.port
       );
 
-      READ_SOCKET.send(fromMsg, source.port, source.address);
+      READ_SOCKET.send(fromMsg, this.source.port, this.source.address);
     });
     this.writeSocket.on("error", (err) => {
       this.#writeError(err);
     });
     this.writeSocket.on("close", () => {
-      console.log("Socket for", source.address, source.port, "closed.");
+      console.log(
+        "Socket for",
+        this.source.address,
+        this.source.port,
+        "closed."
+      );
       CONNECTIONS.delete(source);
     });
 
@@ -168,6 +178,8 @@ class Connection {
    * @param {?Buffer} msg
    */
   send(msg) {
+    this.#resetIdleTimeout();
+
     if (msg && (this.writing || !this.writeSocket)) {
       this.msgBuffer.push(msg);
       return;
@@ -213,7 +225,10 @@ class Connection {
 
   #resetIdleTimeout() {
     clearTimeout(this.idleTimeout);
-    this.idleTimeout = setTimeout(() => idleTimeout(this.source), IDLE_TIMEOUT);
+    this.idleTimeout = setTimeout(
+      () => this.idleTimeoutCallback(),
+      IDLE_TIMEOUT
+    );
   }
 
   /**
@@ -238,14 +253,6 @@ class Connection {
   }
 }
 
-/**
- * @param {Target} source
- */
-function idleTimeout(source) {
-  console.log("Idle timeout for", source.address, source.port);
-  CONNECTIONS.delete(source);
-}
-
 const TARGET_ADDRESS = argv[2];
 const TARGET_PORT = parseInt(argv[3], 10);
 
@@ -274,22 +281,16 @@ READ_SOCKET.on("message", (toMsg, rinfo) => {
   let connection = CONNECTIONS.get(source);
   if (!connection) {
     const socket = createSocket(SOCKET_TYPE);
-    connection = new Connection(
-      source,
-      socket,
-      setTimeout(() => {
-        idleTimeout(source);
-      }, IDLE_TIMEOUT)
-    );
+    connection = new Connection(source, socket, () => {
+      console.log("Idle timeout for", source.address, source.port);
+      CONNECTIONS.delete(source);
+    });
     connection.initTargetSocket();
 
     CONNECTIONS.set(source, connection);
   }
 
-  clearTimeout(connection.idleTimeout);
-  connection.idleTimeout = setTimeout(() => {
-    idleTimeout(source);
-  }, IDLE_TIMEOUT);
+  connection.send(toMsg);
 });
 
 READ_SOCKET.on("error", (err) => {
diff --git a/package.json b/package.json
index c996df5..05a78ca 100644
--- a/package.json
+++ b/package.json
@@ -12,5 +12,6 @@
   },
   "author": "Mikko Ahlroth <mikko@ahlroth.fi>",
   "license": "CC0-1.0",
-  "private": true
+  "private": true,
+  "type": "module"
 }