pranavshuklaa commented on code in PR #1091:
URL: https://github.com/apache/flink-agents/pull/1091#discussion_r4075962528
##########
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:
Fixed this. _open_entry_stream is deleted. _extract_zip_to_dir now uses
zf.open(member) for streaming decompression, reading in 64 KiB chunks. Size
checks run per chunk before any bytes are written to disk. BadZipFile
(bad CRC-32 or truncated data) is now being caught at the outer try level
and re-raised as ValueError("... failed integrity check ..."), so the archive
is still rejected via CRC validation rather than the byte counter. When a size
limit fires mid-read, src.close() is called inside a try/except BadZipFile
block before re-raising the original limit error, suppressing the spurious CRC
error that would otherwise fire on an incomplete read.
--
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]
