diff --git a/onnx_runtime_0_1b_int8/arktts_runtime/runtime.py b/onnx_runtime_0_1b_int8/arktts_runtime/runtime.py index c84e5be..dfafb2f 100644 --- a/onnx_runtime_0_1b_int8/arktts_runtime/runtime.py +++ b/onnx_runtime_0_1b_int8/arktts_runtime/runtime.py @@ -3,7 +3,7 @@ from __future__ import annotations import json import threading from collections.abc import Iterator -from pathlib import Path +from pathlib import Path, PureWindowsPath from typing import Any import numpy as np @@ -22,6 +22,25 @@ ORT_DTYPES: dict[str, Any] = { } + +def _model_path(model_dir: Path, name: str) -> Path: + """Resolve a manifest model path inside an immutable model package.""" + if not isinstance(name, str) or not name or "\x00" in name: + raise ValueError("model path must be a non-empty string") + relative = Path(name) + # Reject foreign-platform absolute/drive paths as well as native ones. + if relative.is_absolute() or PureWindowsPath(name).drive or "\\" in name: + raise ValueError("model path must be package-relative") + if ".." in relative.parts: + raise ValueError("model path must not contain parent traversal") + root = model_dir.resolve(strict=True) + candidate = (root / relative).resolve(strict=True) + if not candidate.is_relative_to(root): + raise ValueError("model path escapes the model package") + if not candidate.is_file(): + raise ValueError("model path must identify a regular file") + return candidate + def _session(path: Path, threads: int | None = None) -> ort.InferenceSession: if not path.is_file(): raise FileNotFoundError(f"ONNX model was not found: {path}") @@ -157,9 +176,13 @@ class ArkTtsRuntime: self.codec_precision, f"codec_decoder_{self.codec_precision}.onnx", ) - self.slow = _session(self.model_dir / slow_name, threads) - self.fast = _session(self.model_dir / fast_name, threads) - self.decoder = _session(self.model_dir / codec_name, threads) + # Validate every manifest path before constructing any ONNX session. + slow_path = _model_path(self.model_dir, slow_name) + fast_path = _model_path(self.model_dir, fast_name) + codec_path = _model_path(self.model_dir, codec_name) + self.slow = _session(slow_path, threads) + self.fast = _session(fast_path, threads) + self.decoder = _session(codec_path, threads) self.slow_inputs = {item.name: item for item in self.slow.get_inputs()} self.fast_inputs = {item.name: item for item in self.fast.get_inputs()}