Merge pull request #8 from NoCLin/feat/base-url-configurable

feat: makes base url configurable

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

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)

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
 )

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)

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)
 

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)