Merge pull request #112 from AikidoSec/AIK-4387

AIK-4387 Make sure route hit reporting works for Java

Author: bitterpanda63

agent_api/src/main/java/dev/aikido/agent_api/storage/Hostnames.java

@@ -24,7 +24,13 @@ public void add(String hostname, int port) {
     }
     public void addArray(HostnameEntry[] hostnameEntries) {
         for (HostnameEntry entry: hostnameEntries) {
-            add(entry.getHostname(), entry.getPort());
+            String key = getKey(entry.getHostname(), entry.getPort());
+            if (map.containsKey(key)) {
+                // Merge hits :
+                map.get(key).incrementHits(entry.getHits());
+            } else {
+                map.put(key, entry);
+            }
         }
     }
     public HostnameEntry[] asArray() {
@@ -53,6 +59,10 @@ public HostnameEntry(String hostname, int port) {
         public void incrementHits() {
             hits++;
         }
+        public void incrementHits(int delta) {
+            hits += delta;
+        }
+
 
         public String getHostname() {
             return hostname;

agent_api/src/test/java/storage/HostnamesTest.java

@@ -105,6 +105,78 @@ public void testExceedMaxEntriesWithMultiplePorts() {
         assertTrue(containsEntry(entries, "test.com", 8080));
         assertTrue(containsEntry(entries, "newsite.com", 3000));
     }
+    @Test
+    public void testAddArrayWithNewEntries() {
+        Hostnames.HostnameEntry[] entriesToAdd = {
+                new Hostnames.HostnameEntry("example.com", 80),
+                new Hostnames.HostnameEntry("test.com", 443)
+        };
+        entriesToAdd[0].incrementHits();
+        entriesToAdd[1].incrementHits();
+
+        hostnames.addArray(entriesToAdd);
+
+        Hostnames.HostnameEntry[] entries = hostnames.asArray();
+        assertEquals(2, entries.length);
+        assertTrue(containsEntry(entries, "example.com", 80));
+        assertTrue(containsEntry(entries, "test.com", 443));
+    }
+
+    @Test
+    public void testAddArrayWithExistingEntries() {
+        hostnames.add("example.com", 80); // Initial entry
+        Hostnames.HostnameEntry[] entriesToAdd = {
+                new Hostnames.HostnameEntry("example.com", 80), // Same entry, should merge hits
+                new Hostnames.HostnameEntry("test.com", 443)    // New entry
+        };
+        entriesToAdd[0].incrementHits();
+        entriesToAdd[1].incrementHits();
+
+        hostnames.addArray(entriesToAdd);
+
+        Hostnames.HostnameEntry[] entries = hostnames.asArray();
+        assertEquals(2, entries.length);
+        assertEquals(2, getHits(entries, "example.com", 80)); // Hits should be 2
+        assertEquals(1, getHits(entries, "test.com", 443));   // Hits should be 1
+    }
+
+    @Test
+    public void testAddArrayExceedMaxEntries() {
+        hostnames.add("example.com", 80);
+        hostnames.add("test.com", 443);
+        hostnames.add("localhost", 3000);
+
+        Hostnames.HostnameEntry[] entriesToAdd = {
+                new Hostnames.HostnameEntry("newsite.com", 8080), // This should cause an eviction
+                new Hostnames.HostnameEntry("example.com", 80)    // Should merge hits
+        };
+        entriesToAdd[0].incrementHits();
+        entriesToAdd[1].incrementHits(10);
+
+        hostnames.addArray(entriesToAdd);
+
+        Hostnames.HostnameEntry[] entries = hostnames.asArray();
+        assertEquals(3, entries.length);
+        assertFalse(containsEntry(entries, "test.com", 443)); // "test.com" should be evicted
+        assertTrue(containsEntry(entries, "localhost", 3000));
+        assertTrue(containsEntry(entries, "newsite.com", 8080));
+        assertEquals(10, getHits(entries, "example.com", 80)); // Hits should be 2
+    }
+
+    @Test
+    public void testAddArrayWithZeroPort() {
+        Hostnames.HostnameEntry[] entriesToAdd = {
+                new Hostnames.HostnameEntry("example.com", 0),
+                new Hostnames.HostnameEntry("test.com", 0)
+        };
+
+        hostnames.addArray(entriesToAdd);
+
+        Hostnames.HostnameEntry[] entries = hostnames.asArray();
+        assertEquals(2, entries.length);
+        assertTrue(containsEntry(entries, "example.com", 0));
+        assertTrue(containsEntry(entries, "test.com", 0));
+    }
 
     private boolean containsEntry(Hostnames.HostnameEntry[] entries, String hostname, int port) {
         for (Hostnames.HostnameEntry entry : entries) {
@@ -114,4 +186,12 @@ private boolean containsEntry(Hostnames.HostnameEntry[] entries, String hostname
         }
         return false;
     }
+    private int getHits(Hostnames.HostnameEntry[] entries, String hostname, int port) {
+        for (Hostnames.HostnameEntry entry : entries) {
+            if (entry.getHostname().equals(hostname) && entry.getPort() == port) {
+                return entry.getHits();
+            }
+        }
+        return 0; // Not found
+    }
 }