[BugFix] Harden memmap path handling

README.md

@@ -192,8 +192,8 @@ utilities.
 
 ```python
 td = TensorDict({"tokens": tokens, "scores": scores}, batch_size=[n])
-td.memmap("/tmp/batch")          # memory-map every leaf
-reloaded = TensorDict.load_memmap("/tmp/batch")
+td.memmap("/path/to/private/batch")  # memory-map every leaf
+reloaded = TensorDict.load_memmap("/path/to/private/batch")
 ```
 
 Memory-mapped TensorDicts are useful for large offline datasets, replay buffers,

docs/source/saving.rst

@@ -47,6 +47,24 @@ However, this approach also has some disadvantages:
   The :class:`~tensordict.NonTensorData` class can be used to represent non-tensor
   data in a regular :class:`~tensordict.TensorDict` instance.
 
+Filesystem considerations
+-------------------------
+
+Keep the save directory (or archive's parent directory) writable only by
+trusted principals. Existing symbolic links are rejected when creating
+individual memory-map files, but a path-based API cannot protect against a
+process that can replace names in the parent directory concurrently. Prefer
+:class:`tempfile.TemporaryDirectory` or another private directory over
+predictable names in a shared ``/tmp`` directory.
+
+Robust key encoding is enabled by default and keeps TensorDict keys within
+their save prefix. ``robust_key=False`` exists only to interoperate with
+legacy layouts; do not use it with untrusted keys or metadata. Individual
+memory-map leaf files overwrite existing regular files by default
+(``existsok=True``); pass ``existsok=False`` to reject a colliding leaf path.
+The metadata file is still refreshed when reusing an existing save directory,
+so ``existsok=False`` does not reserve the directory as a whole.
+
 tensordict's memory-mapped API relies on four core methods:
 :meth:`~tensordict.TensorDictBase.memmap_`, :meth:`~tensordict.TensorDictBase.memmap`,
 :meth:`~tensordict.TensorDictBase.memmap_like` and :meth:`~tensordict.TensorDictBase.load_memmap`.

docs/source/storage.rst

@@ -403,8 +403,8 @@ For **memmap**, non-tensor data is serialised via tensorclass's
    ...     label=NonTensorData(data="cat", batch_size=[4]),
    ...     batch_size=[4],
    ... )
-   >>> td_mm = td.memmap_("/tmp/example")
-   >>> loaded = TensorDict.load_memmap("/tmp/example")  # doctest: +SKIP
+   >>> td_mm = td.memmap_("/path/to/private/example")  # doctest: +SKIP
+   >>> loaded = TensorDict.load_memmap("/path/to/private/example")  # doctest: +SKIP
    >>> loaded["label"].data  # doctest: +SKIP
    'cat'
 

tensordict/_td.py

@@ -67,12 +67,15 @@
     _clone_value,
     _create_segments_from_int,
     _create_segments_from_list,
+    _encode_key_for_filesystem,
     _get_item,
     _get_leaf_tensordict,
+    _get_robust_key_setting_with_warning,
     _get_shape_from_args,
     _getitem_batch_size,
     _index_preserve_data_ptr,
     _infer_size_impl,
+    _is_safe_legacy_key,
     _is_shared,
     _is_unbatched,
     _KEY_ERROR,
@@ -3039,8 +3042,23 @@ def _memmap_(
         for key, value in self.items():
             type_value = type(value)
             if _is_tensor_collection(type_value):
+                if prefix is not None:
+                    effective_robust_key = _get_robust_key_setting_with_warning(
+                        key, robust_key
+                    )
+                    safe_key = _encode_key_for_filesystem(
+                        key, robust=effective_robust_key
+                    )
+                    value_prefix = prefix / safe_key
+                    if value_prefix.is_symlink():
+                        raise RuntimeError(
+                            "Refusing to write a memory-mapped TensorDict through "
+                            f"symlink {value_prefix}."
+                        )
+                else:
+                    value_prefix = None
                 dest._tensordict[key] = value._memmap_(
-                    prefix=prefix / key if prefix is not None else None,
+                    prefix=value_prefix,
                     copy_existing=copy_existing,
                     executor=executor,
                     futures=futures,
@@ -3124,22 +3142,17 @@ def _load_memmap(
         else:
             result = out
 
-        paths = set()
+        paths = []
         for key, entry_metadata in metadata.items():
             if not isinstance(entry_metadata, dict):
                 # there can be other metadata
                 continue
             type_value = entry_metadata.get("type")
             if type_value is not None:
-                paths.add(key)
+                paths.append(key)
                 continue
             dtype = entry_metadata.get("dtype")
             shape = entry_metadata.get("shape")
-            from .utils import (
-                _encode_key_for_filesystem,
-                _get_robust_key_setting_with_warning,
-            )
-
             # Use smart warning for loading that only warns when encoding would differ
             effective_robust_key = _get_robust_key_setting_with_warning(key, robust_key)
 
@@ -3148,7 +3161,11 @@ def _load_memmap(
             memmap_file = prefix / f"{safe_key}.memmap"
 
             # If robust encoding is requested but file doesn't exist, try legacy filename
-            if not memmap_file.exists() and effective_robust_key:
+            if (
+                not memmap_file.exists()
+                and effective_robust_key
+                and _is_safe_legacy_key(key)
+            ):
                 legacy_key = _encode_key_for_filesystem(key, robust=False)
                 legacy_file = prefix / f"{legacy_key}.memmap"
                 if legacy_file.exists():
@@ -3218,35 +3235,64 @@ def _load_memmap(
                 inplace=False,
                 non_blocking=False,
             )
-        # iterate over folders and load them
-        for path in prefix.iterdir():
-            if path.is_dir() and path.parts[-1] in paths:
-                key = path.parts[-1]  # path.parts[len(prefix.parts) :]
-                existing_elt = result._get_str(key, default=None)
-                if existing_elt is not None:
-                    existing_elt.load_memmap_(path)
-                else:
-                    result._set_str(
-                        key,
-                        TensorDict.load_memmap(path, device=device, non_blocking=True),
-                        inplace=False,
-                        validated=False,
-                    )
+        # Load collection directories named by metadata. New saves use robust
+        # encoding; safe single-component legacy names remain readable.
+        for key in paths:
+            effective_robust_key = _get_robust_key_setting_with_warning(key, robust_key)
+            safe_key = _encode_key_for_filesystem(key, robust=effective_robust_key)
+            path = prefix / safe_key
+            if (
+                not path.is_dir()
+                and effective_robust_key
+                and _is_safe_legacy_key(key, is_collection=True)
+            ):
+                legacy_path = prefix / key
+                if legacy_path.is_dir():
+                    path = legacy_path
+            if not path.is_dir():
+                continue
+            existing_elt = result._get_str(key, default=None)
+            if existing_elt is not None:
+                existing_elt.load_memmap_(path, robust_key=robust_key)
+            else:
+                result._set_str(
+                    key,
+                    TensorDict.load_memmap(
+                        path,
+                        device=device,
+                        non_blocking=True,
+                        robust_key=robust_key,
+                    ),
+                    inplace=False,
+                    validated=False,
+                )
         # Archive paths are read-only views inside a zip file: they cannot be
         # used as a target for a subsequent memmap_()/refresh, so only real
         # directories are recorded.
         result._memmap_prefix = prefix if isinstance(prefix, Path) else None
         return result
 
-    def _make_memmap_subtd(self, key):
+    def _make_memmap_subtd(self, key, *, robust_key):
         """Creates a sub-tensordict given a tuple key."""
         result = self
         for key_str in key:
             result_tmp = result._get_str(key_str, default=None)
             if result_tmp is None:
                 result_tmp = result.empty()
                 if result._memmap_prefix is not None:
-                    result_tmp.memmap_(prefix=result._memmap_prefix / key_str)
+                    effective_robust_key = _get_robust_key_setting_with_warning(
+                        key_str, robust_key
+                    )
+                    safe_key = _encode_key_for_filesystem(
+                        key_str, robust=effective_robust_key
+                    )
+                    subtd_prefix = result._memmap_prefix / safe_key
+                    if subtd_prefix.is_symlink():
+                        raise RuntimeError(
+                            "Refusing to write a memory-mapped TensorDict through "
+                            f"symlink {subtd_prefix}."
+                        )
+                    result_tmp.memmap_(prefix=subtd_prefix)
                     metadata = _load_metadata(result._memmap_prefix)
                     _update_metadata(
                         metadata=metadata,
@@ -3276,7 +3322,7 @@ def make_memmap(
 
         key = unravel_key(key)
         if isinstance(key, tuple):
-            last_node = self._make_memmap_subtd(key[:-1])
+            last_node = self._make_memmap_subtd(key[:-1], robust_key=robust_key)
             last_key = key[-1]
         else:
             last_node = self
@@ -3332,7 +3378,7 @@ def make_memmap_from_storage(
 
         key = unravel_key(key)
         if isinstance(key, tuple):
-            last_node = self._make_memmap_subtd(key[:-1])
+            last_node = self._make_memmap_subtd(key[:-1], robust_key=robust_key)
             last_key = key[-1]
         else:
             last_node = self
@@ -3392,7 +3438,7 @@ def make_memmap_from_tensor(
 
         key = unravel_key(key)
         if isinstance(key, tuple):
-            last_node = self._make_memmap_subtd(key[:-1])
+            last_node = self._make_memmap_subtd(key[:-1], robust_key=robust_key)
             last_key = key[-1]
         else:
             last_node = self

tensordict/_utils_key_json.py

@@ -13,6 +13,7 @@
     "_encode_key_for_filesystem",
     "_get_robust_key_setting",
     "_get_robust_key_setting_with_warning",
+    "_is_safe_legacy_key",
     "_json_dumps",
     "get_json_backend",
     "json_dumps",
@@ -28,6 +29,14 @@ def _encode_key_for_filesystem(key: str, *, robust: bool = True) -> str:
     if not robust:
         return key
 
+    # These values do not name a child when used as a path component. A bare
+    # "%" cannot be produced by the encoder (literal percent signs become
+    # "%25"), so it is an unambiguous representation for the empty key.
+    if not key:
+        return "%"
+    if key in (".", ".."):
+        return "%2E" * len(key)
+
     unsafe_chars = set('/<>:"|?*\\ \0%')
     unsafe_chars.update(chr(i) for i in range(32))
     unsafe_chars.add(chr(127))
@@ -42,6 +51,17 @@ def _encode_key_for_filesystem(key: str, *, robust: bool = True) -> str:
     return "".join(encoded_parts)
 
 
+def _is_safe_legacy_key(key: str, *, is_collection: bool = False) -> bool:
+    """Return whether a raw legacy key stays beneath its memmap prefix."""
+    if "/" in key or "\\" in key:
+        return False
+    if len(key) >= 2 and key[0].isalpha() and key[1] == ":":
+        # A drive-relative path such as ``C:payload`` escapes a prefix on
+        # Windows even though it has no slash.
+        return False
+    return not is_collection or key not in ("", ".", "..")
+
+
 def _get_robust_key_setting_with_warning(key: str, robust_key) -> bool:
     """Handle the robust_key parameter after the robust default migration."""
     if robust_key is None:
@@ -58,6 +78,8 @@ def _get_robust_key_setting(robust_key) -> bool:
 
 def _decode_key_from_filesystem(encoded_key: str) -> str:
     """Decode a filesystem-safe key back to the original TensorDict key."""
+    if encoded_key == "%":
+        return ""
     decoded_parts = []
     i = 0
     while i < len(encoded_key):

tensordict/base.py

@@ -79,12 +79,15 @@
     _CloudpickleWrapper,
     _convert_list_to_stack,
     _DTYPE_TO_STR_DTYPE,
+    _encode_key_for_filesystem,
     _GENERIC_NESTED_ERR,
+    _get_robust_key_setting_with_warning,
     _get_shared_executor,
     _is_dataclass as is_dataclass,
     _is_list_tensor_compatible,
     _is_non_tensor,
     _is_number,
+    _is_safe_legacy_key,
     _is_tensorclass,
     _is_unbatched,
     _KEY_ERROR,
@@ -7989,7 +7992,22 @@ def load_memmap(
                         "strings."
                     )
             for part in subpath:
-                prefix = prefix / part
+                effective_robust_key = _get_robust_key_setting_with_warning(
+                    part, robust_key
+                )
+                safe_part = _encode_key_for_filesystem(
+                    part, robust=effective_robust_key
+                )
+                candidate = prefix / safe_part
+                if (
+                    effective_robust_key
+                    and not (candidate / "meta.json").exists()
+                    and _is_safe_legacy_key(part, is_collection=True)
+                ):
+                    legacy_candidate = prefix / part
+                    if (legacy_candidate / "meta.json").exists():
+                        candidate = legacy_candidate
+                prefix = candidate
             if not (prefix / "meta.json").exists():
                 raise ValueError(
                     f"No tensordict found under subpath {'/'.join(subpath)!r} "