pranavshuklaa commented on code in PR #1091:
URL: https://github.com/apache/flink-agents/pull/1091#discussion_r4543829931
##########
runtime/src/test/java/org/apache/flink/agents/runtime/skill/SkillMaterializerTest.java:
##########
@@ -592,4 +599,853 @@ private List<String> getMessages() {
return messages;
}
}
+ // -------------------------------------------------------
+ // Download size cap tests
+ // -------------------------------------------------------
+
+ /**
+ * Server declares a Content-Length larger than the injected cap. The
pre-flight check must
+ * reject before reading any body bytes.
+ */
+ @Test
+ void rejectsDeclaredContentLengthOverCap() throws IOException {
+ long cap = 1024L;
+ SkillMaterializer.Limits limits =
+ new SkillMaterializer.Limits(cap, cap * 10, cap * 100, 1_000);
+
+ long overCap = cap + 1;
+ HttpServer server = HttpServer.create(new
InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext(
+ "/",
+ exchange -> {
+ exchange.getResponseHeaders().add("Content-Length",
String.valueOf(overCap));
+ exchange.sendResponseHeaders(200, 0);
+ exchange.getResponseBody().close();
+ exchange.close();
+ });
+ server.setExecutor(null);
+ server.start();
+ try {
+ int port = server.getAddress().getPort();
+ IOException ex =
+ assertThrows(
+ IOException.class,
+ () ->
+ SkillMaterializer.downloadToTempFile(
+ "http://127.0.0.1:"; + port +
"/skill.zip",
+ 5_000,
+ true,
+ limits));
+ assertTrue(
+ ex.getMessage().contains("exceeding the limit"),
+ "error must mention the limit, got: " + ex.getMessage());
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ /**
+ * Server sends a small {@code Content-Length} header but frames the body
as chunked ({@code
+ * sendResponseHeaders(200, 0)}), so the client reads past the declared
length. The declared
+ * size passes the pre-flight check; only the streaming byte counter can
reject it.
+ *
+ * <p>Uses a 1 KiB injectable cap so the test streams only 1,025 bytes.
+ */
+ @Test
+ void rejectsChunkedBodyOverCapDespiteSmallDeclaredContentLength() throws
IOException {
+ long cap = 1024L;
+ SkillMaterializer.Limits limits =
+ new SkillMaterializer.Limits(cap, cap * 10, cap * 100, 1_000);
+
+ int declaredLength = 100;
+ long actualBytes = cap + 1;
+ HttpServer server = HttpServer.create(new
InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext(
+ "/",
+ exchange -> {
+ exchange.getResponseHeaders()
+ .add("Content-Length",
String.valueOf(declaredLength));
+ exchange.sendResponseHeaders(200, 0);
+ OutputStream body = exchange.getResponseBody();
+ byte[] chunk = new byte[65536];
+ Arrays.fill(chunk, (byte) 'x');
+ long remaining = actualBytes;
+ while (remaining > 0) {
+ int toWrite = (int) Math.min(chunk.length, remaining);
+ try {
+ body.write(chunk, 0, toWrite);
+ body.flush();
+ } catch (IOException ignored) {
+ break;
+ }
+ remaining -= toWrite;
+ }
+ exchange.close();
+ });
+ server.setExecutor(null);
+ server.start();
+ try {
+ int port = server.getAddress().getPort();
+ IOException ex =
+ assertThrows(
+ IOException.class,
+ () ->
+ SkillMaterializer.downloadToTempFile(
+ "http://127.0.0.1:"; + port +
"/skill.zip",
+ 5_000,
+ true,
+ limits));
+ assertTrue(
+ ex.getMessage().contains("exceeded the limit"),
+ "error must mention the limit, got: " + ex.getMessage());
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ /**
+ * Server streams past the cap with no Content-Length header at all. The
byte counter must catch
+ * it.
+ *
+ * <p>Uses a 1 KiB injectable cap so the test streams only 1,025 bytes.
+ */
+ @Test
+ void rejectsStreamWithNoContentLengthAndBodyOverCap() throws IOException {
+ long cap = 1024L;
+ SkillMaterializer.Limits limits =
+ new SkillMaterializer.Limits(cap, cap * 10, cap * 100, 1_000);
+
+ long actualBytes = cap + 1;
+ HttpServer server = HttpServer.create(new
InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext(
+ "/",
+ exchange -> {
+ exchange.sendResponseHeaders(200, 0);
+ OutputStream body = exchange.getResponseBody();
+ byte[] chunk = new byte[65536];
+ Arrays.fill(chunk, (byte) 'x');
+ long remaining = actualBytes;
+ while (remaining > 0) {
+ int toWrite = (int) Math.min(chunk.length, remaining);
+ try {
+ body.write(chunk, 0, toWrite);
+ body.flush();
+ } catch (IOException ignored) {
+ break;
+ }
+ remaining -= toWrite;
+ }
+ exchange.close();
+ });
+ server.setExecutor(null);
+ server.start();
+ try {
+ int port = server.getAddress().getPort();
+ IOException ex =
+ assertThrows(
+ IOException.class,
+ () ->
+ SkillMaterializer.downloadToTempFile(
+ "http://127.0.0.1:"; + port +
"/skill.zip",
+ 5_000,
+ true,
+ limits));
+ assertTrue(
+ ex.getMessage().contains("exceeded the limit"),
+ "error must mention the limit, got: " + ex.getMessage());
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ /**
+ * A body of exactly {@code cap} bytes must succeed (boundary is
inclusive). Uses a 1 KiB
+ * injectable cap so the test streams only 1,024 bytes.
+ */
+ @Test
+ void acceptsBodyExactlyAtDownloadCap() throws IOException {
+ long cap = 1024L;
+ SkillMaterializer.Limits limits =
+ new SkillMaterializer.Limits(cap, cap * 10, cap * 100, 1_000);
+
+ byte[] body = new byte[(int) cap];
+ Arrays.fill(body, (byte) 'z');
+ HttpServer server = startServer(200, body);
+ try {
+ int port = server.getAddress().getPort();
+ Path file =
+ SkillMaterializer.downloadToTempFile(
+ "http://127.0.0.1:"; + port + "/skill.zip", 5_000,
true, limits);
+ try {
+ assertEquals(cap, Files.size(file));
+ } finally {
+ Files.deleteIfExists(file);
+ }
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ /**
+ * A body of {@code cap + 1} bytes must be rejected. Paired with {@link
+ * #acceptsBodyExactlyAtDownloadCap()} to prove the boundary is enforced
at exactly the right
+ * byte.
+ *
+ * <p>The server intentionally omits {@code Content-Length} so the
pre-flight check cannot fire;
+ * only the streaming counter can reject the download. This isolates the
counter under test.
+ */
+ @Test
+ void rejectsBodyOneByteOverDownloadCap() throws IOException {
+ long cap = 1024L;
+ SkillMaterializer.Limits limits =
+ new SkillMaterializer.Limits(cap, cap * 10, cap * 100, 1_000);
+
+ long actualBytes = cap + 1;
Review Comment:
Agreed, I should have consolidated while testing it. Merged
`rejectsBodyOneByteOverDownloadCap` into
`rejectsStreamWithNoContentLengthAndBodyOverCap`, which is now the boundary
pair of `acceptsBodyExactlyAtDownloadCap`. I did the same for the Python twin,
whose assertion now checks `exceeded the limit of 1024 bytes`. I removed
`happyPathExtractionUnchanged`, which is covered by `extractsTopLevelEntries`.
New helpers `downloadCap`, `startStreamingServer`, `assertDownloadRejected` and
`countSkillTempDirs` (and `_download_cap` in Python) replace the repeated setup
and cleanup checks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
