pranavshuklaa commented on code in PR #1091:
URL: https://github.com/apache/flink-agents/pull/1091#discussion_r4075829369
##########
runtime/src/test/java/org/apache/flink/agents/runtime/skill/SkillMaterializerTest.java:
##########
@@ -592,4 +598,670 @@ private List<String> getMessages() {
return messages;
}
}
+ // -------------------------------------------------------
+ // Download size cap tests
+ // -------------------------------------------------------
+
+ /**
+ * Server declares a Content-Length larger than the cap. The pre-flight
check must reject before
+ * reading any body bytes.
+ */
+ @Test
+ void rejectsDeclaredContentLengthOverCap() throws IOException {
+ long overCap = SkillMaterializer.MAX_DOWNLOAD_BYTES + 1;
+ // We serve an empty body but declare a huge Content-Length.
+ // The handler sends the declared length in the header, then closes
immediately.
+ HttpServer server = HttpServer.create(new
InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext(
+ "/",
+ exchange -> {
+ exchange.getResponseHeaders().add("Content-Length",
String.valueOf(overCap));
+ // sendResponseHeaders with -1 means no auto
Content-Length; we set it above.
+ 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));
+ assertTrue(
+ ex.getMessage().contains("exceeding the limit"),
+ "error must mention the limit, got: " + ex.getMessage());
+ // Confirm no temp file was left behind.
+ // (We can't grab the path since the call threw, but we can verify
indirectly
+ // by checking the message does not contain a path — the important
thing is
+ // the exception propagated cleanly. The cleanup assertion below
is the
+ // stronger guarantee tested in cleanupOnDownloadFailure.)
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ /**
+ * Server declares a small (below-cap) Content-Length but actually streams
more bytes. The byte
+ * counter must catch the overage even though the pre-flight passed.
+ */
+ @Test
+ void rejectsUnderstatedContentLengthViaByteCounter() throws IOException {
+ // Declare 100 bytes but stream MAX_DOWNLOAD_BYTES + 1 bytes.
+ int declaredLength = 100;
+ long actualBytes = SkillMaterializer.MAX_DOWNLOAD_BYTES + 1;
+ HttpServer server = HttpServer.create(new
InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext(
+ "/",
+ exchange -> {
+ // Set a small declared size so the pre-flight passes.
+ 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) {
+ // Client closed; stop writing.
+ 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",
+ 30_000,
+ true));
+ 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.
+ */
+ @Test
+ void rejectsStreamWithNoContentLengthAndBodyOverCap() throws IOException {
+ long actualBytes = SkillMaterializer.MAX_DOWNLOAD_BYTES + 1;
+ HttpServer server = HttpServer.create(new
InetSocketAddress("127.0.0.1", 0), 0);
+ server.createContext(
+ "/",
+ exchange -> {
+ // 0 enables chunked transfer without a Content-Length
header.
+ 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",
+ 30_000,
+ true));
+ assertTrue(
+ ex.getMessage().contains("exceeded the limit"),
+ "error must mention the limit, got: " + ex.getMessage());
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ /** A body exactly at the cap (MAX_DOWNLOAD_BYTES bytes) must succeed. */
+ @Test
+ void acceptsBodyExactlyAtDownloadCap() throws IOException {
+ // Using a small cap so the test doesn't actually allocate 512 MiB.
+ // We test the boundary logic by constructing a body of exactly cap
bytes,
+ // where cap here is small. Since MAX_DOWNLOAD_BYTES is a constant we
can't
+ // change per-test, we use a body that is clearly below the cap
instead and
+ // trust the cap+1 tests above cover the boundary.
+ // This test just confirms a normal small download still works
unaffected.
+ byte[] body = new byte[1024];
+ 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);
+ try {
+ assertEquals(1024, Files.size(file));
+ } finally {
+ Files.deleteIfExists(file);
+ }
+ } finally {
+ server.stop(0);
+ }
+ }
Review Comment:
I had pushed the changes and missed on replying to these, sorry for that. To
update the Limits object is now injectable into all materializer methods, so
tests use a 1 KiB cap instead of the 512 MiB production default.
acceptsBodyExactlyAtDownloadCap now sends exactly 1,024 bytes against a
1,024-byte cap and asserts success. rejectsBodyOneByteOverDownloadCap sends
1,025 bytes against the same cap with no Content-Length header, so only the
streaming counter can reject it. The same injectable limits are used across all
download and extraction cap tests,no test allocates more than a few KiB of data.
--
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]
