Reuse HTTP clients and close blocklist responses

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

@@ -11,9 +11,6 @@
 import java.net.http.HttpResponse;
 import java.time.Duration;
 import java.util.Optional;
-import javax.net.ssl.SSLContext;
-import javax.net.ssl.TrustManagerFactory;
-import java.security.KeyStore;
 
 import static dev.aikido.agent_api.background.cloud.SSLContextHelper.createDefaultSSLContext;
 import static dev.aikido.agent_api.helpers.env.Endpoints.getAikidoRealtimeEndpoint;
@@ -23,23 +20,30 @@ public class RealtimeAPI {
     private static final int timeoutInSec = 3; // 3 sec timeout for requests to the realtime endpoint
     private final String endpoint;
     private final Token token;
+    private HttpClient httpClient;
 
     public RealtimeAPI(Token token) {
         endpoint = getAikidoRealtimeEndpoint(token);
         this.token = token;
     }
     public record ConfigResponse(long configUpdatedAt) {}
 
-    public Optional<ConfigResponse> getConfig() {
-        try {
-            HttpClient httpClient = HttpClient.newBuilder()
+    private synchronized HttpClient getHttpClient() throws Exception {
+        if (httpClient == null) {
+            httpClient = HttpClient.newBuilder()
                     .connectTimeout(Duration.ofSeconds(timeoutInSec))
                     .sslContext(createDefaultSSLContext())
                     .build();
+        }
+        return httpClient;
+    }
+
+    public Optional<ConfigResponse> getConfig() {
+        try {
             URI uri = URI.create(endpoint + "config");
             HttpRequest request = createConfigRequest(token.get(), uri);
             // Send the request and get the response
-            HttpResponse<String> httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+            HttpResponse<String> httpResponse = getHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
             return toConfigResponse(httpResponse);
         } catch (Exception e) {
             logger.debug("Error while communicating with cloud: %s", e.getMessage());

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

@@ -11,16 +11,12 @@
 import java.io.InputStreamReader;
 import java.net.HttpURLConnection;
 import java.net.URI;
-import java.net.URL;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.time.Duration;
 import java.util.Optional;
 import java.util.zip.GZIPInputStream;
-import javax.net.ssl.SSLContext;
-import javax.net.ssl.TrustManagerFactory;
-import java.security.KeyStore;
 
 import static dev.aikido.agent_api.background.cloud.SSLContextHelper.createDefaultSSLContext;
 
@@ -29,25 +25,31 @@ public class ReportingApiHTTP extends ReportingApi {
     private final String reportingUrl;
     private final Token token;
     private final Gson gson = new Gson();
+    private HttpClient httpClient;
     public ReportingApiHTTP(String reportingUrl, int timeoutInSec, Token token) {
         // Reporting URL should end with trailing slash for now.
         super(timeoutInSec);
         this.reportingUrl = reportingUrl;
         this.token = token;
     }
 
-    public Optional<APIResponse> fetchNewConfig() {
-        try {
-            HttpClient httpClient = HttpClient.newBuilder()
+    private synchronized HttpClient getHttpClient() throws Exception {
+        if (httpClient == null) {
+            httpClient = HttpClient.newBuilder()
                     .connectTimeout(Duration.ofSeconds(timeoutInSec))
                     .sslContext(createDefaultSSLContext())
                     .build();
+        }
+        return httpClient;
+    }
 
+    public Optional<APIResponse> fetchNewConfig() {
+        try {
             URI uri = URI.create(reportingUrl + "api/runtime/config");
             HttpRequest request = createHttpRequest(Optional.empty(), uri);
 
             // Send the request and get the response
-            HttpResponse<String> httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+            HttpResponse<String> httpResponse = getHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
             logger.trace("Got response for %s: %s", uri.toString(), httpResponse.body());
             return Optional.of(toApiResponse(httpResponse));
         } catch (Exception e) {
@@ -59,16 +61,11 @@ public Optional<APIResponse> fetchNewConfig() {
     @Override
     public Optional<APIResponse> report(APIEvent event) {
         try {
-            HttpClient httpClient = HttpClient.newBuilder()
-                .connectTimeout(Duration.ofSeconds(timeoutInSec))
-                .sslContext(createDefaultSSLContext())
-                .build();
-
             URI uri = URI.create(reportingUrl + "api/runtime/events");
             HttpRequest request = createHttpRequest(Optional.of(event), uri);
 
             // Send the request and get the response
-            HttpResponse<String> httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+            HttpResponse<String> httpResponse = getHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
             logger.trace("Got response for %s: %s", uri.toString(), httpResponse.body());
             return Optional.of(toApiResponse(httpResponse));
         } catch (Exception e) {
@@ -83,11 +80,6 @@ public Optional<APIListsResponse> fetchBlockedLists() {
             return Optional.empty();
         }
         try {
-            HttpClient httpClient = HttpClient.newBuilder()
-                    .connectTimeout(Duration.ofSeconds(timeoutInSec))
-                    .sslContext(createDefaultSSLContext())
-                    .build();
-
             URI uri = URI.create(reportingUrl + "api/runtime/firewall/lists");
             HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                     .uri(uri)
@@ -98,20 +90,17 @@ public Optional<APIListsResponse> fetchBlockedLists() {
             HttpRequest request = requestBuilder.build();
 
             // Send the request and get the response
-            HttpResponse<InputStream> httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofInputStream());
-            if (httpResponse.statusCode() != HttpURLConnection.HTTP_OK) {
-                return Optional.empty();
+            HttpResponse<InputStream> httpResponse = getHttpClient().send(request, HttpResponse.BodyHandlers.ofInputStream());
+            try (InputStream body = httpResponse.body()) {
+                if (httpResponse.statusCode() != HttpURLConnection.HTTP_OK) {
+                    return Optional.empty();
+                }
+
+                boolean gzipped = "gzip".equalsIgnoreCase(httpResponse.headers().firstValue("Content-Encoding").orElse(""));
+                try (InputStream inputStream = gzipped ? new GZIPInputStream(body) : body) {
+                    return Optional.of(gson.fromJson(new InputStreamReader(inputStream), APIListsResponse.class));
+                }
             }
-
-            InputStream inputStream = httpResponse.body();
-            // Check if the response is gzipped
-            if ("gzip".equalsIgnoreCase(httpResponse.headers().firstValue("Content-Encoding").orElse(""))) {
-                inputStream = new GZIPInputStream(inputStream);
-            }
-
-            // Read the response
-            APIListsResponse res = gson.fromJson(new InputStreamReader(inputStream), APIListsResponse.class);
-            return Optional.of(res);
         } catch (Exception e) {
             logger.debug("Failed to fetch blocked lists: %s", e.getMessage());
         }

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

@@ -27,6 +27,7 @@ public void testGetConfigSuccess() throws Exception {
 
         assertTrue(response.isPresent());
         assertEquals(0, response.get().configUpdatedAt());
+        assertEquals(response, realtimeAPI.getConfig());
     }
 
     @SetEnvironmentVariable(key = "AIKIDO_REALTIME_ENDPOINT", value="http://localnothost:2500")

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

@@ -0,0 +1,115 @@
+package background.cloud;
+
+import com.sun.net.httpserver.HttpServer;
+import dev.aikido.agent_api.background.cloud.api.ReportingApiHTTP;
+import dev.aikido.agent_api.background.cloud.api.events.APIEvent;
+import dev.aikido.agent_api.helpers.env.Token;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.io.ByteArrayOutputStream;
+import java.net.InetSocketAddress;
+import java.net.ServerSocket;
+import java.nio.charset.StandardCharsets;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.zip.GZIPOutputStream;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+@Timeout(10)
+public class ReportingAPIResourceTest {
+    @Test
+    public void testReusesConnectionsAcrossCloudRequests() throws Exception {
+        Set<InetSocketAddress> clients = ConcurrentHashMap.newKeySet();
+        HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+        server.createContext("/", exchange -> {
+            clients.add(exchange.getRemoteAddress());
+            exchange.getRequestBody().readAllBytes();
+            byte[] body = "{}".getBytes(StandardCharsets.UTF_8);
+            exchange.sendResponseHeaders(200, body.length);
+            try (var output = exchange.getResponseBody()) {
+                output.write(body);
+            }
+        });
+        server.start();
+        try {
+            ReportingApiHTTP api = new ReportingApiHTTP("http://127.0.0.1:" + server.getAddress().getPort() + "/", 2, new Token("token"));
+            record Event(String type) implements APIEvent {}
+            for (int i = 0; i < 10; i++) {
+                assertTrue(api.fetchNewConfig().isPresent());
+                assertTrue(api.fetchBlockedLists().isPresent());
+                assertTrue(api.report(new Event("heartbeat")).isPresent());
+            }
+            assertEquals(1, clients.size());
+        } finally {
+            server.stop(0);
+        }
+    }
+
+    @Test
+    public void testClosesRejectedResponseWithUnfinishedBody() throws Exception {
+        assertClosesUnfinishedResponse("503 Service Unavailable", "", "no");
+    }
+
+    @Test
+    public void testClosesResponseWhenJsonParsingFails() throws Exception {
+        assertClosesUnfinishedResponse("200 OK", "", "[]     \n");
+    }
+
+    @Test
+    public void testClosesResponseWhenGzipHeaderIsInvalid() throws Exception {
+        assertClosesUnfinishedResponse("200 OK", "Content-Encoding: gzip\r\n", "no");
+    }
+
+    private void assertClosesUnfinishedResponse(String status, String headers, String body) throws Exception {
+        try (ServerSocket server = new ServerSocket(0, 1, java.net.InetAddress.getLoopbackAddress())) {
+            var executor = Executors.newSingleThreadExecutor();
+            try {
+                var disconnected = executor.submit(() -> {
+                    try (var socket = server.accept()) {
+                        socket.setSoTimeout(3000);
+                        var reader = new java.io.BufferedReader(new java.io.InputStreamReader(socket.getInputStream(), StandardCharsets.US_ASCII));
+                        while (!reader.readLine().isEmpty()) {}
+                        var output = socket.getOutputStream();
+                        output.write(("HTTP/1.1 " + status + "\r\nContent-Length: 1048576\r\n" + headers + "\r\n" + body)
+                                .getBytes(StandardCharsets.US_ASCII));
+                        output.flush();
+                        return socket.getInputStream().read();
+                    }
+                });
+                ReportingApiHTTP api = new ReportingApiHTTP("http://localhost:" + server.getLocalPort() + "/", 1, new Token("token"));
+                assertTrue(api.fetchBlockedLists().isEmpty());
+                assertEquals(-1, disconnected.get(4, TimeUnit.SECONDS));
+            } finally {
+                executor.shutdownNow();
+            }
+        }
+    }
+
+    @Test
+    public void testReadsGzippedLists() throws Exception {
+        var bytes = new ByteArrayOutputStream();
+        try (var gzip = new GZIPOutputStream(bytes)) {
+            gzip.write("{\"blockedUserAgents\":\"bot\"}".getBytes(StandardCharsets.UTF_8));
+        }
+        byte[] body = bytes.toByteArray();
+        HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+        server.createContext("/", exchange -> {
+            exchange.getResponseHeaders().add("Content-Encoding", "gzip");
+            exchange.sendResponseHeaders(200, body.length);
+            try (var output = exchange.getResponseBody()) {
+                output.write(body);
+            }
+        });
+        server.start();
+        try {
+            ReportingApiHTTP api = new ReportingApiHTTP("http://127.0.0.1:" + server.getAddress().getPort() + "/", 2, new Token("token"));
+            assertEquals("bot", api.fetchBlockedLists().orElseThrow().blockedUserAgents());
+        } finally {
+            server.stop(0);
+        }
+    }
+}