feat(image_spec): validate container registry names (#2748)

* feat(image_spec): validate Docker registry names Signed-off-by: Kevin Su <[email protected]> * one more test Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * thomas's suggestions Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]>
flyteorg · Sep 17, 2024 · fb55841 · fb55841
1 parent a666246
commit fb55841
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 0 deletions.
diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py
@@ -90,6 +90,13 @@ def __post_init__(self):
         self._is_force_push = os.environ.get(FLYTE_FORCE_PUSH_IMAGE_SPEC, False)  # False by default
         if self.registry:
             self.registry = self.registry.lower()
+            if not validate_container_registry_name(self.registry):
+                raise ValueError(
+                    f"Invalid container registry name: '{self.registry}'.\n Expected formats:\n"
+                    f"- 'localhost:30000' (for local registries)\n"
+                    f"- 'ghcr.io/username' (for GitHub Container Registry)\n"
+                    f"- 'docker.io/username' (for docker hub)\n"
+                )
 
         # If not set, help the user set this option as well, to support the older default behavior where existence
         # of the source root implied that copying of files was needed.
@@ -407,3 +414,12 @@ def _get_builder(cls, builder: str) -> ImageSpecBuilder:
                     f" Please upgrade envd to v0.3.39+."
                 )
         return cls._REGISTRY[builder][0]
+
+
+def validate_container_registry_name(name: str) -> bool:
+    """Validate Docker container registry name."""
+    # Define the regular expression for the registry name
+    registry_pattern = r"^(localhost:\d{1,5}|([a-z\d\._-]+)(:\d{1,5})?)(/[\w\.-]+)*$"
+
+    # Use regex to validate the given name
+    return bool(re.match(registry_pattern, name))
diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py
@@ -215,3 +215,29 @@ def test_update_image_spec_copy_handling():
     update_image_spec_copy_handling(image_spec, ss)
     assert image_spec.source_copy_mode is None
     assert image_spec.source_root is None
+
+
+def test_registry_name():
+    invalid_registry_names = [
+        "invalid:port:50000",
+        "ghcr.io/flyteorg:latest",
+        "flyteorg:latest"
+    ]
+    for invalid_registry_name in invalid_registry_names:
+        with pytest.raises(ValueError, match="Invalid container registry name"):
+            ImageSpec(registry=invalid_registry_name)
+
+    valid_registry_names = [
+        "localhost:30000",
+        "localhost:30000/flyte",
+        "192.168.1.1:30000",
+        "192.168.1.1:30000/myimage",
+        "ghcr.io/flyteorg",
+        "my.registry.com/myimage",
+        "my.registry.com:5000/myimage",
+        "myregistry:5000/myimage",
+        "us-west1-docker.pkg.dev/example.com/my-project/my-repo"
+        "flyteorg",
+    ]
+    for valid_registry_name in valid_registry_names:
+        ImageSpec(registry=valid_registry_name)