Merge pull request #116 from AikidoSec/AIK-4395

Add more test coverage
AikidoSec · Jan 31, 2025 · 2e6a151 · 2e6a151
2 parents d59a197 + 0639b46
commit 2e6a151
Show file tree

Hide file tree

Showing 8 changed files with 474 additions and 0 deletions.
diff --git a/agent_api/build.gradle b/agent_api/build.gradle
@@ -61,6 +61,9 @@ test {
         events "passed", "skipped", "failed"
         exceptionFormat "full"
     }
+    jacoco {
+        enabled = true
+    }
 }
 
 jacocoTestReport {

diff --git a/...t_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/GetBinaryPath.java b/...t_api/src/main/java/dev/aikido/agent_api/vulnerabilities/sql_injection/GetBinaryPath.java
@@ -21,6 +21,8 @@ private static String getFileName() {
             fileName.append("aarch64-"); // Add architecture to file name
         } else if (architecture.contains("64")) {
             fileName.append("x86_64-"); // Add architecture to file name
+        } else {
+            fileName.append("x86_64-"); // Default to x86-64
         }
 
         if (os.contains("win")) {
@@ -29,6 +31,8 @@ private static String getFileName() {
             fileName.append("apple-darwin.dylib"); // macOS
         } else if (os.contains("nix") || os.contains("nux")) {
             fileName.append("unknown-linux-gnu.so"); // Linux
+        } else {
+            fileName.append("unknown-linux-gnu.so"); // Default to linux.
         }
         return fileName.toString();
     }

diff --git a/agent_api/src/test/java/LogLevelTest.java b/agent_api/src/test/java/LogLevelTest.java
@@ -0,0 +1,27 @@
+import dev.aikido.agent_api.helpers.logging.LogLevel;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class LogLevelTest {
+
+    @Test
+    public void testLogLevelValues() {
+        assertEquals(0, LogLevel.TRACE.getLevel());
+        assertEquals(1, LogLevel.DEBUG.getLevel());
+        assertEquals(2, LogLevel.INFO.getLevel());
+        assertEquals(3, LogLevel.WARN.getLevel());
+        assertEquals(4, LogLevel.ERROR.getLevel());
+        assertEquals(5, LogLevel.FATAL.getLevel());
+    }
+
+    @Test
+    public void testLogLevelToString() {
+        assertEquals("TRACE", LogLevel.TRACE.toString());
+        assertEquals("DEBUG", LogLevel.DEBUG.toString());
+        assertEquals("INFO", LogLevel.INFO.toString());
+        assertEquals("WARN", LogLevel.WARN.toString());
+        assertEquals("ERROR", LogLevel.ERROR.toString());
+        assertEquals("FATAL", LogLevel.FATAL.toString());
+    }
+}
diff --git a/agent_api/src/test/java/ShouldBlockRequestTest.java b/agent_api/src/test/java/ShouldBlockRequestTest.java
@@ -1,9 +1,12 @@
 import dev.aikido.agent_api.SetUser;
 import dev.aikido.agent_api.ShouldBlockRequest;
 import dev.aikido.agent_api.background.Endpoint;
+import dev.aikido.agent_api.background.ipc_commands.ShouldRateLimitCommand;
+import dev.aikido.agent_api.background.utilities.ThreadIPCClient;
 import dev.aikido.agent_api.context.Context;
 import dev.aikido.agent_api.context.ContextObject;
 import dev.aikido.agent_api.context.User;
+import dev.aikido.agent_api.ratelimiting.ShouldRateLimit;
 import dev.aikido.agent_api.storage.routes.Routes;
 import dev.aikido.agent_api.thread_cache.ThreadCache;
 import dev.aikido.agent_api.thread_cache.ThreadCacheObject;
@@ -17,7 +20,11 @@
 import java.sql.SQLException;
 import java.util.*;
 
+import static dev.aikido.agent_api.background.utilities.ThreadIPCClientFactory.getDefaultThreadIPCClient;
 import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 import static utils.EmtpyThreadCacheObject.getEmptyThreadCacheObject;
 
 public class ShouldBlockRequestTest {
@@ -209,4 +216,62 @@ public void testThreadClientInvalid2() throws SQLException {
         assertFalse(res1.block());
     }
 
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "valid-token")
+    public void testNoEndpointsConfigured() throws SQLException {
+        // Set up context with a user
+        ContextObject ctx = new SampleContextObject();
+        ctx.setUser(new User("ID3", "Alice", "192.168.1.3", 100));
+        Context.set(ctx);
+
+        // Set up thread cache with no endpoints
+        ThreadCache.set(new ThreadCacheObject(List.of(), Set.of(), Set.of(), new Routes(), Optional.empty()));
+
+        // Call the method
+        var res = ShouldBlockRequest.shouldBlockRequest();
+
+        // Assert that the request is not blocked
+        assertFalse(res.block());
+    }
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "valid-token")
+    public void testBlockedUserWithMultipleEndpoints() throws SQLException {
+        // Set up context with a blocked user
+        ContextObject ctx = new SampleContextObject();
+        ctx.setUser(new User("ID1", "John Doe", "192.168.1.1", 100));
+        Context.set(ctx);
+
+        // Set up thread cache with multiple endpoints and a blocked user
+        ThreadCache.set(new ThreadCacheObject(List.of(
+                new Endpoint("GET", "/api/resource", 1, 1000, Collections.emptyList(), false, false, true),
+                new Endpoint("POST", "/api/resource", 1, 1000, Collections.emptyList(), false, false, true)
+        ), Set.of("ID1"), Set.of(), new Routes(), Optional.empty()));
+
+        // Call the method
+        var res = ShouldBlockRequest.shouldBlockRequest();
+
+        // Assert that the request is blocked due to the user being blocked
+        assertTrue(res.block());
+        assertEquals("user", res.data().trigger());
+        assertEquals("blocked", res.data().type());
+        assertEquals("192.168.1.1", res.data().ip());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "valid-token")
+    public void testNoUserWithEndpoints() throws SQLException {
+        // Set up context without a user
+        Context.set(new SampleContextObject());
+
+        // Set up thread cache with endpoints
+        ThreadCache.set(new ThreadCacheObject(List.of(
+                new Endpoint("GET", "/api/resource", 1, 1000, Collections.emptyList(), false, false, true)
+        ), Set.of(), Set.of(), new Routes(), Optional.empty()));
+
+        // Call the method
+        var res = ShouldBlockRequest.shouldBlockRequest();
+
+        // Assert that the request is not blocked
+        assertFalse(res.block());
+    }
 }
diff --git a/agent_api/src/test/java/helpers/env/BooleanEnvTest.java b/agent_api/src/test/java/helpers/env/BooleanEnvTest.java
@@ -0,0 +1,81 @@
+package helpers.env;
+
+import dev.aikido.agent_api.helpers.env.BooleanEnv;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junitpioneer.jupiter.SetEnvironmentVariable;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class BooleanEnvTest {
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "1")
+    public void testBooleanEnv_WithValueOne() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", false);
+        assertTrue(booleanEnv.getValue());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "true")
+    public void testBooleanEnv_WithValueTrue() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", false);
+        assertTrue(booleanEnv.getValue());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "TRUE")
+    public void testBooleanEnv_WithValueTrueUppercase() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", false);
+        assertTrue(booleanEnv.getValue());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "0")
+    public void testBooleanEnv_WithValueZero() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", true);
+        assertFalse(booleanEnv.getValue());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "false")
+    public void testBooleanEnv_WithValueFalse() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", true);
+        assertFalse(booleanEnv.getValue());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "FALSE")
+    public void testBooleanEnv_WithValueFalseUppercase() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", true);
+        assertFalse(booleanEnv.getValue());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "")
+    public void testBooleanEnv_WithEmptyString() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", true);
+        assertTrue(booleanEnv.getValue());
+    }
+
+    @Test
+    public void testBooleanEnv_WithNullEnvironmentVariable() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", true);
+        assertTrue(booleanEnv.getValue());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "randomString")
+    public void testBooleanEnv_WithRandomString() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", false);
+        assertFalse(booleanEnv.getValue());
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "TEST_BOOLEAN_ENV", value = "TrUe")
+    public void testBooleanEnv_WithMixedCaseTrue() {
+        BooleanEnv booleanEnv = new BooleanEnv("TEST_BOOLEAN_ENV", false);
+        assertTrue(booleanEnv.getValue());
+    }
+}
diff --git a/agent_api/src/test/java/helpers/env/EndpointsTest.java b/agent_api/src/test/java/helpers/env/EndpointsTest.java
@@ -0,0 +1,104 @@
+package helpers.env;
+
+import dev.aikido.agent_api.helpers.env.Endpoints;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.AfterEach;
+import org.junitpioneer.jupiter.SetEnvironmentVariable;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+public class EndpointsTest {
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://custom.aikido.dev")
+    public void testGetAikidoAPIEndpoint_WithCustomEndpoint() {
+        String result = Endpoints.getAikidoAPIEndpoint();
+        assertEquals("https://custom.aikido.dev/", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://custom.aikido.dev")
+    public void testGetAikidoAPIEndpoint_WithCustomEndpointWithoutTrailingSlash() {
+        String result = Endpoints.getAikidoAPIEndpoint();
+        assertEquals("https://custom.aikido.dev/", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "")
+    public void testGetAikidoAPIEndpoint_WithEmptyEnvironmentVariable() {
+        String result = Endpoints.getAikidoAPIEndpoint();
+        assertEquals("https://guard.aikido.dev/", result);
+    }
+
+    @Test
+    public void testGetAikidoAPIEndpoint_WithNullEnvironmentVariable() {
+        // No environment variable set, should return default
+        String result = Endpoints.getAikidoAPIEndpoint();
+        assertEquals("https://guard.aikido.dev/", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_REALTIME_ENDPOINT", value = "https://realtime.aikido.dev")
+    public void testGetAikidoRealtimeEndpoint_WithCustomEndpoint() {
+        String result = Endpoints.getAikidoRealtimeEndpoint();
+        assertEquals("https://realtime.aikido.dev/", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_REALTIME_ENDPOINT", value = "https://realtime.aikido.dev")
+    public void testGetAikidoRealtimeEndpoint_WithCustomEndpointWithoutTrailingSlash() {
+        String result = Endpoints.getAikidoRealtimeEndpoint();
+        assertEquals("https://realtime.aikido.dev/", result);
+    }
+
+    @Test
+    public void testGetAikidoRealtimeEndpoint_WithNullEnvironmentVariable() {
+        // No environment variable set, should return default
+        String result = Endpoints.getAikidoRealtimeEndpoint();
+        assertEquals("https://runtime.aikido.dev/", result);
+    }
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_REALTIME_ENDPOINT", value = "https://realtime.aikido.dev/")
+    public void testGetAikidoRealtimeEndpoint_WithCustomEndpointWithTrailingSlash() {
+        String result = Endpoints.getAikidoRealtimeEndpoint();
+        assertEquals("https://realtime.aikido.dev/", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_REALTIME_ENDPOINT", value = "")
+    public void testGetAikidoRealtimeEndpoint_WithEmptyEnvironmentVariable() {
+        String result = Endpoints.getAikidoRealtimeEndpoint();
+        assertEquals("https://runtime.aikido.dev/", result);
+    }
+
+    // Additional tests
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://missing-slash.aikido.dev")
+    public void testGetAikidoAPIEndpoint_WithMissingSlash() {
+        String result = Endpoints.getAikidoAPIEndpoint();
+        assertEquals("https://missing-slash.aikido.dev/", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://missing-slash.aikido.dev/")
+    public void testGetAikidoAPIEndpoint_WithTrailingSlashAlreadyPresent() {
+        String result = Endpoints.getAikidoAPIEndpoint();
+        assertEquals("https://missing-slash.aikido.dev/", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://another-custom.aikido.dev")
+    public void testGetAikidoAPIEndpoint_WithAnotherCustomEndpoint() {
+        String result = Endpoints.getAikidoAPIEndpoint();
+        assertEquals("https://another-custom.aikido.dev/", result);
+    }
+
+    @Test
+    @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://another-custom.aikido.dev/")
+    public void testGetAikidoAPIEndpoint_WithAnotherCustomEndpointWithTrailingSlash() {
+        String result = Endpoints.getAikidoAPIEndpoint();
+        assertEquals("https://another-custom.aikido.dev/", result);
+    }
+}