Merge pull request #348 from AikidoSec/dont-shadow-ports-in-hostnames

Keep all ports that we connect to

library/agent/Agent.test.ts

@@ -746,10 +746,17 @@ t.test("it sends hostnames and routes along with heartbeat", async () => {
         {
           hostname: "aikido.dev",
           port: 443,
+          hits: 1,
+        },
+        {
+          hostname: "aikido.dev",
+          port: 80,
+          hits: 1,
         },
         {
           hostname: "google.com",
           port: 443,
+          hits: 1,
         },
       ],
       routes: [

library/agent/Hostnames.test.ts

@@ -5,32 +5,104 @@ t.test("it works", async () => {
   const hostnames = new Hostnames(3);
   t.same(hostnames.asArray(), []);
 
-  hostnames.add("aikido.dev", 433);
-  t.same(hostnames.asArray(), [{ hostname: "aikido.dev", port: 433 }]);
+  hostnames.add("aikido.dev", 443);
+  t.same(hostnames.asArray(), [{ hostname: "aikido.dev", port: 443, hits: 1 }]);
 
   hostnames.add("aikido.dev", 80);
-  t.same(hostnames.asArray(), [{ hostname: "aikido.dev", port: 433 }]);
+  t.same(hostnames.asArray(), [
+    { hostname: "aikido.dev", port: 443, hits: 1 },
+    { hostname: "aikido.dev", port: 80, hits: 1 },
+  ]);
 
   hostnames.add("google.com", 80);
   t.same(hostnames.asArray(), [
-    { hostname: "aikido.dev", port: 433 },
-    { hostname: "google.com", port: 80 },
+    { hostname: "aikido.dev", port: 443, hits: 1 },
+    { hostname: "aikido.dev", port: 80, hits: 1 },
+    { hostname: "google.com", port: 80, hits: 1 },
+  ]);
+
+  hostnames.add("google.com", undefined);
+  t.same(hostnames.asArray(), [
+    { hostname: "aikido.dev", port: 80, hits: 1 },
+    { hostname: "google.com", port: 80, hits: 1 },
+    { hostname: "google.com", port: undefined, hits: 1 },
   ]);
 
   hostnames.add("github.com", 80);
   t.same(hostnames.asArray(), [
-    { hostname: "aikido.dev", port: 433 },
-    { hostname: "google.com", port: 80 },
-    { hostname: "github.com", port: 80 },
+    { hostname: "google.com", port: 80, hits: 1 },
+    { hostname: "google.com", port: undefined, hits: 1 },
+    { hostname: "github.com", port: 80, hits: 1 },
   ]);
 
   hostnames.add("jetbrains.com", 80);
   t.same(hostnames.asArray(), [
-    { hostname: "google.com", port: 80 },
-    { hostname: "github.com", port: 80 },
-    { hostname: "jetbrains.com", port: 80 },
+    { hostname: "google.com", port: undefined, hits: 1 },
+    { hostname: "github.com", port: 80, hits: 1 },
+    { hostname: "jetbrains.com", port: 80, hits: 1 },
   ]);
 
   hostnames.clear();
   t.same(hostnames.asArray(), []);
 });
+
+t.test("it respects max size", async () => {
+  const hostnames = new Hostnames(2);
+  hostnames.add("aikido.dev", 1);
+  hostnames.add("aikido.dev", 2);
+
+  t.same(hostnames.asArray(), [
+    { hostname: "aikido.dev", port: 1, hits: 1 },
+    { hostname: "aikido.dev", port: 2, hits: 1 },
+  ]);
+
+  hostnames.add("aikido.dev", 3);
+  hostnames.add("aikido.dev", 4);
+
+  t.same(hostnames.asArray(), [
+    { hostname: "aikido.dev", port: 3, hits: 1 },
+    { hostname: "aikido.dev", port: 4, hits: 1 },
+  ]);
+
+  hostnames.add("google.com", 1);
+
+  t.same(hostnames.asArray(), [
+    { hostname: "aikido.dev", port: 4, hits: 1 },
+    { hostname: "google.com", port: 1, hits: 1 },
+  ]);
+
+  hostnames.add("google.com", 2);
+
+  t.same(hostnames.asArray(), [
+    { hostname: "google.com", port: 1, hits: 1 },
+    { hostname: "google.com", port: 2, hits: 1 },
+  ]);
+});
+
+t.test("it tracks hits", async () => {
+  const hostnames = new Hostnames(3);
+
+  hostnames.add("aikido.dev", 443);
+  hostnames.add("aikido.dev", 443);
+  t.same(hostnames.asArray(), [{ hostname: "aikido.dev", port: 443, hits: 2 }]);
+
+  hostnames.add("aikido.dev", 80);
+  t.same(hostnames.asArray(), [
+    { hostname: "aikido.dev", port: 443, hits: 2 },
+    { hostname: "aikido.dev", port: 80, hits: 1 },
+  ]);
+
+  hostnames.add("google.com", 80);
+  t.same(hostnames.asArray(), [
+    { hostname: "aikido.dev", port: 443, hits: 2 },
+    { hostname: "aikido.dev", port: 80, hits: 1 },
+    { hostname: "google.com", port: 80, hits: 1 },
+  ]);
+
+  hostnames.add("aikido.dev", 443);
+  t.same(hostnames.asArray(), [
+    { hostname: "aikido.dev", port: 443, hits: 3 },
+    { hostname: "aikido.dev", port: 80, hits: 1 },
+    { hostname: "google.com", port: 80, hits: 1 },
+  ]);
+});

library/agent/Hostnames.ts

@@ -1,30 +1,56 @@
+type Ports = Map<number, number>;
+
 export class Hostnames {
-  private map: Map<string, number | undefined> = new Map();
+  private map: Map<string, Ports> = new Map();
 
   constructor(private readonly maxEntries: number = 200) {}
 
-  add(hostname: string, port: number | undefined) {
-    if (this.map.has(hostname)) {
-      return;
+  add(hostname: string, port: number | undefined = -1) {
+    if (!this.map.has(hostname)) {
+      this.map.set(hostname, new Map([[port, 1]]));
+    } else {
+      const ports = this.map.get(hostname) as Ports;
+      if (!ports.has(port)) {
+        ports.set(port, 1);
+      } else {
+        ports.set(port, ports.get(port)! + 1);
+      }
     }
 
-    if (this.map.size >= this.maxEntries) {
+    if (this.length > this.maxEntries) {
       const firstAdded = this.map.keys().next().value;
       if (firstAdded) {
-        this.map.delete(firstAdded);
+        const ports: Ports = this.map.get(firstAdded) as Ports;
+
+        if (ports.size > 1) {
+          const firstPort = ports.keys().next().value;
+          if (firstPort) {
+            ports.delete(firstPort);
+          }
+        } else {
+          this.map.delete(firstAdded);
+        }
       }
     }
+  }
 
-    this.map.set(hostname, port);
+  get length() {
+    return Array.from(this.map.values()).reduce(
+      (total: number, ports: Ports) => total + ports.size,
+      0
+    );
   }
 
   asArray() {
-    return Array.from(this.map.entries()).map(([hostname, port]) => {
-      return {
-        hostname,
-        port,
-      };
-    });
+    return Array.from(this.map.entries()).flatMap(([hostname, ports]) =>
+      Array.from(ports.entries()).map(([port, hits]) => {
+        return {
+          hostname,
+          port: port === -1 ? undefined : port,
+          hits,
+        };
+      })
+    );
   }
 
   clear() {

library/sinks/Fetch.test.ts

@@ -88,14 +88,14 @@ t.test(
     await fetch("http://app.aikido.dev");
 
     t.same(agent.getHostnames().asArray(), [
-      { hostname: "app.aikido.dev", port: 80 },
+      { hostname: "app.aikido.dev", port: 80, hits: 1 },
     ]);
     agent.getHostnames().clear();
 
     await fetch(new URL("https://app.aikido.dev"));
 
     t.same(agent.getHostnames().asArray(), [
-      { hostname: "app.aikido.dev", port: 443 },
+      { hostname: "app.aikido.dev", port: 443, hits: 1 },
     ]);
     agent.getHostnames().clear();
 

library/sinks/HTTPRequest.axios.test.ts

@@ -37,7 +37,7 @@ t.test("it works", { skip: "SSRF redirect check disabled atm" }, async (t) => {
   });
 
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "www.aikido.dev", port: 443 },
+    { hostname: "www.aikido.dev", port: 443, hits: 1 },
   ]);
   agent.getHostnames().clear();
 

library/sinks/HTTPRequest.needle.test.ts

@@ -36,7 +36,7 @@ t.test("it works", { skip: "SSRF redirect check disabled atm" }, async (t) => {
   });
 
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "www.aikido.dev", port: 443 },
+    { hostname: "www.aikido.dev", port: 443, hits: 1 },
   ]);
   agent.getHostnames().clear();
 

library/sinks/HTTPRequest.nodeFetch.test.ts

@@ -36,7 +36,7 @@ t.test("it works", { skip: "SSRF redirect check disabled atm" }, async (t) => {
   });
 
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "www.aikido.dev", port: 443 },
+    { hostname: "www.aikido.dev", port: 443, hits: 1 },
   ]);
   agent.getHostnames().clear();
 

library/sinks/HTTPRequest.test.ts

@@ -69,7 +69,7 @@ t.test("it works", (t) => {
   });
 
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "aikido.dev", port: 80 },
+    { hostname: "aikido.dev", port: 80, hits: 1 },
   ]);
   agent.getHostnames().clear();
 
@@ -78,14 +78,14 @@ t.test("it works", (t) => {
     aikido.end();
   });
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "aikido.dev", port: 443 },
+    { hostname: "aikido.dev", port: 443, hits: 1 },
   ]);
   agent.getHostnames().clear();
 
   const aikido = https.request(new URL("https://aikido.dev"));
   aikido.end();
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "aikido.dev", port: 443 },
+    { hostname: "aikido.dev", port: 443, hits: 1 },
   ]);
   agent.getHostnames().clear();
 
@@ -96,7 +96,7 @@ t.test("it works", (t) => {
   t.same(withoutPort instanceof http.ClientRequest, true);
   withoutPort.end();
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "aikido.dev", port: 443 },
+    { hostname: "aikido.dev", port: 443, hits: 1 },
   ]);
   agent.getHostnames().clear();
 
@@ -107,23 +107,23 @@ t.test("it works", (t) => {
   httpWithoutPort.end();
   t.same(httpWithoutPort instanceof http.ClientRequest, true);
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "aikido.dev", port: 80 },
+    { hostname: "aikido.dev", port: 80, hits: 1 },
   ]);
   agent.getHostnames().clear();
 
   const withPort = https.request({ hostname: "aikido.dev", port: 443 });
   t.same(withPort instanceof http.ClientRequest, true);
   withPort.end();
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "aikido.dev", port: 443 },
+    { hostname: "aikido.dev", port: 443, hits: 1 },
   ]);
   agent.getHostnames().clear();
 
   const withStringPort = https.request({ hostname: "aikido.dev", port: "443" });
   t.same(withStringPort instanceof http.ClientRequest, true);
   withStringPort.end();
   t.same(agent.getHostnames().asArray(), [
-    { hostname: "aikido.dev", port: "443" },
+    { hostname: "aikido.dev", port: "443", hits: 1 },
   ]);
   agent.getHostnames().clear();