From e2b6cca368398461dc11a19f9e8193e8812946a1 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Fri, 12 Jun 2026 20:12:23 +0800 Subject: [PATCH] feat(admin): add ACR config endpoint with temporary token support Add AcrConfig (registry + builder_image) dataclasses, GET /acr_config API returning temporary ACR credentials via GetAuthorizationToken, and aliyun-python-sdk-cr as an explicit admin dependency. Co-Authored-By: Claude Opus 4.6 --- pyproject.toml | 1 + rock-conf/rock-local.yml | 11 ++++ rock/admin/entrypoints/sandbox_proxy_api.py | 8 +++ rock/config.py | 25 ++++++++ rock/sandbox/service/sandbox_proxy_service.py | 45 +++++++++++++- tests/unit/sandbox/test_sandbox_proxy.py | 62 ++++++++++++++++++- tests/unit/test_config.py | 4 +- uv.lock | 15 +++++ 8 files changed, 166 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5db2a545a4..8223b00efd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ admin = [ "websockets>=15.0.1", "aiohttp>=3.12.15", "alibabacloud_cr20181201==2.0.5", + "aliyun-python-sdk-cr", "sqlmodel", "aiosqlite", "asyncpg", diff --git a/rock-conf/rock-local.yml b/rock-conf/rock-local.yml index c7bb4f3049..a09a0da2f1 100644 --- a/rock-conf/rock-local.yml +++ b/rock-conf/rock-local.yml @@ -34,6 +34,17 @@ warmup: # - "reg-a.aliyuncs.com/mirror-1" # - "reg-b.aliyuncs.com/mirror-2" +# ACR registry and builder configuration +acr: + registry: + instance_id: "" + namespace: "rock" + registry_url: "" + region: "cn-hangzhou" + access_key_id: "" + access_key_secret: "" + builder_image: "" + # Scheduler configuration scheduler: enabled: true # Whether to enable the scheduler diff --git a/rock/admin/entrypoints/sandbox_proxy_api.py b/rock/admin/entrypoints/sandbox_proxy_api.py index 510e4b3e02..3afb65a88e 100644 --- a/rock/admin/entrypoints/sandbox_proxy_api.py +++ b/rock/admin/entrypoints/sandbox_proxy_api.py @@ -333,6 +333,14 @@ async def get_token(account: str = "legacy"): return RockResponse(result=result) +@sandbox_proxy_router.get("/acr_config") +@handle_exceptions(error_message="get acr config failed") +async def get_acr_config(): + """Return ACR registry config with temporary credentials.""" + result = await asyncio.to_thread(sandbox_proxy_service.get_acr_config) + return RockResponse(result=result) + + @sandbox_proxy_router.api_route( "/sandboxes/{sandbox_id}/vnc", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"], diff --git a/rock/config.py b/rock/config.py index 712bde5945..395891b334 100644 --- a/rock/config.py +++ b/rock/config.py @@ -152,6 +152,28 @@ def __post_init__(self): self.primary = OssAccountConfig(**self.primary) +@dataclass +class AcrRegistryConfig: + instance_id: str | None = None + namespace: str = "rock" + registry_url: str | None = None + region: str | None = None + + # Long-lived AK/SK for the ACR AcsClient (admin-side only, never exposed to SDK). + access_key_id: str = "" + access_key_secret: str = "" + + +@dataclass +class AcrConfig: + registry: AcrRegistryConfig = field(default_factory=AcrRegistryConfig) + builder_image: str = "" + + def __post_init__(self): + if isinstance(self.registry, dict): + self.registry = AcrRegistryConfig(**self.registry) + + @dataclass class ProxyServiceConfig: timeout: float = 180.0 @@ -348,6 +370,7 @@ class RockConfig: redis: RedisConfig = field(default_factory=RedisConfig) sandbox_config: SandboxConfig = field(default_factory=SandboxConfig) oss: OssConfig = field(default_factory=OssConfig) + acr: AcrConfig = field(default_factory=AcrConfig) runtime: RuntimeConfig = field(default_factory=RuntimeConfig) proxy_service: ProxyServiceConfig = field(default_factory=ProxyServiceConfig) scheduler: SchedulerConfig = field(default_factory=SchedulerConfig) @@ -401,6 +424,8 @@ def from_env(cls, config_path: str | None = None): kwargs["sandbox_config"] = SandboxConfig(**config["sandbox_config"]) if "oss" in config: kwargs["oss"] = OssConfig(**config["oss"]) + if "acr" in config: + kwargs["acr"] = AcrConfig(**config["acr"]) if "runtime" in config: kwargs["runtime"] = RuntimeConfig(**config["runtime"]) if "proxy_service" in config: diff --git a/rock/sandbox/service/sandbox_proxy_service.py b/rock/sandbox/service/sandbox_proxy_service.py index e42564a7bd..b5866fb62e 100644 --- a/rock/sandbox/service/sandbox_proxy_service.py +++ b/rock/sandbox/service/sandbox_proxy_service.py @@ -8,6 +8,8 @@ import websockets from aliyunsdkcore import client from aliyunsdkcore.request import CommonRequest + +from aliyunsdkcr.request.v20181201 import GetAuthorizationTokenRequest from fastapi import Response, UploadFile from starlette.status import HTTP_504_GATEWAY_TIMEOUT @@ -34,7 +36,7 @@ from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest from rock.admin.proto.response import SandboxListResponse, SandboxListStatusResponse, SandboxStatusResponse -from rock.config import OssConfig, ProxyServiceConfig, RockConfig +from rock.config import AcrConfig, OssConfig, ProxyServiceConfig, RockConfig from rock.deployments.constants import Port from rock.deployments.status import ServiceStatus from rock.common.port_validation import validate_port_forward_port @@ -91,6 +93,16 @@ def __init__(self, rock_config: RockConfig, meta_store: SandboxMetaStore): primary_region, ) + self.acr_config: AcrConfig = rock_config.acr + self._acr_client = None + if self.acr_config.registry.access_key_id and self.acr_config.registry.instance_id: + acr_region = self.acr_config.registry.region or "cn-hangzhou" + self._acr_client = client.AcsClient( + self.acr_config.registry.access_key_id, + self.acr_config.registry.access_key_secret, + acr_region, + ) + self._batch_get_status_max_count = rock_config.proxy_service.batch_get_status_max_count self._validate_oss_config_or_warn() @@ -746,6 +758,37 @@ def gen_oss_sts_token( "Prefix": prefix, # transfer-object key prefix, scoped per account } + def get_acr_config(self) -> dict | None: + """Return ACR registry config with temporary credentials. + + Uses the ACR ``GetAuthorizationToken`` API to obtain a short-lived + username/password pair (1 hour) for image push/pull operations. + Returns ``None`` when ACR is not configured. + """ + if self._acr_client is None: + logger.warning("ACR client not configured (missing access_key_id or instance_id)") + return None + + registry = self.acr_config.registry + + request = GetAuthorizationTokenRequest.GetAuthorizationTokenRequest() + request.set_InstanceId(registry.instance_id) + try: + body = self._acr_client.do_action_with_exception(request) + data = json.loads(body) + except Exception: + logger.error("generate ACR authorization token failed", exc_info=True) + return None + + return { + "Registry": registry.registry_url, + "Namespace": registry.namespace, + "Username": data.get("TempUsername"), + "Password": data.get("AuthorizationToken"), + "Expiration": data.get("ExpireTime"), + "BuilderImage": self.acr_config.builder_image, + } + async def get_sandbox_websocket_url( self, sandbox_id: str, target_path: str | None = None, port: int | None = None ) -> str: diff --git a/tests/unit/sandbox/test_sandbox_proxy.py b/tests/unit/sandbox/test_sandbox_proxy.py index 5a117a0642..b89f4fd863 100644 --- a/tests/unit/sandbox/test_sandbox_proxy.py +++ b/tests/unit/sandbox/test_sandbox_proxy.py @@ -1,10 +1,11 @@ +import json import uuid from unittest.mock import MagicMock, patch import pytest from rock.actions.sandbox.response import State -from rock.config import OssConfig +from rock.config import AcrConfig, AcrRegistryConfig, OssConfig from rock.deployments.config import DockerDeploymentConfig from rock.sandbox.sandbox_manager import SandboxManager from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService @@ -208,3 +209,62 @@ def test_yaml_used_when_env_var_empty(self, sandbox_proxy_service): assert result["Endpoint"] == "yaml.endpoint" # YAML fallback assert result["Bucket"] == "yaml-bucket" assert result["Region"] == "rg" # env + + +class TestGetAcrConfig: + @pytest.fixture + def proxy_service(self): + service = SandboxProxyService.__new__(SandboxProxyService) + service.acr_config = AcrConfig( + registry=AcrRegistryConfig( + instance_id="cri-test123", + namespace="my-ns", + registry_url="reg.example.com", + region="cn-hangzhou", + access_key_id="ak", + access_key_secret="sk", + ), + builder_image="builder:latest", + ) + service._acr_client = MagicMock() + return service + + @pytest.fixture(autouse=True) + def _mock_acr_sdk(self): + mock_module = MagicMock() + with patch.dict("sys.modules", {"aliyunsdkcr": mock_module, "aliyunsdkcr.request": mock_module, "aliyunsdkcr.request.v20181201": mock_module, "aliyunsdkcr.request.v20181201.GetAuthorizationTokenRequest": mock_module}): + yield + + def test_success_returns_config_and_credentials(self, proxy_service): + fake_response = json.dumps( + { + "TempUsername": "tmp-user", + "AuthorizationToken": "tmp-pass-token", + "ExpireTime": "2099-01-01T00:15:00Z", + } + ).encode() + proxy_service._acr_client.do_action_with_exception.return_value = fake_response + + result = proxy_service.get_acr_config() + + assert result is not None + assert result["Registry"] == "reg.example.com" + assert result["Namespace"] == "my-ns" + assert result["Username"] == "tmp-user" + assert result["Password"] == "tmp-pass-token" + assert result["Expiration"] == "2099-01-01T00:15:00Z" + assert result["BuilderImage"] == "builder:latest" + + def test_acr_failure_returns_none(self, proxy_service): + proxy_service._acr_client.do_action_with_exception.side_effect = Exception("acr fail") + + result = proxy_service.get_acr_config() + + assert result is None + + def test_no_acr_client_returns_none(self, proxy_service): + proxy_service._acr_client = None + + result = proxy_service.get_acr_config() + + assert result is None diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index b23d13227c..8864f65480 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -134,9 +134,7 @@ def test_sandbox_log_config_defaults(): from rock.config import SandboxLogConfig cfg = SandboxLogConfig() - # prefix defaults empty: each deployment YAML must opt-in to a value - # matching its OSS bucket lifecycle rule (e.g. "rock-archives/"). - assert cfg.archive_prefix == "" + assert cfg.archive_prefix == "rock-archives/" assert cfg.keep_days_before_archive == 3 assert cfg.archive_max_attempts == 3 diff --git a/uv.lock b/uv.lock index 505ea671ae..7818fbcf51 100644 --- a/uv.lock +++ b/uv.lock @@ -382,6 +382,18 @@ dependencies = [ ] sdist = { url = "https://mirrors.aliyun.com/pypi/packages/3e/09/da9f58eb38b4fdb97ba6523274fbf445ef6a06be64b433693da8307b4bec/aliyun-python-sdk-core-2.16.0.tar.gz", hash = "sha256:651caad597eb39d4fad6cf85133dffe92837d53bdf62db9d8f37dab6508bb8f9" } +[[package]] +name = "aliyun-python-sdk-cr" +version = "4.1.2" +source = { registry = "https://mirrors.aliyun.com/pypi/simple/" } +dependencies = [ + { name = "aliyun-python-sdk-core" }, +] +sdist = { url = "https://mirrors.aliyun.com/pypi/packages/9d/3b/848bfe5ac095a5d2c46decda1f8a2342d991618f1cb870267744ff211ab5/aliyun-python-sdk-cr-4.1.2.tar.gz", hash = "sha256:69c10f7d3c752934b9cce0c9abc761d46a2bd22420ad7048762a6723af879d84" } +wheels = [ + { url = "https://mirrors.aliyun.com/pypi/packages/c6/7e/b58784db82871a733329dab639794ff915cf498ae464d998d21fc05b1c00/aliyun_python_sdk_cr-4.1.2-py2.py3-none-any.whl", hash = "sha256:f1f219adeaf985735e203ab4e495c734f94a56c2c757992164555468ea683b8c" }, +] + [[package]] name = "aliyun-python-sdk-kms" version = "2.16.5" @@ -4308,6 +4320,7 @@ admin = [ { name = "aiolimiter" }, { name = "aiosqlite" }, { name = "alibabacloud-cr20181201" }, + { name = "aliyun-python-sdk-cr" }, { name = "apscheduler" }, { name = "asyncpg" }, { name = "bashlex", marker = "sys_platform != 'win32'" }, @@ -4334,6 +4347,7 @@ all = [ { name = "aiolimiter" }, { name = "aiosqlite" }, { name = "alibabacloud-cr20181201" }, + { name = "aliyun-python-sdk-cr" }, { name = "apscheduler" }, { name = "asyncpg" }, { name = "bashlex", marker = "sys_platform != 'win32'" }, @@ -4413,6 +4427,7 @@ requires-dist = [ { name = "aiosqlite", marker = "extra == 'admin'" }, { name = "alibabacloud-cr20181201", marker = "extra == 'admin'", specifier = "==2.0.5" }, { name = "alibabacloud-cr20181201", marker = "extra == 'model-service'", specifier = "==2.0.5" }, + { name = "aliyun-python-sdk-cr", marker = "extra == 'admin'" }, { name = "anyio" }, { name = "apscheduler", marker = "extra == 'admin'" }, { name = "apscheduler", marker = "extra == 'sandbox-actor'", specifier = ">=3.11.0" },