From e30615a7b10339e8854fce4039b1d5635c8f8418 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 18 Aug 2026 18:45:20 -0400 Subject: [PATCH 1/3] fix(container): inspect and strip embedded HEIF media MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The same detect_format contract gap as the top-level router: HEIF-family bytes are reported as "heif", never "heic", but all five embedded-media branches (HTML data: URIs, OOXML zip media, EPUB media) matched "heic" — an embedded HEIC image in a docx/odt/epub was never inspected nor stripped, and its C2PA/XMP markers survived a "cleaned" document. The strip gates also required the action verb "drop", while heif_meta neutralizes in place ("neutralized"/"zeroed") to preserve offsets, so even a routed HEIF would never have been written back. Gates now share _media_strip_succeeded, which accepts both vocabularies and still rejects the no-op case (unchanged bytes). --- .../remove-ai-marks/scripts/container_meta.py | 29 ++++++++++++++----- tests/test_ooxml_xlsx_pptx.py | 25 ++++++++++++++++ 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/skills/remove-ai-marks/scripts/container_meta.py b/skills/remove-ai-marks/scripts/container_meta.py index 72d090d..c3b00f8 100644 --- a/skills/remove-ai-marks/scripts/container_meta.py +++ b/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)) diff --git a/tests/test_ooxml_xlsx_pptx.py b/tests/test_ooxml_xlsx_pptx.py index e5c44a2..d965439 100644 --- a/tests/test_ooxml_xlsx_pptx.py +++ b/tests/test_ooxml_xlsx_pptx.py @@ -271,6 +271,31 @@ def test_docx_embedded_media_cleaning(): assert has_c2pa_after is False +def test_docx_embedded_heic_media_inspected_and_cleaned(): + # Regression: detect_format reports HEIF-family bytes as "heif", but the + # embedded-media branches matched "heic" — embedded HEIF was silently + # skipped by both inspect and clean. + from test_heif_meta import _minimal_heic_with_xmp + + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf: + zf.writestr( + "word/document.xml", + "Doc", + ) + zf.writestr("word/media/image1.heic", _minimal_heic_with_xmp()) + docx_data = buf.getvalue() + + _has_c2pa, has_ai, findings, _ = inspect_docx(docx_data) + assert has_ai is True + assert any("word/media/image1.heic" in f for f in findings) + + cleaned_data, actions = clean_docx(docx_data) + assert any("clean embedded media in word/media/image1.heic" in a for a in actions) + _, has_ai_after, _, _ = inspect_docx(cleaned_data) + assert has_ai_after is False + + def test_clean_container_and_inspect_container_xlsx_pptx(tmp_path): xlsx_path = tmp_path / "budget.xlsx" xlsx_path.write_bytes(_create_synthetic_xlsx()) From a3cfde48e6bfb057284e13b51d941335527e0114 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 18 Aug 2026 18:45:28 -0400 Subject: [PATCH 2/3] fix: server stability and bounded-parser hardening Concurrent hardening pass (master-list sections 2-4), verified green by the full suite: - server: /detect answers ok+kind=unknown with an explanatory report instead of falling through to a failed tempfile workflow; OpenAPI schemas admit the "unknown" kind and the detect report field - image_meta: bound PNG zTXt/iTXt decompression (MAX_PNG_TEXT_BYTES, fail closed on bomb/corrupt streams) via _zlib_decompress_bounded - inspect_file, score_synthid, score_stylometry, text_detectors, synthid_score_server: partial-read and error-path handling - rewrite_text, clean-user-facing-text: alias and validation fixes - compose-check: curl connect/total timeouts - tests: new test_stability_server, test_image_meta_bomb_and_notes, test_inspect_file_partial, test_lightweight_clean_text_aliases plus extensions to binary-guard, markllm-detect, rewrite, stylometry, and text-detector suites --- compose-check.sh | 2 +- .../scripts/clean_text.py | 34 +- skills/remove-ai-marks/scripts/clean_file.py | 30 +- .../scripts/detect_text_watermark.py | 43 ++- skills/remove-ai-marks/scripts/image_meta.py | 44 ++- .../remove-ai-marks/scripts/inspect_file.py | 8 +- .../remove-ai-marks/scripts/rewrite_text.py | 5 +- .../scripts/score_stylometry.py | 10 +- .../remove-ai-marks/scripts/score_synthid.py | 66 ++-- skills/remove-ai-marks/scripts/server.py | 22 +- .../scripts/synthid_score_server.py | 16 +- .../remove-ai-marks/scripts/text_detectors.py | 21 +- tests/test_binary_guard.py | 43 ++- tests/test_image_meta_bomb_and_notes.py | 124 ++++++++ tests/test_inspect_file_partial.py | 99 ++++++ tests/test_lightweight_clean_text_aliases.py | 95 ++++++ tests/test_markllm_detect.py | 61 ++++ tests/test_rewrite_text.py | 26 ++ tests/test_stability_server.py | 291 ++++++++++++++++++ tests/test_stylometry.py | 28 ++ tests/test_text_detectors.py | 85 ++++- 21 files changed, 1089 insertions(+), 64 deletions(-) create mode 100644 tests/test_image_meta_bomb_and_notes.py create mode 100644 tests/test_inspect_file_partial.py create mode 100644 tests/test_lightweight_clean_text_aliases.py create mode 100644 tests/test_stability_server.py diff --git a/compose-check.sh b/compose-check.sh index 7d2fb2f..e313e1c 100755 --- a/compose-check.sh +++ b/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)" diff --git a/skills/clean-user-facing-text/scripts/clean_text.py b/skills/clean-user-facing-text/scripts/clean_text.py index 6fe8708..a02547b 100755 --- a/skills/clean-user-facing-text/scripts/clean_text.py +++ b/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: diff --git a/skills/remove-ai-marks/scripts/clean_file.py b/skills/remove-ai-marks/scripts/clean_file.py index 8ed55bd..ea8efaf 100644 --- a/skills/remove-ai-marks/scripts/clean_file.py +++ b/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: diff --git a/skills/remove-ai-marks/scripts/detect_text_watermark.py b/skills/remove-ai-marks/scripts/detect_text_watermark.py index 84d4dea..36b9c11 100755 --- a/skills/remove-ai-marks/scripts/detect_text_watermark.py +++ b/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) diff --git a/skills/remove-ai-marks/scripts/image_meta.py b/skills/remove-ai-marks/scripts/image_meta.py index cdffc70..f42f1d4 100755 --- a/skills/remove-ai-marks/scripts/image_meta.py +++ b/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 diff --git a/skills/remove-ai-marks/scripts/inspect_file.py b/skills/remove-ai-marks/scripts/inspect_file.py index 2539cab..1d31981 100644 --- a/skills/remove-ai-marks/scripts/inspect_file.py +++ b/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) diff --git a/skills/remove-ai-marks/scripts/rewrite_text.py b/skills/remove-ai-marks/scripts/rewrite_text.py index 4cc9360..41383e8 100644 --- a/skills/remove-ai-marks/scripts/rewrite_text.py +++ b/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 diff --git a/skills/remove-ai-marks/scripts/score_stylometry.py b/skills/remove-ai-marks/scripts/score_stylometry.py index d1f2537..2518641 100755 --- a/skills/remove-ai-marks/scripts/score_stylometry.py +++ b/skills/remove-ai-marks/scripts/score_stylometry.py @@ -34,6 +34,7 @@ DEFAULT_THRESHOLD = 0.65 MIN_SAMPLE_WORDS = 30 FULL_WEIGHT_WORDS = 100 +MAX_SCAN_CHARS = 2_000_000 # High-frequency formulaic transition markers, hedging verbs, and structural # boilerplate commonly overrepresented in AI-generated text across frontier LLMs. @@ -244,13 +245,20 @@ def scan_ai_phrases(text: str) -> list[MarkerMatch]: def score_text_stylometry(text: str, path: str = "") -> StylometryReport: """Run full multi-dimensional stylometric analysis and return a structured report.""" + notes: list[str] = [] + if len(text) > MAX_SCAN_CHARS: + notes.append( + f"Input length {len(text)} exceeds MAX_SCAN_CHARS ({MAX_SCAN_CHARS}); " + f"analysis truncated to first {MAX_SCAN_CHARS} characters" + ) + text = text[:MAX_SCAN_CHARS] + words = extract_words(text) word_count = len(words) sentences = extract_sentences(text) sentence_count = len(sentences) findings: list[str] = [] - notes: list[str] = [] # 1. Length Guard if word_count < MIN_SAMPLE_WORDS: diff --git a/skills/remove-ai-marks/scripts/score_synthid.py b/skills/remove-ai-marks/scripts/score_synthid.py index b5fced5..69274e9 100755 --- a/skills/remove-ai-marks/scripts/score_synthid.py +++ b/skills/remove-ai-marks/scripts/score_synthid.py @@ -23,9 +23,17 @@ import json import os import sys +import threading from pathlib import Path from typing import Any +# Serializes the process-global section of score_file: the sys.path +# insertion of the external checkout, the module imports, the codebook +# load / extraction, and the stdout redirection. The upstream library is +# not thread-safe and redirect_stdout swaps process-wide stdout, so the +# threaded HTTP sidecar must not run two scores concurrently. +_SCORE_LOCK = threading.Lock() + def resolve_upstream(raw: str | None) -> Path | None: if not raw: @@ -74,34 +82,36 @@ def score_file( print(f"codebook not found: {codebook_path}", file=sys.stderr) return 3, None - sys.path.insert(0, str(extraction)) - try: - import cv2 - from robust_extractor import RobustSynthIDExtractor - from synthid_bypass_v4 import SpectralCodebookV4 - except ImportError as e: - print(f"optional scorer dependencies missing: {e}", file=sys.stderr) - return 3, None - - try: - img = cv2.imread(str(path)) - if img is None: - print(f"could not load image: {path}", file=sys.stderr) - return 2, None - rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) - - # Upstream prints progress ("CodebookV4 loaded: ...") straight to - # stdout, which corrupts --json for any caller that parses us - # (image_meta.py json.loads our stdout). Keep stdout ours alone. - with contextlib.redirect_stdout(sys.stderr): - codebook_v4 = SpectralCodebookV4() - codebook_v4.load(str(codebook_path)) - - extractor = RobustSynthIDExtractor() - result = extractor.detect_from_v4_codebook(rgb, codebook_v4, model=model) - except Exception as e: - print(f"scorer error: {e}", file=sys.stderr) - return 1, None + with _SCORE_LOCK: + if str(extraction) not in sys.path: + sys.path.insert(0, str(extraction)) + try: + import cv2 + from robust_extractor import RobustSynthIDExtractor + from synthid_bypass_v4 import SpectralCodebookV4 + except ImportError as e: + print(f"optional scorer dependencies missing: {e}", file=sys.stderr) + return 3, None + + try: + img = cv2.imread(str(path)) + if img is None: + print(f"could not load image: {path}", file=sys.stderr) + return 2, None + rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + # Upstream prints progress ("CodebookV4 loaded: ...") straight to + # stdout, which corrupts --json for any caller that parses us + # (image_meta.py json.loads our stdout). Keep stdout ours alone. + with contextlib.redirect_stdout(sys.stderr): + codebook_v4 = SpectralCodebookV4() + codebook_v4.load(str(codebook_path)) + + extractor = RobustSynthIDExtractor() + result = extractor.detect_from_v4_codebook(rgb, codebook_v4, model=model) + except Exception as e: + print(f"scorer error: {e}", file=sys.stderr) + return 1, None payload = { "available": True, diff --git a/skills/remove-ai-marks/scripts/server.py b/skills/remove-ai-marks/scripts/server.py index d0c673e..d72815c 100755 --- a/skills/remove-ai-marks/scripts/server.py +++ b/skills/remove-ai-marks/scripts/server.py @@ -241,7 +241,9 @@ def _clean_request_schema() -> dict[str, Any]: type="object", properties={ "ok": _schema(type="boolean"), - "kind": _schema(type="string", enum=["text", "image", "container"]), + "kind": _schema( + type="string", enum=["text", "image", "container", "unknown"] + ), "suspicious": _schema(type="boolean"), "report": _schema(type="object"), }, @@ -283,8 +285,11 @@ def _clean_request_schema() -> dict[str, Any]: type="object", properties={ "ok": _schema(type="boolean"), - "kind": _schema(type="string", enum=["text", "image", "container"]), + "kind": _schema( + type="string", enum=["text", "image", "container", "unknown"] + ), "detections": _schema(type="array", items=_schema(type="object")), + "report": _schema(type="object"), }, ) }, @@ -531,6 +536,19 @@ def _handle_inspect(self, data: bytes, name: str, body: dict[str, Any]) -> None: def _handle_detect(self, data: bytes, name: str) -> None: kind = classify_bytes(data, Path(name).suffix) + if kind == "unknown": + self._respond( + HTTPStatus.OK, + { + "ok": True, + "kind": "unknown", + "detections": [], + "report": { + "note": "unrecognized format; use a filename with a known extension", + }, + }, + ) + return with tempfile.TemporaryDirectory(prefix="wm-detect-") as tmp: path = _tmp_path(Path(tmp), name or "input") path.write_bytes(data) diff --git a/skills/remove-ai-marks/scripts/synthid_score_server.py b/skills/remove-ai-marks/scripts/synthid_score_server.py index 33aca1e..932289b 100755 --- a/skills/remove-ai-marks/scripts/synthid_score_server.py +++ b/skills/remove-ai-marks/scripts/synthid_score_server.py @@ -24,6 +24,7 @@ import os import sys import tempfile +import traceback from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path @@ -133,8 +134,19 @@ def do_POST(self) -> None: except OSError as e: self._respond(HTTPStatus.INTERNAL_SERVER_ERROR, {"ok": False, "error": str(e)}) return - code, payload = score_file(path, model=MODEL) - + try: + code, payload = score_file(path, model=MODEL) + except Exception: + # score_file should return (1|2|3, None) instead of raising, + # but an unexpected exception must not kill the request with a + # raw traceback. Log it server-side and answer the fail-soft + # contract image_meta.run_synthid_score expects. + traceback.print_exc(file=sys.stderr) + self._respond( + HTTPStatus.OK, + {"available": False, "error": "scorer error (see sidecar stderr)"}, + ) + return if code == 0 and payload is not None: self._respond(HTTPStatus.OK, payload) elif code == 2: diff --git a/skills/remove-ai-marks/scripts/text_detectors.py b/skills/remove-ai-marks/scripts/text_detectors.py index 98cd487..a392ef4 100644 --- a/skills/remove-ai-marks/scripts/text_detectors.py +++ b/skills/remove-ai-marks/scripts/text_detectors.py @@ -233,7 +233,11 @@ class GeminiSynthIDTextDetector: vendor = "google" def available(self) -> bool: - return bool(os.environ.get("WATERMARKS_GEMINI_API_KEY", "").strip()) + # The DETECT_TEXT_WATERMARK task type is not supported by the Gemini + # generateContent API yet, so this detector cannot run regardless of + # configuration. Never advertise a detector detect() cannot execute; + # flip this once a supported watermark-detection endpoint exists. + return False def detect(self, text: str) -> dict[str, Any]: api_key = os.environ.get("WATERMARKS_GEMINI_API_KEY", "").strip() @@ -366,6 +370,9 @@ def detect(self, text: str) -> dict[str, Any]: cmd += ["--model", self._model] if self._upstream_dir: cmd += ["--upstream-dir", str(Path(upstream).expanduser().resolve())] + rlimit_as = _markllm_rlimit_as() + if rlimit_as is not None: + cmd += ["--rlimit-as", str(rlimit_as)] try: result = run_command( @@ -376,16 +383,18 @@ def detect(self, text: str) -> dict[str, Any]: except ExternalCommandTimeout: report["error"] = "MarkLLM detection timed out" return report + # Use the decoded text views: CommandResult.stderr/stdout are raw + # bytes, and a bytes error field would make the report un-serializable. if result.returncode == 3: - report["error"] = (result.stderr or "").strip()[:400] or "MarkLLM unavailable" + report["error"] = result.stderr_text.strip()[:400] or "MarkLLM unavailable" return report if result.returncode != 0: - report["error"] = (result.stderr or "").strip()[ - :400 - ] or f"MarkLLM exit {result.returncode}" + report["error"] = ( + result.stderr_text.strip()[:400] or f"MarkLLM exit {result.returncode}" + ) return report try: - payload = json.loads(result.stdout or "{}") + payload = json.loads(result.stdout_text or "{}") except json.JSONDecodeError as e: report["error"] = f"bad MarkLLM JSON: {e}" return report diff --git a/tests/test_binary_guard.py b/tests/test_binary_guard.py index ae56700..14363f6 100644 --- a/tests/test_binary_guard.py +++ b/tests/test_binary_guard.py @@ -3,6 +3,7 @@ from __future__ import annotations import io +import json import os import subprocess import sys @@ -169,6 +170,43 @@ def test_clean_file_in_place_as_text_on_docx_leaves_no_backup(tmp_path): assert not (tmp_path / "doc.docx.bak").exists() +def test_clean_file_json_emits_structured_error_for_binary_text(tmp_path): + """guard_binary preflight refusals must surface as structured JSON in --json mode.""" + blob = tmp_path / "binary.txt" + blob.write_bytes(make_docx(tmp_path / "x.docx").read_bytes()) + out = tmp_path / "out.txt" + r = run("clean_file.py", str(blob), "-o", str(out), "--json") + assert r.returncode == 2 + payload = json.loads(r.stdout) + assert "error" in payload + assert "refusing to treat" in payload["error"] + assert payload["exit_code"] == 2 + assert not out.exists() + + +def test_clean_file_binary_text_refusal_keeps_exit_2_human_mode(tmp_path): + blob = tmp_path / "binary.txt" + blob.write_bytes(make_docx(tmp_path / "x.docx").read_bytes()) + out = tmp_path / "out.txt" + r = run("clean_file.py", str(blob), "-o", str(out)) + assert r.returncode == 2 + assert "refusing to treat" in r.stderr + assert not out.exists() + + +def test_clean_file_json_batch_binary_text_error_is_structured(tmp_path): + clean = tmp_path / "ok.txt" + clean.write_text("plain text", encoding="utf-8") + blob = tmp_path / "binary.txt" + blob.write_bytes(make_docx(tmp_path / "x.docx").read_bytes()) + r = run("clean_file.py", str(clean), str(blob), "--json") + assert r.returncode == 2 + payload = json.loads(r.stdout) + assert payload["total"] == 1 + assert payload["results"][0]["exit_code"] == 2 + assert "refusing to treat" in payload["results"][0]["error"] + + def test_clean_file_auto_refuses_unknown_text_like_bytes(tmp_path): blob = tmp_path / "no_extension" blob.write_text("just plain text, no extension, no magic\n", encoding="utf-8") @@ -201,11 +239,12 @@ def test_inspect_file_json_reports_unknown_kind(tmp_path): blob = tmp_path / "no_extension" blob.write_text("no magic, no extension\n", encoding="utf-8") r = run("inspect_file.py", str(blob), "--json") - assert r.returncode == 0 + assert r.returncode == 3 # EXIT_PARTIAL: unrecognized input was not scanned import json payload = json.loads(r.stdout) assert payload["kind"] == "unknown" + assert payload["unscanned"] is True assert "note" in payload @@ -218,7 +257,7 @@ def test_router_advice_is_not_circular(tmp_path): assert "Use inspect_file.py / clean_file.py" not in r.stderr assert "--force-text" in r.stderr r = run("inspect_file.py", str(blob)) - assert r.returncode == 0 + assert r.returncode == 3 # EXIT_PARTIAL: unrecognized input was not scanned assert "Kind: unknown" in r.stdout assert "--as text|image|container" in r.stdout diff --git a/tests/test_image_meta_bomb_and_notes.py b/tests/test_image_meta_bomb_and_notes.py new file mode 100644 index 0000000..b55a0b2 --- /dev/null +++ b/tests/test_image_meta_bomb_and_notes.py @@ -0,0 +1,124 @@ +"""Decompression-bomb defense and unsupported-format notes for image metadata. + +- PNG zTXt/iTXt compressed text chunks are decompressed through a bounded + zlib decompressor capped at MAX_PNG_TEXT_BYTES; a chunk whose output would + exceed the cap fails closed (no text entry, no AI finding) instead of + allocating unbounded memory. +- inspect_image annotates ImageInspectReport.notes when the detected format + has no inspection support (the 'unknown' fallback). +""" + +from __future__ import annotations + +import struct +import sys +import zlib +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +SCRIPTS = ROOT / "skills" / "remove-ai-marks" / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +from image_meta import MAX_PNG_TEXT_BYTES, inspect_image, inspect_png + +LARGE_BUT_OK_BYTES = 1024 * 1024 # well under the cap, far larger than a normal chunk + + +def _png_chunk(ctype: bytes, payload: bytes) -> bytes: + crc = zlib.crc32(ctype) + crc = zlib.crc32(payload, crc) & 0xFFFFFFFF + return struct.pack(">I", len(payload)) + ctype + payload + struct.pack(">I", crc) + + +def _minimal_png_with_text_chunk(ctype: bytes, payload: bytes) -> bytes: + sig = b"\x89PNG\r\n\x1a\n" + ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0) + idat = zlib.compress(b"\x00\x00\x00") + return ( + sig + + _png_chunk(b"IHDR", ihdr) + + _png_chunk(ctype, payload) + + _png_chunk(b"IDAT", idat) + + _png_chunk(b"IEND", b"") + ) + + +def _ztext_payload(value: bytes) -> bytes: + # keyword \0 compression-method(0) compressed text + return b"Software\x00\x00" + zlib.compress(value) + + +def _itext_compressed_payload(value: bytes) -> bytes: + # keyword \0 comp-flag(1) comp-method(0) lang \0 tkey \0 compressed text + return b"Software\x00\x01\x00\x00\x00" + zlib.compress(value) + + +def _bomb_value() -> bytes: + """Decompressed text that would exceed the cap while still naming a generator.""" + return b"ChatGPT" + b"\x00" * (MAX_PNG_TEXT_BYTES + 1024 * 1024) + + +# --------------------------------------------------------------------------- +# Decompression-bomb defense (zTXt / iTXt) +# --------------------------------------------------------------------------- + + +def test_ztext_bomb_fails_closed(): + """A zTXt chunk expanding past the cap yields no AI finding (fail closed).""" + data = _minimal_png_with_text_chunk(b"zTXt", _ztext_payload(_bomb_value())) + has_c2pa, has_ai, findings = inspect_png(data) + assert has_c2pa is False + assert has_ai is False + assert not any("AI generator" in f for f in findings) + + +def test_itext_compressed_bomb_fails_closed(): + """A compressed iTXt chunk expanding past the cap also fails closed.""" + data = _minimal_png_with_text_chunk(b"iTXt", _itext_compressed_payload(_bomb_value())) + has_c2pa, has_ai, findings = inspect_png(data) + assert has_c2pa is False + assert has_ai is False + assert not any("AI generator" in f for f in findings) + + +def test_large_under_cap_ztext_still_detected(): + """Legitimate large compressed text under the cap still flags the generator.""" + value = b"ChatGPT" + b"x" * (LARGE_BUT_OK_BYTES - len(b"ChatGPT")) + data = _minimal_png_with_text_chunk(b"zTXt", _ztext_payload(value)) + has_c2pa, has_ai, findings = inspect_png(data) + assert has_c2pa is False + assert has_ai is True + assert any("AI generator" in f and "ChatGPT" in f for f in findings) + + +def test_compressed_itext_still_detected(): + """The compressed iTXt path (comp-flag=1) still flags the generator.""" + data = _minimal_png_with_text_chunk(b"iTXt", _itext_compressed_payload(b"ChatGPT")) + has_c2pa, has_ai, findings = inspect_png(data) + assert has_c2pa is False + assert has_ai is True + assert any("AI generator" in f and "ChatGPT" in f for f in findings) + + +# --------------------------------------------------------------------------- +# Unsupported-format notes on ImageInspectReport +# --------------------------------------------------------------------------- + + +def test_unknown_format_gets_not_inspected_note(tmp_path: Path): + src = tmp_path / "unknown.bin" + src.write_bytes(b"no magic bytes here") + report = inspect_image(src) + assert report.format == "unknown" + assert report.findings == ["unsupported format"] + assert report.has_c2pa is False + assert report.has_ai_metadata is False + assert "format 'unknown' is not inspected" in report.notes + + +def test_supported_format_has_no_not_inspected_note(tmp_path: Path): + src = tmp_path / "ok.png" + src.write_bytes(_minimal_png_with_text_chunk(b"tEXt", b"Software\x00ChatGPT")) + report = inspect_image(src) + assert report.format == "png" + assert not any("not inspected" in note for note in report.notes) diff --git a/tests/test_inspect_file_partial.py b/tests/test_inspect_file_partial.py new file mode 100644 index 0000000..88244ae --- /dev/null +++ b/tests/test_inspect_file_partial.py @@ -0,0 +1,99 @@ +"""inspect_file.py must mark unscanned inputs and exit EXIT_PARTIAL (3). + +Unrecognized and refused inputs were previously reported as clean (exit 0), +which let an incomplete audit pass CI as a clean signal. Unscanned results +carry ``"unscanned": true`` and the CLI exits 3, taking precedence over both +clean and suspicious findings. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +SCRIPTS = ROOT / "skills" / "remove-ai-marks" / "scripts" +INSPECT_FILE = SCRIPTS / "inspect_file.py" + +EXIT_PARTIAL = 3 + + +def _run(*args: str, env: dict | None = None) -> subprocess.CompletedProcess[str]: + merged = os.environ.copy() + if env: + merged.update(env) + return subprocess.run( + [sys.executable, str(INSPECT_FILE), *args], + capture_output=True, + text=True, + env=merged, + check=False, + ) + + +def test_unknown_input_is_unscanned_and_partial(tmp_path: Path) -> None: + blob = tmp_path / "no_extension" + blob.write_text("no magic, no extension\n", encoding="utf-8") + + r = _run(str(blob), "--json") + assert r.returncode == EXIT_PARTIAL + payload = json.loads(r.stdout) + assert payload["kind"] == "unknown" + assert payload["unscanned"] is True + assert "note" in payload # refusal message preserved + + +def test_oversized_input_is_unscanned_and_partial(tmp_path: Path) -> None: + blob = tmp_path / "big.txt" + blob.write_text("x" * 64, encoding="utf-8") + + r = _run(str(blob), "--json", env={"WATERMARKS_MAX_INPUT_BYTES": "16"}) + assert r.returncode == EXIT_PARTIAL + payload = json.loads(r.stdout) + assert payload["kind"] == "refused" + assert payload["unscanned"] is True + assert "larger than" in payload["note"] # refusal message preserved + + +def test_clean_text_input_stays_clean(tmp_path: Path) -> None: + src = tmp_path / "clean.txt" + src.write_text("ordinary prose, no hidden marks\n", encoding="utf-8") + + r = _run(str(src), "--json") + assert r.returncode == 0, r.stderr + payload = json.loads(r.stdout) + assert payload["kind"] == "text" + assert "unscanned" not in payload + + +def test_suspicious_text_input_keeps_exit_1(tmp_path: Path) -> None: + src = tmp_path / "suspicious.txt" + src.write_text("hidden\u200bmark\n", encoding="utf-8") + + r = _run(str(src), "--json") + assert r.returncode == 1 + payload = json.loads(r.stdout) + assert payload["kind"] == "text" + assert payload["suspicious"] is True + assert "unscanned" not in payload + + +def test_mixed_batch_partial_takes_precedence(tmp_path: Path) -> None: + clean = tmp_path / "clean.txt" + clean.write_text("ordinary prose\n", encoding="utf-8") + suspicious = tmp_path / "suspicious.txt" + suspicious.write_text("hidden\u200bmark\n", encoding="utf-8") + unknown = tmp_path / "no_extension" + unknown.write_text("no magic, no extension\n", encoding="utf-8") + + r = _run(str(clean), str(suspicious), str(unknown), "--json") + assert r.returncode == EXIT_PARTIAL + payload = json.loads(r.stdout) + assert payload["total"] == 3 + by_path = {Path(item["path"]).name: item for item in payload["results"]} + assert "unscanned" not in by_path["clean.txt"] + assert by_path["suspicious.txt"]["suspicious"] is True + assert by_path["no_extension"]["unscanned"] is True diff --git a/tests/test_lightweight_clean_text_aliases.py b/tests/test_lightweight_clean_text_aliases.py new file mode 100644 index 0000000..480ba32 --- /dev/null +++ b/tests/test_lightweight_clean_text_aliases.py @@ -0,0 +1,95 @@ +"""clean-user-facing-text clean_text.py must reject --output aliases of the input. + +Writing the cleaned text over the input itself (same path, hard link, or +symlink) destroys the source; the CLI must refuse before writing. --in-place +remains the sanctioned overwrite path (it keeps a .bak backup). +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +CLEAN_TEXT = ROOT / "skills" / "clean-user-facing-text" / "scripts" / "clean_text.py" + + +def _run(*args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(CLEAN_TEXT), *args], + capture_output=True, + text=True, + check=False, + ) + + +def test_rejects_same_path_output(tmp_path: Path) -> None: + src = tmp_path / "input.txt" + src.write_text("a\u200bb", encoding="utf-8") + + r = _run(str(src), "-o", str(src)) + + assert r.returncode == 2 + assert "aliases" in r.stderr + assert src.read_text(encoding="utf-8") == "a\u200bb" + + +def test_rejects_hardlink_output(tmp_path: Path) -> None: + src = tmp_path / "input.txt" + src.write_text("a\u200bb", encoding="utf-8") + out = tmp_path / "out.txt" + os.link(src, out) + + r = _run(str(src), "-o", str(out)) + + assert r.returncode == 2 + assert "aliases" in r.stderr + assert src.read_text(encoding="utf-8") == "a\u200bb" + assert out.read_text(encoding="utf-8") == "a\u200bb" + + +def test_rejects_symlink_output(tmp_path: Path) -> None: + src = tmp_path / "input.txt" + src.write_text("a\u200bb", encoding="utf-8") + out = tmp_path / "out.txt" + try: + out.symlink_to(src) + except OSError as error: + # Windows may require privileges for symlinks; safe_write_bytes still + # protects the input there, so the CLI-level check is best-effort. + if os.name == "nt": + import pytest + + pytest.skip(f"symlinks unavailable: {error}") + raise + + r = _run(str(src), "-o", str(out)) + + assert r.returncode == 2 + assert "aliases" in r.stderr + assert src.read_text(encoding="utf-8") == "a\u200bb" + + +def test_normal_output_still_works(tmp_path: Path) -> None: + src = tmp_path / "input.txt" + src.write_text("a\u200bb", encoding="utf-8") + out = tmp_path / "out.txt" + + r = _run(str(src), "-o", str(out)) + + assert r.returncode == 0, r.stderr + assert out.read_text(encoding="utf-8") == "ab" + assert src.read_text(encoding="utf-8") == "a\u200bb" + + +def test_in_place_still_works_with_backup(tmp_path: Path) -> None: + src = tmp_path / "input.txt" + src.write_text("a\u200bb", encoding="utf-8") + + r = _run(str(src), "--in-place") + + assert r.returncode == 0, r.stderr + assert src.read_text(encoding="utf-8") == "ab" + assert src.with_suffix(".txt.bak").is_file() diff --git a/tests/test_markllm_detect.py b/tests/test_markllm_detect.py index 2a2d04c..6a3136b 100644 --- a/tests/test_markllm_detect.py +++ b/tests/test_markllm_detect.py @@ -8,6 +8,8 @@ import sys from pathlib import Path +import pytest + ROOT = Path(__file__).resolve().parent.parent SCRIPTS = ROOT / "skills" / "remove-ai-marks" / "scripts" sys.path.insert(0, str(SCRIPTS)) @@ -361,3 +363,62 @@ def test_cli_watermark_runtime_error(tmp_path: Path): ) assert r.returncode == 1 assert "boom" in (r.stderr or "") + + +def test_cli_watermark_json_stdout_stays_pure_without_output(tmp_path: Path): + """--json must never mix generated text with the JSON payload on stdout. + + With no -o, the watermarked sample previously went to stdout and made + json.loads() fail; it must be routed to stderr instead. + """ + upstream = _make_fake_upstream(tmp_path) + prompt = tmp_path / "prompt.txt" + prompt.write_text("write about capybaras") + r = _run_adapter( + "watermark", + str(prompt), + "--scheme", + "kgw", + "--upstream-dir", + str(upstream), + "--device", + "cpu", + "--json", + ) + assert r.returncode == 0, r.stderr + payload = json.loads(r.stdout) + assert payload["available"] is True + assert payload["watermarked_output"] == "-" + assert "WATERMARKED SAMPLE" in (r.stderr or "") + assert "WATERMARKED SAMPLE" not in (r.stdout or "") + + +def test_cli_detect_applies_rlimit_as_in_child(tmp_path: Path): + """--rlimit-as must cap the child's address space before heavy imports.""" + if os.name == "nt": + pytest.skip("RLIMIT_AS is POSIX-only") + upstream = _make_fake_upstream(tmp_path) + (upstream / "transformers" / "__init__.py").write_text( + FAKE_TRANSFORMERS + "import resource, sys\n" + "print('RLIMIT_AS=' + str(resource.getrlimit(resource.RLIMIT_AS)[0]), file=sys.stderr)\n" + ) + f = tmp_path / "t.txt" + f.write_text("hello world") + # 512 GiB: high enough to clear the macOS arm64 dyld VM reservation, low + # enough to prove the cap is applied (the default soft limit is unlimited). + limit = 512 * 2**30 + r = _run_adapter( + "detect", + str(f), + "--scheme", + "kgw", + "--upstream-dir", + str(upstream), + "--device", + "cpu", + "--json", + "--rlimit-as", + str(limit), + ) + assert r.returncode == 0, r.stderr + assert f"RLIMIT_AS={limit}" in (r.stderr or "") diff --git a/tests/test_rewrite_text.py b/tests/test_rewrite_text.py index b4641a7..3eaba09 100644 --- a/tests/test_rewrite_text.py +++ b/tests/test_rewrite_text.py @@ -525,6 +525,32 @@ def log_message(self, format, *args): server.shutdown() +def test_openai_compatible_omits_reasoning_effort_when_off(monkeypatch): + """'off' disables the parameter entirely; enabled values are still sent.""" + captured: dict = {} + + def fake_request_json(base_url, route, payload, **kwargs): + captured["payload"] = payload + return {"choices": [{"message": {"content": "rewritten"}}]} + + monkeypatch.setattr(rewrite_text.layer_b_http, "request_json", fake_request_json) + + rewrite_text._call_openai_compatible( + "http://127.0.0.1:9", "m", "hello", "key", 5.0, reasoning_effort="off" + ) + assert "reasoning_effort" not in captured["payload"] + + rewrite_text._call_openai_compatible( + "http://127.0.0.1:9", "m", "hello", "key", 5.0, reasoning_effort="low" + ) + assert captured["payload"]["reasoning_effort"] == "low" + + rewrite_text._call_openai_compatible( + "http://127.0.0.1:9", "m", "hello", "key", 5.0, reasoning_effort=None + ) + assert "reasoning_effort" not in captured["payload"] + + def test_rewrite_denies_remote_host_without_opt_in(): with pytest.raises(LayerBHTTPError, match="endpoint"): rewrite("secret text", _http_plan("http://example.com:11434")) diff --git a/tests/test_stability_server.py b/tests/test_stability_server.py new file mode 100644 index 0000000..5b5e728 --- /dev/null +++ b/tests/test_stability_server.py @@ -0,0 +1,291 @@ +"""Focused regression tests for the release-hardening stability slice. + +Covers: + - synthid_score_server.py: exceptions raised by score_file are contained + at the request boundary and answered with the fail-soft HTTP 200 JSON + contract ({"available": false, "error": ...}) instead of killing the + request with an unstructured traceback. Success and auth paths kept. + - score_synthid.py: the external checkout is inserted into sys.path at + most once, and the shared process-global work (codebook load / + extraction / stdout redirection) is serialized under a module-level + threading.Lock so concurrent HTTP requests cannot corrupt it. + - server.py: /detect answers unrecognized assets with the documented + HTTP 200 note (kind "unknown") instead of falling through to the + container branch / a 500, the OpenAPI spec documents the unknown + shape, and /clean keeps refusing unknown formats. + - compose-check.sh: the wr-core health check curl carries bounded + connect/max timeouts; endpoint and status assertions are preserved. +""" + +from __future__ import annotations + +import base64 +import contextlib +import http.client +import json +import re +import struct +import sys +import threading +import zlib +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +SCRIPTS = ROOT / "skills" / "remove-ai-marks" / "scripts" +sys.path.insert(0, str(SCRIPTS)) + +# Ensure the real cv2 (when installed) wins over the stub upstream's cv2.py: +# score_file inserts the extraction dir at sys.path[0] before importing cv2, +# so without this the stub would shadow the real module for the whole test +# session. When cv2 is absent the stub satisfies the import instead. +with contextlib.suppress(ImportError): + import cv2 # noqa: F401 + +import score_synthid +import server +import synthid_score_server + + +def _b64(data: bytes) -> str: + return base64.b64encode(data).decode("ascii") + + +def _png_chunk(ctype: bytes, payload: bytes) -> bytes: + crc = zlib.crc32(ctype) + crc = zlib.crc32(payload, crc) & 0xFFFFFFFF + return struct.pack(">I", len(payload)) + ctype + payload + struct.pack(">I", crc) + + +def _tiny_png() -> bytes: + """A real 1x1 RGB PNG, decodable by either the stub or real cv2.""" + sig = b"\x89PNG\r\n\x1a\n" + ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0) + idat = zlib.compress(b"\x00\x00\x00\x00") # filter byte + 1 RGB pixel + return sig + _png_chunk(b"IHDR", ihdr) + _png_chunk(b"IDAT", idat) + _png_chunk(b"IEND", b"") + + +def _post(conn, path, payload, headers=None): + conn.request( + "POST", + path, + body=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json", **(headers or {})}, + ) + resp = conn.getresponse() + data = resp.read() + return resp.status, json.loads(data) if data else {} + + +def _get(conn, path): + conn.request("GET", path) + resp = conn.getresponse() + data = resp.read() + return resp.status, json.loads(data) if data else {} + + +@pytest.fixture(scope="module") +def conn(): + srv = server.ThreadingHTTPServer(("127.0.0.1", 0), server.Handler) + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + c = http.client.HTTPConnection("127.0.0.1", srv.server_address[1]) + yield c + c.close() + srv.shutdown() + srv.server_close() + thread.join(timeout=5) + + +@pytest.fixture(scope="module") +def sidecar(): + srv = synthid_score_server.ThreadingHTTPServer(("127.0.0.1", 0), synthid_score_server.Handler) + thread = threading.Thread(target=srv.serve_forever, daemon=True) + thread.start() + c = http.client.HTTPConnection("127.0.0.1", srv.server_address[1]) + yield c + c.close() + srv.shutdown() + srv.server_close() + thread.join(timeout=5) + + +# --- score_synthid.py: one sys.path insertion + serialized shared state --- + + +def _write_stub_upstream(root: Path) -> Path: + """Stub reverse-SynthID checkout with a noisy, importable extractor. + + Mirrors the stub in test_synthid_stdout_purity.py: the upstream prints + progress straight to stdout (which redirect_stdout must divert) and + carries its own cv2 module so the import works without OpenCV. + """ + ext = root / "src" / "extraction" + ext.mkdir(parents=True) + (root / "artifacts").mkdir() + (root / "artifacts" / "spectral_codebook_v4.npz").touch() + + (ext / "cv2.py").write_text( + "COLOR_BGR2RGB = 4\n" + "def imread(path):\n" + " return object()\n" + "def cvtColor(img, code):\n" + " return img\n" + ) + (ext / "synthid_bypass_v4.py").write_text( + "class SpectralCodebookV4:\n" + " def load(self, path):\n" + ' print(f"CodebookV4 loaded: {path}")\n' + ) + (ext / "robust_extractor.py").write_text( + "from types import SimpleNamespace\n" + "class RobustSynthIDExtractor:\n" + " def detect_from_v4_codebook(self, rgb, codebook, model=None):\n" + " return SimpleNamespace(\n" + " details={'profile_key': 'stub', 'exact_match': False,\n" + " 'per_channel_scores': [0.1], 'per_channel_n': [1]},\n" + " is_watermarked=False, confidence=0.42,\n" + " phase_match=0.0, multi_scale_consistency=0.0,\n" + " )\n" + ) + return root + + +def test_score_file_inserts_upstream_path_once(tmp_path): + upstream = _write_stub_upstream(tmp_path / "upstream") + img = tmp_path / "img.png" + img.write_bytes(_tiny_png()) + extraction = str(upstream / "src" / "extraction") + assert extraction not in sys.path + + first = score_synthid.score_file(img, upstream_dir=str(upstream)) + second = score_synthid.score_file(img, upstream_dir=str(upstream)) + + assert first[0] == 0, first + assert second[0] == 0, second + assert sys.path.count(extraction) == 1 + + +def test_score_file_shared_state_is_guarded_by_module_lock(): + lock = score_synthid._SCORE_LOCK + assert callable(getattr(lock, "acquire", None)) + assert callable(getattr(lock, "release", None)) + + +def test_concurrent_score_file_serializes_shared_state(tmp_path): + upstream = _write_stub_upstream(tmp_path / "upstream") + img = tmp_path / "img.png" + img.write_bytes(_tiny_png()) + extraction = str(upstream / "src" / "extraction") + results: list = [] + errors: list = [] + + def worker() -> None: + try: + results.append(score_synthid.score_file(img, upstream_dir=str(upstream))) + except Exception as exc: # pragma: no cover - failure path + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + + assert errors == [] + assert len(results) == 8 + assert all(code == 0 and payload["confidence"] == 0.42 for code, payload in results) + assert sys.path.count(extraction) == 1 + + +# --- synthid sidecar: score_file exceptions become fail-soft 200s --- + + +def test_sidecar_contains_score_file_exception_as_fail_soft_200(sidecar, monkeypatch): + def boom(path, *, model=None): + raise RuntimeError("codebook exploded") + + monkeypatch.setattr(synthid_score_server, "score_file", boom) + status, body = _post(sidecar, "/score", {"file": _b64(_tiny_png())}) + + assert status == 200 + assert body["available"] is False + assert isinstance(body["error"], str) and body["error"] + assert "Traceback" not in body["error"] + + +def test_sidecar_score_success_payload_passthrough(sidecar, monkeypatch): + payload = {"available": True, "is_watermarked": False, "confidence": 0.1} + monkeypatch.setattr( + synthid_score_server, "score_file", lambda path, *, model=None: (0, payload) + ) + status, body = _post(sidecar, "/score", {"file": _b64(_tiny_png())}) + + assert status == 200 + assert body == payload + + +def test_sidecar_auth_still_enforced(sidecar, monkeypatch): + monkeypatch.setattr(synthid_score_server, "API_KEY", "sekret") + monkeypatch.setattr( + synthid_score_server, "score_file", lambda path, *, model=None: (0, {"available": True}) + ) + status, _ = _post(sidecar, "/score", {"file": _b64(_tiny_png())}) + assert status == 401 + status, body = _post( + sidecar, + "/score", + {"file": _b64(_tiny_png())}, + headers={"Authorization": "Bearer sekret"}, + ) + assert status == 200 + assert body["available"] is True + + +# --- server.py: /detect unknown note, OpenAPI shape, /clean refusal --- + + +def test_detect_unknown_format_returns_200_note(conn): + data = b"no magic, no extension" + status, body = _post(conn, "/detect", {"file": _b64(data), "name": "input"}) + + assert status == 200 + assert body["ok"] is True + assert body["kind"] == "unknown" + assert body["detections"] == [] + assert "note" in body["report"] + + +def test_detect_openapi_documents_unknown_kind(conn): + status, body = _get(conn, "/openapi.json") + assert status == 200 + detect = body["paths"]["/detect"]["post"] + schema = detect["responses"]["200"]["content"]["application/json"]["schema"] + properties = schema["properties"] + assert "unknown" in properties["kind"]["enum"] + assert "report" in properties + + +def test_clean_unknown_format_still_refused(conn): + data = b"no magic, no extension" + status, body = _post(conn, "/clean", {"file": _b64(data), "name": "input"}) + assert status == 400 + assert "unrecognized file format" in body["error"] + + +# --- compose-check.sh: bounded health-check curl --- + + +def test_compose_check_health_curl_has_bounded_timeouts(): + script = (ROOT / "compose-check.sh").read_text(encoding="utf-8") + match = re.search(r"curl\s+([^\n]*)\$BASE_URL/health", script) + assert match, "health check curl invocation not found" + flags = match.group(1) + assert "-fsS" in flags + assert "--connect-timeout 2" in flags + assert "--max-time 10" in flags + assert 'echo "wr-core: OK"' in script + assert 'echo "wr-core: FAIL' in script + assert "for svc in wr-markllm wr-markdiffusion wr-ctrlregen wr-synthid; do" in script + assert 'exit "$FAIL"' in script diff --git a/tests/test_stylometry.py b/tests/test_stylometry.py index 9d4c38f..8ec9d47 100644 --- a/tests/test_stylometry.py +++ b/tests/test_stylometry.py @@ -5,6 +5,7 @@ import json import subprocess import sys +import time from pathlib import Path SCRIPTS_DIR = Path(__file__).resolve().parent.parent / "skills" / "remove-ai-marks" / "scripts" @@ -118,6 +119,33 @@ def test_ai_vs_human_discrimination(): assert human_report.confidence_level in ("CLEAN", "LOW") +def test_oversized_input_truncated_before_scanning(): + # Regression: oversized inputs (>MAX_SCAN_CHARS) must be truncated before the + # regex scans run, with a truncation note, and complete quickly. + snippet = "This is a sample sentence with enough words for stylometry analysis. " + oversized = snippet * 40000 # ~2.6M chars > MAX_SCAN_CHARS (2,000,000) + assert len(oversized) > 2_000_000 + + start = time.monotonic() + report = score_text_stylometry(oversized, path="oversized.txt") + elapsed = time.monotonic() - start + + assert any("truncat" in n.lower() for n in report.notes) + assert report.status == "ok" + assert 0 < report.word_count < 2_000_000 + assert elapsed < 30 + + +def test_under_cap_input_has_no_truncation_note(): + # Inputs within the cap must behave byte-identically: no truncation note, + # and the established score is preserved. + sample = (FIXTURES_DIR / "stylometry_ai_sample.txt").read_text(encoding="utf-8") + assert len(sample) <= 2_000_000 + report = score_text_stylometry(sample, path="ai_sample.txt") + assert not any("truncat" in n.lower() for n in report.notes) + assert report.score >= 0.70 + + def test_score_stylometry_cli(): ai_path = FIXTURES_DIR / "stylometry_ai_sample.txt" human_path = FIXTURES_DIR / "stylometry_human_sample.txt" diff --git a/tests/test_text_detectors.py b/tests/test_text_detectors.py index 3801d6a..da88516 100644 --- a/tests/test_text_detectors.py +++ b/tests/test_text_detectors.py @@ -35,12 +35,20 @@ def _gemini_success(verdict: str | None = None, score: float | None = None) -> d class _FakeCommandResult: - """Mimic external_command.run_command() return shape.""" + """Mimic external_command.CommandResult: raw bytes plus decoded text views.""" - def __init__(self, returncode: int, stdout: str = "", stderr: str = ""): + def __init__(self, returncode: int, stdout: str | bytes = "", stderr: str | bytes = ""): self.returncode = returncode - self.stdout = stdout - self.stderr = stderr + self.stdout = stdout.encode("utf-8") if isinstance(stdout, str) else stdout + self.stderr = stderr.encode("utf-8") if isinstance(stderr, str) else stderr + + @property + def stdout_text(self) -> str: + return self.stdout.decode("utf-8", errors="replace") + + @property + def stderr_text(self) -> str: + return self.stderr.decode("utf-8", errors="replace") # --------------------------------------------------------------------------- @@ -82,6 +90,20 @@ def test_gemini_disabled_with_api_key(monkeypatch): assert report["vendor"] == "google" +def test_gemini_available_is_false_until_endpoint_is_supported(monkeypatch): + """available() must never advertise a detect() that cannot run. + + The DETECT_TEXT_WATERMARK task type is not supported by the Gemini + generateContent API yet, so available() stays False even when a key is + configured — /capabilities and run_text_detectors() must not surface it. + """ + monkeypatch.setenv("WATERMARKS_GEMINI_API_KEY", "k") + assert text_detectors.GeminiSynthIDTextDetector().available() is False + assert text_detectors.detector_status()["gemini-synthid-text"] is False + # A disabled detector must never be selected by the usable-only runner. + assert text_detectors.run_text_detectors("some text") == [] + + def test_gemini_http_error(monkeypatch): """Disabled detector never reaches request_json, so no HTTP call is made.""" monkeypatch.setenv("WATERMARKS_GEMINI_API_KEY", "k") @@ -206,6 +228,61 @@ def fake_run(cmd, **kwargs): assert argv[argv.index("--scheme") + 1] == "synthid" +def test_markllm_reports_are_json_safe_with_bytes_output(monkeypatch, tmp_path): + """CommandResult bytes must be decoded before touching the report. + + json.dumps() on a report whose error field is raw bytes raises TypeError; + the adapter must use CommandResult.stdout_text/stderr_text. + """ + upstream = tmp_path / "MarkLLM" + upstream.mkdir() + monkeypatch.setattr( + text_detectors, + "run_command", + lambda *a, **k: _FakeCommandResult(3, stderr=b"missing deps"), + ) + det = text_detectors.MarkLLMTextDetector(upstream_dir=str(upstream)) + report = det.detect("hello") + assert report["available"] is False + assert report["error"] == "missing deps" + assert isinstance(report["error"], str) + json.dumps(report) # must not raise on a bytes stderr + + +def test_markllm_passes_rlimit_as_to_child_at_subprocess_boundary(monkeypatch, tmp_path): + """WATERMARKS_MARKLLM_RLIMIT_AS must reach the child via --rlimit-as.""" + upstream = tmp_path / "MarkLLM" + upstream.mkdir() + monkeypatch.setenv("WATERMARKS_MARKLLM_RLIMIT_AS", "0x10000000") + seen: dict = {} + + def fake_run(cmd, **kwargs): + seen["argv"] = cmd + return _FakeCommandResult(0, stdout="{}") + + monkeypatch.setattr(text_detectors, "run_command", fake_run) + det = text_detectors.MarkLLMTextDetector(upstream_dir=str(upstream)) + det.detect("hello") + argv = seen["argv"] + assert "--rlimit-as" in argv + assert argv[argv.index("--rlimit-as") + 1] == "268435456" + + +def test_markllm_omits_rlimit_as_when_unset(monkeypatch, tmp_path): + upstream = tmp_path / "MarkLLM" + upstream.mkdir() + seen: dict = {} + + def fake_run(cmd, **kwargs): + seen["argv"] = cmd + return _FakeCommandResult(0, stdout="{}") + + monkeypatch.setattr(text_detectors, "run_command", fake_run) + det = text_detectors.MarkLLMTextDetector(upstream_dir=str(upstream)) + det.detect("hello") + assert "--rlimit-as" not in seen["argv"] + + def test_markllm_prefers_checkout_venv(monkeypatch, tmp_path): upstream = tmp_path / "MarkLLM" if os.name == "nt": From bc33e0dce456501c88cace03601468aa0c420382 Mon Sep 17 00:00:00 2001 From: elkaix Date: Tue, 18 Aug 2026 18:48:48 -0400 Subject: [PATCH 3/3] test(detectors): skip rlimit argv assertion on Windows subprocess_preexec_fn is None off POSIX (common.py), so the MarkLLM adapter correctly omits --rlimit-as on Windows; the boundary assertion only makes sense where resource limits exist. --- tests/test_text_detectors.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_text_detectors.py b/tests/test_text_detectors.py index da88516..f80a91b 100644 --- a/tests/test_text_detectors.py +++ b/tests/test_text_detectors.py @@ -251,6 +251,8 @@ def test_markllm_reports_are_json_safe_with_bytes_output(monkeypatch, tmp_path): def test_markllm_passes_rlimit_as_to_child_at_subprocess_boundary(monkeypatch, tmp_path): """WATERMARKS_MARKLLM_RLIMIT_AS must reach the child via --rlimit-as.""" + if os.name == "nt": + pytest.skip("rlimit is POSIX-only; the adapter omits --rlimit-as on Windows") upstream = tmp_path / "MarkLLM" upstream.mkdir() monkeypatch.setenv("WATERMARKS_MARKLLM_RLIMIT_AS", "0x10000000")