Directory Traversal in zenml-io/zenml

@@ -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 @@
 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 inner_function(*args: Any, **kwargs: Any) -> 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 @@ def stat(self, path: PathType) -> Any:
             The stat descriptor.
         """
 
+    @abstractmethod
     def size(self, path: PathType) -> Optional[int]:
         """Get the size of a file in bytes.
 
@@ -376,30 +435,30 @@ def _register(self) -> None:
         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)

@@ -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)

@@ -23,7 +23,7 @@
 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 @@ def write(self, text: str) -> None:
 
     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"

@@ -156,8 +156,9 @@ def save_visualizations(self, data: Any) -> Dict[str, VisualizationType]:
 
         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")

@@ -26,8 +26,8 @@
     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 @@ def load(self, data_type: Type[Any]) -> Any:
         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 @@ def save(self, data: Any) -> None:
         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 @@ def load(self, data_type: Type[Any]) -> Any:
         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 @@ def save(self, data: Any) -> None:
         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 @@ def save(self, data: Any) -> None:
         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 @@ def save(self, data: Any) -> None:
         # 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"]:

@@ -18,9 +18,9 @@
 
 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 @@ def load(self, data_type: Type[Any]) -> Any:
             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 @@ def load(self, data_type: Type[Any]) -> Any:
 
         # 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 @@ def save(self, data: Any) -> None:
         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 @@ def save(self, data: Any) -> None:
 
         # 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:

@@ -19,8 +19,8 @@
 
 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 @@ def load(self, data_type: Type[Any]) -> "Any":
         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 @@ def load(self, data_type: Type[Any]) -> "Any":
                     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 @@ def save(self, arr: "NDArray[Any]") -> None:
         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 @@ def _save_histogram(self, output_path: str, arr: "NDArray[Any]") -> None:
         """
         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 @@ def _save_image(self, output_path: str, arr: "NDArray[Any]") -> None:
         """
         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(

@@ -18,8 +18,8 @@
 
 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 @@ def load(self, data_type: Type[Any]) -> Union[pd.DataFrame, pd.Series]:
         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 @@ def load(self, data_type: Type[Any]) -> Union[pd.DataFrame, pd.Series]:
                     "'`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 @@ def save(self, df: Union[pd.DataFrame, pd.Series]) -> None:
         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 @@ def save_visualizations(
         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}
 

@@ -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 @@ def load(self, data_type: Type[Any]) -> BaseService:
         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 @@ def save(self, service: BaseService) -> None:
         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(

@@ -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 @@ def load(self, data_type: Type[STRUCTURED_STRINGS]) -> STRUCTURED_STRINGS:
         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 @@ def save(self, data: STRUCTURED_STRINGS) -> None:
         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(

@@ -17,7 +17,7 @@
 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,16 +72,17 @@ 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(
             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
 
@@ -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)

@@ -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)

@@ -262,7 +262,7 @@ def 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(

@@ -26,6 +26,7 @@
     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

@@ -16,6 +16,7 @@
 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 @@ def load(self, data_type: Type[CustomType]) -> Optional[CustomType]:
         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

@@ -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)

@@ -16,6 +16,7 @@
 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)