diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java b/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java index d08e070a..02419abe 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/background/BackgroundProcess.java @@ -38,7 +38,7 @@ public void run() { } logger.trace("Starting agent thread: %s", Thread.currentThread()); - ReportingApiHTTP api = new ReportingApiHTTP(getAikidoAPIEndpoint(), API_TIMEOUT, token); + ReportingApiHTTP api = new ReportingApiHTTP(getAikidoAPIEndpoint(token), API_TIMEOUT, token); RealtimeAPI realtimeApi = new RealtimeAPI(token); // make api calls on start diff --git a/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/Endpoints.java b/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/Endpoints.java index 190c5136..463689ed 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/Endpoints.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/Endpoints.java @@ -2,7 +2,7 @@ public final class Endpoints { private Endpoints() {} - public static String getAikidoAPIEndpoint() { + public static String getAikidoAPIEndpoint(Token token) { String endpoint = System.getenv("AIKIDO_ENDPOINT"); if (endpoint != null && !endpoint.isEmpty()) { if (!endpoint.endsWith("/")) { @@ -11,8 +11,13 @@ public static String getAikidoAPIEndpoint() { return endpoint; } - // Default option : - return "https://guard.aikido.dev/"; + String region = token != null ? token.getRegion() : "EU"; + return switch (region) { + case "US" -> "https://guard.us.aikido.dev/"; + case "ME" -> "https://guard.me.aikido.dev/"; + case "AU" -> "https://guard.au.aikido.dev/"; + default -> "https://guard.aikido.dev/"; + }; } public static String getAikidoRealtimeEndpoint() { diff --git a/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/Token.java b/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/Token.java index ea73868b..bace6fdf 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/Token.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/helpers/env/Token.java @@ -25,6 +25,27 @@ public static Token fromEnv() { return new Token(token); } + /** + * Extracts the region from a runtime token. + * New format: AIK_RUNTIME_{sys_group_id}_{service_id}_{region}_{random} + * Old format: AIK_RUNTIME_{sys_group_id}_{service_id}_{random} + */ + public static String extractRegionFromToken(String token) { + if (token == null || !token.startsWith("AIK_RUNTIME_")) { + return "EU"; + } + String tokenWithoutPrefix = token.substring("AIK_RUNTIME_".length()); + String[] parts = tokenWithoutPrefix.split("_"); + if (parts.length == 4) { + return parts[2]; + } + return "EU"; + } + + public String getRegion() { + return extractRegionFromToken(token); + } + /** * Hashes the token with SHA-256 and returns the hashed bytes in a hex representation (alphanum) */ diff --git a/agent_api/src/test/java/helpers/TokenTest.java b/agent_api/src/test/java/helpers/TokenTest.java index def95dae..66096ed5 100644 --- a/agent_api/src/test/java/helpers/TokenTest.java +++ b/agent_api/src/test/java/helpers/TokenTest.java @@ -135,4 +135,42 @@ void constructor_shouldAcceptTokenWithMixedCase() { Token token = new Token(mixedCaseToken); assertEquals(mixedCaseToken, token.get()); } + + @Test + void getRegion_shouldReturnRegion_forNewFormatToken() { + Token token = new Token("AIK_RUNTIME_1_2_US_random"); + assertEquals("US", token.getRegion()); + } + + @Test + void getRegion_shouldReturnEU_forOldFormatToken() { + Token token = new Token("AIK_RUNTIME_1_2_random"); + assertEquals("EU", token.getRegion()); + } + + @Test + void getRegion_shouldReturnEU_forNonRuntimeToken() { + Token token = new Token("some-other-token"); + assertEquals("EU", token.getRegion()); + } + + @Test + void extractRegionFromToken_shouldReturnEU_whenTokenIsNull() { + assertEquals("EU", Token.extractRegionFromToken(null)); + } + + @Test + void extractRegionFromToken_shouldReturnEU_whenTokenIsEmpty() { + assertEquals("EU", Token.extractRegionFromToken("")); + } + + @Test + void extractRegionFromToken_shouldReturnRegion_forME() { + assertEquals("ME", Token.extractRegionFromToken("AIK_RUNTIME_1_2_ME_random")); + } + + @Test + void extractRegionFromToken_shouldReturnRegion_forAU() { + assertEquals("AU", Token.extractRegionFromToken("AIK_RUNTIME_1_2_AU_random")); + } } \ No newline at end of file diff --git a/agent_api/src/test/java/helpers/env/EndpointsTest.java b/agent_api/src/test/java/helpers/env/EndpointsTest.java index 5bfb5884..2d65f9cb 100644 --- a/agent_api/src/test/java/helpers/env/EndpointsTest.java +++ b/agent_api/src/test/java/helpers/env/EndpointsTest.java @@ -1,6 +1,7 @@ package helpers.env; import dev.aikido.agent_api.helpers.env.Endpoints; +import dev.aikido.agent_api.helpers.env.Token; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.api.Assertions; @@ -14,28 +15,28 @@ public class EndpointsTest { @Test @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://custom.aikido.dev") public void testGetAikidoAPIEndpoint_WithCustomEndpoint() { - String result = Endpoints.getAikidoAPIEndpoint(); + String result = Endpoints.getAikidoAPIEndpoint(null); assertEquals("https://custom.aikido.dev/", result); } @Test @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://custom.aikido.dev") public void testGetAikidoAPIEndpoint_WithCustomEndpointWithoutTrailingSlash() { - String result = Endpoints.getAikidoAPIEndpoint(); + String result = Endpoints.getAikidoAPIEndpoint(null); assertEquals("https://custom.aikido.dev/", result); } @Test @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "") public void testGetAikidoAPIEndpoint_WithEmptyEnvironmentVariable() { - String result = Endpoints.getAikidoAPIEndpoint(); + String result = Endpoints.getAikidoAPIEndpoint(null); assertEquals("https://guard.aikido.dev/", result); } @Test public void testGetAikidoAPIEndpoint_WithNullEnvironmentVariable() { // No environment variable set, should return default - String result = Endpoints.getAikidoAPIEndpoint(); + String result = Endpoints.getAikidoAPIEndpoint(null); assertEquals("https://guard.aikido.dev/", result); } @@ -77,28 +78,78 @@ public void testGetAikidoRealtimeEndpoint_WithEmptyEnvironmentVariable() { @Test @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://missing-slash.aikido.dev") public void testGetAikidoAPIEndpoint_WithMissingSlash() { - String result = Endpoints.getAikidoAPIEndpoint(); + String result = Endpoints.getAikidoAPIEndpoint(null); 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(); + String result = Endpoints.getAikidoAPIEndpoint(null); 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(); + String result = Endpoints.getAikidoAPIEndpoint(null); 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(); + String result = Endpoints.getAikidoAPIEndpoint(null); assertEquals("https://another-custom.aikido.dev/", result); } + + @Test + public void testGetAikidoAPIEndpoint_WithUSRegionToken() { + Token token = new Token("AIK_RUNTIME_1_2_US_random"); + String result = Endpoints.getAikidoAPIEndpoint(token); + assertEquals("https://guard.us.aikido.dev/", result); + } + + @Test + public void testGetAikidoAPIEndpoint_WithMERegionToken() { + Token token = new Token("AIK_RUNTIME_1_2_ME_random"); + String result = Endpoints.getAikidoAPIEndpoint(token); + assertEquals("https://guard.me.aikido.dev/", result); + } + + @Test + public void testGetAikidoAPIEndpoint_WithAURegionToken() { + Token token = new Token("AIK_RUNTIME_1_2_AU_random"); + String result = Endpoints.getAikidoAPIEndpoint(token); + assertEquals("https://guard.au.aikido.dev/", result); + } + + @Test + public void testGetAikidoAPIEndpoint_WithEURegionToken() { + Token token = new Token("AIK_RUNTIME_1_2_EU_random"); + String result = Endpoints.getAikidoAPIEndpoint(token); + assertEquals("https://guard.aikido.dev/", result); + } + + @Test + public void testGetAikidoAPIEndpoint_WithOldFormatToken() { + Token token = new Token("AIK_RUNTIME_1_2_random"); + String result = Endpoints.getAikidoAPIEndpoint(token); + assertEquals("https://guard.aikido.dev/", result); + } + + @Test + public void testGetAikidoAPIEndpoint_WithNonRuntimeToken() { + Token token = new Token("some-other-token"); + String result = Endpoints.getAikidoAPIEndpoint(token); + assertEquals("https://guard.aikido.dev/", result); + } + + @Test + @SetEnvironmentVariable(key = "AIKIDO_ENDPOINT", value = "https://custom.aikido.dev") + public void testGetAikidoAPIEndpoint_CustomEndpointTakesPrecedenceOverRegion() { + Token token = new Token("AIK_RUNTIME_1_2_US_random"); + String result = Endpoints.getAikidoAPIEndpoint(token); + assertEquals("https://custom.aikido.dev/", result); + } }