|
8 | 8 | import websockets |
9 | 9 | from aliyunsdkcore import client |
10 | 10 | from aliyunsdkcore.request import CommonRequest |
| 11 | + |
| 12 | +from aliyunsdkcr.request.v20181201 import GetAuthorizationTokenRequest |
11 | 13 | from fastapi import Response, UploadFile |
12 | 14 | from starlette.status import HTTP_504_GATEWAY_TIMEOUT |
13 | 15 |
|
|
34 | 36 | from rock.admin.proto.request import SandboxReadFileRequest as ReadFileRequest |
35 | 37 | from rock.admin.proto.request import SandboxWriteFileRequest as WriteFileRequest |
36 | 38 | 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 |
38 | 40 | from rock.deployments.constants import Port |
39 | 41 | from rock.deployments.status import ServiceStatus |
40 | 42 | from rock.common.port_validation import validate_port_forward_port |
@@ -91,6 +93,16 @@ def __init__(self, rock_config: RockConfig, meta_store: SandboxMetaStore): |
91 | 93 | primary_region, |
92 | 94 | ) |
93 | 95 |
|
| 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 | + |
94 | 106 | self._batch_get_status_max_count = rock_config.proxy_service.batch_get_status_max_count |
95 | 107 | self._validate_oss_config_or_warn() |
96 | 108 |
|
@@ -746,6 +758,37 @@ def gen_oss_sts_token( |
746 | 758 | "Prefix": prefix, # transfer-object key prefix, scoped per account |
747 | 759 | } |
748 | 760 |
|
| 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 | + |
749 | 792 | async def get_sandbox_websocket_url( |
750 | 793 | self, sandbox_id: str, target_path: str | None = None, port: int | None = None |
751 | 794 | ) -> str: |
|
0 commit comments