fix(utils): prevent path traversal in encode_images image filenames

[andesyteoss]

Summary

omniparse.utils.encode_images writes each parsed image to disk using the dict key from the parser as the full file path, then deletes it:

for i, (filename, image) in enumerate(images.items()):
    image.save(filename, "PNG")
    ...
    os.remove(filename)

filename is not a sanitized basename — it is whatever key the underlying document parser (marker) emits for an embedded image. If any code path allows a document to influence that key (image resource names, XObject names in PDFs, embedded media names in PPTX/DOCX), a value like ../../etc/omniparse_pwn or an absolute path becomes:

1. An arbitrary-location PNG write with the server process's privileges (image.save(filename, "PNG")).
2. An arbitrary-location file delete (os.remove(filename)) — the same path is unlinked at the end.

encode_images is reached from unauthenticated FastAPI routes (omniparse/documents/router.py: /pdf, /ppt, /docs, and the image router), so the trigger is an unauthenticated document upload.

• CWE-22: Improper Limitation of a Pathname to a Restricted Directory (Path Traversal)
• Affected function: omniparse.utils.encode_images (omniparse/utils.py)
• Callers: omniparse/documents/__init__.py lines 57/114/169, omniparse/image/__init__.py line 96
• Routes: document_router.post("/pdf" | "/ppt" | "/docs") in omniparse/documents/router.py; image router
• Data flow: HTTP upload → marker/parser → images dict → filename key → image.save(filename) / os.remove(filename)

Fix

Two changes in encode_images, both minimal and behavior-preserving for legitimate inputs:

1. Strip to basename — safe_name = os.path.basename(filename) or f"image_{i}.png". Any ../ segments or absolute-path prefixes are discarded before the value is ever handed to the filesystem.
2. Write to a sandboxed temp directory — tempfile.TemporaryDirectory() gives a fresh, private directory per call. The image is saved inside it, read back, and the directory (with all contents) is removed on context exit. This eliminates the separate os.remove(filename) call, which was itself a second traversal sink, and also removes a pre-existing bug where the file was written into the process CWD.

The image_name recorded on the response document is now the sanitized basename, which is what downstream consumers already treat as a display name.

Tests

• python -c "from omniparse.utils import encode_images" — imports cleanly.
• Unit smoke: constructed a fake images dict with keys ../../evil.png, /tmp/absolute.png, and a normal img_0.png, plus a stub PIL-like object with a .save(path, fmt) method. Before the patch, save was invoked with the raw traversal path; after the patch, save is only ever invoked with a path inside the TemporaryDirectory, and the recorded image_name is evil.png / absolute.png / img_0.png respectively. No file is left on disk after the call.
• git diff main..HEAD --stat → 1 file changed, 17 insertions, 12 deletions. No API surface change.

Adversarial review

Before submitting we tried to disprove this. The main "is it really exploitable?" question is whether the marker parser actually propagates document-controlled bytes into the image dict keys. We could not fully confirm that against marker's source in this environment, so the exploit precondition (crafted document yields a traversal-shaped key) is not proven end-to-end here. However: (a) the sink is unambiguously unsafe — image.save(filename) and os.remove(filename) on an unvalidated string is a path-traversal primitive regardless of the current parser's behavior; (b) the routes are unauthenticated (no Depends(...) auth guard on document_router or the image router), so any exploitability is remote and pre-auth; (c) the fix is defense-in-depth that costs nothing — os.path.basename + tempfile.TemporaryDirectory is standard practice and does not change legitimate output. Even in the worst case where today's marker never emits traversal keys, this patch prevents a future parser change or a new caller from silently reintroducing the sink.

Proof of concept

Reproduce the sink behavior directly against the pre-patch function (no marker needed — this demonstrates that the function trusts its input):

# repro.py — run against the code at main (pre-patch)
import os, sys
sys.path.insert(0, ".")
from PIL import Image
from omniparse.utils import encode_images
from omniparse.models import responseDocument

target = "/tmp/omniparse_pwn.png"       # attacker-chosen path
if os.path.exists(target): os.remove(target)

img = Image.new("RGB", (2, 2), "red")
images = {"../../../../../../tmp/omniparse_pwn.png": img}   # traversal key
doc = responseDocument(text="")
encode_images(images, doc)

print("file written to attacker path:", os.path.exists(target))
# Pre-patch: file was created at the attacker path before being unlinked,
# proving the arbitrary-write + arbitrary-delete primitives.
# Post-patch: the save target is confined to a private TemporaryDirectory.

End-to-end trigger shape against a running server (unauthenticated):

curl -X POST -F "file=@crafted.pdf" http://localhost:8000/parse_document/pdf

where crafted.pdf is a document whose parser output contains an image resource name shaped like ../../tmp/x.png. The server-side effect is image.save("../../tmp/x.png", "PNG") followed by os.remove("../../tmp/x.png"), both executed with the omniparse process's privileges.

Impact

• Arbitrary file write of a valid PNG to any path the server user can write to (config directories, cron drop-in directories, web roots, etc.).
• Arbitrary file delete of any path the server user can unlink.
• Unauthenticated remote trigger via the document/image parsing endpoints.

Happy to iterate on the patch shape if you'd prefer a different sanitization strategy (e.g. secure_filename from werkzeug, or enforcing a UUID name). Thanks for maintaining omniparse.

---

Summary by CodeRabbit

• Bug Fixes
  • Improved image processing security by preventing uploaded or parsed images from being written to attacker-controlled file paths.
  • Image files are now handled in temporary storage and cleaned up automatically.
  • Preserved safe image names when attaching encoded images to documents.

[coderabbitai]

📝 Walkthrough

Walkthrough

encode_images now sanitizes image names and stores generated PNGs in a temporary directory before base64 encoding, removing direct writes and cleanup against caller-provided paths.

Changes

Image storage hardening

Layer / File(s)	Summary
Safe image encoding workflow omniparse/utils.py	encode_images derives a basename-safe image name, writes PNG data in a TemporaryDirectory, reads and base64-encodes the temporary file, and attaches it to inputDocument under the sanitized name.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: preventing path traversal in encode_images image filenames.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

omniparse/utils.py (1)

16-26: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Pass the image directly here. responseDocument.add_image already accepts PIL.Image, so the temp-file/base64 round-trip is redundant; removing it also lets you drop import tempfile.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@omniparse/utils.py` around lines 16 - 26, Update the image-handling flow
around inputDocument.add_image to pass the PIL.Image object directly instead of
saving it to a temporary PNG, reading bytes, and base64-encoding them. Remove
the now-unused tempfile import and any related temporary-file conversion logic,
while preserving the existing image_name value.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@omniparse/utils.py`:
- Around line 16-26: Update the image-handling flow around
inputDocument.add_image to pass the PIL.Image object directly instead of saving
it to a temporary PNG, reading bytes, and base64-encoding them. Remove the
now-unused tempfile import and any related temporary-file conversion logic,
while preserving the existing image_name value.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 615b8a71-656e-4ec3-8634-9bab4e32027c

📥 Commits

Reviewing files that changed from the base of the PR and between 9d1ae83 and 213b67f.

📒 Files selected for processing (1)

• omniparse/utils.py