diff --git a/products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/AgentlessConfigurationSource.java b/products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/AgentlessConfigurationSource.java index f4fc35cce30..7dc0c74f36b 100644 --- a/products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/AgentlessConfigurationSource.java +++ b/products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/AgentlessConfigurationSource.java @@ -200,6 +200,13 @@ private boolean fetchAndApply() { synchronized (lifecycleLock) { return !closed && apply(response); } + } catch (final UfcResponseBodyReader.ResponseTooLargeException e) { + if (!closed) { + ratelimitedLogger.warn( + "Feature Flagging agentless endpoint response exceeded the {}-byte safety limit", + e.limitBytes); + } + return false; } catch (final IOException e) { if (!closed) { ratelimitedLogger.warn( @@ -363,7 +370,8 @@ public UfcHttpResponse fetch(final HttpUrl endpoint, final Config config, final if (etag != null) { headers.put("If-None-Match", etag); } - // Leave Accept-Encoding unset so OkHttp negotiates gzip and transparently decompresses it. + // Set this explicitly so OkHttp keeps the compressed stream available for byte limits. + headers.put("Accept-Encoding", "gzip"); final Request request = prepareRequest(endpoint, headers, config, isDatadogManagedEndpoint(endpoint, config)) .get() @@ -401,7 +409,11 @@ public HttpRetryPolicy create() { final int status = response.code(); final String responseEtag = response.header("ETag"); try (ResponseBody responseBody = response.body()) { - final byte[] body = responseBody != null ? responseBody.bytes() : null; + final byte[] body = + status == HttpURLConnection.HTTP_OK && responseBody != null + ? UfcResponseBodyReader.read( + responseBody, response.header("Content-Encoding")) + : null; return new UfcHttpResponse(status, responseEtag, body); } }); @@ -454,6 +466,7 @@ static final class AgentlessRetryPolicy extends HttpRetryPolicy { @Override public boolean shouldRetry(final Exception exception) { return exception instanceof IOException + && !(exception instanceof UfcResponseBodyReader.ResponseTooLargeException) && !Thread.currentThread().isInterrupted() && reserveRetry(); } diff --git a/products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/UfcResponseBodyReader.java b/products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/UfcResponseBodyReader.java new file mode 100644 index 00000000000..f5a08d796b1 --- /dev/null +++ b/products/feature-flagging/feature-flagging-lib/src/main/java/com/datadog/featureflag/UfcResponseBodyReader.java @@ -0,0 +1,130 @@ +package com.datadog.featureflag; + +import java.io.FilterInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.zip.GZIPInputStream; +import okhttp3.ResponseBody; + +/** Reads UFC HTTP responses without allowing input size or gzip expansion to grow without limit. */ +final class UfcResponseBodyReader { + + private static final int MAX_RESPONSE_BYTES = 10 << 20; + static final int MAX_COMPRESSED_BYTES = MAX_RESPONSE_BYTES; + static final int MAX_DECOMPRESSED_BYTES = MAX_RESPONSE_BYTES; + private static final int CHUNK_SIZE = 8 << 10; + + private UfcResponseBodyReader() {} + + static byte[] read(final ResponseBody responseBody, final String contentEncoding) + throws IOException { + final boolean gzip = contentEncoding != null && "gzip".equalsIgnoreCase(contentEncoding.trim()); + final String inputKind = gzip ? "compressed" : "decompressed"; + final long contentLength = responseBody.contentLength(); + if (contentLength > MAX_RESPONSE_BYTES) { + throw new ResponseTooLargeException(inputKind, MAX_RESPONSE_BYTES); + } + if (!gzip && contentLength == 0) { + return new byte[0]; + } + + final InputStream responseStream = responseBody.byteStream(); + if (gzip) { + try (InputStream compressed = + new LimitedInputStream(responseStream, MAX_COMPRESSED_BYTES, "compressed"); + InputStream decompressed = new GZIPInputStream(compressed)) { + return readBounded(decompressed, MAX_DECOMPRESSED_BYTES, "decompressed"); + } + } + return readBounded(responseStream, MAX_DECOMPRESSED_BYTES, "decompressed"); + } + + private static byte[] readBounded( + final InputStream input, final int limitBytes, final String kind) throws IOException { + final List chunks = new ArrayList<>(); + int totalBytes = 0; + boolean endOfInput = false; + + while (totalBytes < limitBytes && !endOfInput) { + final int chunkSize = Math.min(CHUNK_SIZE, limitBytes - totalBytes); + byte[] chunk = new byte[chunkSize]; + int chunkBytes = 0; + while (chunkBytes < chunkSize) { + final int read = input.read(chunk, chunkBytes, chunkSize - chunkBytes); + if (read == -1) { + endOfInput = true; + break; + } + chunkBytes += read; + } + if (chunkBytes > 0) { + if (chunkBytes < chunk.length) { + chunk = Arrays.copyOf(chunk, chunkBytes); + } + chunks.add(chunk); + totalBytes += chunkBytes; + } + } + + if (!endOfInput && input.read() != -1) { + throw new ResponseTooLargeException(kind, limitBytes); + } + if (chunks.isEmpty()) { + return new byte[0]; + } + if (chunks.size() == 1) { + return chunks.get(0); + } + + final byte[] body = new byte[totalBytes]; + int offset = 0; + for (final byte[] chunk : chunks) { + System.arraycopy(chunk, 0, body, offset, chunk.length); + offset += chunk.length; + } + return body; + } + + static final class ResponseTooLargeException extends IOException { + final int limitBytes; + + ResponseTooLargeException(final String kind, final int limitBytes) { + super("Feature Flagging " + kind + " response exceeds " + limitBytes + " bytes"); + this.limitBytes = limitBytes; + } + } + + static final class LimitedInputStream extends FilterInputStream { + private final int limitBytes; + private final String kind; + private int readBytes; + + LimitedInputStream(final InputStream input, final int limitBytes, final String kind) { + super(input); + this.limitBytes = limitBytes; + this.kind = kind; + } + + @Override + public int read() throws IOException { + final int value = super.read(); + if (value != -1 && ++readBytes > limitBytes) { + throw new ResponseTooLargeException(kind, limitBytes); + } + return value; + } + + @Override + public int read(final byte[] bytes, final int offset, final int length) throws IOException { + final int allowed = Math.min(length, limitBytes - readBytes + 1); + final int read = super.read(bytes, offset, allowed); + if (read > 0 && (readBytes += read) > limitBytes) { + throw new ResponseTooLargeException(kind, limitBytes); + } + return read; + } + } +} diff --git a/products/feature-flagging/feature-flagging-lib/src/test/java/com/datadog/featureflag/AgentlessConfigurationSourceTest.java b/products/feature-flagging/feature-flagging-lib/src/test/java/com/datadog/featureflag/AgentlessConfigurationSourceTest.java index b02bd05642b..4e023268a40 100644 --- a/products/feature-flagging/feature-flagging-lib/src/test/java/com/datadog/featureflag/AgentlessConfigurationSourceTest.java +++ b/products/feature-flagging/feature-flagging-lib/src/test/java/com/datadog/featureflag/AgentlessConfigurationSourceTest.java @@ -265,7 +265,7 @@ void handlesGzipAndKeepsLastKnownGoodWhenNextResponseIsTruncated() throws Except } @Test - void downloadsAndAppliesLargeUfcWithoutPayloadLimit() throws Exception { + void downloadsAndAppliesLargeUfcWithinPayloadLimit() throws Exception { final int flagCount = 5_000; final String largeConfig = largeConfig(flagCount); assertTrue(largeConfig.getBytes(UTF_8).length > 500_000); @@ -299,6 +299,40 @@ void downloadsAndAppliesLargeUfcWithoutPayloadLimit() throws Exception { } } + @Test + void oversizedKnownLengthIsNotRetried() throws Exception { + final List requests = new ArrayList<>(); + final byte[] oversized = new byte[UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES + 1]; + final AgentlessConfigurationSource.OkHttpUfcHttpClient client = + scriptedClient( + requests, + delay -> {}, + () -> 1.0, + new AgentlessConfigurationSource.UfcHttpResponse(200, "etag-bad", oversized)); + + assertThrows( + UfcResponseBodyReader.ResponseTooLargeException.class, + () -> client.fetch(HttpUrl.get("http://localhost" + CONFIG_PATH), config(), "etag-good")); + assertEquals(1, requests.size()); + } + + @Test + void oversizedResponseKeepsLastKnownGoodConfigurationAndEtag() { + final FakeClient client = + new FakeClient( + response(200, "etag-good", emptyConfig()), + new UfcResponseBodyReader.ResponseTooLargeException( + "decompressed", UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES)); + final AgentlessConfigurationSource service = service(client); + FeatureFlaggingGateway.addConfigListener(listener); + + assertTrue(service.pollOnce()); + assertFalse(service.pollOnce()); + + verify(listener).accept(any(ServerConfiguration.class)); + assertEquals("etag-good", client.requests.get(1).etag); + } + @Test void realHttpClientAllowsMissingEtagAndEmptyResponseBody() throws Exception { try (JavaTestHttpServer server = @@ -322,7 +356,7 @@ void realHttpClientAllowsMissingEtagAndEmptyResponseBody() throws Exception { assertEquals(HttpURLConnection.HTTP_NO_CONTENT, response.status); assertNull(response.etag); - assertEquals(0, response.body.length); + assertNull(response.body); assertNull(server.getLastRequest().getHeader("If-None-Match")); } finally { httpClient.dispatcher().executorService().shutdownNow(); @@ -830,6 +864,10 @@ void retryPolicyRejectsNonIoAndInterruptedIoFailures() { retryPolicy(new AtomicBoolean()); assertFalse(policy.shouldRetry(new IllegalStateException("not an I/O failure"))); + assertFalse( + policy.shouldRetry( + new UfcResponseBodyReader.ResponseTooLargeException( + "decompressed", UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES))); Thread.currentThread().interrupt(); try { diff --git a/products/feature-flagging/feature-flagging-lib/src/test/java/com/datadog/featureflag/UfcResponseBodyReaderTest.java b/products/feature-flagging/feature-flagging-lib/src/test/java/com/datadog/featureflag/UfcResponseBodyReaderTest.java new file mode 100644 index 00000000000..40d294a7eee --- /dev/null +++ b/products/feature-flagging/feature-flagging-lib/src/test/java/com/datadog/featureflag/UfcResponseBodyReaderTest.java @@ -0,0 +1,268 @@ +package com.datadog.featureflag; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.util.Arrays; +import java.util.Random; +import java.util.zip.GZIPOutputStream; +import okhttp3.MediaType; +import okhttp3.ResponseBody; +import okio.BufferedSource; +import okio.Okio; +import org.junit.jupiter.api.Test; + +class UfcResponseBodyReaderTest { + + @Test + void acceptsIdentityResponseAtLimit() throws Exception { + final byte[] body = + UfcResponseBodyReader.read( + responseBody( + new RepeatingInputStream(UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES), + UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES), + null); + + assertEquals(UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES, body.length); + } + + @Test + void acceptsIdentityContentEncoding() throws Exception { + final byte[] expected = "identity UFC".getBytes("UTF-8"); + + final byte[] actual = + UfcResponseBodyReader.read( + responseBody(new ByteArrayInputStream(expected), expected.length), "identity"); + + assertArrayEquals(expected, actual); + } + + @Test + void acceptsKnownEmptyIdentityResponseWithoutReading() throws Exception { + assertArrayEquals(new byte[0], UfcResponseBodyReader.read(unreadableResponseBody(0), null)); + } + + @Test + void acceptsChunkedEmptyIdentityResponse() throws Exception { + assertArrayEquals( + new byte[0], + UfcResponseBodyReader.read(responseBody(new ByteArrayInputStream(new byte[0]), -1), null)); + } + + @Test + void rejectsKnownIdentityLengthBeforeReading() { + final ResponseBody body = + unreadableResponseBody(UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES + 1L); + + final UfcResponseBodyReader.ResponseTooLargeException failure = + assertThrows( + UfcResponseBodyReader.ResponseTooLargeException.class, + () -> UfcResponseBodyReader.read(body, null)); + + assertEquals(UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES, failure.limitBytes); + } + + @Test + void rejectsChunkedIdentityResponseAboveLimit() { + final ResponseBody body = + responseBody( + new RepeatingInputStream(UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES + 1L), -1); + + assertThrows( + UfcResponseBodyReader.ResponseTooLargeException.class, + () -> UfcResponseBodyReader.read(body, null)); + } + + @Test + void rejectsBodyThatExceedsMisleadingContentLength() { + final ResponseBody body = + responseBody( + new RepeatingInputStream(UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES + 1L), 1); + + assertThrows( + UfcResponseBodyReader.ResponseTooLargeException.class, + () -> UfcResponseBodyReader.read(body, null)); + } + + @Test + void rejectsKnownCompressedLengthBeforeReading() { + final ResponseBody body = + unreadableResponseBody(UfcResponseBodyReader.MAX_COMPRESSED_BYTES + 1L); + + final UfcResponseBodyReader.ResponseTooLargeException failure = + assertThrows( + UfcResponseBodyReader.ResponseTooLargeException.class, + () -> UfcResponseBodyReader.read(body, "gzip")); + + assertEquals(UfcResponseBodyReader.MAX_COMPRESSED_BYTES, failure.limitBytes); + } + + @Test + void rejectsChunkedCompressedResponseAboveLimit() throws Exception { + final byte[] compressed = gzipRandom(UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES); + assertTrue(compressed.length > UfcResponseBodyReader.MAX_COMPRESSED_BYTES); + + final UfcResponseBodyReader.ResponseTooLargeException failure = + assertThrows( + UfcResponseBodyReader.ResponseTooLargeException.class, + () -> + UfcResponseBodyReader.read( + responseBody(new ByteArrayInputStream(compressed), -1), "gzip")); + + assertTrue(failure.getMessage().contains("compressed")); + } + + @Test + void rejectsGzipExpansionAboveLimit() throws Exception { + final byte[] compressed = gzipRepeated(UfcResponseBodyReader.MAX_DECOMPRESSED_BYTES + 1L); + + assertThrows( + UfcResponseBodyReader.ResponseTooLargeException.class, + () -> + UfcResponseBodyReader.read( + responseBody(new ByteArrayInputStream(compressed), compressed.length), "gzip")); + } + + @Test + void decodesGzipWithinLimits() throws Exception { + final byte[] expected = "bounded UFC".getBytes("UTF-8"); + final byte[] compressed = gzip(expected); + + final byte[] actual = + UfcResponseBodyReader.read( + responseBody(new ByteArrayInputStream(compressed), compressed.length), " GZIP "); + + assertArrayEquals(expected, actual); + } + + @Test + void limitsSingleByteReads() throws Exception { + final InputStream input = + new UfcResponseBodyReader.LimitedInputStream( + new ByteArrayInputStream(new byte[] {1, 2}), 1, "compressed"); + + assertEquals(1, input.read()); + assertThrows(UfcResponseBodyReader.ResponseTooLargeException.class, input::read); + } + + @Test + void allowsEndOfInputAfterBulkReads() throws Exception { + final InputStream input = + new UfcResponseBodyReader.LimitedInputStream( + new ByteArrayInputStream(new byte[0]), 1, "compressed"); + + assertEquals(-1, input.read(new byte[1])); + } + + private static ResponseBody responseBody(final InputStream input, final long contentLength) { + final BufferedSource source = Okio.buffer(Okio.source(input)); + return new ResponseBody() { + @Override + public MediaType contentType() { + return MediaType.get("application/json"); + } + + @Override + public long contentLength() { + return contentLength; + } + + @Override + public BufferedSource source() { + return source; + } + }; + } + + private static ResponseBody unreadableResponseBody(final long contentLength) { + return new ResponseBody() { + @Override + public MediaType contentType() { + return MediaType.get("application/json"); + } + + @Override + public long contentLength() { + return contentLength; + } + + @Override + public BufferedSource source() { + throw new AssertionError("response body must not be read"); + } + }; + } + + private static byte[] gzip(final byte[] value) throws IOException { + final ByteArrayOutputStream output = new ByteArrayOutputStream(); + try (GZIPOutputStream gzip = new GZIPOutputStream(output)) { + gzip.write(value); + } + return output.toByteArray(); + } + + private static byte[] gzipRepeated(final long size) throws IOException { + final ByteArrayOutputStream output = new ByteArrayOutputStream(); + final byte[] chunk = new byte[8 << 10]; + Arrays.fill(chunk, (byte) 'x'); + long remaining = size; + try (GZIPOutputStream gzip = new GZIPOutputStream(output)) { + while (remaining > 0) { + final int count = (int) Math.min(chunk.length, remaining); + gzip.write(chunk, 0, count); + remaining -= count; + } + } + return output.toByteArray(); + } + + private static byte[] gzipRandom(final long size) throws IOException { + final ByteArrayOutputStream output = new ByteArrayOutputStream(); + final byte[] chunk = new byte[8 << 10]; + final Random random = new Random(123456789L); + long remaining = size; + try (GZIPOutputStream gzip = new GZIPOutputStream(output)) { + while (remaining > 0) { + random.nextBytes(chunk); + final int count = (int) Math.min(chunk.length, remaining); + gzip.write(chunk, 0, count); + remaining -= count; + } + } + return output.toByteArray(); + } + + private static final class RepeatingInputStream extends InputStream { + private long remaining; + + private RepeatingInputStream(final long remaining) { + this.remaining = remaining; + } + + @Override + public int read() { + if (remaining == 0) { + return -1; + } + remaining--; + return 'x'; + } + + @Override + public int read(final byte[] bytes, final int offset, final int length) { + if (remaining == 0) { + return -1; + } + final int count = (int) Math.min(length, remaining); + Arrays.fill(bytes, offset, offset + count, (byte) 'x'); + remaining -= count; + return count; + } + } +}