diff --git a/opa-services/src/main/java/io/github/open_policy_agent/opa/plugins/BundleDownloader.java b/opa-services/src/main/java/io/github/open_policy_agent/opa/plugins/BundleDownloader.java index 99af2444..7d529351 100644 --- a/opa-services/src/main/java/io/github/open_policy_agent/opa/plugins/BundleDownloader.java +++ b/opa-services/src/main/java/io/github/open_policy_agent/opa/plugins/BundleDownloader.java @@ -2,24 +2,21 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.InputStream; import java.net.URI; import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; -import java.nio.ByteBuffer; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; import java.nio.file.attribute.FileTime; import java.time.Duration; import java.util.HashSet; -import java.util.List; import java.util.Locale; import java.util.OptionalLong; import java.util.Set; import java.util.concurrent.CompletableFuture; -import java.util.concurrent.CompletionStage; -import java.util.concurrent.Flow; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledThreadPoolExecutor; @@ -57,6 +54,12 @@ public abstract class BundleDownloader { protected long lastModifiedTime = 0; protected long maxSizeBytes = Config.BundleConfig.DEFAULT_MAX_SIZE_BYTES; + // Fallback when no service config is available; matches the response_header_timeout_seconds + // default in Config.ServiceConfig. + private static final int DEFAULT_RESPONSE_HEADER_TIMEOUT_SECONDS = 10; + + private static final int READ_BUFFER_BYTES = 8192; + private static final Set ALLOWED_CONTENT_TYPES = Set.of( "application/vnd.openpolicyagent.bundles", @@ -189,48 +192,63 @@ public CompletableFuture startPolling(ScheduledExecutorService scheduler) ? polling.getMaxDelaySeconds() : 120; - scheduler.schedule(this::downloadBundle, 0, TimeUnit.SECONDS); - scheduleNextPoll(scheduler, minDelay, maxDelay); + scheduleDownload(scheduler, minDelay, maxDelay, 0); return initialActivation; } - // Re-schedules the next download with a uniformly random delay in [minDelay, maxDelay], - // matching Go-OPA's jittered polling. ScheduledExecutorService has no built-in jitter, so the - // task chains itself. RejectedExecutionException after a shutdown breaks the chain cleanly. - private void scheduleNextPoll(ScheduledExecutorService scheduler, int minDelay, int maxDelay) { - long delay = - minDelay >= maxDelay - ? minDelay - : ThreadLocalRandom.current().nextLong(minDelay, (long) maxDelay + 1); + // Schedules one download after delaySeconds; + // once that download completes, schedules the next poll with a fresh jittered delay in [minDelay, maxDelay]. + // Because the next poll is chained off the download's completion — not fired on an independent timer — + // at most one download is ever in flight, matching Go-OPA's jittered polling. + // ScheduledExecutorService has no built-in jitter, so the task chains itself; a + // RejectedExecutionException after shutdown ends the chain cleanly. + private void scheduleDownload( + ScheduledExecutorService scheduler, int minDelay, int maxDelay, long delaySeconds) { try { scheduler.schedule( - () -> { - try { - downloadBundle(); - } catch (Exception e) { - // downloadBundle() handles its own logging; swallow so the chain keeps polling. - // Only Exception is caught here — Errors (OOM, etc.) propagate and let the - // executor's uncaught-exception handler tear down the pool, which is the right - // outcome for unrecoverable conditions. - } finally { - scheduleNextPoll(scheduler, minDelay, maxDelay); - } - }, - delay, + () -> + downloadBundle() + .whenComplete( + (result, throwable) -> + scheduleDownload( + scheduler, + minDelay, + maxDelay, + nextDelaySeconds(minDelay, maxDelay))), + delaySeconds, TimeUnit.SECONDS); } catch (RejectedExecutionException stopped) { // Scheduler was shut down; let the chain end. } } + private static long nextDelaySeconds(int minDelay, int maxDelay) { + return minDelay >= maxDelay + ? minDelay + : ThreadLocalRandom.current().nextLong(minDelay, (long) maxDelay + 1); + } + + // Per-request download timeout. + private Duration requestTimeout() { + int seconds = + authService != null + ? authService.getResponseHeaderTimeoutSeconds() + : DEFAULT_RESPONSE_HEADER_TIMEOUT_SECONDS; + return Duration.ofSeconds(seconds > 0 ? seconds : DEFAULT_RESPONSE_HEADER_TIMEOUT_SECONDS); + } + /** * Download the bundle from the configured service and resource. * *

Handles HTTP/HTTPS downloads with ETag caching, file:// URIs, and filesystem paths. Calls * {@link #activateBundle(byte[])} when new bundle data is available. + * + * @return a future that completes when this download attempt — including any asynchronous + * activation — has finished. The poll scheduler awaits it before scheduling the next poll, so + * downloads never overlap. */ - protected void downloadBundle() { + protected CompletableFuture downloadBundle() { try { Config.ServiceConfig serviceConfig = manager.getConfig().getService(service); if (serviceConfig == null) { @@ -239,7 +257,7 @@ protected void downloadBundle() { initialActivation.completeExceptionally( new RuntimeException("Service '" + service + "' not found")); } - return; + return CompletableFuture.completedFuture(null); } String baseUrl = serviceConfig.getUrl(); @@ -257,16 +275,17 @@ protected void downloadBundle() { // Handle file:// URIs if ("file".equalsIgnoreCase(uri.getScheme())) { handleFileDownload(Paths.get(uri)); - return; + return CompletableFuture.completedFuture(null); } // Handle HTTP/HTTPS URIs - handleHttpDownload(uri); + return handleHttpDownload(uri); } else { // It's a file path (relative or absolute) Path basePath = Paths.get(baseUrl); Path filePath = basePath.resolve(resource); handleFileDownload(filePath); + return CompletableFuture.completedFuture(null); } } catch (Exception e) { @@ -274,6 +293,7 @@ protected void downloadBundle() { if (!initialActivation.isDone()) { initialActivation.completeExceptionally(e); } + return CompletableFuture.completedFuture(null); } } @@ -307,11 +327,14 @@ private void handleFileDownload(Path filePath) throws IOException { * Handle downloading from an HTTP/HTTPS URI. * * @param uri the URI to download from + * @return a future that completes when the response has been handled. The actual outcome is + * recorded on {@code initialActivation}. */ - private void handleHttpDownload(URI uri) { + private CompletableFuture handleHttpDownload(URI uri) { HttpRequest.Builder requestBuilder = HttpRequest.newBuilder() .uri(uri) + .timeout(requestTimeout()) .header("Accept", "application/vnd.openpolicyagent.bundles") .GET(); @@ -325,61 +348,85 @@ private void handleHttpDownload(URI uri) { } HttpRequest request = requestBuilder.build(); - httpClient - .sendAsync(request, sizeLimitedBodyHandler(maxSizeBytes)) - .whenComplete( + + CompletableFuture> exchange = + httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofInputStream()); + return exchange + .handle( (response, throwable) -> { if (throwable != null) { - Throwable cause = - throwable instanceof java.util.concurrent.CompletionException - ? throwable.getCause() - : throwable; - manager.getLogger().error("Bundle '%s': Download error: %s", name, cause.getMessage()); - if (!initialActivation.isDone()) { - initialActivation.completeExceptionally(cause); - } - return; + exchange.cancel(true); } + handleHttpResponse(response, throwable); + return null; + }); + } - if (response.statusCode() == 304) { - manager.getLogger().debug("Bundle '%s': Not modified (ETag match)", name); - if (!initialActivation.isDone()) { - initialActivation.complete(null); - } - return; - } + // Handles a completed HTTP exchange: records success or failure on initialActivation. Runs on the + // HTTP client's executor thread, but the poll chain never starts the next download until the + // future returned by handleHttpDownload completes, so activateBundle is never called concurrently. + private void handleHttpResponse(HttpResponse response, Throwable throwable) { + if (throwable != null) { + Throwable cause = + throwable instanceof java.util.concurrent.CompletionException + ? throwable.getCause() + : throwable; + manager.getLogger().error("Bundle '%s': Download error: %s", name, cause.getMessage()); + if (!initialActivation.isDone()) { + initialActivation.completeExceptionally(cause); + } + return; + } - if (response.statusCode() == 200) { - String contentType = response.headers().firstValue("Content-Type").orElse(""); - if (!isAcceptableContentType(contentType)) { - String errorMsg = "Unexpected Content-Type: '" + contentType + "'"; - manager.getLogger().error("Bundle '%s': %s", name, errorMsg); - if (!initialActivation.isDone()) { - initialActivation.completeExceptionally(new RuntimeException(errorMsg)); - } - return; - } - response.headers().firstValue("ETag").ifPresent(newEtag -> this.etag = newEtag); - try { - activateBundle(response.body()); - } catch (Exception e) { - manager.getLogger().error("Bundle '%s': Activation failed: %s", name, e.getMessage()); - if (!initialActivation.isDone()) { - initialActivation.completeExceptionally(e); - } - return; - } - if (!initialActivation.isDone()) { - initialActivation.complete(null); - } - } else { - String errorMsg = "Download failed with status " + response.statusCode(); - manager.getLogger().error("Bundle '%s': %s", name, errorMsg); - if (!initialActivation.isDone()) { - initialActivation.completeExceptionally(new RuntimeException(errorMsg)); - } - } - }); + if (response.statusCode() == 304) { + manager.getLogger().debug("Bundle '%s': Not modified (ETag match)", name); + if (!initialActivation.isDone()) { + initialActivation.complete(null); + } + return; + } + + if (response.statusCode() == 200) { + String contentType = response.headers().firstValue("Content-Type").orElse(""); + if (!isAcceptableContentType(contentType)) { + String errorMsg = "Unexpected Content-Type: '" + contentType + "'"; + manager.getLogger().error("Bundle '%s': %s", name, errorMsg); + if (!initialActivation.isDone()) { + initialActivation.completeExceptionally(new RuntimeException(errorMsg)); + } + return; + } + OptionalLong contentLength = response.headers().firstValueAsLong("Content-Length"); + byte[] body; + try { + body = readBodyWithLimit(response.body(), contentLength); + } catch (Exception e) { + manager.getLogger().error("Bundle '%s': Download error: %s", name, e.getMessage()); + if (!initialActivation.isDone()) { + initialActivation.completeExceptionally(e); + } + return; + } + response.headers().firstValue("ETag").ifPresent(newEtag -> this.etag = newEtag); + try { + activateBundle(body); + } catch (Exception e) { + manager.getLogger().error("Bundle '%s': Activation failed: %s", name, e.getMessage()); + if (!initialActivation.isDone()) { + initialActivation.completeExceptionally(e); + } + return; + } + if (!initialActivation.isDone()) { + initialActivation.complete(null); + } + } else { + String errorMsg = "Download failed with status " + response.statusCode(); + manager.getLogger().error("Bundle '%s': %s", name, errorMsg); + if (!initialActivation.isDone()) { + initialActivation.completeExceptionally(new RuntimeException(errorMsg)); + } + } } /** @@ -395,101 +442,51 @@ private static boolean isAcceptableContentType(String contentType) { } /** - * Returns a BodyHandler that rejects over-sized responses. - * - *

If Content-Length is present and already exceeds the limit, the body is rejected upfront - * without reading it. Otherwise, a {@link SizeLimitedByteArraySubscriber} streams the body and - * aborts the moment the cumulative byte count exceeds the limit — so a server that omits - * Content-Length (or uses chunked encoding) cannot force the client to buffer an unbounded body. + * Reads the response body into memory, enforcing {@code maxSizeBytes} both upfront (via + * Content-Length) and as the bytes arrive. * + *

The limit is capped at {@link Integer#MAX_VALUE} because the body is materialized into a + * {@code byte[]}, whose maximum length is ~2 GB. */ - private static HttpResponse.BodyHandler sizeLimitedBodyHandler(long maxBytes) { - long cappedMax = Math.min(maxBytes, Integer.MAX_VALUE); - return responseInfo -> { - OptionalLong contentLength = responseInfo.headers().firstValueAsLong("Content-Length"); - if (contentLength.isPresent() && contentLength.getAsLong() > cappedMax) { - return new SizeLimitedByteArraySubscriber( - cappedMax, - "Content-Length " - + contentLength.getAsLong() - + " exceeds limit of " - + cappedMax - + " bytes"); - } - return new SizeLimitedByteArraySubscriber(cappedMax); - }; + private byte[] readBodyWithLimit(InputStream body, OptionalLong contentLength) + throws IOException { + long limit = Math.min(maxSizeBytes, Integer.MAX_VALUE); + try (InputStream in = body) { + rejectIfContentLengthExceedsLimit(contentLength, limit); + return readUpTo(in, limit); + } } /** - * Accumulates the response body while enforcing a size limit as bytes arrive. The instant the cumulative byte - * count exceeds {@code limit}, it cancels the subscription. + * Rejects a response whose advertised Content-Length already exceeds the limit, so an oversized + * body is never read at all. Servers are turned away before a single byte is transferred. */ - private static final class SizeLimitedByteArraySubscriber - implements HttpResponse.BodySubscriber { - private final long limit; - private final String rejectImmediatelyMessage; - private final CompletableFuture result = new CompletableFuture<>(); - private final ByteArrayOutputStream buffer = new ByteArrayOutputStream(); - private long total = 0; - private Flow.Subscription subscription; - - SizeLimitedByteArraySubscriber(long limit) { - this(limit, null); - } - - SizeLimitedByteArraySubscriber(long limit, String rejectImmediatelyMessage) { - this.limit = limit; - this.rejectImmediatelyMessage = rejectImmediatelyMessage; - } - - @Override - public CompletionStage getBody() { - return result; - } - - @Override - public void onSubscribe(Flow.Subscription subscription) { - this.subscription = subscription; - if (rejectImmediatelyMessage != null) { - subscription.cancel(); - result.completeExceptionally(new BundleSizeLimitException(rejectImmediatelyMessage)); - return; - } - subscription.request(Long.MAX_VALUE); + private static void rejectIfContentLengthExceedsLimit(OptionalLong contentLength, long limit) { + if (contentLength.isPresent() && contentLength.getAsLong() > limit) { + throw new BundleSizeLimitException( + "Content-Length " + contentLength.getAsLong() + " exceeds limit of " + limit + " bytes"); } + } - @Override - public void onNext(List items) { - if (result.isDone()) { - return; - } - for (ByteBuffer item : items) { - total += item.remaining(); - } + /** + * Reads the stream into a byte array, aborting the moment the cumulative size exceeds the limit. + * A server that omits Content-Length (or uses chunked encoding) therefore cannot force the client + * to buffer an unbounded body: the read stops mid-stream and the caller closes the stream, which + * cancels the exchange so the remaining bytes are never transferred. + */ + private static byte[] readUpTo(InputStream in, long limit) throws IOException { + ByteArrayOutputStream buffer = new ByteArrayOutputStream(); + byte[] chunk = new byte[READ_BUFFER_BYTES]; + long total = 0; + int read; + while ((read = in.read(chunk)) != -1) { + total += read; if (total > limit) { - subscription.cancel(); - result.completeExceptionally( - new BundleSizeLimitException("Response body exceeds limit of " + limit + " bytes")); - return; - } - for (ByteBuffer item : items) { - byte[] chunk = new byte[item.remaining()]; - item.get(chunk); - buffer.write(chunk, 0, chunk.length); - } - } - - @Override - public void onError(Throwable throwable) { - result.completeExceptionally(throwable); - } - - @Override - public void onComplete() { - if (!result.isDone()) { - result.complete(buffer.toByteArray()); + throw new BundleSizeLimitException("Response body exceeds limit of " + limit + " bytes"); } + buffer.write(chunk, 0, read); } + return buffer.toByteArray(); } static class BundleSizeLimitException extends RuntimeException { diff --git a/opa-services/src/test/java/io/github/open_policy_agent/opa/plugins/BundleDownloaderSecurityTest.java b/opa-services/src/test/java/io/github/open_policy_agent/opa/plugins/BundleDownloaderSecurityTest.java index 7c0c852e..276acf11 100644 --- a/opa-services/src/test/java/io/github/open_policy_agent/opa/plugins/BundleDownloaderSecurityTest.java +++ b/opa-services/src/test/java/io/github/open_policy_agent/opa/plugins/BundleDownloaderSecurityTest.java @@ -1,5 +1,6 @@ package io.github.open_policy_agent.opa.plugins; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -16,8 +17,13 @@ import java.net.Socket; import java.util.Collections; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.zip.GZIPOutputStream; import org.apache.commons.compress.archivers.tar.TarArchiveEntry; import org.apache.commons.compress.archivers.tar.TarArchiveOutputStream; @@ -32,12 +38,20 @@ class BundleDownloaderSecurityTest { private HttpServer server; private Store store; + private ExecutorService serverExecutor; + private ScheduledExecutorService scheduler; @AfterEach void tearDown() { + if (scheduler != null) { + scheduler.shutdownNow(); + } if (server != null) { server.stop(0); } + if (serverExecutor != null) { + serverExecutor.shutdownNow(); + } } @Test @@ -306,6 +320,151 @@ void download_chunkedBodyExceedsLimit_failsMidStream() throws Exception { assertTrue(ex.getCause().getMessage().contains("exceeds limit")); } + @Test + void polling_serverAcceptsButNeverReplies_timesOutAndKeepsPolling() throws Exception { + CountDownLatch connections = new CountDownLatch(2); + ServerSocket silentServer = new ServerSocket(0); + int port = silentServer.getLocalPort(); + + Thread accepter = new Thread(() -> { + try { + while (!Thread.currentThread().isInterrupted()) { + silentServer.accept(); + connections.countDown(); + } + } catch (IOException ignored) { + } + }); + accepter.setDaemon(true); + accepter.start(); + + try { + Config config = new Config(); + config.setServices(Collections.singletonMap("test-service", + new Config.ServiceConfig() + .setName("test-service") + .setUrl("http://localhost:" + port) + .setResponseHeaderTimeoutSeconds(1))); + config.setBundles(Collections.singletonMap("authz", + new Config.BundleConfig() + .setService("test-service") + .setResource("/bundle.tar.gz") + .setPolling( + new Config.PollingConfig().setMinDelaySeconds(0).setMaxDelaySeconds(0)))); + + Logger logger = new Logger.StandardLogger(); + PluginManager manager = + new PluginManager.Builder() + .withId("test") + .withStore(new InMem()) + .withConfig(config) + .withLogger(logger) + .build(); + + ServicePlugin servicePlugin = (ServicePlugin) new ServicePlugin().initialize(manager); + manager.registerPlugin("services", servicePlugin); + servicePlugin.start(); + + BundlePlugin bundlePlugin = (BundlePlugin) new BundlePlugin().initialize(manager); + manager.registerPlugin("bundles", bundlePlugin); + try { + bundlePlugin.start(); + + assertTrue( + connections.await(10, TimeUnit.SECONDS), + "polling stopped after a stalled request — expected the request timeout to fail the " + + "exchange so the poll chain retries"); + } finally { + bundlePlugin.stop(); + } + } finally { + silentServer.close(); + accepter.interrupt(); + } + } + + @Test + void polling_slowActivationExceedingInterval_neverRunsConcurrently() throws Exception { + byte[] bundleData = createValidBundle(); + + serverExecutor = Executors.newFixedThreadPool(8); + server = HttpServer.create(new InetSocketAddress("localhost", 0), 0); + server.setExecutor(serverExecutor); + server.createContext("/bundle.tar.gz", exchange -> { + exchange.getResponseHeaders().add("Content-Type", "application/vnd.openpolicyagent.bundles"); + exchange.sendResponseHeaders(200, bundleData.length); + exchange.getResponseBody().write(bundleData); + exchange.close(); + }); + server.start(); + + Config config = new Config(); + config.setServices(Collections.singletonMap("test-service", + new Config.ServiceConfig() + .setName("test-service") + .setUrl("http://localhost:" + server.getAddress().getPort()))); + + Logger logger = new Logger.StandardLogger(); + PluginManager manager = + new PluginManager.Builder() + .withId("test") + .withStore(new InMem()) + .withConfig(config) + .withLogger(logger) + .build(); + + ConcurrencyTrackingDownloader downloader = (ConcurrencyTrackingDownloader) new ConcurrencyTrackingDownloader("authz", manager) + .setService("test-service") + .setResource("/bundle.tar.gz") + .setPolling(new Config.PollingConfig().setMinDelaySeconds(0).setMaxDelaySeconds(0)); + + scheduler = BundleDownloader.newPollScheduler("concurrency-test"); + downloader.startPolling(scheduler); + + // Let many poll cycles run back-to-back. + Thread.sleep(800); + scheduler.shutdownNow(); + scheduler.awaitTermination(2, TimeUnit.SECONDS); + + assertTrue( + downloader.activations.get() > 1, + "expected multiple poll cycles to have run, but only " + downloader.activations.get() + + " did"); + assertEquals( + 1, + downloader.maxInFlight.get(), + "activateBundle ran concurrently — max simultaneous activations was " + + downloader.maxInFlight.get()); + } + + /** + * A BundleDownloader whose activation is deliberately slow and records whether any two + * activations ever overlap, so a test can assert that polling stays serialized. + */ + private static final class ConcurrencyTrackingDownloader extends BundleDownloader { + final AtomicInteger inFlight = new AtomicInteger(); + final AtomicInteger maxInFlight = new AtomicInteger(); + final AtomicInteger activations = new AtomicInteger(); + + ConcurrencyTrackingDownloader(String name, PluginManager manager) { + super(name, manager, null); + } + + @Override + protected void activateBundle(byte[] bundleData) { + int now = inFlight.incrementAndGet(); + maxInFlight.accumulateAndGet(now, Math::max); + try { + Thread.sleep(50); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + inFlight.decrementAndGet(); + activations.incrementAndGet(); + } + } + } + private int startRawChunkedServer(int chunkSize, int totalChunks) throws IOException { ServerSocket serverSocket = new ServerSocket(0); int port = serverSocket.getLocalPort();