From baeeed66806cb406bf2f85e714a33f31834f0110 Mon Sep 17 00:00:00 2001
From: nate nowack <thrast36@gmail.com>
Date: Mon, 8 Jul 2024 13:47:34 -0500
Subject: [PATCH] Fix docker worker volumes (#14516)

---
 .../prefect-docker/prefect_docker/worker.py   | 27 +++++++++++++-
 .../prefect-docker/tests/test_worker.py       | 37 +++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/src/integrations/prefect-docker/prefect_docker/worker.py b/src/integrations/prefect-docker/prefect_docker/worker.py
index 6490378aeb73..c914167472ac 100644
--- a/src/integrations/prefect-docker/prefect_docker/worker.py
+++ b/src/integrations/prefect-docker/prefect_docker/worker.py
@@ -62,7 +62,32 @@ class ImagePullPolicy(enum.Enum):
     NEVER = "Never"
 
 
-VolumeStr = Annotated[str, AfterValidator(lambda v: ":" in v)]
+def assert_volume_str(v: str) -> str:
+    """Assert that a string is a valid Docker volume string."""
+    if not isinstance(v, str):
+        raise ValueError("Volume must be a string")
+
+    # Regex pattern for valid Docker volume strings, including Windows paths
+    pattern = r"^([a-zA-Z]:\\|/)?[^:]+:(/)?[^:]+(:ro|:rw)?$"
+
+    match = re.match(pattern, v)
+    if not match:
+        raise ValueError(f"Invalid volume string: {v!r}")
+
+    _, _, mode = match.groups()
+
+    # Check for empty parts
+    if ":" not in v or v.startswith(":") or v.endswith(":"):
+        raise ValueError(f"Volume string contains empty part: {v!r}")
+
+    # If there's a mode, it must be 'ro' or 'rw'
+    if mode and mode not in (":ro", ":rw"):
+        raise ValueError(f"Invalid volume mode: {mode!r}. Must be ':ro' or ':rw'")
+
+    return v
+
+
+VolumeStr = Annotated[str, AfterValidator(assert_volume_str)]
 
 
 class DockerWorkerJobConfiguration(BaseJobConfiguration):
diff --git a/src/integrations/prefect-docker/tests/test_worker.py b/src/integrations/prefect-docker/tests/test_worker.py
index f9917e08f639..800128e7aa32 100644
--- a/src/integrations/prefect-docker/tests/test_worker.py
+++ b/src/integrations/prefect-docker/tests/test_worker.py
@@ -15,7 +15,9 @@
     CONTAINER_LABELS,
     DockerWorker,
     DockerWorkerJobConfiguration,
+    VolumeStr,
 )
+from pydantic import TypeAdapter, ValidationError
 
 from prefect.client.schemas import FlowRun
 from prefect.events import RelatedResource
@@ -262,6 +264,41 @@ async def test_uses_volumes_setting(
     assert "c:d" in call_volumes
 
 
+@pytest.mark.parametrize(
+    "volume_str",
+    [
+        "a:b",
+        "/host/path:/container/path",
+        "named_volume:/app/data",
+        "/home/user:/home/docker:ro",
+        "/home/user:/home/docker:rw",
+        "C:\\path\\on\\windows:/path/in/container",
+        "\\\\host\\share:/path/in/container",
+    ],
+)
+def test_valid_volume_strings(volume_str):
+    assert TypeAdapter(VolumeStr).validate_python(volume_str) == volume_str
+
+
+@pytest.mark.parametrize(
+    "volume_str",
+    [
+        "invalid_volume",
+        ":missing_host",
+        "missing_container:",
+        "/double:/colon:/path",
+        "/path:/path:invalid_mode",
+        ":/:",
+        " : : ",
+        "/host:/container:rw:extra",
+        "",  # empty string
+    ],
+)
+def test_invalid_volume_strings(volume_str):
+    with pytest.raises(ValidationError, match="Invalid volume"):
+        TypeAdapter(VolumeStr).validate_python(volume_str)
+
+
 async def test_uses_privileged_setting(
     mock_docker_client, flow_run, default_docker_worker_job_configuration
 ):