[K8s] Fix validating label keys (#5670)

mlrun · Jun 2, 2024 · cadaf2b · cadaf2b
1 parent 76482a4
commit cadaf2b
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 14 deletions.
diff --git a/mlrun/k8s_utils.py b/mlrun/k8s_utils.py
@@ -141,17 +141,6 @@ def verify_label_key(key: str):
     if not key:
         raise mlrun.errors.MLRunInvalidArgumentError("label key cannot be empty")
 
-    mlrun.utils.helpers.verify_field_regex(
-        f"project.metadata.labels.'{key}'",
-        key,
-        mlrun.utils.regex.k8s_character_limit,
-    )
-
-    if key.startswith("k8s.io/") or key.startswith("kubernetes.io/"):
-        raise mlrun.errors.MLRunInvalidArgumentError(
-            "Labels cannot start with 'k8s.io/' or 'kubernetes.io/'"
-        )
-
     parts = key.split("/")
     if len(parts) == 1:
         name = parts[0]
@@ -173,12 +162,22 @@ def verify_label_key(key: str):
             "Label key can only contain one '/'"
         )
 
+    mlrun.utils.helpers.verify_field_regex(
+        f"project.metadata.labels.'{key}'",
+        name,
+        mlrun.utils.regex.k8s_character_limit,
+    )
     mlrun.utils.helpers.verify_field_regex(
         f"project.metadata.labels.'{key}'",
         name,
         mlrun.utils.regex.qualified_name,
     )
 
+    if key.startswith("k8s.io/") or key.startswith("kubernetes.io/"):
+        raise mlrun.errors.MLRunInvalidArgumentError(
+            "Labels cannot start with 'k8s.io/' or 'kubernetes.io/'"
+        )
+
 
 def verify_label_value(value, label_key):
     mlrun.utils.helpers.verify_field_regex(

diff --git a/mlrun/utils/helpers.py b/mlrun/utils/helpers.py
@@ -181,8 +181,12 @@ def verify_field_regex(
             )
             if mode == mlrun.common.schemas.RegexMatchModes.all:
                 if raise_on_failure:
+                    if len(field_name) > max_chars:
+                        field_name = field_name[:max_chars] + "...truncated"
+                    if len(field_value) > max_chars:
+                        field_value = field_value[:max_chars] + "...truncated"
                     raise mlrun.errors.MLRunInvalidArgumentError(
-                        f"Field '{field_name[:max_chars]}' is malformed. '{field_value[:max_chars]}' "
+                        f"Field '{field_name}' is malformed. '{field_value}' "
                         f"does not match required pattern: {pattern}"
                     )
                 return False

diff --git a/tests/test_k8s_utils.py b/tests/test_k8s_utils.py
@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from contextlib import nullcontext as does_not_raise
+
 import pytest
 
-from mlrun.k8s_utils import sanitize_label_value
+import mlrun.errors
+import mlrun.k8s_utils
 
 
 @pytest.mark.parametrize(
@@ -30,4 +33,58 @@
     ],
 )
 def test_sanitize_label_value(value: str, expected: str):
-    assert sanitize_label_value(value) == expected
+    assert mlrun.k8s_utils.sanitize_label_value(value) == expected
+
+
+@pytest.mark.parametrize(
+    "label_key, exception",
+    [
+        # valid
+        ("a/" + "b" * 63, does_not_raise()),
+        ("a" * 253 + "/b", does_not_raise()),
+        ("a" * 253 + "/" + "b" * 63, does_not_raise()),
+        ("my-key", does_not_raise()),
+        ("a/b", does_not_raise()),
+        ("prefix/valid-key", does_not_raise()),
+        ("prefix.with.dots/valid-key", does_not_raise()),
+        ("prefix-with-dashes/valid-key", does_not_raise()),
+        # preserved
+        ("k8s.io/a", pytest.raises(mlrun.errors.MLRunInvalidArgumentError)),
+        ("kubernetes.io/a", pytest.raises(mlrun.errors.MLRunInvalidArgumentError)),
+        # prefix too long
+        (
+            "toolong" + "a" * 248 + "/key",
+            pytest.raises(mlrun.errors.MLRunInvalidArgumentError),
+        ),
+        # name too long
+        (
+            "a/" + "b" * 64,
+            pytest.raises(mlrun.errors.MLRunInvalidArgumentError),
+        ),
+        # prefix has invalid character - '_'
+        (
+            "prefix_with_underscores/valid-key",
+            pytest.raises(mlrun.errors.MLRunInvalidArgumentError),
+        ),
+        # invalid prefix
+        (
+            "invalid-prefix-.com/key",
+            pytest.raises(mlrun.errors.MLRunInvalidArgumentError),
+        ),
+        # trailing slash in key
+        ("invalid-prefix/key/", pytest.raises(mlrun.errors.MLRunInvalidArgumentError)),
+        # leading slash in key
+        ("/invalid-key", pytest.raises(mlrun.errors.MLRunInvalidArgumentError)),
+        # empty key
+        ("", pytest.raises(mlrun.errors.MLRunInvalidArgumentError)),
+        # trailing dash
+        ("invalid-key-", pytest.raises(mlrun.errors.MLRunInvalidArgumentError)),
+        # trailing underscore
+        ("invalid-key_", pytest.raises(mlrun.errors.MLRunInvalidArgumentError)),
+        # trailing dot
+        ("invalid-key.", pytest.raises(mlrun.errors.MLRunInvalidArgumentError)),
+    ],
+)
+def test_verify_label_key(label_key, exception):
+    with exception:
+        mlrun.k8s_utils.verify_label_key(label_key)