Add timeout for API requests to Core

agent_api/src/main/java/dev/aikido/agent_api/background/RealtimeTask.java

@@ -34,7 +34,7 @@ public void run() {
             if(configLastUpdatedAt.isEmpty() || configLastUpdatedAt.get() < configAccordingToCloudUpdatedAt) {
                 // Config was updated
                 configLastUpdatedAt = Optional.of(configAccordingToCloudUpdatedAt); // Store new time of last update
-                Optional<APIResponse> newConfig = connectionManager.getApi().fetchNewConfig(connectionManager.getToken(), /* Timeout in seconds: */ 3);
+                Optional<APIResponse> newConfig = connectionManager.getApi().fetchNewConfig(connectionManager.getToken());
                 newConfig.ifPresent(connectionManager.getConfig()::updateConfig);
 
                 // Fetch blocked lists from separate API route :

agent_api/src/main/java/dev/aikido/agent_api/background/cloud/CloudConnectionManager.java

@@ -34,7 +34,7 @@ public class CloudConnectionManager {
     private final Hostnames hostnames;
 
     public CloudConnectionManager(boolean block, Token token, String serverless) {
-        this(block, token, serverless, new ReportingApiHTTP(getAikidoAPIEndpoint()));
+        this(block, token, serverless, new ReportingApiHTTP(getAikidoAPIEndpoint(), timeout));
     }
     public CloudConnectionManager(boolean block, Token token, String serverless, ReportingApi api) {
         this.config = new ServiceConfiguration(block, serverless);
@@ -54,7 +54,7 @@ public void onStart() {
         config.storeBlockedListsRes(api.fetchBlockedLists(token));
     }
     public void reportEvent(APIEvent event, boolean updateConfig) {
-        Optional<APIResponse> res = this.api.report(this.token, event, timeout);
+        Optional<APIResponse> res = this.api.report(this.token, event);
         if (res.isPresent() && updateConfig) {
             config.updateConfig(res.get());
         }

agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApi.java

@@ -7,6 +7,12 @@
 import java.util.Optional;
 
 public abstract class ReportingApi {
+    public int timeoutInSec;
+
+    public ReportingApi(int timeoutInSec) {
+        this.timeoutInSec = timeoutInSec;
+    }
+
     /**
      * Converts results into an API response object.
      *
@@ -20,10 +26,8 @@ public abstract class ReportingApi {
      *
      * @param token        The authentication token.
      * @param event        The event to report.
-     * @param timeoutInSec The timeout in seconds.
-     * @return
      */
-    public abstract Optional<APIResponse> report(String token, APIEvent event, int timeoutInSec);
+    public abstract Optional<APIResponse> report(String token, APIEvent event);
 
     public record APIListsResponse(List<ListsResponseEntry> blockedIPAddresses, String blockedUserAgents) {}
     public record ListsResponseEntry(String source, String description, List<String> ips) {}

agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApiHTTP.java

@@ -23,17 +23,21 @@ public class ReportingApiHTTP extends ReportingApi {
     private final Logger logger = LogManager.getLogger(ReportingApiHTTP.class);
     private final String reportingUrl;
     private final Gson gson = new Gson();
-    public ReportingApiHTTP(String reportingUrl) {
+    public ReportingApiHTTP(String reportingUrl, int timeoutInSec) {
         // Reporting URL should end with trailing slash for now.
+        super(timeoutInSec);
         this.reportingUrl = reportingUrl;
     }
-    public Optional<APIResponse> fetchNewConfig(String token, int timeoutInSec) {
+
+    public Optional<APIResponse> fetchNewConfig(String token) {
         try {
             HttpClient httpClient = HttpClient.newBuilder()
                     .connectTimeout(Duration.ofSeconds(timeoutInSec))
                     .build();
+
             URI uri = URI.create(reportingUrl + "api/runtime/config");
             HttpRequest request = createHttpRequest(Optional.empty(), token, uri);
+
             // Send the request and get the response
             HttpResponse<String> httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
             return Optional.of(toApiResponse(httpResponse));
@@ -42,14 +46,17 @@ public Optional<APIResponse> fetchNewConfig(String token, int timeoutInSec) {
         }
         return Optional.empty();
     }
+
     @Override
-    public Optional<APIResponse> report(String token, APIEvent event, int timeoutInSec) {
+    public Optional<APIResponse> report(String token, APIEvent event) {
         try {
             HttpClient httpClient = HttpClient.newBuilder()
                 .connectTimeout(Duration.ofSeconds(timeoutInSec))
                 .build();
+
             URI uri = URI.create(reportingUrl + "api/runtime/events");
             HttpRequest request = createHttpRequest(Optional.of(event), token, uri);
+
             // Send the request and get the response
             HttpResponse<String> httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
             return Optional.of(toApiResponse(httpResponse));
@@ -105,9 +112,10 @@ public APIResponse toApiResponse(HttpResponse<String> res) {
         }
         return getUnsuccessfulAPIResponse("unknown_error");
     }
-    private static HttpRequest createHttpRequest(Optional<APIEvent> event, String token, URI uri) {
+    private HttpRequest createHttpRequest(Optional<APIEvent> event, String token, URI uri) {
         HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
             .uri(uri) // Change to your target URL
+            .timeout(Duration.ofSeconds(timeoutInSec))
             .header("Content-Type", "application/json") // Set Content-Type header
             .header("Authorization", token); // Set Authorization header
         if (event.isPresent()) {

agent_api/src/test/java/background/RealtimeTaskTest.java

@@ -46,7 +46,7 @@ void testRunWithUpdatedConfig() {
         RealtimeAPI.ConfigResponse configResponse = mock(RealtimeAPI.ConfigResponse.class);
         // Mock the API response for fetching new config
         APIResponse apiResponse = mock(APIResponse.class);
-        when(reportingApiHTTP.fetchNewConfig(token, 3)).thenReturn(Optional.of(apiResponse));
+        when(reportingApiHTTP.fetchNewConfig(token)).thenReturn(Optional.of(apiResponse));
 
         // Act
         realtimeTask.run();
@@ -67,7 +67,7 @@ void testRunWithNoConfigUpdate() {
         RealtimeAPI.ConfigResponse configResponse = mock(RealtimeAPI.ConfigResponse.class);
         // Mock the API response for fetching new config
         APIResponse apiResponse = mock(APIResponse.class);
-        when(reportingApiHTTP.fetchNewConfig(token, 3)).thenReturn(Optional.of(apiResponse));
+        when(reportingApiHTTP.fetchNewConfig(token)).thenReturn(Optional.of(apiResponse));
 
         // Act
         realtimeTask.run();
@@ -89,7 +89,7 @@ void testRunWithEmptyResponse() {
         // Arrange
         String token = "test-token";
         when(connectionManager.getToken()).thenReturn(token);
-        when(reportingApiHTTP.fetchNewConfig(token, 3)).thenReturn(Optional.empty());
+        when(reportingApiHTTP.fetchNewConfig(token)).thenReturn(Optional.empty());
         // Act
         realtimeTask.run();
 

agent_api/src/test/java/background/cloud/CloudConnectionManagerTest.java

@@ -37,21 +37,21 @@ void setUp() {
     @Test
     void testOnStartReportsEvent() {
         // Arrange
-        when(mockApi.report(anyString(), any(APIEvent.class), anyInt())).thenReturn(Optional.of(mock(APIResponse.class)));
+        when(mockApi.report(anyString(), any(APIEvent.class))).thenReturn(Optional.of(mock(APIResponse.class)));
 
         // Act
         cloudConnectionManager.onStart();
 
         // Assert
         ArgumentCaptor<APIEvent> eventCaptor = ArgumentCaptor.forClass(APIEvent.class);
-        verify(mockApi).report(eq("token"), eventCaptor.capture(), eq(10));
+        verify(mockApi).report(eq("token"), eventCaptor.capture());
     }
 
     @Test
     void testReportEventUpdatesConfigWhenResponseIsPresent() {
         // Arrange
         APIResponse mockResponse = mock(APIResponse.class);
-        when(mockApi.report(anyString(), any(APIEvent.class), anyInt())).thenReturn(Optional.of(mockResponse));
+        when(mockApi.report(anyString(), any(APIEvent.class))).thenReturn(Optional.of(mockResponse));
 
         // Act
         cloudConnectionManager.reportEvent(Started.get(cloudConnectionManager), true);
@@ -63,14 +63,14 @@ void testReportEventUpdatesConfigWhenResponseIsPresent() {
     @Test
     void testReportEventDoesNotUpdateConfigWhenResponseIsNotPresent() {
         // Arrange
-        when(mockApi.report(anyString(), any(APIEvent.class), anyInt())).thenReturn(Optional.empty());
+        when(mockApi.report(anyString(), any(APIEvent.class))).thenReturn(Optional.empty());
 
         // Act
         cloudConnectionManager.reportEvent(Started.get(cloudConnectionManager), true);
 
         // Assert
         // No interaction with config update
-        verify(mockApi).report(anyString(), any(APIEvent.class), anyInt());
+        verify(mockApi).report(anyString(), any(APIEvent.class));
     }
 
     @Test

agent_api/src/test/java/background/cloud/ReportingAPITest.java

@@ -17,21 +17,36 @@ public class ReportingAPITest {
     ReportingApiHTTP api;
     @BeforeEach
     public void setup() {
-        api = new ReportingApiHTTP("http://localhost:5000/");
+        api = new ReportingApiHTTP("http://localhost:5000/", 2);
+    }
+
+    @Test
+    public void testTimeoutValid() {
+        api = new ReportingApiHTTP("http://localhost:5000/delayed/2/", 3); // Allowed delay
+        Optional<APIResponse> res = api.fetchNewConfig("token");
+        assertTrue(res.isPresent());
+        assertTrue(res.get().block());
+        assertEquals(1, res.get().endpoints().size());
+    }
+    @Test
+    public void testTimeoutInvalid() {
+        api = new ReportingApiHTTP("http://localhost:5000/delayed/4/", 3); // Allowed delay
+        Optional<APIResponse> res = api.fetchNewConfig("token");
+        assertFalse(res.isPresent());
     }
 
     @Test
     public void testFetchNewConfig() {
-        Optional<APIResponse> res = api.fetchNewConfig("token", 2);
+        Optional<APIResponse> res = api.fetchNewConfig("token");
         assertTrue(res.isPresent());
         assertTrue(res.get().block());
         assertEquals(1, res.get().endpoints().size());
     }
     @Test
     @StdIo
     public void testFetchNewConfigInvalidEndpoint(StdOut out) {
-        this.api = new ReportingApiHTTP("http://unknown.app.here:1234/");
-        Optional<APIResponse> res = api.fetchNewConfig("token", 2);
+        this.api = new ReportingApiHTTP("http://unknown.app.here:1234/", 2);
+        Optional<APIResponse> res = api.fetchNewConfig("token");
         assertEquals(Optional.empty(), res);
         assertTrue(
                 out.capturedString().contains("DEBUG dev.aikido.agent_api.background.cloud.api.ReportingApiHTTP: Error while fetching new config from cloud"),
@@ -48,7 +63,7 @@ public void testListsResponseWithTokenNull() {
     @Test
     @StdIo
     public void testListsResponseWithWrongEndpoint(StdOut out) {
-        this.api = new ReportingApiHTTP("http://unknown.app.here:1234/");
+        this.api = new ReportingApiHTTP("http://unknown.app.here:1234/", 2);
         Optional<ReportingApiHTTP.APIListsResponse> res = api.fetchBlockedLists("token");
         assertEquals(Optional.empty(), res);
         assertTrue(

end2end/server/mock_aikido_core.py

@@ -69,6 +69,10 @@ def post_events():
         events.append(request.get_json())
     return jsonify(responses["config"])
 
+@app.route('/delayed/<int:delay>/api/runtime/config')
+def delayed_route(delay):
+    time.sleep(delay)
+    return jsonify(responses["config"])
 
 @app.route('/mock/config', methods=['POST'])
 def mock_set_config():