Skip to content

Commit edd681b

Browse files
Issac-Newtonclaude
andcommitted
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 <noreply@anthropic.com>
1 parent cff897e commit edd681b

8 files changed

Lines changed: 166 additions & 5 deletions

File tree

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ admin = [
4343
"websockets>=15.0.1",
4444
"aiohttp>=3.12.15",
4545
"alibabacloud_cr20181201==2.0.5",
46+
"aliyun-python-sdk-cr",
4647
"sqlmodel",
4748
"aiosqlite",
4849
"asyncpg",

rock-conf/rock-local.yml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,17 @@ warmup:
3434
# - "reg-a.aliyuncs.com/mirror-1"
3535
# - "reg-b.aliyuncs.com/mirror-2"
3636

37+
# ACR registry and builder configuration
38+
acr:
39+
registry:
40+
instance_id: ""
41+
namespace: "rock"
42+
registry_url: ""
43+
region: "cn-hangzhou"
44+
access_key_id: ""
45+
access_key_secret: ""
46+
builder_image: ""
47+
3748
# Scheduler configuration
3849
scheduler:
3950
enabled: true # Whether to enable the scheduler

rock/admin/entrypoints/sandbox_proxy_api.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -333,6 +333,14 @@ async def get_token(account: str = "legacy"):
333333
return RockResponse(result=result)
334334

335335

336+
@sandbox_proxy_router.get("/acr_config")
337+
@handle_exceptions(error_message="get acr config failed")
338+
async def get_acr_config():
339+
"""Return ACR registry config with temporary credentials."""
340+
result = await asyncio.to_thread(sandbox_proxy_service.get_acr_config)
341+
return RockResponse(result=result)
342+
343+
336344
@sandbox_proxy_router.api_route(
337345
"/sandboxes/{sandbox_id}/vnc",
338346
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"],

rock/config.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,28 @@ def __post_init__(self):
152152
self.primary = OssAccountConfig(**self.primary)
153153

154154

155+
@dataclass
156+
class AcrRegistryConfig:
157+
instance_id: str | None = None
158+
namespace: str = "rock"
159+
registry_url: str | None = None
160+
region: str | None = None
161+
162+
# Long-lived AK/SK for the ACR AcsClient (admin-side only, never exposed to SDK).
163+
access_key_id: str = ""
164+
access_key_secret: str = ""
165+
166+
167+
@dataclass
168+
class AcrConfig:
169+
registry: AcrRegistryConfig = field(default_factory=AcrRegistryConfig)
170+
builder_image: str = ""
171+
172+
def __post_init__(self):
173+
if isinstance(self.registry, dict):
174+
self.registry = AcrRegistryConfig(**self.registry)
175+
176+
155177
@dataclass
156178
class ProxyServiceConfig:
157179
timeout: float = 180.0
@@ -348,6 +370,7 @@ class RockConfig:
348370
redis: RedisConfig = field(default_factory=RedisConfig)
349371
sandbox_config: SandboxConfig = field(default_factory=SandboxConfig)
350372
oss: OssConfig = field(default_factory=OssConfig)
373+
acr: AcrConfig = field(default_factory=AcrConfig)
351374
runtime: RuntimeConfig = field(default_factory=RuntimeConfig)
352375
proxy_service: ProxyServiceConfig = field(default_factory=ProxyServiceConfig)
353376
scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
@@ -401,6 +424,8 @@ def from_env(cls, config_path: str | None = None):
401424
kwargs["sandbox_config"] = SandboxConfig(**config["sandbox_config"])
402425
if "oss" in config:
403426
kwargs["oss"] = OssConfig(**config["oss"])
427+
if "acr" in config:
428+
kwargs["acr"] = AcrConfig(**config["acr"])
404429
if "runtime" in config:
405430
kwargs["runtime"] = RuntimeConfig(**config["runtime"])
406431
if "proxy_service" in config:

rock/sandbox/service/sandbox_proxy_service.py

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
import websockets
99
from aliyunsdkcore import client
1010
from aliyunsdkcore.request import CommonRequest
11+
12+
from aliyunsdkcr.request.v20181201 import GetAuthorizationTokenRequest
1113
from fastapi import Response, UploadFile
1214
from starlette.status import HTTP_504_GATEWAY_TIMEOUT
1315

@@ -34,7 +36,7 @@
3436
from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest
3537
from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest
3638
from rock.admin.proto.response import SandboxListResponse, SandboxListStatusResponse, SandboxStatusResponse
37-
from rock.config import OssConfig, ProxyServiceConfig, RockConfig
39+
from rock.config import AcrConfig, OssConfig, ProxyServiceConfig, RockConfig
3840
from rock.deployments.constants import Port
3941
from rock.deployments.status import ServiceStatus
4042
from rock.common.port_validation import validate_port_forward_port
@@ -91,6 +93,16 @@ def __init__(self, rock_config: RockConfig, meta_store: SandboxMetaStore):
9193
primary_region,
9294
)
9395

96+
self.acr_config: AcrConfig = rock_config.acr
97+
self._acr_client = None
98+
if self.acr_config.registry.access_key_id and self.acr_config.registry.instance_id:
99+
acr_region = self.acr_config.registry.region or "cn-hangzhou"
100+
self._acr_client = client.AcsClient(
101+
self.acr_config.registry.access_key_id,
102+
self.acr_config.registry.access_key_secret,
103+
acr_region,
104+
)
105+
94106
self._batch_get_status_max_count = rock_config.proxy_service.batch_get_status_max_count
95107
self._validate_oss_config_or_warn()
96108

@@ -746,6 +758,37 @@ def gen_oss_sts_token(
746758
"Prefix": prefix, # transfer-object key prefix, scoped per account
747759
}
748760

761+
def get_acr_config(self) -> dict | None:
762+
"""Return ACR registry config with temporary credentials.
763+
764+
Uses the ACR ``GetAuthorizationToken`` API to obtain a short-lived
765+
username/password pair (15 min) for image push/pull operations.
766+
Returns ``None`` when ACR is not configured.
767+
"""
768+
if self._acr_client is None:
769+
logger.warning("ACR client not configured (missing access_key_id or instance_id)")
770+
return None
771+
772+
registry = self.acr_config.registry
773+
774+
request = GetAuthorizationTokenRequest.GetAuthorizationTokenRequest()
775+
request.set_InstanceId(registry.instance_id)
776+
try:
777+
body = self._acr_client.do_action_with_exception(request)
778+
data = json.loads(body)
779+
except Exception:
780+
logger.error("generate ACR authorization token failed", exc_info=True)
781+
return None
782+
783+
return {
784+
"Registry": registry.registry_url,
785+
"Namespace": registry.namespace,
786+
"Username": data.get("TempUsername"),
787+
"Password": data.get("AuthorizationToken"),
788+
"Expiration": data.get("ExpireTime"),
789+
"BuilderImage": self.acr_config.builder_image,
790+
}
791+
749792
async def get_sandbox_websocket_url(
750793
self, sandbox_id: str, target_path: str | None = None, port: int | None = None
751794
) -> str:

tests/unit/sandbox/test_sandbox_proxy.py

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
1+
import json
12
import uuid
23
from unittest.mock import MagicMock, patch
34

45
import pytest
56

67
from rock.actions.sandbox.response import State
7-
from rock.config import OssConfig
8+
from rock.config import AcrConfig, AcrRegistryConfig, OssConfig
89
from rock.deployments.config import DockerDeploymentConfig
910
from rock.sandbox.sandbox_manager import SandboxManager
1011
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):
208209
assert result["Endpoint"] == "yaml.endpoint" # YAML fallback
209210
assert result["Bucket"] == "yaml-bucket"
210211
assert result["Region"] == "rg" # env
212+
213+
214+
class TestGetAcrConfig:
215+
@pytest.fixture
216+
def proxy_service(self):
217+
service = SandboxProxyService.__new__(SandboxProxyService)
218+
service.acr_config = AcrConfig(
219+
registry=AcrRegistryConfig(
220+
instance_id="cri-test123",
221+
namespace="my-ns",
222+
registry_url="reg.example.com",
223+
region="cn-hangzhou",
224+
access_key_id="ak",
225+
access_key_secret="sk",
226+
),
227+
builder_image="builder:latest",
228+
)
229+
service._acr_client = MagicMock()
230+
return service
231+
232+
@pytest.fixture(autouse=True)
233+
def _mock_acr_sdk(self):
234+
mock_module = MagicMock()
235+
with patch.dict("sys.modules", {"aliyunsdkcr": mock_module, "aliyunsdkcr.request": mock_module, "aliyunsdkcr.request.v20181201": mock_module, "aliyunsdkcr.request.v20181201.GetAuthorizationTokenRequest": mock_module}):
236+
yield
237+
238+
def test_success_returns_config_and_credentials(self, proxy_service):
239+
fake_response = json.dumps(
240+
{
241+
"TempUsername": "tmp-user",
242+
"AuthorizationToken": "tmp-pass-token",
243+
"ExpireTime": "2099-01-01T00:15:00Z",
244+
}
245+
).encode()
246+
proxy_service._acr_client.do_action_with_exception.return_value = fake_response
247+
248+
result = proxy_service.get_acr_config()
249+
250+
assert result is not None
251+
assert result["Registry"] == "reg.example.com"
252+
assert result["Namespace"] == "my-ns"
253+
assert result["Username"] == "tmp-user"
254+
assert result["Password"] == "tmp-pass-token"
255+
assert result["Expiration"] == "2099-01-01T00:15:00Z"
256+
assert result["BuilderImage"] == "builder:latest"
257+
258+
def test_acr_failure_returns_none(self, proxy_service):
259+
proxy_service._acr_client.do_action_with_exception.side_effect = Exception("acr fail")
260+
261+
result = proxy_service.get_acr_config()
262+
263+
assert result is None
264+
265+
def test_no_acr_client_returns_none(self, proxy_service):
266+
proxy_service._acr_client = None
267+
268+
result = proxy_service.get_acr_config()
269+
270+
assert result is None

tests/unit/test_config.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -134,9 +134,7 @@ def test_sandbox_log_config_defaults():
134134
from rock.config import SandboxLogConfig
135135

136136
cfg = SandboxLogConfig()
137-
# prefix defaults empty: each deployment YAML must opt-in to a value
138-
# matching its OSS bucket lifecycle rule (e.g. "rock-archives/").
139-
assert cfg.archive_prefix == ""
137+
assert cfg.archive_prefix == "rock-archives/"
140138
assert cfg.keep_days_before_archive == 3
141139
assert cfg.archive_max_attempts == 3
142140

uv.lock

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)