CVE-2024-2083: Insufficient Information (NVD-CWE-noinfo)
value: A directory traversal vulnerability exists in the zenml-io/zenml repository, specifically within the /api/v1/steps endpoint. Attackers can exploit this vulnerability by manipulating the 'logs' URI path in the request to fetch arbitrary file content, bypassing intended access restrictions. The vulnerability arises due to the lack of validation for directory traversal patterns, allowing attackers to access files outside of the restricted directory.
CWE description: There is insufficient information about the issue to classify it; details are unkown or unspecified.
Published: 2024-04-16T00:15Z
Repository: https://github.com/zenml-io/zenml
Language: Python
Files changed: src/zenml/artifact_stores/base_artifact_store.py, src/zenml/artifacts/utils.py, src/zenml/logging/step_logging.py, src/zenml/materializers/base_materializer.py, src/zenml/materializers/built_in_materializer.py, src/zenml/materializers/cloudpickle_materializer.py, src/zenml/materializers/numpy_materializer.py, src/zenml/materializers/pandas_materializer.py, src/zenml/materializers/service_materializer.py, src/zenml/materializers/structured_string_materializer.py, src/zenml/orchestrators/output_utils.py, tests/integration/functional/artifacts/test_base_artifact_store.py, tests/integration/functional/artifacts/test_utils.py, tests/unit/artifacts/test_utils.py, tests/unit/materializers/test_built_in_materializer.py, tests/unit/materializers/test_cloudpickle_materializer.py, tests/unit/test_general.py
Commit date: 2024-03-05 09:02:33 +0100
Commit message: Improve Artifact Store isolation (#2490)

* dir traversal issue

* Auto-update of Starter template

* Auto-update of NLP template

* reroute artifacts and logs via AS

* reroute materializers via AS

* simplify to one deco

* fix materializer tests

* allow local download

* Auto-update of E2E template

* fix test issues

* rework based on comments

* fix bugs

* lint

* Candidate (#2493)

Co-authored-by: Stefan Nica <stefan@zenml.io>

* darglint

---------

Co-authored-by: GitHub Actions <actions@github.com>
Co-authored-by: Stefan Nica <stefan@zenml.io>
Diff: diff --git a/src/zenml/artifact_stores/base_artifact_store.py b/src/zenml/artifact_stores/base_artifact_store.py
index fbadbaaae..1a1f775b3 100644
--- a/src/zenml/artifact_stores/base_artifact_store.py
+++ b/src/zenml/artifact_stores/base_artifact_store.py
@@ -11,12 +11,14 @@
 #  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.
 """The base interface to extend the ZenML artifact store."""
 
+import inspect
 import textwrap
 from abc import abstractmethod
+from pathlib import Path
 from typing import (
     Any,
     Callable,
     ClassVar,
     Dict,
@@ -42,71 +44,127 @@ from zenml.utils import io_utils
 logger = get_logger(__name__)
 
 PathType = Union[bytes, str]
 
 
-def _sanitize_potential_path(potential_path: Any) -> Any:
-    """Sanitizes the input if it is a path.
+class _sanitize_paths:
+    """Sanitizes path inputs before calling the original function.
 
-    If the input is a **remote** path, this function replaces backslash path
-    separators by forward slashes.
+    Extra decoration layer is needed to pass in fixed artifact store root
+    path for static methods that are called on filesystems directly.
 
     Args:
-        potential_path: Value that potentially refers to a (remote) path.
+        func: The function to decorate.
+        fixed_root_path: The fixed artifact store root path.
+        is_static: Whether the function is static or not.
 
     Returns:
-        The original input or a sanitized version of it in case of a remote
-        path.
+        Function that calls the input function with sanitized path inputs.
     """
-    if isinstance(potential_path, bytes):
-        path = fileio.convert_to_str(potential_path)
-    elif isinstance(potential_path, str):
-        path = potential_path
-    else:
-        # Neither string nor bytes, this is not a path
-        return potential_path
 
-    if io_utils.is_remote(path):
-        # If we have a remote path, replace windows path separators with
-        # slashes
-        import ntpath
-        import posixpath
+    def __init__(self, func: Callable[..., Any], fixed_root_path: str) -> None:
+        """Initializes the decorator.
 
-        path = path.replace(ntpath.sep, posixpath.sep)
+        Args:
+            func: The function to decorate.
+            fixed_root_path: The fixed artifact store root path.
+        """
+        self.func = func
+        self.fixed_root_path = fixed_root_path
 
-    return path
+        self.path_args: List[int] = []
+        self.path_kwargs: List[str] = []
+        for i, param in enumerate(
+            inspect.signature(self.func).parameters.values()
+        ):
+            if param.annotation == PathType:
+                self.path_kwargs.append(param.name)
+                if param.default == inspect.Parameter.empty:
+                    self.path_args.append(i)
 
+    def _validate_path(self, path: str) -> None:
+        """Validates a path.
 
-def _sanitize_paths(_func: Callable[..., Any]) -> Callable[..., Any]:
-    """Sanitizes path inputs before calling the original function.
+        Args:
+            path: The path to validate.
 
-    Args:
-        _func: The function for which to sanitize the inputs.
+        Raises:
+            FileNotFoundError: If the path is outside of the artifact store
+                bounds.
+        """
+        if not path.startswith(self.fixed_root_path):
+            raise FileNotFoundError(
+                f"File `{path}` is outside of "
+                f"artifact store bounds `{self.fixed_root_path}`"
+            )
 
-    Returns:
-        Function that calls the input function with sanitized path inputs.
-    """
+    def _sanitize_potential_path(self, potential_path: Any) -> Any:
+        """Sanitizes the input if it is a path.
+
+        If the input is a **remote** path, this function replaces backslash path
+        separators by forward slashes.
 
-    def inner_function(*args: Any, **kwargs: Any) -> Any:
-        """Inner function.
+        Args:
+            potential_path: Value that potentially refers to a (remote) path.
+
+        Returns:
+            The original input or a sanitized version of it in case of a remote
+            path.
+        """
+        if isinstance(potential_path, bytes):
+            path = fileio.convert_to_str(potential_path)
+        elif isinstance(potential_path, str):
+            path = potential_path
+        else:
+            # Neither string nor bytes, this is not a path
+            return potential_path
+
+        if io_utils.is_remote(path):
+            # If we have a remote path, replace windows path separators with
+            # slashes
+            import ntpath
+            import posixpath
+
+            path = path.replace(ntpath.sep, posixpath.sep)
+            self._validate_path(path)
+        else:
+            self._validate_path(str(Path(path).absolute().resolve()))
+
+        return path
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Decorator function that sanitizes paths before calling the original function.
 
         Args:
             *args: Positional args.
             **kwargs: Keyword args.
 
         Returns:
             Output of the input function called with sanitized paths.
         """
-        args = tuple(_sanitize_potential_path(arg) for arg in args)
+        # verify if `self` is part of the args
+        has_self = bool(args and isinstance(args[0], BaseArtifactStore))
+
+        # sanitize inputs for relevant args and kwargs, keep rest unchanged
+        args = tuple(
+            self._sanitize_potential_path(
+                arg,
+            )
+            if i + has_self in self.path_args
+            else arg
+            for i, arg in enumerate(args)
+        )
         kwargs = {
-            key: _sanitize_potential_path(value)
+            key: self._sanitize_potential_path(
+                value,
+            )
+            if key in self.path_kwargs
+            else value
             for key, value in kwargs.items()
         }
 
-        return _func(*args, **kwargs)
-
-    return inner_function
+        return self.func(*args, **kwargs)
 
 
 class BaseArtifactStoreConfig(StackComponentConfig):
     """Config class for `BaseArtifactStore`."""
 
@@ -321,10 +379,11 @@ class BaseArtifactStore(StackComponent):
 
         Returns:
             The stat descriptor.
         """
 
+    @abstractmethod
     def size(self, path: PathType) -> Optional[int]:
         """Get the size of a file in bytes.
 
         Args:
             path: The path to the file.
@@ -374,34 +433,34 @@ class BaseArtifactStore(StackComponent):
         """Create and register a filesystem within the filesystem registry."""
         from zenml.io.filesystem import BaseFilesystem
         from zenml.io.filesystem_registry import default_filesystem_registry
         from zenml.io.local_filesystem import LocalFilesystem
 
+        overloads: Dict[str, Any] = {
+            "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
+        }
+        for abc_method in inspect.getmembers(BaseArtifactStore):
+            if getattr(abc_method[1], "__isabstractmethod__", False):
+                sanitized_method = _sanitize_paths(
+                    getattr(self, abc_method[0]), self.path
+                )
+                # prepare overloads for filesystem methods
+                overloads[abc_method[0]] = staticmethod(sanitized_method)
+
+                # decorate artifact store methods
+                setattr(
+                    self,
+                    abc_method[0],
+                    sanitized_method,
+                )
+
         # Local filesystem is always registered, no point in doing it again.
         if isinstance(self, LocalFilesystem):
             return
 
         filesystem_class = type(
-            self.__class__.__name__,
-            (BaseFilesystem,),
-            {
-                "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
-                "open": staticmethod(_sanitize_paths(self.open)),
-                "copyfile": staticmethod(_sanitize_paths(self.copyfile)),
-                "exists": staticmethod(_sanitize_paths(self.exists)),
-                "glob": staticmethod(_sanitize_paths(self.glob)),
-                "isdir": staticmethod(_sanitize_paths(self.isdir)),
-                "listdir": staticmethod(_sanitize_paths(self.listdir)),
-                "makedirs": staticmethod(_sanitize_paths(self.makedirs)),
-                "mkdir": staticmethod(_sanitize_paths(self.mkdir)),
-                "remove": staticmethod(_sanitize_paths(self.remove)),
-                "rename": staticmethod(_sanitize_paths(self.rename)),
-                "rmtree": staticmethod(_sanitize_paths(self.rmtree)),
-                "size": staticmethod(_sanitize_paths(self.size)),
-                "stat": staticmethod(_sanitize_paths(self.stat)),
-                "walk": staticmethod(_sanitize_paths(self.walk)),
-            },
+            self.__class__.__name__, (BaseFilesystem,), overloads
         )
 
         default_filesystem_registry.register(filesystem_class)
 
 
diff --git a/src/zenml/artifacts/utils.py b/src/zenml/artifacts/utils.py
index ae00b743f..39b9bf636 100644
--- a/src/zenml/artifacts/utils.py
+++ b/src/zenml/artifacts/utils.py
@@ -150,21 +150,21 @@ def save_artifact(
     if not uri:
         uri = os.path.join("custom_artifacts", name, str(uuid4()))
     if not uri.startswith(artifact_store.path):
         uri = os.path.join(artifact_store.path, uri)
 
-    if manual_save and fileio.exists(uri):
+    if manual_save and artifact_store.exists(uri):
         # This check is only necessary for manual saves as we already check
         # it when creating the directory for step output artifacts
         other_artifacts = client.list_artifact_versions(uri=uri, size=1)
         if other_artifacts and (other_artifact := other_artifacts[0]):
             raise RuntimeError(
                 f"Cannot save new artifact {name} version to URI "
                 f"{uri} because the URI is already used by artifact "
                 f"{other_artifact.name} (version {other_artifact.version})."
             )
-    fileio.makedirs(uri)
+    artifact_store.makedirs(uri)
 
     # Find and initialize the right materializer class
     if isinstance(materializer, type):
         materializer_class = materializer
     elif materializer:
@@ -750,19 +750,22 @@ def _load_file_from_artifact_store(
         The loaded file.
 
     Raises:
         DoesNotExistException: If the file does not exist in the artifact store.
         NotImplementedError: If the artifact store cannot open the file.
+        IOError: If the artifact store rejects the request.
     """
     try:
         with artifact_store.open(uri, mode) as text_file:
             return text_file.read()
     except FileNotFoundError:
         raise DoesNotExistException(
             f"File '{uri}' does not exist in artifact store "
             f"'{artifact_store.name}'."
         )
+    except IOError as e:
+        raise e
     except Exception as e:
         logger.exception(e)
         link = "https://docs.zenml.io/stacks-and-components/component-guide/artifact-stores/custom#enabling-artifact-visualizations-with-custom-artifact-stores"
         raise NotImplementedError(
             f"File '{uri}' could not be loaded because the underlying artifact "
@@ -817,11 +820,12 @@ def load_model_from_metadata(model_uri: str) -> Any:
 
     Returns:
         The ML model object loaded into memory.
     """
     # Load the model from its metadata
-    with fileio.open(
+    artifact_store = Client().active_stack.artifact_store
+    with artifact_store.open(
         os.path.join(model_uri, MODEL_METADATA_YAML_FILE_NAME), "r"
     ) as f:
         metadata = read_yaml(f.name)
     data_type = metadata["datatype"]
     materializer = metadata["materializer"]
diff --git a/src/zenml/logging/step_logging.py b/src/zenml/logging/step_logging.py
index 178cf388c..ea5a566eb 100644
--- a/src/zenml/logging/step_logging.py
+++ b/src/zenml/logging/step_logging.py
@@ -21,11 +21,11 @@ from contextvars import ContextVar
 from types import TracebackType
 from typing import Any, Callable, List, Optional, Type
 from uuid import uuid4
 
 from zenml.artifact_stores import BaseArtifactStore
-from zenml.io import fileio
+from zenml.client import Client
 from zenml.logger import get_logger
 from zenml.logging import (
     STEP_LOGS_STORAGE_INTERVAL_SECONDS,
     STEP_LOGS_STORAGE_MAX_MESSAGES,
 )
@@ -62,30 +62,31 @@ def prepare_logs_uri(
         log_key: The unique identification key of the log file.
 
     Returns:
         The URI of the logs file.
     """
+    artifact_store = Client().active_stack.artifact_store
     if log_key is None:
         log_key = str(uuid4())
 
     logs_base_uri = os.path.join(
         artifact_store.path,
         step_name,
         "logs",
     )
 
     # Create the dir
-    if not fileio.exists(logs_base_uri):
-        fileio.makedirs(logs_base_uri)
+    if not artifact_store.exists(logs_base_uri):
+        artifact_store.makedirs(logs_base_uri)
 
     # Delete the file if it already exists
     logs_uri = os.path.join(logs_base_uri, f"{log_key}.log")
-    if fileio.exists(logs_uri):
+    if artifact_store.exists(logs_uri):
         logger.warning(
             f"Logs file {logs_uri} already exists! Removing old log file..."
         )
-        fileio.remove(logs_uri)
+        artifact_store.remove(logs_uri)
     return logs_uri
 
 
 class StepLogsStorage:
     """Helper class which buffers and stores logs to a given URI."""
@@ -133,16 +134,17 @@ class StepLogsStorage:
             ):
                 self.save_to_file()
 
     def save_to_file(self) -> None:
         """Method to save the buffer to the given URI."""
+        artifact_store = Client().active_stack.artifact_store
         if not self.disabled:
             try:
                 self.disabled = True
 
                 if self.buffer:
-                    with fileio.open(self.logs_uri, "a") as file:
+                    with artifact_store.open(self.logs_uri, "a") as file:
                         for message in self.buffer:
                             file.write(
                                 remove_ansi_escape_codes(message) + "\n"
                             )
 
diff --git a/src/zenml/materializers/base_materializer.py b/src/zenml/materializers/base_materializer.py
index 82ef92037..9c6b8a1ed 100644
--- a/src/zenml/materializers/base_materializer.py
+++ b/src/zenml/materializers/base_materializer.py
@@ -154,12 +154,13 @@ class BaseMaterializer(metaclass=BaseMaterializerMeta):
         When overriding this method, make sure to save all visualizations to
         files within `self.uri`.
 
         Example:
         ```
+        artifact_store = Client().active_stack.artifact_store
         visualization_uri = os.path.join(self.uri, "visualization.html")
-        with fileio.open(visualization_uri, "w") as f:
+        with artifact_store.open(visualization_uri, "w") as f:
             f.write("<html><body>data</body></html>")
 
         visualization_uri_2 = os.path.join(self.uri, "visualization.png")
         data.save_as_png(visualization_uri_2)
 
diff --git a/src/zenml/materializers/built_in_materializer.py b/src/zenml/materializers/built_in_materializer.py
index de09c4292..65acf157c 100644
--- a/src/zenml/materializers/built_in_materializer.py
+++ b/src/zenml/materializers/built_in_materializer.py
@@ -24,12 +24,12 @@ from typing import (
     Tuple,
     Type,
     Union,
 )
 
+from zenml.client import Client
 from zenml.enums import ArtifactType
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.materializer_registry import materializer_registry
 from zenml.utils import source_utils, yaml_utils
 
@@ -133,20 +133,22 @@ class BytesMaterializer(BaseMaterializer):
             data_type: The type of the data to read.
 
         Returns:
             The data read.
         """
-        with fileio.open(self.data_path, "rb") as file_:
+        artifact_store = Client().active_stack.artifact_store
+        with artifact_store.open(self.data_path, "rb") as file_:
             return file_.read()
 
     def save(self, data: Any) -> None:
         """Save a bytes object to file.
 
         Args:
             data: The data to store.
         """
-        with fileio.open(self.data_path, "wb") as file_:
+        artifact_store = Client().active_stack.artifact_store
+        with artifact_store.open(self.data_path, "wb") as file_:
             file_.write(data)
 
 
 def _all_serializable(iterable: Iterable[Any]) -> bool:
     """For an iterable, check whether all of its elements are JSON-serializable.
@@ -280,21 +282,22 @@ class BuiltInContainerMaterializer(BaseMaterializer):
             The data read.
 
         Raises:
             RuntimeError: If the data was not found.
         """
+        artifact_store = Client().active_stack.artifact_store
         # If the data was not serialized, there must be metadata present.
-        if not fileio.exists(self.data_path) and not fileio.exists(
-            self.metadata_path
-        ):
+        if not artifact_store.exists(
+            self.data_path
+        ) and not artifact_store.exists(self.metadata_path):
             raise RuntimeError(
                 f"Materialization of type {data_type} failed. Expected either"
                 f"{self.data_path} or {self.metadata_path} to exist."
             )
 
         # If the data was serialized as JSON, deserialize it.
-        if fileio.exists(self.data_path):
+        if artifact_store.exists(self.data_path):
             outputs = yaml_utils.read_json(self.data_path)
 
         # Otherwise, use the metadata to reconstruct the data as a list.
         else:
             metadata = yaml_utils.read_json(self.metadata_path)
@@ -353,10 +356,11 @@ class BuiltInContainerMaterializer(BaseMaterializer):
             data: The built-in container object to materialize.
 
         Raises:
             Exception: If any exception occurs, it is raised after cleanup.
         """
+        artifact_store = Client().active_stack.artifact_store
         # tuple and set: handle as list.
         if isinstance(data, tuple) or isinstance(data, set):
             data = list(data)
 
         # If the data is serializable, just write it into a single JSON file.
@@ -373,11 +377,11 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         metadata: List[Dict[str, str]] = []
         materializers: List[BaseMaterializer] = []
         try:
             for i, element in enumerate(data):
                 element_path = os.path.join(self.uri, str(i))
-                fileio.mkdir(element_path)
+                artifact_store.mkdir(element_path)
                 type_ = type(element)
                 materializer_class = materializer_registry[type_]
                 materializer = materializer_class(uri=element_path)
                 materializers.append(materializer)
                 metadata.append(
@@ -396,15 +400,15 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                 materializer.validate_type_compatibility(type(element))
                 materializer.save(element)
         # If an error occurs, delete all created files.
         except Exception as e:
             # Delete metadata
-            if fileio.exists(self.metadata_path):
-                fileio.remove(self.metadata_path)
+            if artifact_store.exists(self.metadata_path):
+                artifact_store.remove(self.metadata_path)
             # Delete all elements that were already saved.
             for entry in metadata:
-                fileio.rmtree(entry["path"])
+                artifact_store.rmtree(entry["path"])
             raise e
 
     def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:
         """Extract metadata from the given built-in container object.
 
diff --git a/src/zenml/materializers/cloudpickle_materializer.py b/src/zenml/materializers/cloudpickle_materializer.py
index 856375cd6..a84606b10 100644
--- a/src/zenml/materializers/cloudpickle_materializer.py
+++ b/src/zenml/materializers/cloudpickle_materializer.py
@@ -16,13 +16,13 @@
 import os
 from typing import Any, ClassVar, Tuple, Type
 
 import cloudpickle
 
+from zenml.client import Client
 from zenml.enums import ArtifactType
 from zenml.environment import Environment
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.utils.io_utils import (
     read_file_contents_as_string,
     write_file_contents_as_string,
@@ -57,10 +57,11 @@ class CloudpickleMaterializer(BaseMaterializer):
 
         Returns:
             The loaded artifact data.
         """
         # validate python version
+        artifact_store = Client().active_stack.artifact_store
         source_python_version = self._load_python_version()
         current_python_version = Environment().python_version()
         if source_python_version != current_python_version:
             logger.warning(
                 f"Your artifact was materialized under Python version "
@@ -70,11 +71,11 @@ class CloudpickleMaterializer(BaseMaterializer):
                 "versions. Attempting to load anyway..."
             )
 
         # load data
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
-        with fileio.open(filepath, "rb") as fid:
+        with artifact_store.open(filepath, "rb") as fid:
             data = cloudpickle.load(fid)
         return data
 
     def _load_python_version(self) -> str:
         """Loads the Python version that was used to materialize the artifact.
@@ -91,10 +92,11 @@ class CloudpickleMaterializer(BaseMaterializer):
         """Saves an artifact to a cloudpickle file.
 
         Args:
             data: The data to save.
         """
+        artifact_store = Client().active_stack.artifact_store
         # Log a warning if this materializer was not explicitly specified for
         # the given data type.
         if type(self) == CloudpickleMaterializer:
             logger.warning(
                 f"No materializer is registered for type `{type(data)}`, so "
@@ -109,11 +111,11 @@ class CloudpickleMaterializer(BaseMaterializer):
         # save python version for validation on loading
         self._save_python_version()
 
         # save data
         filepath = os.path.join(self.uri, DEFAULT_FILENAME)
-        with fileio.open(filepath, "wb") as fid:
+        with artifact_store.open(filepath, "wb") as fid:
             cloudpickle.dump(data, fid)
 
     def _save_python_version(self) -> None:
         """Saves the Python version used to materialize the artifact."""
         filepath = os.path.join(self.uri, DEFAULT_PYTHON_VERSION_FILENAME)
diff --git a/src/zenml/materializers/numpy_materializer.py b/src/zenml/materializers/numpy_materializer.py
index 38bea87c2..e2eef1935 100644
--- a/src/zenml/materializers/numpy_materializer.py
+++ b/src/zenml/materializers/numpy_materializer.py
@@ -17,12 +17,12 @@ import os
 from collections import Counter
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
 import numpy as np
 
+from zenml.client import Client
 from zenml.enums import ArtifactType, VisualizationType
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.metadata.metadata_types import DType, MetadataType
 
 if TYPE_CHECKING:
@@ -55,16 +55,17 @@ class NumpyMaterializer(BaseMaterializer):
             ImportError: If pyarrow is not installed.
 
         Returns:
             The numpy array.
         """
+        artifact_store = Client().active_stack.artifact_store
         numpy_file = os.path.join(self.uri, NUMPY_FILENAME)
 
-        if fileio.exists(numpy_file):
-            with fileio.open(numpy_file, "rb") as f:
+        if artifact_store.exists(numpy_file):
+            with artifact_store.open(numpy_file, "rb") as f:
                 return np.load(f, allow_pickle=True)
-        elif fileio.exists(os.path.join(self.uri, DATA_FILENAME)):
+        elif artifact_store.exists(os.path.join(self.uri, DATA_FILENAME)):
             logger.warning(
                 "A legacy artifact was found. "
                 "This artifact was created with an older version of "
                 "ZenML. You can still use it, but it will be "
                 "converted to the new format on the next materialization."
@@ -79,11 +80,11 @@ class NumpyMaterializer(BaseMaterializer):
                 # Read numpy array from parquet file
                 shape_dict = yaml_utils.read_json(
                     os.path.join(self.uri, SHAPE_FILENAME)
                 )
                 shape_tuple = tuple(shape_dict.values())
-                with fileio.open(
+                with artifact_store.open(
                     os.path.join(self.uri, DATA_FILENAME), "rb"
                 ) as f:
                     input_stream = pa.input_stream(f)
                     data = pq.read_table(input_stream)
                 vals = getattr(data.to_pandas(), DATA_VAR).values
@@ -100,11 +101,14 @@ class NumpyMaterializer(BaseMaterializer):
         """Writes a np.ndarray to the artifact store as a `.npy` file.
 
         Args:
             arr: The numpy array to write.
         """
-        with fileio.open(os.path.join(self.uri, NUMPY_FILENAME), "wb") as f:
+        artifact_store = Client().active_stack.artifact_store
+        with artifact_store.open(
+            os.path.join(self.uri, NUMPY_FILENAME), "wb"
+        ) as f:
             np.save(f, arr)
 
     def save_visualizations(
         self, arr: "NDArray[Any]"
     ) -> Dict[str, VisualizationType]:
@@ -153,12 +157,13 @@ class NumpyMaterializer(BaseMaterializer):
             output_path: The path to save the histogram to.
             arr: The numpy array of which to save the histogram.
         """
         import matplotlib.pyplot as plt
 
+        artifact_store = Client().active_stack.artifact_store
         plt.hist(arr)
-        with fileio.open(output_path, "wb") as f:
+        with artifact_store.open(output_path, "wb") as f:
             plt.savefig(f)
         plt.close()
 
     def _save_image(self, output_path: str, arr: "NDArray[Any]") -> None:
         """Saves a numpy array as an image.
@@ -167,11 +172,12 @@ class NumpyMaterializer(BaseMaterializer):
             output_path: The path to save the image to.
             arr: The numpy array to save.
         """
         from matplotlib.image import imsave
 
-        with fileio.open(output_path, "wb") as f:
+        artifact_store = Client().active_stack.artifact_store
+        with artifact_store.open(output_path, "wb") as f:
             imsave(f, arr)
 
     def extract_metadata(
         self, arr: "NDArray[Any]"
     ) -> Dict[str, "MetadataType"]:
diff --git a/src/zenml/materializers/pandas_materializer.py b/src/zenml/materializers/pandas_materializer.py
index c25f988ec..d36e1469d 100644
--- a/src/zenml/materializers/pandas_materializer.py
+++ b/src/zenml/materializers/pandas_materializer.py
@@ -16,12 +16,12 @@
 import os
 from typing import Any, ClassVar, Dict, Tuple, Type, Union
 
 import pandas as pd
 
+from zenml.client import Client
 from zenml.enums import ArtifactType, VisualizationType
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.metadata.metadata_types import DType, MetadataType
 
 logger = get_logger(__name__)
@@ -75,24 +75,25 @@ class PandasMaterializer(BaseMaterializer):
             ImportError: If pyarrow or fastparquet is not installed.
 
         Returns:
             The pandas dataframe or series.
         """
-        if fileio.exists(self.parquet_path):
+        artifact_store = Client().active_stack.artifact_store
+        if artifact_store.exists(self.parquet_path):
             if self.pyarrow_exists:
-                with fileio.open(self.parquet_path, mode="rb") as f:
+                with artifact_store.open(self.parquet_path, mode="rb") as f:
                     df = pd.read_parquet(f)
             else:
                 raise ImportError(
                     "You have an old version of a `PandasMaterializer` "
                     "data artifact stored in the artifact store "
                     "as a `.parquet` file, which requires `pyarrow` "
                     "for reading, You can install `pyarrow` by running "
                     "'`pip install pyarrow fastparquet`'."
                 )
         else:
-            with fileio.open(self.csv_path, mode="rb") as f:
+            with artifact_store.open(self.csv_path, mode="rb") as f:
                 df = pd.read_csv(f, index_col=0, parse_dates=True)
 
         # validate the type of the data.
         def is_dataframe_or_series(
             df: Union[pd.DataFrame, pd.Series],
@@ -120,18 +121,19 @@ class PandasMaterializer(BaseMaterializer):
         """Writes a pandas dataframe or series to the specified filename.
 
         Args:
             df: The pandas dataframe or series to write.
         """
+        artifact_store = Client().active_stack.artifact_store
         if isinstance(df, pd.Series):
             df = df.to_frame(name="series")
 
         if self.pyarrow_exists:
-            with fileio.open(self.parquet_path, mode="wb") as f:
+            with artifact_store.open(self.parquet_path, mode="wb") as f:
                 df.to_parquet(f, compression=COMPRESSION_TYPE)
         else:
-            with fileio.open(self.csv_path, mode="wb") as f:
+            with artifact_store.open(self.csv_path, mode="wb") as f:
                 df.to_csv(f, index=True)
 
     def save_visualizations(
         self, df: Union[pd.DataFrame, pd.Series]
     ) -> Dict[str, VisualizationType]:
@@ -141,13 +143,14 @@ class PandasMaterializer(BaseMaterializer):
             df: The pandas dataframe or series to visualize.
 
         Returns:
             A dictionary of visualization URIs and their types.
         """
+        artifact_store = Client().active_stack.artifact_store
         describe_uri = os.path.join(self.uri, "describe.csv")
         describe_uri = describe_uri.replace("\\", "/")
-        with fileio.open(describe_uri, mode="wb") as f:
+        with artifact_store.open(describe_uri, mode="wb") as f:
             df.describe().to_csv(f)
         return {describe_uri: VisualizationType.CSV}
 
     def extract_metadata(
         self, df: Union[pd.DataFrame, pd.Series]
diff --git a/src/zenml/materializers/service_materializer.py b/src/zenml/materializers/service_materializer.py
index a33f51c1b..7659cabe0 100644
--- a/src/zenml/materializers/service_materializer.py
+++ b/src/zenml/materializers/service_materializer.py
@@ -14,12 +14,12 @@
 """Implementation of a materializer to read and write ZenML service instances."""
 
 import os
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Type
 
+from zenml.client import Client
 from zenml.enums import ArtifactType
-from zenml.io import fileio
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.services.service import BaseService
 from zenml.services.service_registry import ServiceRegistry
 
 if TYPE_CHECKING:
@@ -44,12 +44,13 @@ class ServiceMaterializer(BaseMaterializer):
             data_type: The type of the data to read.
 
         Returns:
             A ZenML service instance.
         """
+        artifact_store = Client().active_stack.artifact_store
         filepath = os.path.join(self.uri, SERVICE_CONFIG_FILENAME)
-        with fileio.open(filepath, "r") as f:
+        with artifact_store.open(filepath, "r") as f:
             service = ServiceRegistry().load_service_from_json(f.read())
         return service
 
     def save(self, service: BaseService) -> None:
         """Writes a ZenML service.
@@ -58,12 +59,13 @@ class ServiceMaterializer(BaseMaterializer):
         are serialized and saved as an artifact.
 
         Args:
             service: A ZenML service instance.
         """
+        artifact_store = Client().active_stack.artifact_store
         filepath = os.path.join(self.uri, SERVICE_CONFIG_FILENAME)
-        with fileio.open(filepath, "w") as f:
+        with artifact_store.open(filepath, "w") as f:
             f.write(service.json(indent=4))
 
     def extract_metadata(
         self, service: BaseService
     ) -> Dict[str, "MetadataType"]:
diff --git a/src/zenml/materializers/structured_string_materializer.py b/src/zenml/materializers/structured_string_materializer.py
index dacda8748..933831645 100644
--- a/src/zenml/materializers/structured_string_materializer.py
+++ b/src/zenml/materializers/structured_string_materializer.py
@@ -14,12 +14,12 @@
 """Implementation of HTMLString materializer."""
 
 import os
 from typing import Dict, Type, Union
 
+from zenml.client import Client
 from zenml.enums import ArtifactType, VisualizationType
-from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.types import CSVString, HTMLString, MarkdownString
 
 logger = get_logger(__name__)
@@ -45,20 +45,22 @@ class StructuredStringMaterializer(BaseMaterializer):
             data_type: The type of the data to read.
 
         Returns:
             The loaded data.
         """
-        with fileio.open(self._get_filepath(data_type), "r") as f:
+        artifact_store = Client().active_stack.artifact_store
+        with artifact_store.open(self._get_filepath(data_type), "r") as f:
             return data_type(f.read())
 
     def save(self, data: STRUCTURED_STRINGS) -> None:
         """Save data as an HTML or Markdown file.
 
         Args:
             data: The data to save as an HTML or Markdown file.
         """
-        with fileio.open(self._get_filepath(type(data)), "w") as f:
+        artifact_store = Client().active_stack.artifact_store
+        with artifact_store.open(self._get_filepath(type(data)), "w") as f:
             f.write(data)
 
     def save_visualizations(
         self, data: STRUCTURED_STRINGS
     ) -> Dict[str, VisualizationType]:
diff --git a/src/zenml/orchestrators/output_utils.py b/src/zenml/orchestrators/output_utils.py
index 21f8b74cc..8bd18e09a 100644
--- a/src/zenml/orchestrators/output_utils.py
+++ b/src/zenml/orchestrators/output_utils.py
@@ -15,11 +15,11 @@
 
 import os
 from typing import TYPE_CHECKING, Dict, Sequence
 from uuid import uuid4
 
-from zenml.io import fileio
+from zenml.client import Client
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
     from zenml.artifact_stores import BaseArtifactStore
     from zenml.config.step_configurations import Step
@@ -70,28 +70,30 @@ def prepare_output_artifact_uris(
         RuntimeError: If an artifact URI already exists.
 
     Returns:
         A dictionary mapping output names to artifact URIs.
     """
+    artifact_store = stack.artifact_store
     output_artifact_uris: Dict[str, str] = {}
     for output_name in step.config.outputs.keys():
         artifact_uri = generate_artifact_uri(
             artifact_store=stack.artifact_store,
             step_run=step_run,
             output_name=output_name,
         )
-        if fileio.exists(artifact_uri):
+        if artifact_store.exists(artifact_uri):
             raise RuntimeError("Artifact already exists")
-        fileio.makedirs(artifact_uri)
+        artifact_store.makedirs(artifact_uri)
         output_artifact_uris[output_name] = artifact_uri
     return output_artifact_uris
 
 
 def remove_artifact_dirs(artifact_uris: Sequence[str]) -> None:
     """Removes the artifact directories.
 
     Args:
         artifact_uris: URIs of the artifacts to remove the directories for.
     """
+    artifact_store = Client().active_stack.artifact_store
     for artifact_uri in artifact_uris:
-        if fileio.isdir(artifact_uri):
-            fileio.rmtree(artifact_uri)
+        if artifact_store.isdir(artifact_uri):
+            artifact_store.rmtree(artifact_uri)
diff --git a/tests/integration/functional/artifacts/test_base_artifact_store.py b/tests/integration/functional/artifacts/test_base_artifact_store.py
new file mode 100644
index 000000000..b6f6cc024
--- /dev/null
+++ b/tests/integration/functional/artifacts/test_base_artifact_store.py
@@ -0,0 +1,65 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  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 os
+from pathlib import Path
+
+import pytest
+
+from zenml.client import Client
+
+
+def test_files_outside_of_artifact_store_are_not_reachable_by_it(
+    clean_client: "Client",
+):
+    """Tests that no operations outside of bounds of artifact store could happen."""
+    a_s = clean_client.active_stack.artifact_store
+
+    outside_dir = Path(a_s.path) / ".."
+    outside_file = str(outside_dir / "tmp.file")
+    try:
+        # create a file outside of artifact store
+        with open(outside_file, "w") as f:
+            f.write("test")
+        # try to open it via artifact store interface
+        with pytest.raises(FileNotFoundError):
+            a_s.open(outside_file, "r")
+        # try to copy it via artifact store interface
+        with pytest.raises(FileNotFoundError):
+            a_s.copyfile(outside_file, ".", "r")
+    except Exception as e:
+        raise e
+    finally:
+        os.remove(outside_file)
+
+    inside_file = str(Path(a_s.path) / "tmp.file")
+    try:
+        # create a file inside of artifact store
+        with open(inside_file, "w") as f:
+            f.write("test")
+        # try to open it via artifact store interface
+        assert a_s.open(inside_file, "r").read() == "test"
+        # try to copy it via artifact store interface
+        inside_file2 = str(Path(a_s.path) / "tmp2.file")
+        a_s.copyfile(inside_file, inside_file2, "r")
+        # try to open it via artifact store interface
+        assert open(inside_file2, "r").read() == "test"
+        # try to copy it via artifact store interface, but with target outside of bounds
+        with pytest.raises(FileNotFoundError):
+            a_s.copyfile(inside_file, ".", "r")
+    except Exception as e:
+        raise e
+    finally:
+        os.remove(inside_file)
+        os.remove(inside_file2)
diff --git a/tests/integration/functional/artifacts/test_utils.py b/tests/integration/functional/artifacts/test_utils.py
index 6571f1592..2d1a49ed2 100644
--- a/tests/integration/functional/artifacts/test_utils.py
+++ b/tests/integration/functional/artifacts/test_utils.py
@@ -260,11 +260,11 @@ def test_log_artifact_metadata_raises_error_if_output_name_unclear(
     with pytest.raises(ValueError):
         artifact_metadata_logging_pipeline()
 
 
 def test_download_artifact_files_from_response(
-    tmp_path, clean_client_with_run
+    tmp_path, clean_client_with_run: "Client"
 ):
     """Test that we can download artifact files from an artifact version."""
     artifact: ArtifactResponse = clean_client_with_run.get_artifact(
         name_id_or_prefix="connected_two_step_pipeline::step_1::output"
     )
diff --git a/tests/unit/artifacts/test_utils.py b/tests/unit/artifacts/test_utils.py
index 187055e10..6f6689e1f 100644
--- a/tests/unit/artifacts/test_utils.py
+++ b/tests/unit/artifacts/test_utils.py
@@ -24,10 +24,11 @@ from zenml.artifacts.utils import (
     _load_artifact_from_uri,
     load_artifact_from_response,
     load_model_from_metadata,
     save_model_metadata,
 )
+from zenml.client import Client
 from zenml.constants import MODEL_METADATA_YAML_FILE_NAME
 from zenml.materializers.numpy_materializer import NUMPY_FILENAME
 from zenml.models import ArtifactVersionResponse, Page
 
 
@@ -63,16 +64,18 @@ def test_save_model_metadata(model_artifact):
         assert "materializer" in file_contents
         assert model_artifact.materializer in file_contents
 
 
 @pytest.fixture
-def model_metadata_dir(model_artifact):
+def model_metadata_dir(model_artifact, clean_client: "Client"):
     # Save the model metadata to a temporary file
     file_path = save_model_metadata(model_artifact)
 
     # Move the file to a temporary directory
-    temp_dir = tempfile.mkdtemp()
+    temp_dir = tempfile.mkdtemp(
+        dir=clean_client.active_stack.artifact_store.path
+    )
     shutil.move(
         file_path, os.path.join(temp_dir, MODEL_METADATA_YAML_FILE_NAME)
     )
 
     # Yield the temporary directory
@@ -117,13 +120,15 @@ def test_load_artifact_from_response(mocker, model_artifact):
     # Ensure the _load_artifact_from_uri function is called
     mocker_load_artifact.assert_called_once()
 
 
 @pytest.fixture
-def numpy_file_uri():
+def numpy_file_uri(clean_client: "Client"):
     # Create a temporary file to save the numpy array
-    temp_dir = tempfile.mkdtemp()
+    temp_dir = tempfile.mkdtemp(
+        dir=clean_client.active_stack.artifact_store.path
+    )
     numpy_file = os.path.join(temp_dir, NUMPY_FILENAME)
 
     # Save a numpy array to the temporary file
     arr = np.array([1, 2, 3, 4, 5])
     np.save(numpy_file, arr)
diff --git a/tests/unit/materializers/test_built_in_materializer.py b/tests/unit/materializers/test_built_in_materializer.py
index 451cbe324..90d822a58 100644
--- a/tests/unit/materializers/test_built_in_materializer.py
+++ b/tests/unit/materializers/test_built_in_materializer.py
@@ -14,10 +14,11 @@
 import os
 from tempfile import TemporaryDirectory
 from typing import Optional, Type
 
 from tests.unit.test_general import _test_materializer
+from zenml.client import Client
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.built_in_materializer import (
     BuiltInContainerMaterializer,
 )
 
@@ -188,11 +189,13 @@ class CustomTypeMaterializer(BaseMaterializer):
     def load(self, data_type: Type[CustomType]) -> Optional[CustomType]:
         """Load the data."""
         return data_type()
 
 
-def test_container_materializer_for_custom_types(mocker):
+def test_container_materializer_for_custom_types(
+    mocker, clean_client: "Client"
+):
     """Test container materializer for custom types.
 
     This ensures that:
     - The container materializer can handle custom types.
     - Custom types are loaded as the correct type.
@@ -200,11 +203,13 @@ def test_container_materializer_for_custom_types(mocker):
         materializer registry when the container is loaded.
     """
     from zenml.materializers.materializer_registry import materializer_registry
 
     example = [CustomType(), CustomSubType()]
-    with TemporaryDirectory() as artifact_uri:
+    with TemporaryDirectory(
+        dir=clean_client.active_stack.artifact_store.path
+    ) as artifact_uri:
         materializer = BuiltInContainerMaterializer(uri=artifact_uri)
 
         # Container materializer should find materializer for both elements in
         # the default materializer registry.
         materializer.save(example)
diff --git a/tests/unit/materializers/test_cloudpickle_materializer.py b/tests/unit/materializers/test_cloudpickle_materializer.py
index 829dfa6af..98213c8c0 100644
--- a/tests/unit/materializers/test_cloudpickle_materializer.py
+++ b/tests/unit/materializers/test_cloudpickle_materializer.py
@@ -58,11 +58,13 @@ def test_cloudpickle_materializer_is_not_registered(clean_client):
 
 
 def test_cloudpickle_materializer_can_load_pickle(clean_client):
     """Test that the cloudpickle materializer can load regular pickle."""
     my_object = Unmaterializable()
-    with TemporaryDirectory() as artifact_uri:
+    with TemporaryDirectory(
+        dir=clean_client.active_stack.artifact_store.path
+    ) as artifact_uri:
         artifact_filepath = os.path.join(artifact_uri, DEFAULT_FILENAME)
         with open(artifact_filepath, "wb") as f:
             pickle.dump(my_object, f)
         materializer = CloudpickleMaterializer(uri=artifact_uri)
         loaded_object = materializer.load(data_type=Unmaterializable)
diff --git a/tests/unit/test_general.py b/tests/unit/test_general.py
index d55d3ee2c..5176acb8a 100644
--- a/tests/unit/test_general.py
+++ b/tests/unit/test_general.py
@@ -14,10 +14,11 @@
 
 import os
 from tempfile import TemporaryDirectory
 from typing import Any, Callable, Optional, Type
 
+from zenml.client import Client
 from zenml.constants import ENV_ZENML_DEBUG
 from zenml.enums import VisualizationType
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.materializer_registry import materializer_registry
 from zenml.metadata.metadata_types import MetadataTypeTuple
@@ -74,11 +75,12 @@ def _test_materializer(
         step_output_type = type(step_output)
 
     if materializer_class is None:
         materializer_class = materializer_registry[step_output_type]
 
-    with TemporaryDirectory() as artifact_uri:
+    artifact_store_uri = Client().active_stack.artifact_store.path
+    with TemporaryDirectory(dir=artifact_store_uri) as artifact_uri:
         materializer = materializer_class(uri=artifact_uri)
         existing_files = os.listdir(artifact_uri)
 
         # Assert that materializer saves something to disk
         materializer.save(step_output)

Vulnerable code: def _sanitize_potential_path(potential_path: Any) -> Any:
    """Sanitizes the input if it is a path.
    If the input is a **remote** path, this function replaces backslash path
    separators by forward slashes.
        potential_path: Value that potentially refers to a (remote) path.
        The original input or a sanitized version of it in case of a remote
        path.
    if isinstance(potential_path, bytes):
        path = fileio.convert_to_str(potential_path)
    elif isinstance(potential_path, str):
        path = potential_path
    else:
        # Neither string nor bytes, this is not a path
        return potential_path
    if io_utils.is_remote(path):
        # If we have a remote path, replace windows path separators with
        # slashes
        import ntpath
        import posixpath
        path = path.replace(ntpath.sep, posixpath.sep)
    return path
def _sanitize_paths(_func: Callable[..., Any]) -> Callable[..., Any]:
    """Sanitizes path inputs before calling the original function.
    Args:
        _func: The function for which to sanitize the inputs.
    Returns:
        Function that calls the input function with sanitized path inputs.
    """
    def inner_function(*args: Any, **kwargs: Any) -> Any:
        """Inner function.
        args = tuple(_sanitize_potential_path(arg) for arg in args)
            key: _sanitize_potential_path(value)
        return _func(*args, **kwargs)

    return inner_function
            self.__class__.__name__,
            (BaseFilesystem,),
            {
                "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
                "open": staticmethod(_sanitize_paths(self.open)),
                "copyfile": staticmethod(_sanitize_paths(self.copyfile)),
                "exists": staticmethod(_sanitize_paths(self.exists)),
                "glob": staticmethod(_sanitize_paths(self.glob)),
                "isdir": staticmethod(_sanitize_paths(self.isdir)),
                "listdir": staticmethod(_sanitize_paths(self.listdir)),
                "makedirs": staticmethod(_sanitize_paths(self.makedirs)),
                "mkdir": staticmethod(_sanitize_paths(self.mkdir)),
                "remove": staticmethod(_sanitize_paths(self.remove)),
                "rename": staticmethod(_sanitize_paths(self.rename)),
                "rmtree": staticmethod(_sanitize_paths(self.rmtree)),
                "size": staticmethod(_sanitize_paths(self.size)),
                "stat": staticmethod(_sanitize_paths(self.stat)),
                "walk": staticmethod(_sanitize_paths(self.walk)),
            },
    if manual_save and fileio.exists(uri):
    fileio.makedirs(uri)
    with fileio.open(
from zenml.io import fileio
    if not fileio.exists(logs_base_uri):
        fileio.makedirs(logs_base_uri)
    if fileio.exists(logs_uri):
        fileio.remove(logs_uri)
                    with fileio.open(self.logs_uri, "a") as file:
        with fileio.open(visualization_uri, "w") as f:
from zenml.io import fileio
        with fileio.open(self.data_path, "rb") as file_:
        with fileio.open(self.data_path, "wb") as file_:
        if not fileio.exists(self.data_path) and not fileio.exists(
            self.metadata_path
        ):
        if fileio.exists(self.data_path):
                fileio.mkdir(element_path)
            if fileio.exists(self.metadata_path):
                fileio.remove(self.metadata_path)
                fileio.rmtree(entry["path"])
from zenml.io import fileio
        with fileio.open(filepath, "rb") as fid:
        with fileio.open(filepath, "wb") as fid:
from zenml.io import fileio
        if fileio.exists(numpy_file):
            with fileio.open(numpy_file, "rb") as f:
        elif fileio.exists(os.path.join(self.uri, DATA_FILENAME)):
                with fileio.open(
        with fileio.open(os.path.join(self.uri, NUMPY_FILENAME), "wb") as f:
        with fileio.open(output_path, "wb") as f:
        with fileio.open(output_path, "wb") as f:
from zenml.io import fileio
        if fileio.exists(self.parquet_path):
                with fileio.open(self.parquet_path, mode="rb") as f:
            with fileio.open(self.csv_path, mode="rb") as f:
            with fileio.open(self.parquet_path, mode="wb") as f:
            with fileio.open(self.csv_path, mode="wb") as f:
        with fileio.open(describe_uri, mode="wb") as f:
from zenml.io import fileio
        with fileio.open(filepath, "r") as f:
        with fileio.open(filepath, "w") as f:
from zenml.io import fileio
        with fileio.open(self._get_filepath(data_type), "r") as f:
        with fileio.open(self._get_filepath(type(data)), "w") as f:
from zenml.io import fileio
        if fileio.exists(artifact_uri):
        fileio.makedirs(artifact_uri)
        if fileio.isdir(artifact_uri):
            fileio.rmtree(artifact_uri)
    tmp_path, clean_client_with_run
def model_metadata_dir(model_artifact):
    temp_dir = tempfile.mkdtemp()
def numpy_file_uri():
    temp_dir = tempfile.mkdtemp()
def test_container_materializer_for_custom_types(mocker):
    with TemporaryDirectory() as artifact_uri:
    with TemporaryDirectory() as artifact_uri:
    with TemporaryDirectory() as artifact_uri:
Fixed code: import inspect
from pathlib import Path
class _sanitize_paths:
    """Sanitizes path inputs before calling the original function.
    Extra decoration layer is needed to pass in fixed artifact store root
    path for static methods that are called on filesystems directly.
        func: The function to decorate.
        fixed_root_path: The fixed artifact store root path.
        is_static: Whether the function is static or not.
        Function that calls the input function with sanitized path inputs.
    def __init__(self, func: Callable[..., Any], fixed_root_path: str) -> None:
        """Initializes the decorator.
        Args:
            func: The function to decorate.
            fixed_root_path: The fixed artifact store root path.
        """
        self.func = func
        self.fixed_root_path = fixed_root_path
        self.path_args: List[int] = []
        self.path_kwargs: List[str] = []
        for i, param in enumerate(
            inspect.signature(self.func).parameters.values()
        ):
            if param.annotation == PathType:
                self.path_kwargs.append(param.name)
                if param.default == inspect.Parameter.empty:
                    self.path_args.append(i)
    def _validate_path(self, path: str) -> None:
        """Validates a path.
        Args:
            path: The path to validate.
        Raises:
            FileNotFoundError: If the path is outside of the artifact store
                bounds.
        """
        if not path.startswith(self.fixed_root_path):
            raise FileNotFoundError(
                f"File `{path}` is outside of "
                f"artifact store bounds `{self.fixed_root_path}`"
            )
    def _sanitize_potential_path(self, potential_path: Any) -> Any:
        """Sanitizes the input if it is a path.

        If the input is a **remote** path, this function replaces backslash path
        separators by forward slashes.
        Args:
            potential_path: Value that potentially refers to a (remote) path.

        Returns:
            The original input or a sanitized version of it in case of a remote
            path.
        """
        if isinstance(potential_path, bytes):
            path = fileio.convert_to_str(potential_path)
        elif isinstance(potential_path, str):
            path = potential_path
        else:
            # Neither string nor bytes, this is not a path
            return potential_path

        if io_utils.is_remote(path):
            # If we have a remote path, replace windows path separators with
            # slashes
            import ntpath
            import posixpath

            path = path.replace(ntpath.sep, posixpath.sep)
            self._validate_path(path)
        else:
            self._validate_path(str(Path(path).absolute().resolve()))

        return path

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Decorator function that sanitizes paths before calling the original function.
        # verify if `self` is part of the args
        has_self = bool(args and isinstance(args[0], BaseArtifactStore))

        # sanitize inputs for relevant args and kwargs, keep rest unchanged
        args = tuple(
            self._sanitize_potential_path(
                arg,
            )
            if i + has_self in self.path_args
            else arg
            for i, arg in enumerate(args)
        )
            key: self._sanitize_potential_path(
                value,
            )
            if key in self.path_kwargs
            else value
        return self.func(*args, **kwargs)
    @abstractmethod
        overloads: Dict[str, Any] = {
            "SUPPORTED_SCHEMES": self.config.SUPPORTED_SCHEMES,
        }
        for abc_method in inspect.getmembers(BaseArtifactStore):
            if getattr(abc_method[1], "__isabstractmethod__", False):
                sanitized_method = _sanitize_paths(
                    getattr(self, abc_method[0]), self.path
                )
                # prepare overloads for filesystem methods
                overloads[abc_method[0]] = staticmethod(sanitized_method)

                # decorate artifact store methods
                setattr(
                    self,
                    abc_method[0],
                    sanitized_method,
                )

            self.__class__.__name__, (BaseFilesystem,), overloads
    if manual_save and artifact_store.exists(uri):
    artifact_store.makedirs(uri)
        IOError: If the artifact store rejects the request.
    except IOError as e:
        raise e
    artifact_store = Client().active_stack.artifact_store
    with artifact_store.open(
from zenml.client import Client
    artifact_store = Client().active_stack.artifact_store
    if not artifact_store.exists(logs_base_uri):
        artifact_store.makedirs(logs_base_uri)
    if artifact_store.exists(logs_uri):
        artifact_store.remove(logs_uri)
        artifact_store = Client().active_stack.artifact_store
                    with artifact_store.open(self.logs_uri, "a") as file:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(visualization_uri, "w") as f:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(self.data_path, "rb") as file_:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(self.data_path, "wb") as file_:
        artifact_store = Client().active_stack.artifact_store
        if not artifact_store.exists(
            self.data_path
        ) and not artifact_store.exists(self.metadata_path):
        if artifact_store.exists(self.data_path):
        artifact_store = Client().active_stack.artifact_store
                artifact_store.mkdir(element_path)
            if artifact_store.exists(self.metadata_path):
                artifact_store.remove(self.metadata_path)
                artifact_store.rmtree(entry["path"])
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(filepath, "rb") as fid:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(filepath, "wb") as fid:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        if artifact_store.exists(numpy_file):
            with artifact_store.open(numpy_file, "rb") as f:
        elif artifact_store.exists(os.path.join(self.uri, DATA_FILENAME)):
                with artifact_store.open(
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(
            os.path.join(self.uri, NUMPY_FILENAME), "wb"
        ) as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(output_path, "wb") as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(output_path, "wb") as f:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        if artifact_store.exists(self.parquet_path):
                with artifact_store.open(self.parquet_path, mode="rb") as f:
            with artifact_store.open(self.csv_path, mode="rb") as f:
        artifact_store = Client().active_stack.artifact_store
            with artifact_store.open(self.parquet_path, mode="wb") as f:
            with artifact_store.open(self.csv_path, mode="wb") as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(describe_uri, mode="wb") as f:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(filepath, "r") as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(filepath, "w") as f:
from zenml.client import Client
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(self._get_filepath(data_type), "r") as f:
        artifact_store = Client().active_stack.artifact_store
        with artifact_store.open(self._get_filepath(type(data)), "w") as f:
from zenml.client import Client
    artifact_store = stack.artifact_store
        if artifact_store.exists(artifact_uri):
        artifact_store.makedirs(artifact_uri)
    artifact_store = Client().active_stack.artifact_store
        if artifact_store.isdir(artifact_uri):
            artifact_store.rmtree(artifact_uri)
#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at:
#
#       https://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  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 os
from pathlib import Path

import pytest

from zenml.client import Client


def test_files_outside_of_artifact_store_are_not_reachable_by_it(
    clean_client: "Client",
):
    """Tests that no operations outside of bounds of artifact store could happen."""
    a_s = clean_client.active_stack.artifact_store

    outside_dir = Path(a_s.path) / ".."
    outside_file = str(outside_dir / "tmp.file")
    try:
        # create a file outside of artifact store
        with open(outside_file, "w") as f:
            f.write("test")
        # try to open it via artifact store interface
        with pytest.raises(FileNotFoundError):
            a_s.open(outside_file, "r")
        # try to copy it via artifact store interface
        with pytest.raises(FileNotFoundError):
            a_s.copyfile(outside_file, ".", "r")
    except Exception as e:
        raise e
    finally:
        os.remove(outside_file)

    inside_file = str(Path(a_s.path) / "tmp.file")
    try:
        # create a file inside of artifact store
        with open(inside_file, "w") as f:
            f.write("test")
        # try to open it via artifact store interface
        assert a_s.open(inside_file, "r").read() == "test"
        # try to copy it via artifact store interface
        inside_file2 = str(Path(a_s.path) / "tmp2.file")
        a_s.copyfile(inside_file, inside_file2, "r")
        # try to open it via artifact store interface
        assert open(inside_file2, "r").read() == "test"
        # try to copy it via artifact store interface, but with target outside of bounds
        with pytest.raises(FileNotFoundError):
            a_s.copyfile(inside_file, ".", "r")
    except Exception as e:
        raise e
    finally:
        os.remove(inside_file)
        os.remove(inside_file2)
    tmp_path, clean_client_with_run: "Client"
from zenml.client import Client
def model_metadata_dir(model_artifact, clean_client: "Client"):
    temp_dir = tempfile.mkdtemp(
        dir=clean_client.active_stack.artifact_store.path
    )
def numpy_file_uri(clean_client: "Client"):
    temp_dir = tempfile.mkdtemp(
        dir=clean_client.active_stack.artifact_store.path
    )
from zenml.client import Client
def test_container_materializer_for_custom_types(
    mocker, clean_client: "Client"
):
    with TemporaryDirectory(
        dir=clean_client.active_stack.artifact_store.path
    ) as artifact_uri:
    with TemporaryDirectory(
        dir=clean_client.active_stack.artifact_store.path
    ) as artifact_uri:
from zenml.client import Client
    artifact_store_uri = Client().active_stack.artifact_store.path
    with TemporaryDirectory(dir=artifact_store_uri) as artifact_uri: