From 96db0a1e076b15befebf42746ad4905a65b889a7 Mon Sep 17 00:00:00 2001 From: Issac-Newton <1556820213@qq.com> Date: Tue, 16 Jun 2026 16:10:07 +0800 Subject: [PATCH] feat(admin): add ACR credentials endpoint with temporary token support Move ACR credentials from sandbox proxy router to a dedicated image router at POST /apis/envs/image/v1/generate_registry_credentials. Rename get_acr_credentials to generate_acr_credentials to better reflect the temporary token generation semantics. Co-Authored-By: Claude Opus 4.6 --- pyproject.toml | 1 + rock-conf/rock-local.yml | 7 +++ rock/admin/entrypoints/image_api.py | 23 ++++++++ rock/admin/main.py | 3 ++ rock/config.py | 11 ++++ rock/sandbox/service/sandbox_proxy_service.py | 39 +++++++++++++- tests/unit/sandbox/test_sandbox_proxy.py | 53 ++++++++++++++++++- uv.lock | 15 ++++++ 8 files changed, 150 insertions(+), 2 deletions(-) create mode 100644 rock/admin/entrypoints/image_api.py 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..9fd5ae977a 100644 --- a/rock-conf/rock-local.yml +++ b/rock-conf/rock-local.yml @@ -34,6 +34,13 @@ warmup: # - "reg-a.aliyuncs.com/mirror-1" # - "reg-b.aliyuncs.com/mirror-2" +# ACR credentials (admin-side only, never exposed to SDK) +acr: + instance_id: "" + region: "cn-hangzhou" + access_key_id: "" + access_key_secret: "" + # Scheduler configuration scheduler: enabled: true # Whether to enable the scheduler diff --git a/rock/admin/entrypoints/image_api.py b/rock/admin/entrypoints/image_api.py new file mode 100644 index 0000000000..e94c9715f6 --- /dev/null +++ b/rock/admin/entrypoints/image_api.py @@ -0,0 +1,23 @@ +import asyncio + +from fastapi import APIRouter + +from rock.actions import RockResponse +from rock.common.exception import handle_exceptions +from rock.sandbox.service.sandbox_proxy_service import SandboxProxyService + +image_router = APIRouter() +image_service: SandboxProxyService + + +def set_image_service(service: SandboxProxyService): + global image_service + image_service = service + + +@image_router.post("/generate_registry_credentials") +@handle_exceptions(error_message="generate registry credentials failed") +async def generate_registry_credentials(): + """Return ACR registry credentials with temporary token.""" + result = await asyncio.to_thread(image_service.generate_acr_credentials) + return RockResponse(result=result) diff --git a/rock/admin/main.py b/rock/admin/main.py index a3c7537f40..05224ffd39 100644 --- a/rock/admin/main.py +++ b/rock/admin/main.py @@ -21,6 +21,7 @@ from rock.admin.core.scheduler_task_table import SchedulerTaskTable from rock.admin.entrypoints.admin_ops_api import admin_ops_router, set_ops_service from rock.admin.entrypoints.sandbox_api import sandbox_router, set_sandbox_manager +from rock.admin.entrypoints.image_api import image_router, set_image_service from rock.admin.entrypoints.sandbox_proxy_api import sandbox_proxy_router, set_sandbox_proxy_service from rock.admin.entrypoints.warmup_api import set_warmup_service, warmup_router from rock.admin.gem.api import gem_router, set_env_service @@ -199,6 +200,7 @@ async def lifespan(app: FastAPI): else: sandbox_manager = SandboxProxyService(rock_config=rock_config, meta_store=meta_store) set_sandbox_proxy_service(sandbox_manager) + set_image_service(sandbox_manager) logger.info("rock-admin start") @@ -304,6 +306,7 @@ def main(): app.include_router(admin_ops_router, prefix="/apis/envs/sandbox/v1/ops", tags=["admin-ops"]) else: app.include_router(sandbox_proxy_router, prefix="/apis/envs/sandbox/v1", tags=["sandbox"]) + app.include_router(image_router, prefix="/apis/envs/image/v1", tags=["image"]) app.include_router(warmup_router, prefix="/apis/envs/sandbox/v1", tags=["warmup"]) app.include_router(gem_router, prefix="/apis/v1/envs/gem", tags=["gem"]) diff --git a/rock/config.py b/rock/config.py index 712bde5945..68e456cae6 100644 --- a/rock/config.py +++ b/rock/config.py @@ -152,6 +152,14 @@ def __post_init__(self): self.primary = OssAccountConfig(**self.primary) +@dataclass +class AcrConfig: + instance_id: str | None = None + region: str = "cn-hangzhou" + access_key_id: str = "" + access_key_secret: str = "" + + @dataclass class ProxyServiceConfig: timeout: float = 180.0 @@ -348,6 +356,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 +410,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..ec80fe4e1c 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,15 @@ 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.access_key_id and self.acr_config.instance_id: + self._acr_client = client.AcsClient( + self.acr_config.access_key_id, + self.acr_config.access_key_secret, + self.acr_config.region, + ) + self._batch_get_status_max_count = rock_config.proxy_service.batch_get_status_max_count self._validate_oss_config_or_warn() @@ -746,6 +757,32 @@ def gen_oss_sts_token( "Prefix": prefix, # transfer-object key prefix, scoped per account } + def generate_acr_credentials(self) -> dict | None: + """Return ACR temporary credentials for image push/pull. + + 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 + + request = GetAuthorizationTokenRequest.GetAuthorizationTokenRequest() + request.set_InstanceId(self.acr_config.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 { + "Username": data.get("TempUsername"), + "Password": data.get("AuthorizationToken"), + "Expiration": data.get("ExpireTime"), + } + 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..7fee4f144c 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, 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,53 @@ 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 TestGetAcrCredentials: + @pytest.fixture + def proxy_service(self): + service = SandboxProxyService.__new__(SandboxProxyService) + service.acr_config = AcrConfig( + instance_id="cri-test123", + access_key_id="ak", + access_key_secret="sk", + ) + 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_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.generate_acr_credentials() + + assert result is not None + assert result["Username"] == "tmp-user" + assert result["Password"] == "tmp-pass-token" + assert result["Expiration"] == "2099-01-01T00:15:00Z" + + 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.generate_acr_credentials() + + assert result is None + + def test_no_acr_client_returns_none(self, proxy_service): + proxy_service._acr_client = None + + result = proxy_service.generate_acr_credentials() + + assert result is None 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" },