wenjin272 commented on code in PR #1091:
URL: https://github.com/apache/flink-agents/pull/1091#discussion_r4464305124
##########
runtime/src/main/java/org/apache/flink/agents/runtime/skill/repository/SkillMaterializer.java:
##########
@@ -203,19 +448,16 @@ public static Materialized
extractClasspathFromJars(List<URL> jarUrls, String re
if (jarUrls == null || jarUrls.isEmpty()) {
throw new IllegalArgumentException("jarUrls must be non-empty");
}
+
Path extractDir = Files.createTempDirectory(TEMP_DIR_PREFIX);
Thread hook = registerCleanup(extractDir);
String prefix = resourcePrefix.endsWith("/") ? resourcePrefix :
resourcePrefix + "/";
- Materialized materialized = new Materialized(extractDir, hook);
- try {
- for (URL jarUrl : jarUrls) {
- copyJarEntries(jarUrl, prefix, extractDir);
- }
- return materialized;
- } catch (IOException | RuntimeException | Error e) {
- materialized.close();
- throw e;
+
+ for (URL jarUrl : jarUrls) {
+ copyJarEntries(jarUrl, prefix, extractDir);
}
+
+ return new Materialized(extractDir, hook);
Review Comment:
Could we restore the original exception cleanup here? If a later JAR fails,
previously extracted files and the shutdown hook remain until JVM exit, and the
caller never receives a handle to close. Please also add a regression test with
a valid first JAR followed by a failing one.
##########
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:
Non-blocking: this duplicates
`rejectsStreamWithNoContentLengthAndBodyOverCap`: both send 1,025 bytes without
Content-Length against a 1,024-byte cap. Could we consolidate them and the
equivalent Python duplicates? The new Java happy-path extraction test also
repeats existing coverage. Small helpers could reduce repeated server setup and
cleanup assertions without removing distinct scenarios.
##########
python/flink_agents/runtime/skill/tests/test_materialize.py:
##########
@@ -309,3 +400,532 @@ def
test_logs_sanitized_effective_url_for_same_protocol_redirect(
assert "redirect-fragment" not in warning
finally:
path.unlink(missing_ok=True)
+
+
+# ---------------------------------------------------------------------------
+# Helpers for size-cap tests
+# ---------------------------------------------------------------------------
+
+
+def _make_streaming_server(
+ declared_content_length: int | None, bytes_to_stream: int
+) -> tuple[str, HTTPServer]:
+ class _StreamingHandler(BaseHTTPRequestHandler):
+ def do_GET(self) -> None:
+ self.send_response(200)
+ if declared_content_length is not None:
+ self.send_header("Content-Length",
str(declared_content_length))
+ self.end_headers()
+ chunk = b"x" * 65536
+ remaining = bytes_to_stream
+ while remaining > 0:
+ to_write = min(len(chunk), remaining)
+ try:
+ self.wfile.write(chunk[:to_write])
+ self.wfile.flush()
+ except OSError:
+ break
+ remaining -= to_write
+
+ def log_message(self, *_args: object) -> None:
+ pass
+
+ server = HTTPServer(("127.0.0.1", 0), _StreamingHandler)
+ port = server.server_address[1]
+ threading.Thread(target=server.serve_forever, daemon=True).start()
+ return f"http://127.0.0.1:{port}";, server
+
+
+def _make_zip_with_patched_declared_sizes(
+ zip_path: Path, entries: dict[str, bytes], declared_size: int
+) -> None:
+ r"""Write a ZIP and patch the declared uncompressed size in both the local
file
+ header (LFH, ``PK\x03\x04``) and the central directory header (CDH,
``PK\x01\x02``)
+ for every entry.
+
+ Patching only the CDH is insufficient for the per-entry streaming test:
Python's
+ ``zipfile`` reads the LFH declared size to bound how many bytes it
decompresses, so
+ leaving the LFH intact causes a CRC mismatch before the byte counter is
reached.
+ Patching both headers lets the full real payload flow through the
streaming counter.
Review Comment:
Nit: this explanation is outdated. The current forged-size test expects CRC
validation to reject the archive, while byte counters are tested through
`_bounded_copy`. Could we update the comment to reflect that?
##########
api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java:
##########
@@ -208,6 +208,37 @@ public enum ConditionEvaluationFailureStrategy {
public static final ConfigOption<Integer> TRACE_LOG_MAX_DEPTH =
new ConfigOption<>("trace-log.standard.max-depth", Integer.class,
5);
+ /**
+ * The maximum number of bytes accepted from a single skill archive
download over HTTP(S). Must
+ * be strictly positive; there is no unlimited setting. Defaults to 64 MiB.
+ */
+ public static final ConfigOption<Long> SKILL_SOURCE_URL_MAX_DOWNLOAD_BYTES
=
+ new ConfigOption<>(
+ "skill.source.url.max-download-bytes", Long.class, 64L *
1024 * 1024);
+
+ /**
+ * The maximum uncompressed size in bytes of a single entry when a skill
archive is extracted.
+ * Must be strictly positive; there is no unlimited setting. Defaults to
64 MiB.
+ */
+ public static final ConfigOption<Long>
SKILL_SOURCE_URL_MAX_EXTRACT_ENTRY_BYTES =
+ new ConfigOption<>(
+ "skill.source.url.max-extract-entry-bytes", Long.class,
64L * 1024 * 1024);
Review Comment:
Could we group these four options in a dedicated `SkillOptions` in the
skills API module, in both Java and Python?
Also, the extraction limits apply to local ZIP sources, so `url` is
misleading for those three keys. I suggest keeping
`skill.source.url.max-download-bytes` and using
`skill.source.max-extract-{entry-bytes,total-bytes,entries}` for extraction.
Please update the docs and configuration parity checks accordingly.
##########
python/flink_agents/runtime/skill/tests/test_url_repository.py:
##########
@@ -208,3 +212,34 @@ def test_404_error(self, zip_server: "tuple[str,
type[_ZipHandler]]") -> None:
with pytest.raises(HTTPError) as exc_info:
URLSkillRepository(f"{url}?token=top-secret",
allow_insecure_http=True)
assert "top-secret" not in str(exc_info.value)
+
+
+class TestURLSkillRepositoryLimits:
+ def test_download_uses_given_limits(
+ self, zip_server: "tuple[str, type[_ZipHandler]]"
+ ) -> None:
+ url, handler = zip_server
+ size = len(handler.zip_bytes)
+ limits = MaterializerLimits(
+ max_download_bytes=size - 1,
+ max_extract_entry_bytes=DEFAULT_LIMITS.max_extract_entry_bytes,
+ max_extract_total_bytes=DEFAULT_LIMITS.max_extract_total_bytes,
+ max_extract_entries=DEFAULT_LIMITS.max_extract_entries,
+ )
+ with pytest.raises(ValueError, match=f"limit of {size - 1} bytes"):
+ URLSkillRepository(url, allow_insecure_http=True, limits=limits)
+ # Control: the same archive loads with the default limits.
+ assert URLSkillRepository(url, allow_insecure_http=True).get_skills()
Review Comment:
Nit: please close the repository created by this successful control case, so
its temporary directory is released immediately. The filesystem repository
control case needs the same cleanup.
--
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]
