Fix embedded HEIF media handling; land section 2-4 hardening

compose-check.sh

@@ -11,7 +11,7 @@ BASE_URL="${WATERMARKS_SERVICE_URL:-http://127.0.0.1:8765}"
 COMPOSE=(docker compose --profile harness --profile heavy)
 FAIL=0
 
-if curl -fsS "$BASE_URL/health" >/dev/null 2>&1; then
+if curl -fsS --connect-timeout 2 --max-time 10 "$BASE_URL/health" >/dev/null 2>&1; then
   echo "wr-core: OK"
 else
   echo "wr-core: FAIL (no /health at $BASE_URL)"

skills/clean-user-facing-text/scripts/clean_text.py

@@ -10,8 +10,24 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 
-from common import backup_path, cleaned_path, eprint, read_text_input, write_text_output  # noqa: E402
-from text_unicode import clean_text  # noqa: E402
+from common import (
+    backup_path,
+    cleaned_path,
+    eprint,
+    read_text_input,
+    write_text_output,
+)
+from text_unicode import clean_text
+
+
+def _paths_alias(left: Path, right: Path) -> bool:
+    """Return whether two paths name the same file (path, hard link, symlink)."""
+    if left == right:
+        return True
+    try:
+        return left.samefile(right)
+    except OSError:
+        return False
 
 
 def main() -> int:
@@ -63,11 +79,23 @@ def main() -> int:
             eprint("--in-place requires a file path")
             return 2
         src = Path(args.path)
-        bak = backup_path(src)
+        backup_path(src)
         out = str(src)
     elif out is None and args.path not in (None, "-"):
         out = str(cleaned_path(Path(args.path)))
 
+    # --output must never overwrite the input itself (same path, hard link,
+    # or symlink). --in-place is the sanctioned overwrite path and keeps a
+    # .bak backup, so it is exempt. safe_write_bytes hardens the actual write.
+    if (
+        not args.in_place
+        and out not in (None, "-")
+        and args.path not in (None, "-")
+        and _paths_alias(Path(args.path), Path(out))
+    ):
+        eprint(f"refusing to overwrite input: --output {out} aliases {args.path}")
+        return 2
+
     write_text_output(cleaned, out)
 
     if args.stats:

skills/remove-ai-marks/scripts/clean_file.py

@@ -244,6 +244,14 @@ def main() -> int:
             eprint(f"error on {error.path}: {error}")
         return ExitCode.RESIDUAL_OR_ERROR.value
     except ValueError as error:
+        # Preflight refusals (unrecognized format, binary-as-text, oversized
+        # input, output collisions) must not vanish as a bare exit code when
+        # --json is requested: emit the same structured shape a batch consumer
+        # can parse, while keeping the human message and the usage-error exit.
+        if args.json:
+            entry = {"error": str(error), "exit_code": ExitCode.USAGE_ERROR.value}
+            payload = {"total": 1, "results": [entry]} if batch else entry
+            print(json.dumps(payload, indent=2, ensure_ascii=False))
         eprint(f"invalid output selection: {error}")
         return ExitCode.USAGE_ERROR.value
     if args.dry_run:
@@ -329,12 +337,22 @@ def _plan_work(
         if kind == "text" and not args.force_text:
             with item.path.open("rb") as source:
                 head = source.read(8192)
-            guard_binary(
-                head,
-                str(item.path),
-                allow_binary=args.force_text,
-                advice=ROUTER_ADVICE,
-            )
+            try:
+                guard_binary(
+                    head,
+                    str(item.path),
+                    allow_binary=args.force_text,
+                    advice=ROUTER_ADVICE,
+                )
+            except SystemExit as error:
+                # guard_binary raises SystemExit(2) directly, which would blow
+                # past the preflight error handling below and swallow the
+                # structured JSON report in --json mode. Convert to ValueError
+                # so main() reports it like every other preflight refusal while
+                # preserving the usage-error exit mapping.
+                raise ValueError(
+                    f"refusing to treat {item.path} as text: binary content"
+                ) from error
         try:
             plan = _build_clean_plan(args, dest, kind)
         except Exception as error:

skills/remove-ai-marks/scripts/container_meta.py

@@ -198,6 +198,19 @@ def _blob_hits(blob: bytes) -> tuple[bool, bool, list[str]]:
 )
 
 
+def _media_strip_succeeded(sub_actions: list[str], cleaned: bytes, raw: bytes) -> bool:
+    """True when a media stripper changed bytes while reporting a removal.
+
+    Most raster strippers "drop" chunks/segments; heif_meta neutralizes in
+    place ("neutralized"/"zeroed") to preserve offsets, so accept both
+    vocabularies. The no-op case always returns the input bytes unchanged.
+    """
+    if cleaned == raw:
+        return False
+    verbs = ("drop", "neutraliz", "zero")
+    return any(verb in action.lower() for action in sub_actions for verb in verbs)
+
+
 def _inspect_embedded_data_uris(text: str) -> tuple[bool, bool, list[str]]:
     has_c2pa = False
     has_ai = False
@@ -236,7 +249,7 @@ def _inspect_embedded_data_uris(text: str) -> tuple[bool, bool, list[str]]:
             sub_c2pa, sub_ai, sub_findings = inspect_webp(data)
         elif fmt == "avif":
             sub_c2pa, sub_ai, sub_findings = inspect_avif(data)
-        elif fmt == "heic":
+        elif fmt == "heif":
             sub_c2pa, sub_ai, sub_findings = inspect_heic(data)
         elif "svg" in mime or data.lstrip().startswith(b"<"):
             sub_c2pa, sub_ai, sub_findings, _ = inspect_svg(data)
@@ -296,14 +309,14 @@ def _replace_uri(m: re.Match[str]) -> str:
                 cleaned_bytes, sub_actions = strip_webp(data, strip_all_metadata=strip_all_metadata)
             elif fmt == "avif":
                 cleaned_bytes, sub_actions = strip_avif(data, strip_all=strip_all_metadata)
-            elif fmt == "heic":
+            elif fmt == "heif":
                 cleaned_bytes, sub_actions = strip_heic(data, strip_all=strip_all_metadata)
             elif "svg" in mime.lower() or data.lstrip().startswith(b"<"):
                 cleaned_bytes, sub_actions = clean_svg(data)
         except Exception:
             return full_match
 
-        if not any("drop" in a.lower() for a in sub_actions) or cleaned_bytes == data:
+        if not _media_strip_succeeded(sub_actions, cleaned_bytes, data):
             return full_match
 
         actions.append(f"cleaned embedded data:image/{mime} ({', '.join(sub_actions[:2])})")
@@ -820,7 +833,7 @@ def _inspect_ooxml_zip(data: bytes, fmt: str) -> tuple[bool, bool, list[str], di
                         sub_c2pa, sub_ai, sub_findings = inspect_webp(raw)
                     elif img_fmt == "avif":
                         sub_c2pa, sub_ai, sub_findings = inspect_avif(raw)
-                    elif img_fmt == "heic":
+                    elif img_fmt == "heif":
                         sub_c2pa, sub_ai, sub_findings = inspect_heic(raw)
                     elif img_fmt == "gif":
                         sub_c2pa, sub_ai, sub_findings = inspect_gif(raw)
@@ -1025,7 +1038,7 @@ def _scrub_ooxml_zip(
                         cleaned_bytes, sub_actions = strip_webp(raw, strip_all_metadata=True)
                     elif img_fmt == "avif":
                         cleaned_bytes, sub_actions = strip_avif(raw, strip_all=True)
-                    elif img_fmt == "heic":
+                    elif img_fmt == "heif":
                         cleaned_bytes, sub_actions = strip_heic(raw, strip_all=True)
                     elif img_fmt == "gif":
                         cleaned_bytes, sub_actions = strip_gif(raw, strip_all_metadata=True)
@@ -1037,7 +1050,7 @@ def _scrub_ooxml_zip(
                         cleaned_bytes, sub_actions = clean_svg(raw)
                 except Exception:  # noqa: S110
                     pass
-                if any("drop" in a.lower() for a in sub_actions) and cleaned_bytes != raw:
+                if _media_strip_succeeded(sub_actions, cleaned_bytes, raw):
                     actions.append(f"clean embedded media in {name} ({', '.join(sub_actions[:2])})")
                     raw = cleaned_bytes
                 kept.append((info, raw))
@@ -1494,7 +1507,7 @@ def clean_epub(data: bytes, *, also_layer_a_text: bool = True) -> tuple[bytes, l
                         cleaned, sub_actions = strip_webp(raw, strip_all_metadata=True)
                     elif img_fmt == "avif":
                         cleaned, sub_actions = strip_avif(raw, strip_all=True)
-                    elif img_fmt == "heic":
+                    elif img_fmt == "heif":
                         cleaned, sub_actions = strip_heic(raw, strip_all=True)
                     elif img_fmt == "gif":
                         cleaned, sub_actions = strip_gif(raw, strip_all_metadata=True)
@@ -1506,7 +1519,7 @@ def clean_epub(data: bytes, *, also_layer_a_text: bool = True) -> tuple[bytes, l
                         cleaned, sub_actions = clean_svg(raw)
                 except Exception:  # noqa: S110
                     pass
-                if any("drop" in a.lower() for a in sub_actions) and cleaned != raw:
+                if _media_strip_succeeded(sub_actions, cleaned, raw):
                     actions.append(f"clean embedded media in {name} ({', '.join(sub_actions[:2])})")
                     raw = cleaned
                 kept.append((info, raw))

skills/remove-ai-marks/scripts/detect_text_watermark.py

@@ -226,7 +226,11 @@ def _cmd_watermark(args: argparse.Namespace, upstream: Path, alg: str) -> int:
 
     wm_out = "-" if args.watermarked_output is None else args.watermarked_output
     if wm_out == "-":
-        sys.stdout.write(watermarked)
+        # Never mix generated text with the --json payload on stdout: a batch
+        # consumer json.loads()ing stdout would choke on the sample. In JSON
+        # mode the sample goes to stderr; non-JSON mode keeps the CLI contract
+        # of writing the sample to stdout.
+        (sys.stderr if args.json else sys.stdout).write(watermarked)
     else:
         atomic_write_text(Path(wm_out), watermarked)
     if unwatermarked is not None:
@@ -264,6 +268,12 @@ def _add_common(p: argparse.ArgumentParser) -> None:
         default=None,
         help="MarkLLM checkout root (default: $MARKLLM_DIR)",
     )
+    p.add_argument(
+        "--rlimit-as",
+        type=int,
+        default=None,
+        help=argparse.SUPPRESS,
+    )
     p.add_argument(
         "--scheme",
         required=True,
@@ -298,6 +308,36 @@ def _add_common(p: argparse.ArgumentParser) -> None:
     )
 
 
+def _apply_rlimit_as(bytes_: int | None) -> None:
+    """Cap this process's address space before any heavy import (POSIX only).
+
+    Runs as the first thing the child does, so the limit is in force for the
+    whole MarkLLM harness (torch, transformers, ...) — equivalent to a
+    preexec_fn-set rlimit, applied at the subprocess boundary. The adapter
+    (text_detectors.MarkLLMTextDetector) passes --rlimit-as when
+    WATERMARKS_MARKLLM_RLIMIT_AS is configured. Failures degrade silently,
+    matching common.subprocess_rlimits().
+    """
+    if bytes_ is None:
+        return
+    try:
+        import resource
+    except ImportError:
+        return
+    try:
+        resource.setrlimit(resource.RLIMIT_AS, (bytes_, bytes_))
+    except ValueError:
+        # macOS rejects lowering the hard limit while the current soft limit
+        # is still RLIM_INFINITY; lower the soft limit first, then the hard.
+        try:
+            resource.setrlimit(resource.RLIMIT_AS, (bytes_, resource.RLIM_INFINITY))
+            resource.setrlimit(resource.RLIMIT_AS, (bytes_, bytes_))
+        except (OSError, ValueError):
+            pass
+    except OSError:
+        pass
+
+
 def main() -> int:
     p = argparse.ArgumentParser(description=__doc__)
     sub = p.add_subparsers(dest="cmd", required=True)
@@ -330,6 +370,7 @@ def main() -> int:
     wm.set_defaults(handler=_cmd_watermark)
 
     args = p.parse_args()
+    _apply_rlimit_as(args.rlimit_as)
 
     raw_upstream = args.upstream_dir or os.environ.get("MARKLLM_DIR")
     upstream = resolve_upstream(str(raw_upstream) if raw_upstream else None)

skills/remove-ai-marks/scripts/image_meta.py

@@ -28,6 +28,11 @@
 
 OPTIONAL_TOOL_OUTPUT_LIMIT = 2 * 1024 * 1024
 SYNTHID_OUTPUT_LIMIT = 2 * 1024 * 1024
+#: Hard cap on the decompressed size of a PNG zTXt/iTXt text chunk.
+#: Compressed chunks that would expand past this are treated as
+#: decompression bombs: fail closed (no text entry) rather than allocating
+#: unbounded output.
+MAX_PNG_TEXT_BYTES = 8 * 1024 * 1024
 
 PNG_SIG = b"\x89PNG\r\n\x1a\n"
 JPEG_SOI = b"\xff\xd8"
@@ -178,6 +183,34 @@ def _contains_any(blob: bytes, needles: tuple[bytes, ...]) -> list[str]:
     return found
 
 
+def _zlib_decompress_bounded(data: bytes, max_bytes: int = MAX_PNG_TEXT_BYTES) -> bytes | None:
+    """Decompress a zlib stream with an explicit output cap.
+
+    Returns the decompressed bytes when the stream is valid and its output
+    fits in *max_bytes*; returns None for corrupt/truncated streams and when
+    the output would exceed the cap. None means "no usable text": callers
+    fail closed (skip the chunk) instead of allocating unbounded memory.
+    """
+    decompressor = zlib.decompressobj()
+    out = bytearray()
+    remaining = data
+    while True:
+        budget = max_bytes + 1 - len(out)
+        if budget <= 0:
+            return None
+        try:
+            out += decompressor.decompress(remaining, budget)
+        except zlib.error:
+            return None
+        if len(out) > max_bytes:
+            return None
+        if decompressor.eof:
+            return bytes(out)
+        if not decompressor.unconsumed_tail:
+            return None
+        remaining = decompressor.unconsumed_tail
+
+
 def _png_text_entries(payload: bytes, ctype: bytes) -> list[tuple[str, str]]:
     """Parse a PNG text-chunk payload into (key, value) pairs.
 
@@ -199,9 +232,8 @@ def _png_text_entries(payload: bytes, ctype: bytes) -> list[tuple[str, str]]:
         key, sep, rest = payload.partition(b"\x00")
         if not sep or len(rest) < 2:
             return entries
-        try:
-            text = zlib.decompress(rest[1:])
-        except zlib.error:
+        text = _zlib_decompress_bounded(rest[1:])
+        if text is None:
             return entries
         entries.append(
             (
@@ -222,9 +254,8 @@ def _png_text_entries(payload: bytes, ctype: bytes) -> list[tuple[str, str]]:
         if not sep3:
             return entries
         if comp_flag == 1:
-            try:
-                text = zlib.decompress(text)
-            except zlib.error:
+            text = _zlib_decompress_bounded(text)
+            if text is None:
                 return entries
         entries.append(
             (
@@ -568,6 +599,7 @@ def inspect_image(
         has_c2pa, has_ai, findings = inspect_tiff(data)
     else:
         has_c2pa, has_ai, findings = False, False, ["unsupported format"]
+        notes.append(f"format '{fmt}' is not inspected")
 
     tools = run_optional_tools(path)
     # Elevate flags from tools

skills/remove-ai-marks/scripts/inspect_file.py

@@ -15,7 +15,7 @@
 sys.path.insert(0, str(Path(__file__).resolve().parent))
 from asset_kind import SUPPORTED_EXTENSIONS, classify_asset
 from batch_inputs import select_inputs
-from common import MAX_INPUT_BYTES, emit_json, eprint, read_text_input
+from common import EXIT_PARTIAL, MAX_INPUT_BYTES, emit_json, eprint, read_text_input
 from container_meta import inspect_container
 from image_meta import inspect_image
 from inspect_soft_binding import inspect_soft_binding
@@ -51,6 +51,10 @@ def main() -> int:
         emit_json({"total": len(results), "results": results} if batch else results[0])
     elif batch:
         eprint(f"inspected {len(results)} file(s)")
+    # An incomplete audit is the more important CI signal: any input that was
+    # not scanned (unrecognized or refused) outranks both clean and suspicious.
+    if any(r.get("unscanned") for r in results):
+        return EXIT_PARTIAL
     return 0 if all(not r.get("suspicious", False) for r in results) else 1
 
 
@@ -61,6 +65,7 @@ def _inspect_single(path: Path, args) -> dict:
             "path": str(path),
             "note": f"input larger than {MAX_INPUT_BYTES} bytes",
             "suspicious": False,
+            "unscanned": True,
         }
     kind = classify_asset(path, forced_kind=args.force_type)
     if kind == "text":
@@ -111,6 +116,7 @@ def _inspect_single(path: Path, args) -> dict:
             "path": str(path),
             "note": note,
             "suspicious": False,
+            "unscanned": True,
         }
 
     report = inspect_container(path)

skills/remove-ai-marks/scripts/rewrite_text.py

@@ -393,7 +393,10 @@ def _call_openai_compatible(
         "messages": [{"role": "user", "content": prompt}],
         "temperature": temperature,
     }
-    if reasoning_effort:
+    # "off" is a documented sentinel that omits the parameter entirely
+    # (generic OpenAI-compatible servers may reject it); every other value,
+    # including "none", is a real reasoning-effort request and is sent.
+    if reasoning_effort and reasoning_effort != "off":
         payload["reasoning_effort"] = reasoning_effort
     if disable_thinking:
         # Supported by Qwen/Transformers-compatible servers; opt-in so generic