wenjin272 commented on code in PR #1091:
URL: https://github.com/apache/flink-agents/pull/1091#discussion_r4044294137
##########
runtime/src/main/java/org/apache/flink/agents/runtime/skill/repository/SkillMaterializer.java:
##########
@@ -58,6 +61,104 @@ public final class SkillMaterializer {
private static final int JAR_URL_PREFIX_LEN = "jar:".length();
+ // --- Size caps for download and extraction (issue #1072) ---
+
+ /**
+ * All four resource limits for a single materializer operation, grouped
so tests can inject
+ * small thresholds without touching production defaults and so future
config wiring has a
+ * single object to populate from {@code SkillMaterializerOptions}.
+ *
+ * <p>Production callers use {@link #DEFAULT}; tests construct a small
instance and pass it to
+ * the overloads of {@link SkillMaterializer#downloadToTempFile} and {@link
+ * SkillMaterializer#extractZipSafely} that accept a {@code Limits}
argument.
+ */
+ public static final class Limits {
+ /** Maximum number of bytes accepted from a single HTTP download. */
+ public final long maxDownloadBytes;
+
+ /**
+ * Maximum uncompressed size of any single entry during zip
extraction. Enforced against
+ * actual bytes written, not the attacker-controlled declared size.
+ */
+ public final long maxExtractEntryBytes;
+
+ /**
+ * Maximum cumulative uncompressed bytes written across all entries
during a single zip
+ * extraction. Enforced against actual bytes written.
+ */
+ public final long maxExtractTotalBytes;
+
+ /** Maximum number of entries permitted in a single zip archive. */
+ public final int maxExtractEntries;
+
+ /**
+ * Conservative production defaults. These are intentionally lower
than the original issue
+ * #1072 values to limit disk consumption per materialization,
especially when multiple
+ * skills are materialized concurrently. Deployments that need larger
archives should raise
+ * them explicitly via a {@code Limits} instance.
+ *
+ * <p>TODO: wire these through {@code AgentConfigOptions} / {@code
SkillMaterializerOptions}
+ * so deployments can override them from the YAML config without code
changes (follow-up
+ * PR).
+ */
+ public static final Limits DEFAULT =
+ new Limits(
+ 64L * 1024 * 1024, // 64 MiB download
+ 64L * 1024 * 1024, // 64 MiB per entry
+ 256L * 1024 * 1024, // 256 MiB total extraction
+ 1_000); // entries
+
+ /**
+ * Construct a {@code Limits} instance. All values must be strictly
positive.
+ *
+ * @throws IllegalArgumentException if any value is not strictly
positive.
+ */
+ public Limits(
+ long maxDownloadBytes,
+ long maxExtractEntryBytes,
+ long maxExtractTotalBytes,
+ int maxExtractEntries) {
+ if (maxDownloadBytes <= 0
+ || maxExtractEntryBytes <= 0
+ || maxExtractTotalBytes <= 0
+ || maxExtractEntries <= 0) {
+ throw new IllegalArgumentException("All Limits values must be
strictly positive");
+ }
+ this.maxDownloadBytes = maxDownloadBytes;
+ this.maxExtractEntryBytes = maxExtractEntryBytes;
+ this.maxExtractTotalBytes = maxExtractTotalBytes;
+ this.maxExtractEntries = maxExtractEntries;
+ }
+ }
+
+ /**
+ * Backward-compatible alias for {@link Limits#DEFAULT#maxDownloadBytes}.
+ *
+ * @deprecated Use {@link Limits#DEFAULT} or inject a {@link Limits}
instance.
+ */
+ public static final long MAX_DOWNLOAD_BYTES =
Limits.DEFAULT.maxDownloadBytes;
Review Comment:
Thanks for lowering the defaults and grouping the limits into one object.
However, `Limits` / `MaterializerLimits` currently only provides an injection
point for tests: all production repository paths still use the defaults, while
the actual configuration wiring is deferred by the TODO. I think this
configuration surface needs to be completed in this PR, with the four positive
limits exposed consistently through the Java, Python, and YAML APIs and
propagated through the real materialization path. Also, the `MAX_*` aliases
were introduced in this still-unmerged PR, so there is no released API
compatibility to preserve. Could we remove these already-deprecated public
constants and use the configured limits object directly? This would also
eliminate the current compiler warnings caused by Javadoc `@deprecated` tags
without corresponding `@Deprecated` annotations.
##########
python/flink_agents/runtime/skill/repository/_materialize.py:
##########
@@ -160,37 +218,179 @@ def copy_dir_to_temp(src_dir: Path) -> Materialized:
return materialized
-def extract_zip_safely(zip_path: Path) -> Materialized:
+def extract_zip_safely(
+ zip_path: Path, *, limits: MaterializerLimits = DEFAULT_LIMITS
+) -> Materialized:
"""Extract a zip into a fresh temp dir, returning a :class:`Materialized`.
Each entry is validated against zip-slip. ``close()`` the returned handle
to free the dir eagerly; an atexit cleanup is the fallback.
Args:
zip_path: Path to the zip file to extract.
+ limits: Resource limits to enforce during extraction. Defaults to
+ :data:`DEFAULT_LIMITS`. Pass a small :class:`MaterializerLimits`
+ instance in tests to avoid allocating large fixtures.
Returns:
A :class:`Materialized` handle owning the extraction directory.
Raises:
- ValueError: if any zip entry resolves outside the extraction directory.
+ ValueError: if any zip entry resolves outside the extraction directory,
+ or if any size or entry-count limit is exceeded.
"""
extract_dir = Path(tempfile.mkdtemp(prefix=_TEMP_DIR_PREFIX)).resolve()
# Construct the handle before validation so the (empty) tempdir is always
reclaimed,
# even if validation raises.
materialized = Materialized(extract_dir)
+ try:
+ _extract_zip_to_dir(zip_path, extract_dir, limits)
+ except Exception:
+ materialized.close()
+ raise
+ return materialized
+
+def _validate_zip_members(
+ members: list, extract_dir: Path, limits: MaterializerLimits
+) -> None:
+ if len(members) > limits.max_extract_entries:
+ msg = (
+ f"Skill archive contains {len(members)} entries, "
+ f"exceeding the limit of {limits.max_extract_entries}"
+ )
+ raise ValueError(msg)
+
+ for member in members:
+ target = (extract_dir / member.filename).resolve()
+ if not target.is_relative_to(extract_dir):
+ msg = f"Unsafe zip entry: {member.filename}"
+ raise ValueError(msg)
+
+ total_declared = 0
+ for member in members:
+ if member.is_dir():
+ continue
+ declared = member.file_size
+ if declared > limits.max_extract_entry_bytes:
+ msg = (
+ f"Skill archive entry '{member.filename}' declared size
{declared} "
+ f"exceeds the per-entry limit of
{limits.max_extract_entry_bytes} bytes"
+ )
+ raise ValueError(msg)
+ if declared > 0:
+ total_declared += declared
+ if total_declared > limits.max_extract_total_bytes:
+ msg = (
+ f"Skill archive declared total uncompressed size {total_declared} "
+ f"exceeds the limit of {limits.max_extract_total_bytes} bytes"
+ )
+ raise ValueError(msg)
+
+def _open_entry_stream(zip_path: Path, member: zipfile.ZipInfo) ->
io.RawIOBase:
+ with zip_path.open("rb") as f:
+ f.seek(member.header_offset)
+ f.read(4) # local file header signature
+ f.read(2) # version needed
+ f.read(2) # general purpose bit flag
+ f.read(2) # compression method
+ f.read(2) # last mod file time
+ f.read(2) # last mod file date
+ f.read(4) # crc-32
+ f.read(4) # compressed size
+ f.read(4) # uncompressed size
+ fn_len = struct.unpack("<H", f.read(2))[0]
+ ex_len = struct.unpack("<H", f.read(2))[0]
+ f.read(fn_len + ex_len) # filename + extra field
+ compressed = f.read(member.compress_size)
+
+ if member.compress_type == zipfile.ZIP_STORED:
+ return io.BytesIO(compressed)
+ if member.compress_type == zipfile.ZIP_DEFLATED:
+ return io.BytesIO(zlib.decompress(compressed, -15))
Review Comment:
`zlib.decompress()` expands the entire entry in memory before the size
checks run, so a zip bomb can cause OOM before being rejected. It also bypasses
CRC validation. Could we keep `zf.open(member)` for streaming extraction and
test the byte counter through a separate bounded-copy helper?
--
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]
