feat(admin): add ACR credentials endpoint (#1110)

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",

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

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)

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"])
 

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:

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:

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