diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/RealtimeAPI.java b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/RealtimeAPI.java index ea60c0c9..118a52f7 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/RealtimeAPI.java +++ b/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,6 +20,7 @@ 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); @@ -30,16 +28,22 @@ public RealtimeAPI(Token token) { } public record ConfigResponse(long configUpdatedAt) {} - public Optional 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 getConfig() { + try { URI uri = URI.create(endpoint + "config"); HttpRequest request = createConfigRequest(token.get(), uri); // Send the request and get the response - HttpResponse httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + HttpResponse httpResponse = getHttpClient().send(request, HttpResponse.BodyHandlers.ofString()); return toConfigResponse(httpResponse); } catch (Exception e) { logger.debug("Error while communicating with cloud: %s", e.getMessage()); diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApiHTTP.java b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApiHTTP.java index 7cd29aa5..3a0be419 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApiHTTP.java +++ b/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,6 +25,7 @@ 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); @@ -36,18 +33,23 @@ public ReportingApiHTTP(String reportingUrl, int timeoutInSec, Token token) { this.token = token; } - public Optional 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 fetchNewConfig() { + try { URI uri = URI.create(reportingUrl + "api/runtime/config"); HttpRequest request = createHttpRequest(Optional.empty(), uri); // Send the request and get the response - HttpResponse httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + HttpResponse 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 fetchNewConfig() { @Override public Optional 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 httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + HttpResponse 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 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 fetchBlockedLists() { HttpRequest request = requestBuilder.build(); // Send the request and get the response - HttpResponse httpResponse = httpClient.send(request, HttpResponse.BodyHandlers.ofInputStream()); - if (httpResponse.statusCode() != HttpURLConnection.HTTP_OK) { - return Optional.empty(); + HttpResponse 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()); } diff --git a/agent_api/src/test/java/background/cloud/RealtimeAPITest.java b/agent_api/src/test/java/background/cloud/RealtimeAPITest.java index 9adc57d2..6380d301 100644 --- a/agent_api/src/test/java/background/cloud/RealtimeAPITest.java +++ b/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") diff --git a/agent_api/src/test/java/background/cloud/ReportingAPIResourceTest.java b/agent_api/src/test/java/background/cloud/ReportingAPIResourceTest.java new file mode 100644 index 00000000..a2dd4595 --- /dev/null +++ b/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 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); + } + } +}