CVE-2024-2083 (published 2024-04-16T00:15Z)

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.


Weakness information:
cwe_id: NVD-CWE-noinfo
cwe_name: Insufficient Information
description: There is insufficient information about the issue to classify it; details are unkown or unspecified.
extended_description: Insufficient Information

Fix information:
repo_url: https://github.com/zenml-io/zenml
commit_details: repo_url: https://github.com/zenml-io/zenml
author: Andrei Vishniakov
author_date: 2024-03-05 09:02:33+01:00
msg: 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>
repository: repo_url: https://github.com/zenml-io/zenml
repo_name: zenml-io/zenml
description: ZenML 🙏: Build portable, production-ready MLOps pipelines. https://zenml.io.
repo_language: Python
owner: zenml-io
file_changes: filename: base_artifact_store.py
change_type: ModificationType.MODIFY
diff: @@ -13,8 +13,10 @@
 #  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,
@@ -44,50 +46,93 @@ 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.
@@ -96,15 +141,28 @@ def _sanitize_paths(_func: Callable[..., Any]) -> Callable[..., Any]:
         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):
@@ -323,6 +381,7 @@ class BaseArtifactStore(StackComponent):
             The stat descriptor.
         """
 
+    @abstractmethod
     def size(self, path: PathType) -> Optional[int]:
         """Get the size of a file in bytes.
 
@@ -376,30 +435,30 @@ class BaseArtifactStore(StackComponent):
         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)

programming_language: Python
filename: utils.py
change_type: ModificationType.MODIFY
diff: @@ -152,7 +152,7 @@ def save_artifact(
     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)
@@ -162,7 +162,7 @@ def save_artifact(
                 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):
@@ -752,6 +752,7 @@ def _load_file_from_artifact_store(
     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:
@@ -761,6 +762,8 @@ def _load_file_from_artifact_store(
             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"
@@ -819,7 +822,8 @@ def load_model_from_metadata(model_uri: str) -> Any:
         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)

programming_language: Python
filename: step_logging.py
change_type: ModificationType.MODIFY
diff: @@ -23,7 +23,7 @@ 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,
@@ -64,6 +64,7 @@ def prepare_logs_uri(
     Returns:
         The URI of the logs file.
     """
+    artifact_store = Client().active_stack.artifact_store
     if log_key is None:
         log_key = str(uuid4())
 
@@ -74,16 +75,16 @@ def prepare_logs_uri(
     )
 
     # 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
 
 
@@ -135,12 +136,13 @@ class StepLogsStorage:
 
     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"

programming_language: Python
filename: base_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -156,8 +156,9 @@ class BaseMaterializer(metaclass=BaseMaterializerMeta):
 
         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")

programming_language: Python
filename: built_in_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -26,8 +26,8 @@ from typing import (
     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
@@ -135,7 +135,8 @@ class BytesMaterializer(BaseMaterializer):
         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:
@@ -144,7 +145,8 @@ class BytesMaterializer(BaseMaterializer):
         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)
 
 
@@ -282,17 +284,18 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         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.
@@ -355,6 +358,7 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         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)
@@ -375,7 +379,7 @@ class BuiltInContainerMaterializer(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)
@@ -398,11 +402,11 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         # 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"]:

programming_language: Python
filename: cloudpickle_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -18,9 +18,9 @@ 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 (
@@ -59,6 +59,7 @@ class CloudpickleMaterializer(BaseMaterializer):
             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:
@@ -72,7 +73,7 @@ class CloudpickleMaterializer(BaseMaterializer):
 
         # 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
 
@@ -93,6 +94,7 @@ class CloudpickleMaterializer(BaseMaterializer):
         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:
@@ -111,7 +113,7 @@ class CloudpickleMaterializer(BaseMaterializer):
 
         # 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:

programming_language: Python
filename: numpy_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -19,8 +19,8 @@ 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
@@ -57,12 +57,13 @@ class NumpyMaterializer(BaseMaterializer):
         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 "
@@ -81,7 +82,7 @@ class NumpyMaterializer(BaseMaterializer):
                     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)
@@ -102,7 +103,10 @@ class NumpyMaterializer(BaseMaterializer):
         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(
@@ -155,8 +159,9 @@ class NumpyMaterializer(BaseMaterializer):
         """
         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()
 
@@ -169,7 +174,8 @@ class NumpyMaterializer(BaseMaterializer):
         """
         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(

programming_language: Python
filename: pandas_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -18,8 +18,8 @@ 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
@@ -77,9 +77,10 @@ class PandasMaterializer(BaseMaterializer):
         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(
@@ -90,7 +91,7 @@ class PandasMaterializer(BaseMaterializer):
                     "'`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.
@@ -122,14 +123,15 @@ class PandasMaterializer(BaseMaterializer):
         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(
@@ -143,9 +145,10 @@ class PandasMaterializer(BaseMaterializer):
         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}
 

programming_language: Python
filename: service_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -16,8 +16,8 @@
 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
@@ -46,8 +46,9 @@ class ServiceMaterializer(BaseMaterializer):
         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
 
@@ -60,8 +61,9 @@ class ServiceMaterializer(BaseMaterializer):
         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(

programming_language: Python
filename: structured_string_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -16,8 +16,8 @@
 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
@@ -47,7 +47,8 @@ class StructuredStringMaterializer(BaseMaterializer):
         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:
@@ -56,7 +57,8 @@ class StructuredStringMaterializer(BaseMaterializer):
         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(

programming_language: Python
filename: output_utils.py
change_type: ModificationType.MODIFY
diff: @@ -17,7 +17,7 @@ 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:
@@ -72,6 +72,7 @@ def prepare_output_artifact_uris(
     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(
@@ -79,9 +80,9 @@ def prepare_output_artifact_uris(
             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
 
@@ -92,6 +93,7 @@ def remove_artifact_dirs(artifact_uris: Sequence[str]) -> None:
     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)

programming_language: Python
filename: test_base_artifact_store.py
change_type: ModificationType.ADD
diff: @@ -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)

programming_language: Python
filename: test_utils.py
change_type: ModificationType.MODIFY
diff: @@ -262,7 +262,7 @@ def test_log_artifact_metadata_raises_error_if_output_name_unclear(
 
 
 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(

programming_language: Python
filename: test_utils.py
change_type: ModificationType.MODIFY
diff: @@ -26,6 +26,7 @@ from zenml.artifacts.utils import (
     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
@@ -65,12 +66,14 @@ def test_save_model_metadata(model_artifact):
 
 
 @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)
     )
@@ -119,9 +122,11 @@ def test_load_artifact_from_response(mocker, model_artifact):
 
 
 @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

programming_language: Python
filename: test_built_in_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -16,6 +16,7 @@ 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,
@@ -190,7 +191,9 @@ class CustomTypeMaterializer(BaseMaterializer):
         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:
@@ -202,7 +205,9 @@ def test_container_materializer_for_custom_types(mocker):
     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

programming_language: Python
filename: test_cloudpickle_materializer.py
change_type: ModificationType.MODIFY
diff: @@ -60,7 +60,9 @@ 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)

programming_language: Python
filename: test_general.py
change_type: ModificationType.MODIFY
diff: @@ -16,6 +16,7 @@ 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
@@ -76,7 +77,8 @@ def _test_materializer(
     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)
 

programming_language: Python