diff --git a/src/mirrorsrun/config.py b/src/mirrorsrun/config.py index a446145..47446b2 100644 --- a/src/mirrorsrun/config.py +++ b/src/mirrorsrun/config.py @@ -10,3 +10,24 @@ assert SCHEME in ["http", "https"] CACHE_DIR = os.environ.get("CACHE_DIR", "/app/cache/") EXTERNAL_HOST_ARIA2 = f"aria2.{BASE_DOMAIN}" EXTERNAL_URL_ARIA2 = f"{SCHEME}://{EXTERNAL_HOST_ARIA2}/aria2/index.html" + +BASE_URL_PYTORCH = os.environ.get("BASE_URL_PYTORCH", "https://download.pytorch.org") +BASE_URL_DOCKERHUB = os.environ.get( + "BASE_URL_DOCKERHUB", "https://registry-1.docker.io" +) +BASE_URL_NPM = os.environ.get("BASE_URL_NPM", "https://registry.npmjs.org") +BASE_URL_PYPI = os.environ.get("BASE_URL_PYPI", "https://pypi.org") +BASE_URL_PYPI_FILES = os.environ.get( + "BASE_URL_PYPI_FILES", "https://files.pythonhosted.org" +) + +BASE_URL_ALPINE = os.environ.get("BASE_URL_ALPINE", "https://dl-cdn.alpinelinux.org") +BASE_URL_UBUNTU = os.environ.get("BASE_URL_UBUNTU", "http://archive.ubuntu.com") +BASE_URL_UBUNTU_PORTS = os.environ.get( + "BASE_URL_UBUNTU_PORTS", "http://ports.ubuntu.com" +) + +BASE_URL_K8S = os.environ.get("BASE_URL_K8S", "https://registry.k8s.io") +BASE_URL_QUAY = os.environ.get("BASE_URL_QUAY", "https://quay.io") +BASE_URL_GHCR = os.environ.get("BASE_URL_GHCR", "https://ghcr.io") +BASE_URL_NVCR = os.environ.get("BASE_URL_NVCR", "https://nvcr.io") diff --git a/src/mirrorsrun/sites/common.py b/src/mirrorsrun/sites/common.py index 89616cb..2e4dc3a 100644 --- a/src/mirrorsrun/sites/common.py +++ b/src/mirrorsrun/sites/common.py @@ -1,5 +1,6 @@ from starlette.requests import Request +from mirrorsrun.config import BASE_URL_ALPINE, BASE_URL_UBUNTU, BASE_URL_UBUNTU_PORTS from mirrorsrun.proxy.direct import direct_proxy from starlette.responses import Response @@ -9,10 +10,10 @@ async def common(request: Request): if path == "/": return if path.startswith("/alpine"): - return await direct_proxy(request, "https://dl-cdn.alpinelinux.org" + path) + return await direct_proxy(request, BASE_URL_ALPINE + path) if path.startswith("/ubuntu/"): - return await direct_proxy(request, "http://archive.ubuntu.com" + path) + return await direct_proxy(request, BASE_URL_UBUNTU + path) if path.startswith("/ubuntu-ports/"): - return await direct_proxy(request, "http://ports.ubuntu.com" + path) + return await direct_proxy(request, BASE_URL_UBUNTU_PORTS + path) return Response("Not Found", status_code=404) diff --git a/src/mirrorsrun/sites/docker.py b/src/mirrorsrun/sites/docker.py index ba3c6d3..25581f2 100644 --- a/src/mirrorsrun/sites/docker.py +++ b/src/mirrorsrun/sites/docker.py @@ -1,6 +1,13 @@ import logging import re +from mirrorsrun.config import ( + BASE_URL_DOCKERHUB, + BASE_URL_K8S, + BASE_URL_QUAY, + BASE_URL_GHCR, + BASE_URL_NVCR, +) from mirrorsrun.proxy.direct import direct_proxy from mirrorsrun.proxy.file_cache import try_file_based_cache from starlette.requests import Request @@ -8,6 +15,7 @@ from starlette.responses import Response logger = logging.getLogger(__name__) + HEADER_AUTH_KEY = "www-authenticate" mirror_root_realm_mapping = {} @@ -130,16 +138,11 @@ def dockerhub_name_mapper(name): return name -k8s = build_docker_registry_handler( - "https://registry.k8s.io", -) -quay = build_docker_registry_handler( - "https://quay.io", -) -ghcr = build_docker_registry_handler( - "https://ghcr.io", -) -nvcr = build_docker_registry_handler("https://nvcr.io") +k8s = build_docker_registry_handler(BASE_URL_K8S) +quay = build_docker_registry_handler(BASE_URL_QUAY) +ghcr = build_docker_registry_handler(BASE_URL_GHCR) +nvcr = build_docker_registry_handler(BASE_URL_NVCR) + dockerhub = build_docker_registry_handler( - "https://registry-1.docker.io", name_mapper=dockerhub_name_mapper + BASE_URL_DOCKERHUB, name_mapper=dockerhub_name_mapper ) diff --git a/src/mirrorsrun/sites/npm.py b/src/mirrorsrun/sites/npm.py index 76c7213..303eba5 100644 --- a/src/mirrorsrun/sites/npm.py +++ b/src/mirrorsrun/sites/npm.py @@ -1,11 +1,10 @@ from starlette.requests import Request +from mirrorsrun.config import BASE_URL_NPM from mirrorsrun.proxy.direct import direct_proxy -BASE_URL = "https://registry.npmjs.org" - async def npm(request: Request): path = request.url.path - return await direct_proxy(request, BASE_URL + path) + return await direct_proxy(request, BASE_URL_NPM + path) diff --git a/src/mirrorsrun/sites/pypi.py b/src/mirrorsrun/sites/pypi.py index c643386..0e2d765 100644 --- a/src/mirrorsrun/sites/pypi.py +++ b/src/mirrorsrun/sites/pypi.py @@ -3,12 +3,10 @@ import re from starlette.requests import Request from starlette.responses import Response +from mirrorsrun.config import BASE_URL_PYPI, BASE_URL_PYPI_FILES from mirrorsrun.proxy.direct import direct_proxy from mirrorsrun.proxy.file_cache import try_file_based_cache -pypi_file_base_url = "https://files.pythonhosted.org" -pypi_base_url = "https://pypi.org" - def pypi_replace(request: Request, response: Response) -> Response: is_detail_page = re.search(r"/simple/([^/]+)/", request.url.path) is not None @@ -18,7 +16,7 @@ def pypi_replace(request: Request, response: Response) -> Response: if is_detail_page: mirror_url = f"{request.url.scheme}://{request.url.netloc}" content = response.body - content = content.replace(pypi_file_base_url.encode(), mirror_url.encode()) + content = content.replace(BASE_URL_PYPI_FILES.encode(), mirror_url.encode()) response.body = content del response.headers["content-length"] del response.headers["content-encoding"] @@ -32,9 +30,10 @@ async def pypi(request: Request) -> Response: path = "/simple/" if path.startswith("/simple/"): - target_url = pypi_base_url + path + # FIXME: join + target_url = BASE_URL_PYPI + path elif path.startswith("/packages/"): - target_url = pypi_file_base_url + path + target_url = BASE_URL_PYPI_FILES + path else: return Response(content="Not Found", status_code=404) diff --git a/src/mirrorsrun/sites/torch.py b/src/mirrorsrun/sites/torch.py index 7ac13c6..2f0fd1e 100644 --- a/src/mirrorsrun/sites/torch.py +++ b/src/mirrorsrun/sites/torch.py @@ -1,11 +1,10 @@ from starlette.requests import Request from starlette.responses import Response +from mirrorsrun.config import BASE_URL_PYTORCH from mirrorsrun.proxy.file_cache import try_file_based_cache from mirrorsrun.proxy.direct import direct_proxy -BASE_URL = "https://download.pytorch.org" - async def torch(request: Request): path = request.url.path @@ -16,7 +15,7 @@ async def torch(request: Request): if path == "/whl": path = "/whl/" - target_url = BASE_URL + path + target_url = BASE_URL_PYTORCH + path if path.endswith(".whl") or path.endswith(".tar.gz"): return await try_file_based_cache(request, target_url)