diff --git a/src/diffusers/utils/hub_utils.py b/src/diffusers/utils/hub_utils.py index b5eb9ab2e17f..561d27392fad 100644 --- a/src/diffusers/utils/hub_utils.py +++ b/src/diffusers/utils/hub_utils.py @@ -392,6 +392,13 @@ def _get_checkpoint_shard_files( index = json.loads(f.read()) original_shard_filenames = sorted(set(index["weight_map"].values())) + for shard_filename in original_shard_filenames: + if os.path.basename(shard_filename) != shard_filename: + raise ValueError( + f"The shard filename {shard_filename!r} in the checkpoint index contains a path separator or a " + "parent-directory reference, which is not allowed. Shard filenames must be plain filenames located " + "in the model directory." + ) sharded_metadata = index["metadata"] sharded_metadata["all_checkpoint_keys"] = list(index["weight_map"].keys()) sharded_metadata["weight_map"] = index["weight_map"].copy() diff --git a/tests/others/test_hub_utils.py b/tests/others/test_hub_utils.py index de6ab198e501..c897afd26db9 100644 --- a/tests/others/test_hub_utils.py +++ b/tests/others/test_hub_utils.py @@ -12,11 +12,17 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import json +import os import unittest from pathlib import Path from tempfile import TemporaryDirectory -from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card +from diffusers.utils.hub_utils import ( + _get_checkpoint_shard_files, + load_or_create_model_card, + populate_model_card, +) class CreateModelCardTest(unittest.TestCase): @@ -27,3 +33,36 @@ def test_generate_model_card_with_library_name(self): model_card = load_or_create_model_card(file_path) populate_model_card(model_card) assert model_card.data.library_name == "foo" + + +class GetCheckpointShardFilesTest(unittest.TestCase): + def _write_index(self, model_dir, shard_filename): + index = {"metadata": {"total_size": 1}, "weight_map": {"w": shard_filename}} + index_filename = os.path.join(model_dir, "diffusion_pytorch_model.safetensors.index.json") + with open(index_filename, "w") as f: + json.dump(index, f) + return index_filename + + def test_rejects_parent_directory_traversal(self): + with TemporaryDirectory() as tmpdir: + model_dir = os.path.join(tmpdir, "model") + os.makedirs(model_dir) + index_filename = self._write_index(model_dir, "../secret/SECRET.safetensors") + with self.assertRaises(ValueError): + _get_checkpoint_shard_files(model_dir, index_filename) + + def test_rejects_absolute_path(self): + with TemporaryDirectory() as tmpdir: + model_dir = os.path.join(tmpdir, "model") + os.makedirs(model_dir) + index_filename = self._write_index(model_dir, os.path.join(tmpdir, "secret", "SECRET.safetensors")) + with self.assertRaises(ValueError): + _get_checkpoint_shard_files(model_dir, index_filename) + + def test_rejects_subdirectory_component(self): + with TemporaryDirectory() as tmpdir: + model_dir = os.path.join(tmpdir, "model") + os.makedirs(model_dir) + index_filename = self._write_index(model_dir, "sub/shard.safetensors") + with self.assertRaises(ValueError): + _get_checkpoint_shard_files(model_dir, index_filename)