Fix runner download from Gitea CDN

controller.py

@@ -26,6 +26,7 @@ RUNNER_URL = (
     "https://dl.gitea.com/gitea-runner/3.5.0/gitea-runner-3.5.0-linux-amd64"
 )
 RUNNER_SHA256 = "16dc5c89ad5ac5496369642f01ed29c9bb130a2512e16a8cc3b400d9d0bfea73"
+RUNNER_USER_AGENT = "ces-autodl-controller/1.0"
 RUNNER_LABELS = ["ubuntu-latest:host", "ces-ci:host", "gpu:host"]
 MODEL_FILES = [
     "config.json",
@@ -256,7 +257,8 @@ def ensure_runner(paths: Paths) -> Path:
     if paths.runner_cache.is_file() and sha256(paths.runner_cache) == RUNNER_SHA256:
         return paths.runner_cache
     paths.runner_cache.parent.mkdir(parents=True, exist_ok=True)
-    with urllib.request.urlopen(RUNNER_URL, timeout=300) as response, tempfile.NamedTemporaryFile(
+    request = urllib.request.Request(RUNNER_URL, headers={"User-Agent": RUNNER_USER_AGENT})
+    with urllib.request.urlopen(request, timeout=300) as response, tempfile.NamedTemporaryFile(
         dir=paths.runner_cache.parent, delete=False
     ) as temporary:
         while chunk := response.read(1024 * 1024):

tests/test_controller.py

@@ -1,13 +1,40 @@
+import hashlib
+import io
 import json
 import os
 import tempfile
 import unittest
 from pathlib import Path
+from unittest import mock
 
 import controller
 
 
 class ControllerTests(unittest.TestCase):
+    def test_runner_download_uses_explicit_user_agent(self):
+        payload = b"test-runner"
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            paths = controller.Paths(
+                root / "model",
+                "/model",
+                "/artifacts",
+                root / "artifacts",
+                root / "key",
+                root / "known_hosts",
+                root / "state.json",
+                root / "runner",
+            )
+            with mock.patch.object(
+                controller, "RUNNER_SHA256", hashlib.sha256(payload).hexdigest()
+            ), mock.patch.object(
+                controller.urllib.request, "urlopen", return_value=io.BytesIO(payload)
+            ) as urlopen:
+                controller.ensure_runner(paths)
+
+            request = urlopen.call_args.args[0]
+            self.assertEqual(request.get_header("User-agent"), controller.RUNNER_USER_AGENT)
+
     def test_list_images_accepts_null_for_an_empty_account(self):
         client = controller.AutoDLClient("test-token")
         client.request = lambda method, path, body: {"list": None}