From 95279609d2c2662ad1d86e438613ff382146a0ee Mon Sep 17 00:00:00 2001 From: rookiestar28 Date: Wed, 8 Apr 2026 00:46:47 +0800 Subject: [PATCH] fix(security): harden path boundary handling --- services/checkpoints.py | 26 ++++++++++++++--- services/integrity.py | 45 ++++++++++++++++++++++------- services/model_manager_transfer.py | 15 +++++++--- tests/test_checkpoints.py | 6 ++++ tests/test_model_manager_service.py | 40 +++++++++++++++++++++++++ tests/test_r77_integrity.py | 4 +++ 6 files changed, 117 insertions(+), 19 deletions(-) diff --git a/services/checkpoints.py b/services/checkpoints.py index 1da1ecd..c39482f 100644 --- a/services/checkpoints.py +++ b/services/checkpoints.py @@ -19,6 +19,7 @@ except ImportError: DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data") from .integrity import IntegrityError, load_verified, save_verified +from .safe_io import PathTraversalError, resolve_under_root logger = logging.getLogger("ComfyUI-OpenClaw.services.checkpoints") @@ -31,10 +32,21 @@ def _ensure_dir(): os.makedirs(CHECKPOINTS_DIR, exist_ok=True) +def _normalize_checkpoint_id(checkpoint_id: str) -> str: + text = str(checkpoint_id or "").strip() + try: + return str(uuid.UUID(text)) + except (ValueError, AttributeError, TypeError) as exc: + raise ValueError("invalid checkpoint id") from exc + + def _get_paths(checkpoint_id: str) -> Tuple[str, str]: """Return (meta_path, payload_path) for a given ID.""" - base = os.path.join(CHECKPOINTS_DIR, checkpoint_id) - return f"{base}.meta.json", f"{base}.workflow.json" + cid = _normalize_checkpoint_id(checkpoint_id) + return ( + resolve_under_root(CHECKPOINTS_DIR, f"{cid}.meta.json"), + resolve_under_root(CHECKPOINTS_DIR, f"{cid}.workflow.json"), + ) def list_checkpoints() -> List[Dict[str, Any]]: @@ -67,7 +79,10 @@ def list_checkpoints() -> List[Dict[str, Any]]: def get_checkpoint(checkpoint_id: str) -> Optional[Dict[str, Any]]: """Get full checkpoint data (meta + workflow).""" - meta_path, payload_path = _get_paths(checkpoint_id) + try: + meta_path, payload_path = _get_paths(checkpoint_id) + except (ValueError, PathTraversalError): + return None if not os.path.exists(meta_path) or not os.path.exists(payload_path): return None @@ -172,7 +187,10 @@ def create_checkpoint( def delete_checkpoint(checkpoint_id: str) -> bool: """Delete a checkpoint.""" - meta_path, payload_path = _get_paths(checkpoint_id) + try: + meta_path, payload_path = _get_paths(checkpoint_id) + except (ValueError, PathTraversalError): + return False deleted = False if os.path.exists(meta_path): diff --git a/services/integrity.py b/services/integrity.py index 3f03688..b88cf59 100644 --- a/services/integrity.py +++ b/services/integrity.py @@ -13,6 +13,8 @@ import tempfile from dataclasses import asdict, dataclass from typing import Any, Dict, Optional, Union +from .safe_io import PathTraversalError, resolve_under_root + logger = logging.getLogger("ComfyUI-OpenClaw.services.integrity") @@ -35,6 +37,23 @@ class IntegrityError(Exception): pass +def _normalize_verified_path(path: Union[str, os.PathLike[str]]) -> str: + raw_path = os.fspath(path) + if not raw_path: + raise IntegrityError("verified file path is empty") + + abs_path = os.path.abspath(raw_path) + dir_name = os.path.dirname(abs_path) + file_name = os.path.basename(abs_path) + if not dir_name or not file_name or file_name in {".", ".."}: + raise IntegrityError(f"Invalid verified file path: {raw_path}") + + try: + return resolve_under_root(dir_name, file_name, follow_symlinks=False) + except PathTraversalError as exc: + raise IntegrityError(f"Invalid verified file path: {raw_path}") from exc + + def canonical_dumps(data: Any) -> bytes: """ Serialize data to canonical JSON (sorted keys, no whitespace). @@ -58,6 +77,7 @@ def save_verified(path: str, data: Dict[str, Any], version: int = 1) -> None: Save data wrapped in an integrity envelope. Atomic write. """ + safe_path = _normalize_verified_path(path) data_hash = calculate_hash(data) envelope = IntegrityEnvelope( version=version, data=data, hash=data_hash, algo="sha256" @@ -67,11 +87,11 @@ def save_verified(path: str, data: Dict[str, Any], version: int = 1) -> None: try: content = json.dumps(asdict(envelope), indent=2) except Exception as e: - logger.error(f"Failed to serialize integrity envelope for {path}: {e}") + logger.error(f"Failed to serialize integrity envelope for {safe_path}: {e}") raise # Atomic write - dir_name = os.path.dirname(os.path.abspath(path)) + dir_name = os.path.dirname(safe_path) os.makedirs(dir_name, exist_ok=True) fd, tmp_path = tempfile.mkstemp(dir=dir_name, text=True) @@ -82,9 +102,9 @@ def save_verified(path: str, data: Dict[str, Any], version: int = 1) -> None: os.fsync(fd) # Renaissance-style atomic rename - os.replace(tmp_path, path) + os.replace(tmp_path, safe_path) except Exception as e: - logger.error(f"Failed to save verified file {path}: {e}") + logger.error(f"Failed to save verified file {safe_path}: {e}") if os.path.exists(tmp_path): os.remove(tmp_path) raise @@ -101,14 +121,15 @@ def load_verified( Raises IntegrityError if hash mismatch or malformed. """ - if not os.path.exists(path): - raise FileNotFoundError(f"File not found: {path}") + safe_path = _normalize_verified_path(path) + if not os.path.exists(safe_path): + raise FileNotFoundError(f"File not found: {safe_path}") try: - with open(path, "r", encoding="utf-8") as f: + with open(safe_path, "r", encoding="utf-8") as f: raw = json.load(f) except json.JSONDecodeError as e: - raise IntegrityError(f"Corrupt JSON file {path}: {e}") + raise IntegrityError(f"Corrupt JSON file {safe_path}: {e}") # Check if it's an envelope if isinstance(raw, dict) and "hash" in raw and "data" in raw and "version" in raw: @@ -118,7 +139,9 @@ def load_verified( computed_hash = calculate_hash(stored_data) if computed_hash != stored_hash: - raise IntegrityError(f"Integrity check failed for {path} (hash mismatch)") + raise IntegrityError( + f"Integrity check failed for {safe_path} (hash mismatch)" + ) # Verify version if needed # We can implement version migration logic here if multiple envelope versions exist @@ -128,9 +151,9 @@ def load_verified( # Legacy Fallback if migrate: logger.info( - f"R77: Loaded legacy file {path}, integrity check skipped (pending migration)." + f"R77: Loaded legacy file {safe_path}, integrity check skipped (pending migration)." ) # For legacy files, we assume the whole content is the data. return raw - raise IntegrityError(f"File {path} is not a valid integrity envelope") + raise IntegrityError(f"File {safe_path} is not a valid integrity envelope") diff --git a/services/model_manager_transfer.py b/services/model_manager_transfer.py index 6253e33..eb47a0f 100644 --- a/services/model_manager_transfer.py +++ b/services/model_manager_transfer.py @@ -11,7 +11,7 @@ import tempfile import time import urllib.request import uuid -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import Any, Dict, List, Optional from .request_contracts import ( @@ -435,13 +435,20 @@ def import_downloaded_model( if actual != expected or computed != expected: raise manager._error("sha256_mismatch", f"expected {expected}, got {actual}") manager._validate_provenance(task.provenance) + install_root = Path(manager.install_root).resolve() subdir = manager._sanitize_subdir(destination_subdir or task.destination_subdir) fname = manager._sanitize_filename(filename or task.filename) - rel_target = f"{subdir}/{fname}" + rel_target = PurePosixPath(subdir) / fname # IMPORTANT: keep root-bounded resolution; plain joins re-enable traversal risks. abs_target = Path( - manager._resolve_install_target(str(manager.install_root), rel_target) + manager._resolve_install_target(str(install_root), rel_target.as_posix()) ) + try: + safe_rel_target = abs_target.relative_to(install_root).as_posix() + except ValueError as exc: + raise manager._error( + "invalid_destination", "resolved import target escapes install root" + ) from exc abs_target.parent.mkdir(parents=True, exist_ok=True) fd, tmp = tempfile.mkstemp( prefix=f".{abs_target.name}.tmp.", dir=str(abs_target.parent), text=False @@ -479,7 +486,7 @@ def import_downloaded_model( "sha256": expected, "size_bytes": abs_target.stat().st_size if abs_target.exists() else None, "provenance": dict(task.provenance), - "installation_path": rel_target.replace("\\", "/"), + "installation_path": safe_rel_target, "tenant_id": task.tenant_id, "installed_at": time.time(), "tags": safe_tags, diff --git a/tests/test_checkpoints.py b/tests/test_checkpoints.py index ce1d8a6..86d0943 100644 --- a/tests/test_checkpoints.py +++ b/tests/test_checkpoints.py @@ -95,6 +95,12 @@ class TestCheckpointsService(unittest.TestCase): with self.assertRaises(ValueError): create_checkpoint("Valid Name", {}, long_desc) + def test_invalid_checkpoint_ids_fail_closed(self): + self.assertIsNone(get_checkpoint("../escape")) + self.assertIsNone(get_checkpoint("not-a-uuid")) + self.assertFalse(delete_checkpoint("../escape")) + self.assertFalse(delete_checkpoint("not-a-uuid")) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_model_manager_service.py b/tests/test_model_manager_service.py index 70c926b..47a3cba 100644 --- a/tests/test_model_manager_service.py +++ b/tests/test_model_manager_service.py @@ -15,6 +15,7 @@ from services.model_manager import ( ModelManager, ModelManagerError, ) +from services.safe_io import PathTraversalError class _FakeResponse: @@ -228,6 +229,45 @@ class TestModelManagerService(unittest.TestCase): self.manager.import_downloaded_model(task_id=task.task_id) self.assertEqual(ctx.exception.code, "sha256_mismatch") + def test_resolve_install_target_rejects_escape(self): + with self.assertRaises(PathTraversalError): + self.manager._resolve_install_target(str(self.install_root), "../escape.bin") + + def test_import_records_resolved_relative_installation_path(self): + payload = b"model-bytes" + digest = hashlib.sha256(payload).hexdigest() + staged_dir = self.manager.staging_dir / "task-safe-path" + staged_dir.mkdir(parents=True, exist_ok=True) + staged_file = staged_dir / "model.safetensors" + staged_file.write_bytes(payload) + task = DownloadTask( + task_id="task-safe-path", + model_id="model-safe-path", + name="Model Safe Path", + model_type="checkpoint", + source="catalog", + source_label="Catalog", + download_url="https://example.com/model.safetensors", + destination_subdir="checkpoints//nested", + filename="model.safetensors", + expected_sha256=digest, + provenance={ + "publisher": "OpenClaw", + "license": "OpenRAIL", + "source_url": "https://example.com/model", + }, + tenant_id="default", + state="completed", + staged_path=str(staged_file), + computed_sha256=digest, + ) + self.manager._tasks[task.task_id] = task + rec = self.manager.import_downloaded_model(task_id=task.task_id) + self.assertEqual( + rec["installation_path"], "checkpoints/nested/model.safetensors" + ) + self.assertTrue((self.install_root / rec["installation_path"]).exists()) + def test_list_download_tasks_delta_cursor_contract(self): first = DownloadTask( task_id="task-1", diff --git a/tests/test_r77_integrity.py b/tests/test_r77_integrity.py index 9212675..76b552f 100644 --- a/tests/test_r77_integrity.py +++ b/tests/test_r77_integrity.py @@ -102,6 +102,10 @@ class TestR77Integrity(unittest.TestCase): with self.assertRaises(IntegrityError): load_verified(path, migrate=False) + def test_save_verified_rejects_invalid_leaf_target(self): + with self.assertRaises(IntegrityError): + save_verified("", {"k": "v"}) + if __name__ == "__main__": unittest.main()