diff --git a/main.py b/main.py index 251975a..c669171 100644 --- a/main.py +++ b/main.py @@ -24,6 +24,9 @@ BASE_DIR = Path("/usr/src/app/taobin_project") SERVICE_NAME = os.getenv("SERVICE_NAME") PREFIX = os.getenv("PREFIX", "/taobin-image") +STRIP_PREFIXES = tuple( + dict.fromkeys(prefix for prefix in (PREFIX, "/taobin-image") if prefix) +) GIT_REPO_SERVER_URL = os.getenv("GIT_REPO_SERVER_URL") FRONTEND_NOTIFY_URL = os.getenv("FRONTEND_NOTIFY_URL") @@ -44,8 +47,13 @@ ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp"} async def strip_prefix_middleware(request: Request, call_next): path = request.scope["path"] - if PREFIX and path.startswith(PREFIX): - new_path = path[len(PREFIX):] + matched_prefix = next( + (prefix for prefix in STRIP_PREFIXES if path == prefix or path.startswith(prefix + "/")), + None + ) + + if matched_prefix: + new_path = path[len(matched_prefix):] if not new_path: new_path = "/"