From e24b75e7d8d662573ba7789a3eefffe29baa21cf Mon Sep 17 00:00:00 2001
From: Brandon Gillis <brandon.gillis@epitech.eu>
Date: Mon, 4 Dec 2023 13:36:32 +0100
Subject: [PATCH] fix(#124): Improve localhost handling, add cache for ipv6 &
 ipv4 check

---
 .../src/ipc/network/axios-instance.js         | 54 ++++++++++++-------
 1 file changed, 34 insertions(+), 20 deletions(-)

diff --git a/packages/bruno-electron/src/ipc/network/axios-instance.js b/packages/bruno-electron/src/ipc/network/axios-instance.js
index 2894ee008..dcc57a07e 100644
--- a/packages/bruno-electron/src/ipc/network/axios-instance.js
+++ b/packages/bruno-electron/src/ipc/network/axios-instance.js
@@ -1,6 +1,11 @@
 const URL = require('url');
 const Socket = require('net').Socket;
 const axios = require('axios');
+const connectionCache = new Map(); // Cache to store checkConnection() results
+
+const LOCAL_IPV6 = '::1';
+const LOCAL_IPV4 = '127.0.0.1';
+const LOCALHOST = 'localhost';
 
 const getTld = (hostname) => {
   if (!hostname) {
@@ -12,19 +17,28 @@ const getTld = (hostname) => {
 
 const checkConnection = (host, port) =>
   new Promise((resolve) => {
-    const socket = new Socket();
+    const key = `${host}:${port}`;
+    const cachedResult = connectionCache.get(key);
+
+    if (cachedResult !== undefined) {
+      resolve(cachedResult);
+    } else {
+      const socket = new Socket();
 
-    socket.once('connect', () => {
-      socket.end();
-      resolve(true);
-    });
+      socket.once('connect', () => {
+        socket.end();
+        connectionCache.set(key, true); // Cache successful connection
+        resolve(true);
+      });
 
-    socket.once('error', () => {
-      resolve(false);
-    });
+      socket.once('error', () => {
+        connectionCache.set(key, false); // Cache failed connection
+        resolve(false);
+      });
 
-    // Try to connect to the host and port
-    socket.connect(port, host);
+      // Try to connect to the host and port
+      socket.connect(port, host);
+    }
   });
 
 /**
@@ -43,16 +57,16 @@ function makeAxiosInstance() {
     // Resolve all *.localhost to localhost and check if it should use IPv6 or IPv4
     // RFC: 6761 section 6.3 (https://tools.ietf.org/html/rfc6761#section-6.3)
     // @see https://github.com/usebruno/bruno/issues/124
-    // temporarily disabling the fix (- Anoop)
-    // if (getTld(url.hostname) === 'localhost') {
-    //   config.headers.Host = url.hostname; // Put original hostname in Host
-
-    //   const portNumber = Number(url.port) || (url.protocol.includes('https') ? 443 : 80);
-    //   const useIpv6 = await checkConnection('::1', portNumber);
-    //   url.hostname = useIpv6 ? '::1' : '127.0.0.1';
-    //   delete url.host; // Clear hostname cache
-    //   config.url = URL.format(url);
-    // }
+    if (getTld(url.hostname) === LOCALHOST || url.hostname === LOCAL_IPV4 || url.hostname === LOCAL_IPV6) {
+      // use custom DNS lookup for localhost
+      config.lookup = (hostname, options, callback) => {
+        const portNumber = Number(url.port) || (url.protocol.includes('https') ? 443 : 80);
+        checkConnection(LOCAL_IPV6, portNumber).then((useIpv6) => {
+          const ip = useIpv6 ? LOCAL_IPV6 : LOCAL_IPV4;
+          callback(null, ip, useIpv6 ? 6 : 4);
+        });
+      };
+    }
 
     config.headers['request-start-time'] = Date.now();
     return config;