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 5baaa1f6..51bd22e5 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 @@ -59,7 +59,7 @@ public void run() { // one time check to report initial stats scheduler.schedule(new HeartbeatTask(api, true), 60, TimeUnit.SECONDS); - if (token != null && FeatureFlags.AIKIDO_FEATURE_SSE.isEnabled()) { + if (token != null && (FeatureFlags.AIKIDO_FEATURE_SSE.isEnabled() || ServiceConfigStore.isRealtimeUpdatesEnabled())) { new RealtimeSSETask(new RealtimeSSEAPI(token), api).start(); } } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/APIResponse.java b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/APIResponse.java index 30d86232..18573f28 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/APIResponse.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/APIResponse.java @@ -16,6 +16,23 @@ public record APIResponse( List domains, boolean receivedAnyStats, boolean block, - List excludedUserIdsFromRateLimiting + List excludedUserIdsFromRateLimiting, + List enabledFeatures ) { + public APIResponse( + boolean success, + String error, + long configUpdatedAt, + List endpoints, + List blockedUserIds, + List allowedIPAddresses, + boolean blockNewOutgoingRequests, + List domains, + boolean receivedAnyStats, + boolean block, + List excludedUserIdsFromRateLimiting + ) { + this(success, error, configUpdatedAt, endpoints, blockedUserIds, allowedIPAddresses, blockNewOutgoingRequests, + domains, receivedAnyStats, block, excludedUserIdsFromRateLimiting, List.of()); + } } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfigStore.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfigStore.java index 472a4f72..df972aed 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfigStore.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfigStore.java @@ -41,6 +41,15 @@ public static boolean isBlockedUserAgent(String userAgent) { } } + public static boolean isRealtimeUpdatesEnabled() { + mutex.readLock().lock(); + try { + return config.isRealtimeUpdatesEnabled(); + } finally { + mutex.readLock().unlock(); + } + } + public static void updateFromAPIResponse(APIResponse apiResponse) { mutex.writeLock().lock(); try { diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfiguration.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfiguration.java index 681d2525..51fd4066 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfiguration.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfiguration.java @@ -25,6 +25,7 @@ public class ServiceConfiguration { private IPList bypassedIPs = new IPList(); private HashSet blockedUserIDs = new HashSet<>(); private HashSet excludedUserIdsFromRateLimiting = new HashSet<>(); + private HashSet enabledFeatures = new HashSet<>(); private List endpoints = new ArrayList<>(); private OutboundDomains outboundDomains = new OutboundDomains(); @@ -47,6 +48,9 @@ public void updateConfig(APIResponse apiResponse) { if (apiResponse.excludedUserIdsFromRateLimiting() != null) { this.excludedUserIdsFromRateLimiting = new HashSet<>(apiResponse.excludedUserIdsFromRateLimiting()); } + if (apiResponse.enabledFeatures() != null) { + this.enabledFeatures = new HashSet<>(apiResponse.enabledFeatures()); + } if (apiResponse.endpoints() != null) { this.endpoints = apiResponse.endpoints(); } @@ -58,6 +62,10 @@ public boolean isBlockingEnabled() { return blockingEnabled; } + public boolean isRealtimeUpdatesEnabled() { + return enabledFeatures.contains("realtime_updates"); + } + public void setBlocking(boolean block) { this.blockingEnabled = block; } diff --git a/agent_api/src/test/java/storage/ServiceConfigStoreTest.java b/agent_api/src/test/java/storage/ServiceConfigStoreTest.java index c997addc..2bb8cf59 100644 --- a/agent_api/src/test/java/storage/ServiceConfigStoreTest.java +++ b/agent_api/src/test/java/storage/ServiceConfigStoreTest.java @@ -58,4 +58,18 @@ public void testShouldBlockOutgoingRequestUnknownWhenBlockNewEnabled() { )); assertTrue(ServiceConfigStore.shouldBlockOutgoingRequest("unknown.com")); } + + @Test + public void testIsRealtimeUpdatesEnabledFalseByDefault() { + assertFalse(ServiceConfigStore.isRealtimeUpdatesEnabled()); + } + + @Test + public void testIsRealtimeUpdatesEnabledTrueWhenEnabled() { + ServiceConfigStore.updateFromAPIResponse(new APIResponse( + true, null, 0L, null, null, null, + false, null, true, false, null, List.of("realtime_updates") + )); + assertTrue(ServiceConfigStore.isRealtimeUpdatesEnabled()); + } } diff --git a/agent_api/src/test/java/storage/ServiceConfigurationTest.java b/agent_api/src/test/java/storage/ServiceConfigurationTest.java index cd4801e9..def666c7 100644 --- a/agent_api/src/test/java/storage/ServiceConfigurationTest.java +++ b/agent_api/src/test/java/storage/ServiceConfigurationTest.java @@ -665,4 +665,48 @@ public void testIsIpBlockedWithAllowedIPsAndBlockedIPs() { assertTrue(resultBlocked.blocked()); assertFalse(resultNotAllowedLocal.blocked()); } + + @Test + public void testIsRealtimeUpdatesEnabledDefaultsFalse() { + assertFalse(serviceConfiguration.isRealtimeUpdatesEnabled()); + } + + @Test + public void testIsRealtimeUpdatesEnabledTrueWhenPresent() { + APIResponse apiResponse = new APIResponse( + true, null, 0L, null, null, null, + false, null, true, false, null, List.of("realtime_updates") + ); + serviceConfiguration.updateConfig(apiResponse); + + assertTrue(serviceConfiguration.isRealtimeUpdatesEnabled()); + } + + @Test + public void testIsRealtimeUpdatesEnabledFalseWhenRemoved() { + APIResponse enabled = new APIResponse( + true, null, 0L, null, null, null, + false, null, true, false, null, List.of("realtime_updates") + ); + serviceConfiguration.updateConfig(enabled); + assertTrue(serviceConfiguration.isRealtimeUpdatesEnabled()); + + APIResponse disabled = new APIResponse( + true, null, 0L, null, null, null, + false, null, true, false, null, List.of() + ); + serviceConfiguration.updateConfig(disabled); + assertFalse(serviceConfiguration.isRealtimeUpdatesEnabled()); + } + + @Test + public void testIsRealtimeUpdatesEnabledFalseForUnrelatedFeature() { + APIResponse apiResponse = new APIResponse( + true, null, 0L, null, null, null, + false, null, true, false, null, List.of("some_other_feature") + ); + serviceConfiguration.updateConfig(apiResponse); + + assertFalse(serviceConfiguration.isRealtimeUpdatesEnabled()); + } }