fix(novita): resolve whole ARG names and stop leaking the sandbox

src/openenv/core/containers/runtime/novita_provider.py

@@ -62,6 +62,7 @@
 _ARG_DEFAULT_RE = re.compile(
     r"^\s*ARG\s+(?P<name>\w+)=(?P<value>\S+)\s*$", re.IGNORECASE
 )
+_ARG_REFERENCE_RE = re.compile(r"\$(?:\{(?P<braced>\w+)\}|(?P<bare>\w+))")
 
 
 def _strip_mount_flags(content: str) -> str:
@@ -116,6 +117,23 @@ def _strip_mount_flags(content: str) -> str:
     return "\n".join(out)
 
 
+def _substitute_args(reference: str, arg_defaults: Dict[str, str]) -> str:
+    """Expand ``$NAME`` and ``${NAME}`` in one pass against *arg_defaults*.
+
+    Substituting name by name would let a shorter ARG consume the prefix of a
+    longer one: with ``ARG BASE`` declared before ``ARG BASE_IMAGE``, an
+    unbraced ``FROM $BASE_IMAGE`` became ``<base value>_IMAGE``. Matching the
+    whole name in a single pass makes the result independent of declaration
+    order. Names with no default are left verbatim, as before.
+    """
+
+    def replace(match: "re.Match[str]") -> str:
+        name = match.group("braced") or match.group("bare")
+        return arg_defaults.get(name, match.group(0))
+
+    return _ARG_REFERENCE_RE.sub(replace, reference)
+
+
 def _resolve_from_references(content: str) -> str:
     """Substitute ``ARG`` defaults into ``FROM`` lines and drop ``--platform``.
 
@@ -130,9 +148,8 @@ def _resolve_from_references(content: str) -> str:
       as part of the name. Novita builds for its own platform, so the flag is
       dropped rather than propagated.
 
-    Substitutions are applied longest-name-first so a short unbraced prefix
-    cannot corrupt a longer sibling (e.g. ``$BASE`` must not rewrite
-    ``$BASE_IMAGE`` / ``${BASE_IMAGE}`` into ``python:3.12_IMAGE``).
+    Substitutions match whole names in one pass, leaving undeclared names
+    intact even when a declared ARG is a prefix of the reference.
     """
     arg_defaults: Dict[str, str] = {}
     for line in content.split("\n"):
@@ -142,11 +159,6 @@ def _resolve_from_references(content: str) -> str:
         if match:
             arg_defaults[match.group("name")] = match.group("value")
 
-    # Longest names first: `$BASE` must not rewrite `$BASE_IMAGE` mid-token.
-    arg_items = sorted(
-        arg_defaults.items(), key=lambda item: len(item[0]), reverse=True
-    )
-
     out: List[str] = []
     for line in content.split("\n"):
         match = _FROM_RE.match(line)
@@ -156,9 +168,7 @@ def _resolve_from_references(content: str) -> str:
 
         reference = match.group("rest").strip()
         reference = re.sub(r"^(--platform=\S+\s*)+", "", reference).strip()
-        for name, value in arg_items:
-            reference = reference.replace(f"${{{name}}}", value)
-            reference = reference.replace(f"${name}", value)
+        reference = _substitute_args(reference, arg_defaults)
         out.append(f"FROM {reference}")
     return "\n".join(out)
 

tests/test_core/test_novita_provider.py

@@ -902,6 +902,64 @@ def test_arg_after_first_from_not_in_scope(self):
         out = _resolve_from_references("FROM python:3.11\nARG X=alpine\nRUN echo hi\n")
         assert "FROM python:3.11" in out
 
+    @pytest.mark.parametrize("reference", ["$BASE_IMAGE", "${BASE_IMAGE}"])
+    def test_shorter_arg_does_not_consume_a_longer_name(self, reference):
+        from openenv.core.containers.runtime.novita_provider import (
+            _resolve_from_references,
+        )
+
+        # ARG BASE is declared first, so name-by-name substitution used to turn
+        # `FROM $BASE_IMAGE` into `python:3.12_IMAGE`.
+        out = _resolve_from_references(
+            "ARG BASE=python:3.12\n"
+            "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n"
+            f"FROM {reference} AS builder\n"
+        )
+        assert "FROM ghcr.io/huggingface/openenv-base:latest AS builder" in out
+
+    def test_substitution_is_independent_of_declaration_order(self):
+        from openenv.core.containers.runtime.novita_provider import (
+            _resolve_from_references,
+        )
+
+        declarations = [
+            "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest",
+            "ARG BASE=python:3.12",
+        ]
+        outputs = {
+            _resolve_from_references("\n".join([*order, "FROM $BASE_IMAGE\n"]))
+            for order in (declarations, list(reversed(declarations)))
+        }
+        assert outputs == {
+            "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n"
+            "ARG BASE=python:3.12\n"
+            "FROM ghcr.io/huggingface/openenv-base:latest\n",
+            "ARG BASE=python:3.12\n"
+            "ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest\n"
+            "FROM ghcr.io/huggingface/openenv-base:latest\n",
+        }
+
+    @pytest.mark.parametrize(
+        "reference", ["$UNSET_IMAGE", "$BASE_IMAGE", "${BASE_IMAGE}"]
+    )
+    def test_undeclared_arg_reference_is_left_verbatim(self, reference):
+        from openenv.core.containers.runtime.novita_provider import (
+            _resolve_from_references,
+        )
+
+        out = _resolve_from_references(f"ARG BASE=python:3.12\nFROM {reference}\n")
+        assert f"FROM {reference}" in out
+
+    def test_substituted_value_is_not_expanded_again(self):
+        from openenv.core.containers.runtime.novita_provider import (
+            _resolve_from_references,
+        )
+
+        out = _resolve_from_references(
+            "ARG REGISTRY=example.com\nARG IMAGE=$REGISTRY/app\nFROM $IMAGE\n"
+        )
+        assert "FROM $REGISTRY/app" in out
+
     def test_flatten_drops_same_path_copy(self):
         from openenv.core.containers.runtime.novita_provider import _flatten_multistage