diff --git a/app.py b/app.py index fece2f1..ce814c0 100644 --- a/app.py +++ b/app.py @@ -7,6 +7,7 @@ and a global lock. All model handling lives in ``bgfilter.service.PipelineManage from __future__ import annotations import asyncio +import hmac import io import logging import os @@ -16,7 +17,7 @@ from contextlib import asynccontextmanager from pathlib import Path import numpy as np -from fastapi import FastAPI, File, Form, HTTPException, Response, UploadFile +from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Response, UploadFile from PIL import Image, ImageOps, UnidentifiedImageError from bgfilter.config import load_settings, override_settings @@ -48,6 +49,23 @@ MAX_IMAGE_PIXELS = _int_env("BGFILTER_MAX_IMAGE_PIXELS", 4_194_304) # Pillow's separate decompression-bomb ceiling (which would otherwise raise its own error). Image.MAX_IMAGE_PIXELS = None +# Optional app-level auth. When BGFILTER_AUTH_TOKEN is set, protected endpoints +# require the X-Genarrative-Image-Token header to equal it (constant-time compare) -- +# the same header the production nginx uses, so callers only swap the URL + token. +# When unset the check is a no-op (open), matching pre-auth behaviour and leaving +# auth to an upstream proxy (e.g. nginx in production). +AUTH_TOKEN = os.environ.get("BGFILTER_AUTH_TOKEN", "").strip() + + +async def require_auth(x_genarrative_image_token: str | None = Header(default=None)): + if not AUTH_TOKEN: + return + if not x_genarrative_image_token or not hmac.compare_digest( + x_genarrative_image_token, AUTH_TOKEN + ): + raise HTTPException(status_code=401, detail="missing or invalid auth token") + + _manager: PipelineManager | None = None _lock = asyncio.Lock() @@ -79,13 +97,18 @@ async def lifespan(app: FastAPI): settings = _build_settings() preload = os.environ.get("BGFILTER_PRELOAD", "1").strip().lower() not in ("0", "false", "no", "") logger.info( - "starting bgfilter service: device=%s default_seg=%s cross_check=%s preload=%s max_pixels=%d", + "starting bgfilter service: device=%s default_seg=%s cross_check=%s preload=%s max_pixels=%d auth=%s", settings.model.device, settings.segmentation.backend, settings.cross_check.enabled, preload, MAX_IMAGE_PIXELS, + "enabled" if AUTH_TOKEN else "DISABLED(open)", ) + if not AUTH_TOKEN: + logger.warning( + "BGFILTER_AUTH_TOKEN not set: /remove-background is OPEN (no app-level auth)" + ) _manager = PipelineManager(settings, preload=preload) yield @@ -179,6 +202,7 @@ async def remove_background( screen_color: str | None = Form(None), seg_model: str | None = Form(None), cross_check: str | None = Form(None), + _auth: None = Depends(require_auth), ): if _manager is None: raise HTTPException(status_code=503, detail="service not ready")