From 286e29db59580f1d8df93449318ef4d91123768e Mon Sep 17 00:00:00 2001
From: Yash Mehta <yash.mehta@owasp.org>
Date: Thu, 17 Aug 2023 16:20:52 +0530
Subject: [PATCH] feat: Add support for using existing http(s).Agent

---
 lib/index.js      | 18 ++++++++++---
 package-lock.json | 12 ++++-----
 test/index.js     | 67 ++++++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 83 insertions(+), 14 deletions(-)

diff --git a/lib/index.js b/lib/index.js
index 6ebb448..98c5a08 100644
--- a/lib/index.js
+++ b/lib/index.js
@@ -18,10 +18,12 @@ const checkIp = (ip) => {
   return true;
 };
 
-const manageConnection = (url) => {
-  const httpAgent = new http.Agent();
-  const httpsAgent = new https.Agent();
-  const agent = url.startsWith('https') ? httpsAgent : httpAgent;
+// prevent memory leak
+const ACTIVE = Symbol('active');
+
+const requestFilterHandler = (agent)=>{
+  if (agent[ACTIVE]) return agent;
+  agent[ACTIVE] = true;
   const {createConnection} = agent;
   agent.createConnection = function(options, func) {
     const {host: address} = options;
@@ -40,4 +42,12 @@ const manageConnection = (url) => {
   return agent;
 };
 
+const manageConnection = (url) => {
+  const httpAgent = new http.Agent();
+  const httpsAgent = new https.Agent();
+  const agent = url.startsWith('https') ? httpsAgent : httpAgent;
+  return requestFilterHandler(agent);
+};
+
 module.exports = (url) => manageConnection(url);
+module.exports.requestFilterHandler = (agent) => requestFilterHandler(agent);
diff --git a/package-lock.json b/package-lock.json
index c1a2857..012b2bd 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -2378,9 +2378,9 @@
       }
     },
     "node_modules/flatted": {
-      "version": "3.1.0",
-      "resolved": "https://registry.npmjs.org/flatted/-/flatted-3.1.0.tgz",
-      "integrity": "sha512-tW+UkmtNg/jv9CSofAKvgVcO7c2URjhTdW1ZTkcAritblu8tajiYy7YisnIflEwtKssCtOxpnBRoCB7iap0/TA==",
+      "version": "3.1.1",
+      "resolved": "https://registry.npmjs.org/flatted/-/flatted-3.1.1.tgz",
+      "integrity": "sha512-zAoAQiudy+r5SvnSw3KJy5os/oRJYHzrzja/tBDqrZtNhUw8bt6y8OBzMWcjWr+8liV8Eb6yOhw8WZ7VFZ5ZzA==",
       "dev": true
     },
     "node_modules/follow-redirects": {
@@ -10600,9 +10600,9 @@
       }
     },
     "flatted": {
-      "version": "3.1.0",
-      "resolved": "https://registry.npmjs.org/flatted/-/flatted-3.1.0.tgz",
-      "integrity": "sha512-tW+UkmtNg/jv9CSofAKvgVcO7c2URjhTdW1ZTkcAritblu8tajiYy7YisnIflEwtKssCtOxpnBRoCB7iap0/TA==",
+      "version": "3.1.1",
+      "resolved": "https://registry.npmjs.org/flatted/-/flatted-3.1.1.tgz",
+      "integrity": "sha512-zAoAQiudy+r5SvnSw3KJy5os/oRJYHzrzja/tBDqrZtNhUw8bt6y8OBzMWcjWr+8liV8Eb6yOhw8WZ7VFZ5ZzA==",
       "dev": true
     },
     "follow-redirects": {
diff --git a/test/index.js b/test/index.js
index 12aad86..1bccf67 100644
--- a/test/index.js
+++ b/test/index.js
@@ -1,4 +1,7 @@
 const ssrfFilter = require('../lib/index.js');
+const {requestFilterHandler} = require('../lib/index.js');
+const http = require('http');
+const https = require('https');
 const axios = require('axios');
 const fs = require('fs');
 const expect = require('chai').expect;
@@ -7,7 +10,11 @@ const allowedUrlsFile = `${__dirname}/allowedUrls.txt`;
 let blockUrls;
 let allowedUrls;
 
-// Test: Blocked URLs
+
+const httpAgent = new http.Agent();
+const httpsAgent = new https.Agent();
+
+// ssrfFilter: Test: Blocked URLs
 try {
   blockUrls = JSON.parse(fs.readFileSync(blockUrlsFile));
 } catch (err) {
@@ -17,11 +24,11 @@ try {
 blockUrls.forEach((url)=>{
   it(`${url} is Blocked`, async () => {
     let check = 0;
+    // eslint-disable-next-line max-len
     const response = await axios.get(url, {httpAgent: ssrfFilter(url),
       httpsAgent: ssrfFilter(url)})
         .then((response) => {
           check = 1;
-          console.log(response);
         })
         .catch((error) => {
           check = 0;
@@ -33,7 +40,7 @@ blockUrls.forEach((url)=>{
   });
 });
 
-// Test: Allowed URLs
+// ssrfFilter: Test: Allowed URLs
 try {
   allowedUrls = JSON.parse(fs.readFileSync(allowedUrlsFile));
 } catch (err) {
@@ -58,7 +65,7 @@ allowedUrls.forEach((url)=>{
   });
 });
 
-// Test: DNS Rebind
+// ssrfFilter: Test: DNS Rebind
 it(`Test DNS Rebind`, async () => {
   let check = 0;
   const url = 'http://s-35.185.206.165-127.0.0.1-'+ new Date().valueOf() +'-rr-e.d.rebind.it';
@@ -80,3 +87,55 @@ it(`Test DNS Rebind`, async () => {
       });
   expect(response).to.equal(1);
 });
+
+// requestFilterHandler: Test: Blocked URLs
+try {
+  blockUrls = JSON.parse(fs.readFileSync(blockUrlsFile));
+} catch (err) {
+  console.log(err);
+}
+
+blockUrls.forEach((url)=>{
+  it(`${url} is Blocked`, async () => {
+    let check = 0;
+    // eslint-disable-next-line max-len
+    const response = await axios.get(url, {httpAgent: requestFilterHandler(httpAgent),
+      httpsAgent: requestFilterHandler(httpsAgent)})
+        .then((response) => {
+          check = 1;
+        })
+        .catch((error) => {
+          check = 0;
+        })
+        .then(() => {
+          return check;
+        });
+    expect(response).to.equal(0);
+  });
+});
+
+// requestFilterHandler: Test: Allowed URLs
+try {
+  allowedUrls = JSON.parse(fs.readFileSync(allowedUrlsFile));
+} catch (err) {
+  console.log(err);
+}
+
+allowedUrls.forEach((url)=>{
+  it(`${url} is Allowed`, async () => {
+    let check = 0;
+    // eslint-disable-next-line max-len
+    const response = await axios.get(url, {httpAgent: requestFilterHandler(httpAgent),
+      httpsAgent: requestFilterHandler(httpsAgent)})
+        .then((response) => {
+          check = 1;
+        })
+        .catch((error) => {
+          check = 0;
+        })
+        .then(() => {
+          return check;
+        });
+    expect(response).to.equal(1);
+  });
+});