perf(feature-flagging): bound UFC response memory

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();
     }

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<byte[]> 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;
+    }
+  }
+}

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<okhttp3.Request> 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 {