diff --git a/moshi/moshi/models/lm.py b/moshi/moshi/models/lm.py index a64d736b4..276680524 100644 --- a/moshi/moshi/models/lm.py +++ b/moshi/moshi/models/lm.py @@ -976,7 +976,7 @@ def load_voice_prompt(self, voice_prompt: str): def load_voice_prompt_embeddings(self, path: str): self.voice_prompt = path - state = torch.load(path) + state = torch.load(path, weights_only=True) self.voice_prompt_audio = None self.voice_prompt_embeddings = state["embeddings"].to(self.lm_model.device) diff --git a/moshi/moshi/models/loaders.py b/moshi/moshi/models/loaders.py index d38dcc3bb..4b1b28989 100644 --- a/moshi/moshi/models/loaders.py +++ b/moshi/moshi/models/loaders.py @@ -157,7 +157,7 @@ def get_mimi(filename: str | Path, if _is_safetensors(filename): load_model(model, filename) else: - pkg = torch.load(filename, "cpu") + pkg = torch.load(filename, map_location="cpu", weights_only=True) model.load_state_dict(pkg["model"]) model.set_num_codebooks(8) return model @@ -214,7 +214,7 @@ def get_moshi_lm( else: # torch checkpoint with open(filename, "rb") as f: - state_dict = torch.load(f, map_location="cpu") + state_dict = torch.load(f, map_location="cpu", weights_only=True) # Patch 1: expand depformer self_attn weights if needed model_sd = model.state_dict() for name, tensor in list(state_dict.items()): @@ -292,7 +292,7 @@ def _get_moshi_lm_with_offload( state_dict = load_file(filename, device="cpu") else: with open(filename, "rb") as f: - state_dict = torch.load(f, map_location="cpu") + state_dict = torch.load(f, map_location="cpu", weights_only=True) # Apply weight patches (same as non-offload path) model_sd = model.state_dict() diff --git a/moshi/moshi/offline.py b/moshi/moshi/offline.py index f690620d2..0c01afc37 100644 --- a/moshi/moshi/offline.py +++ b/moshi/moshi/offline.py @@ -122,6 +122,18 @@ def decode_tokens_to_pcm(mimi: MimiModel, other_mimi: MimiModel, lm_gen: LMGen, return pcm +def _safe_tar_extract(tar: tarfile.TarFile, path: str | Path) -> None: + """Extract tar contents with path traversal protection.""" + dest = os.path.realpath(str(path)) + for member in tar.getmembers(): + member_path = os.path.realpath(os.path.join(dest, member.name)) + if not member_path.startswith(dest + os.sep) and member_path != dest: + raise RuntimeError( + f"Refusing to extract {member.name!r}: would write outside {dest}" + ) + tar.extractall(path=path) + + def _get_voice_prompt_dir(voice_prompt_dir: Optional[str], hf_repo: str) -> Optional[str]: """ If voice_prompt_dir is None: @@ -142,7 +154,7 @@ def _get_voice_prompt_dir(voice_prompt_dir: Optional[str], hf_repo: str) -> Opti if not voices_dir.exists(): log("info", f"extracting {voices_tgz} to {voices_dir}") with tarfile.open(voices_tgz, "r:gz") as tar: - tar.extractall(path=voices_tgz.parent) + _safe_tar_extract(tar, voices_tgz.parent) if not voices_dir.exists(): raise RuntimeError("voices.tgz did not contain a 'voices/' directory") diff --git a/moshi/moshi/server.py b/moshi/moshi/server.py index 771f491da..f8a9dd15d 100644 --- a/moshi/moshi/server.py +++ b/moshi/moshi/server.py @@ -43,7 +43,6 @@ import sentencepiece import sphn import torch -import random from .client_utils import make_log, colorize from .models import loaders, MimiModel, LMModel, LMGen @@ -152,7 +151,13 @@ async def handle_chat(self, request): voice_prompt_filename = request.query["voice_prompt"] requested_voice_prompt_path = None if voice_prompt_filename is not None: - requested_voice_prompt_path = os.path.join(self.voice_prompt_dir, voice_prompt_filename) + # Prevent path traversal: strip directory components from client input + safe_filename = os.path.basename(voice_prompt_filename) + if safe_filename != voice_prompt_filename: + raise ValueError( + f"Invalid voice prompt filename: {voice_prompt_filename!r}" + ) + requested_voice_prompt_path = os.path.join(self.voice_prompt_dir, safe_filename) # If the voice prompt file does not exist, find a valid (s0) voiceprompt file in the directory if requested_voice_prompt_path is None or not os.path.exists(requested_voice_prompt_path): raise FileNotFoundError( @@ -168,7 +173,7 @@ async def handle_chat(self, request): else: self.lm_gen.load_voice_prompt(voice_prompt_path) self.lm_gen.text_prompt_tokens = self.text_tokenizer.encode(wrap_with_system_tags(request.query["text_prompt"])) if len(request.query["text_prompt"]) > 0 else None - seed = int(request["seed"]) if "seed" in request.query else None + seed = int(request.query["seed"]) if "seed" in request.query else None async def recv_loop(): nonlocal close @@ -309,6 +314,18 @@ async def is_alive(): return ws +def _safe_tar_extract(tar: tarfile.TarFile, path: str | Path) -> None: + """Extract tar contents with path traversal protection.""" + dest = os.path.realpath(str(path)) + for member in tar.getmembers(): + member_path = os.path.realpath(os.path.join(dest, member.name)) + if not member_path.startswith(dest + os.sep) and member_path != dest: + raise RuntimeError( + f"Refusing to extract {member.name!r}: would write outside {dest}" + ) + tar.extractall(path=path) + + def _get_voice_prompt_dir(voice_prompt_dir: Optional[str], hf_repo: str) -> Optional[str]: """ If voice_prompt_dir is None: @@ -330,7 +347,7 @@ def _get_voice_prompt_dir(voice_prompt_dir: Optional[str], hf_repo: str) -> Opti if not voices_dir.exists(): logger.info(f"extracting {voices_tgz} to {voices_dir}") with tarfile.open(voices_tgz, "r:gz") as tar: - tar.extractall(path=voices_tgz.parent) + _safe_tar_extract(tar, voices_tgz.parent) if not voices_dir.exists(): raise RuntimeError("voices.tgz did not contain a 'voices/' directory") @@ -346,7 +363,7 @@ def _get_static_path(static: Optional[str]) -> Optional[str]: dist = dist_tgz.parent / "dist" if not dist.exists(): with tarfile.open(dist_tgz, "r:gz") as tar: - tar.extractall(path=dist_tgz.parent) + _safe_tar_extract(tar, dist_tgz.parent) return str(dist) elif static != "none": # When set to the "none" string, we don't serve any static content. @@ -465,7 +482,7 @@ async def handle_root(_): logger.info(f"serving static content from {static_path}") app.router.add_get("/", handle_root) app.router.add_static( - "/", path=static_path, follow_symlinks=True, name="static" + "/", path=static_path, follow_symlinks=False, name="static" ) protocol = "http" ssl_context = None