tenseleyflow/documentlanguagemodel / aab55d5

Browse files

fix(export): validate ollama model name before subprocess (audit-04 F5)

Authored by espadonne
SHA
aab55d558f97a1b939b02e3177e5f835724f74e6
Parents
e7925ff
Tree
aa6742e

2 changed files

StatusFile+-
M src/dlm/export/ollama/register.py 35 0
M tests/unit/export/ollama/test_register.py 86 0
src/dlm/export/ollama/register.pymodified
@@ -19,6 +19,7 @@ Modelfile (`./base.gguf`, `./adapter.gguf`) resolve. 10-minute timeout.
1919
 from __future__ import annotations
2020
2121
 import os
22
+import re
2223
 import subprocess  # nosec B404
2324
 from pathlib import Path
2425
@@ -29,6 +30,39 @@ from dlm.store.lock import exclusive
2930
 _DEFAULT_TIMEOUT_SECONDS = 600.0
3031
 _DEFAULT_LOCK_TIMEOUT_SECONDS = 120.0
3132
33
+# Conservative match for `[namespace/]name[:tag]` per Ollama's model-name
34
+# grammar — alphanumerics, `.`, `-`, `_`, with optional single namespace
35
+# segment and optional tag. Rejects shell metacharacters, whitespace,
36
+# leading dots, and traversal (`..`) before Ollama ever sees them so we
37
+# fail fast with a useful error instead of a mid-registry subprocess crash.
38
+_NAME_SEGMENT = r"[a-zA-Z0-9][a-zA-Z0-9._-]*"
39
+_NAME_RE = re.compile(rf"^(?:{_NAME_SEGMENT}/)?{_NAME_SEGMENT}(?::{_NAME_SEGMENT})?$")
40
+_MAX_NAME_LEN = 128
41
+
42
+
43
+def _validate_name(name: str) -> None:
44
+    if not name or len(name) > _MAX_NAME_LEN:
45
+        raise OllamaCreateError(
46
+            stdout="",
47
+            stderr=(
48
+                f"ollama model name is empty or exceeds {_MAX_NAME_LEN} chars; "
49
+                "use `[namespace/]name[:tag]` with alphanumerics, `.`, `-`, `_`."
50
+            ),
51
+        )
52
+    if ".." in name:
53
+        raise OllamaCreateError(
54
+            stdout="",
55
+            stderr=f"ollama model name {name!r} contains `..` (path traversal).",
56
+        )
57
+    if not _NAME_RE.match(name):
58
+        raise OllamaCreateError(
59
+            stdout="",
60
+            stderr=(
61
+                f"ollama model name {name!r} does not match "
62
+                "`[namespace/]name[:tag]` with alphanumerics, `.`, `-`, `_`."
63
+            ),
64
+        )
65
+
3266
3367
 def ollama_lock_path(dlm_home: Path | None = None) -> Path:
3468
     """Return `~/.dlm/ollama.lock` (override via `dlm_home` for tests).
@@ -58,6 +92,7 @@ def ollama_create(
5892
5993
     `binary` / `dlm_home` are test hooks.
6094
     """
95
+    _validate_name(name)
6196
     exe = binary or locate_ollama()
6297
     lock_path = ollama_lock_path(dlm_home)
6398
tests/unit/export/ollama/test_register.pymodified
@@ -130,6 +130,92 @@ class TestOllamaCreate:
130130
         assert lock_observed["existed"]
131131
132132
133
+class TestNameValidation:
134
+    """`name` is validated before subprocess runs — audit-04 F5."""
135
+
136
+    @pytest.mark.parametrize(
137
+        "name",
138
+        [
139
+            "dlm",
140
+            "dlm-01test",
141
+            "dlm-01test:v0001",
142
+            "user/dlm",
143
+            "user/dlm:latest",
144
+            "a.b.c",
145
+            "A",
146
+        ],
147
+    )
148
+    def test_valid_names_pass(self, tmp_path: Path, name: str) -> None:
149
+        exe = tmp_path / "ollama"
150
+        exe.write_text("")
151
+        modelfile = tmp_path / "Modelfile"
152
+        modelfile.write_text("x")
153
+        with patch(
154
+            "dlm.export.ollama.register.subprocess.run",
155
+            return_value=_ok_proc(),
156
+        ):
157
+            ollama_create(
158
+                name=name,
159
+                modelfile_path=modelfile,
160
+                cwd=tmp_path,
161
+                binary=exe,
162
+                dlm_home=tmp_path,
163
+            )
164
+
165
+    @pytest.mark.parametrize(
166
+        "name",
167
+        [
168
+            "",
169
+            "foo; rm -rf /",
170
+            "foo bar",
171
+            "foo\nbar",
172
+            "../evil",
173
+            "../../etc/passwd",
174
+            "$(whoami)",
175
+            "`whoami`",
176
+            ".hidden",
177
+            "-dashstart",
178
+            "a/b/c",  # only one namespace segment allowed
179
+            "a:b:c",  # only one tag segment allowed
180
+        ],
181
+    )
182
+    def test_invalid_names_raise_before_subprocess(self, tmp_path: Path, name: str) -> None:
183
+        exe = tmp_path / "ollama"
184
+        exe.write_text("")
185
+        modelfile = tmp_path / "Modelfile"
186
+        modelfile.write_text("x")
187
+        with (
188
+            patch("dlm.export.ollama.register.subprocess.run") as mock_run,
189
+            pytest.raises(OllamaCreateError),
190
+        ):
191
+            ollama_create(
192
+                name=name,
193
+                modelfile_path=modelfile,
194
+                cwd=tmp_path,
195
+                binary=exe,
196
+                dlm_home=tmp_path,
197
+            )
198
+        assert mock_run.call_count == 0
199
+
200
+    def test_overlong_name_rejected(self, tmp_path: Path) -> None:
201
+        exe = tmp_path / "ollama"
202
+        exe.write_text("")
203
+        modelfile = tmp_path / "Modelfile"
204
+        modelfile.write_text("x")
205
+        with (
206
+            patch("dlm.export.ollama.register.subprocess.run") as mock_run,
207
+            pytest.raises(OllamaCreateError, match="exceeds"),
208
+        ):
209
+            ollama_create(
210
+                name="a" * 200,
211
+                modelfile_path=modelfile,
212
+                cwd=tmp_path,
213
+                binary=exe,
214
+                dlm_home=tmp_path,
215
+            )
216
+        assert mock_run.call_count == 0
217
+
218
+
133219
 class TestLockContention:
134220
     def test_held_lock_blocks_second_caller(self, tmp_path: Path) -> None:
135221
         """A second caller times out while the store-lock holds the file."""