129 lines
5.2 KiB
Python
129 lines
5.2 KiB
Python
"""Resolve verified public-image mirrors without changing Docker's global registry settings."""
|
||
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
from urllib.parse import urlsplit
|
||
|
||
from docker_command import safe_line
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
DEFAULT_MIRROR_PREFIX = "m.daocloud.io"
|
||
MANIFEST_TIMEOUT = 30
|
||
PULL_TIMEOUT = 180
|
||
PUBLIC_REPOSITORIES = {
|
||
"docker.io/library/python", "docker.io/library/postgres", "docker.io/library/redis",
|
||
"ghcr.io/astral-sh/uv",
|
||
}
|
||
|
||
|
||
class MirrorError(RuntimeError):
|
||
pass
|
||
|
||
|
||
def build_images(environment: dict[str, str]) -> dict[str, str]:
|
||
defaults = dict(re.findall(
|
||
r"^ARG (PYTHON_IMAGE|UV_IMAGE)=(.+)$", (ROOT / "aiprovider/Dockerfile").read_text(), re.M
|
||
))
|
||
return {key: environment.get(key) or value for key, value in defaults.items()}
|
||
|
||
|
||
def mirror_prefix(environment: dict[str, str]) -> str:
|
||
prefix = environment.get("PLANET_DOCKER_MIRROR_PREFIX", DEFAULT_MIRROR_PREFIX).rstrip("/")
|
||
if not re.fullmatch(r"[a-zA-Z0-9.-]+(?::[0-9]+)?(?:/[a-zA-Z0-9._-]+)*", prefix):
|
||
raise MirrorError("备用镜像地址须为 registry 或 registry/path，不得包含凭据或 URL 参数")
|
||
return prefix
|
||
|
||
|
||
def canonical_image(image: str) -> str:
|
||
first = image.split("/")[0]
|
||
if "/" not in image or not ("." in first or ":" in first or first == "localhost"):
|
||
image = "docker.io/" + ("library/" if "/" not in image else "") + image
|
||
if image.startswith("index.docker.io/"):
|
||
image = "docker.io/" + image.removeprefix("index.docker.io/")
|
||
if image.startswith("docker.io/") and image.count("/") == 1:
|
||
image = "docker.io/library/" + image.removeprefix("docker.io/")
|
||
return image
|
||
|
||
|
||
def mirror_image(image: str, environment: dict[str, str]) -> str:
|
||
canonical = canonical_image(image)
|
||
repository = canonical.split("@")[0].rsplit("/", 1)
|
||
repository = repository[0] + "/" + repository[1].split(":")[0]
|
||
if repository not in PUBLIC_REPOSITORIES:
|
||
raise MirrorError("仅为 Planet 的公开基础镜像回退；自定义或私有镜像须配置自身仓库")
|
||
return mirror_prefix(environment) + "/" + canonical
|
||
|
||
|
||
def mirror_registry_urls(environment: dict[str, str]) -> list[str]:
|
||
for image in build_images(environment).values():
|
||
mirror_image(image, environment)
|
||
return ["https://" + urlsplit("https://" + mirror_prefix(environment)).netloc + "/v2/"]
|
||
|
||
|
||
def run(command: list[str], timeout: int = MANIFEST_TIMEOUT) -> str:
|
||
result = subprocess.run(command, capture_output=True, text=True, timeout=timeout, check=False)
|
||
if result.returncode:
|
||
detail = safe_line(result.stderr.strip() or result.stdout.strip())
|
||
raise MirrorError(detail[-600:] or "备用镜像命令失败")
|
||
return result.stdout
|
||
|
||
|
||
def resolve_image(image: str, environment: dict[str, str]) -> str:
|
||
mirrored = mirror_image(image, environment)
|
||
manifest = json.loads(run([
|
||
"docker", "buildx", "imagetools", "inspect", mirrored, "--format", "{{json .Manifest}}"
|
||
]))
|
||
digest = manifest.get("digest") if isinstance(manifest, dict) else None
|
||
if not isinstance(digest, str) or not re.fullmatch(r"sha256:[0-9a-f]{64}", digest):
|
||
raise MirrorError("备用镜像未返回有效的 manifest digest")
|
||
if "@" in image and digest != image.rsplit("@", 1)[1]:
|
||
raise MirrorError("备用镜像 digest 与指定镜像不一致")
|
||
repository = mirrored.split("@")[0]
|
||
if ":" in repository.rsplit("/", 1)[1]:
|
||
repository = repository.rsplit(":", 1)[0]
|
||
return repository + "@" + digest
|
||
|
||
|
||
def pull_services(services: list[str], environment: dict[str, str]) -> None:
|
||
if not services or any(service not in {"postgres", "redis"} for service in services):
|
||
raise MirrorError("备用拉取只处理本次请求的 PostgreSQL / Redis 服务")
|
||
configuration = json.loads(run(["docker", "compose", "config", "--format", "json"]))
|
||
for service in services:
|
||
image = configuration["services"][service]["image"]
|
||
local = subprocess.run(
|
||
["docker", "image", "inspect", image], capture_output=True, timeout=MANIFEST_TIMEOUT
|
||
)
|
||
if local.returncode == 0:
|
||
continue
|
||
resolved = resolve_image(image, environment)
|
||
print(f"备用源拉取 {service}", flush=True)
|
||
# Docker verifies layer digests while pulling the immutable manifest reference.
|
||
run(["docker", "pull", resolved], PULL_TIMEOUT)
|
||
run(["docker", "tag", resolved, image])
|
||
|
||
|
||
def main() -> None:
|
||
environment = dict(os.environ)
|
||
if sys.argv[1:] == ["build"]:
|
||
resolved = {key: resolve_image(value, environment)
|
||
for key, value in build_images(environment).items()}
|
||
for key, value in resolved.items():
|
||
print(f"{key}={value}")
|
||
elif sys.argv[1:2] == ["pull"]:
|
||
pull_services(sys.argv[2:], environment)
|
||
else:
|
||
raise MirrorError("未知备用镜像操作")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
main()
|
||
except (MirrorError, OSError, ValueError, KeyError, subprocess.TimeoutExpired) as error:
|
||
detail = str(error) if isinstance(error, MirrorError) else "镜像验证或拉取未完成"
|
||
sys.stderr.write("PLANET_DOCKER_MIRROR_FAILED: " + detail + "\n")
|
||
sys.exit(1)
|