Stable Diffusion GUI
{{ERROR}} - + + + Sign in with GitHub + @@ -202,33 +192,121 @@ async def login_page(error: str | None = None) -> HTMLResponse: return HTMLResponse(content=html) -@router.post("/comfy/login") -async def login_submit(username: str = Form(...), password: str = Form(...)) -> Response: - """Handle login form submission.""" - if not COMFYUI_USER or not COMFYUI_PASS: +@router.get("/comfy/auth/github") +async def github_auth(request: Request) -> Response: + """Redirect to GitHub OAuth.""" + if not _is_auth_configured(): return RedirectResponse( - url="/comfy/login?error=Authentication+not+configured", + url="/comfy/login?error=GitHub+OAuth+not+configured", status_code=status.HTTP_303_SEE_OTHER, ) - if username == COMFYUI_USER and password == COMFYUI_PASS: - token = _create_session_token(username) - response = RedirectResponse(url="/comfy/", status_code=status.HTTP_303_SEE_OTHER) - response.set_cookie( - key="comfy_session", - value=token, - max_age=SESSION_MAX_AGE, - httponly=True, - samesite="lax", - ) - return response + # Generate state for CSRF protection + state = secrets.token_urlsafe(32) + _oauth_states[state] = time.time() + # Clean up old states (older than 10 minutes) + cutoff = time.time() - 600 + for s in list(_oauth_states.keys()): + if _oauth_states[s] < cutoff: + del _oauth_states[s] + + # Build GitHub OAuth URL + params = { + "client_id": GITHUB_CLIENT_ID, + "redirect_uri": str(request.url_for("github_callback")), + "scope": "read:user", + "state": state, + } + query = "&".join(f"{k}={v}" for k, v in params.items()) return RedirectResponse( - url="/comfy/login?error=Invalid+username+or+password", + url=f"https://github.com/login/oauth/authorize?{query}", status_code=status.HTTP_303_SEE_OTHER, ) +@router.get("/comfy/auth/callback") +async def github_callback(request: Request, code: str | None = None, state: str | None = None) -> Response: + """Handle GitHub OAuth callback.""" + # Verify state + if not state or state not in _oauth_states: + return RedirectResponse( + url="/comfy/login?error=Invalid+OAuth+state", + status_code=status.HTTP_303_SEE_OTHER, + ) + del _oauth_states[state] + + if not code: + return RedirectResponse( + url="/comfy/login?error=No+authorization+code", + status_code=status.HTTP_303_SEE_OTHER, + ) + + # Exchange code for access token + async with httpx.AsyncClient() as client: + token_response = await client.post( + "https://github.com/login/oauth/access_token", + data={ + "client_id": GITHUB_CLIENT_ID, + "client_secret": GITHUB_CLIENT_SECRET, + "code": code, + }, + headers={"Accept": "application/json"}, + ) + token_data = token_response.json() + + if "error" in token_data: + return RedirectResponse( + url=f"/comfy/login?error={token_data.get('error_description', 'OAuth+error')}", + status_code=status.HTTP_303_SEE_OTHER, + ) + + access_token = token_data.get("access_token") + if not access_token: + return RedirectResponse( + url="/comfy/login?error=No+access+token", + status_code=status.HTTP_303_SEE_OTHER, + ) + + # Get user info + async with httpx.AsyncClient() as client: + user_response = await client.get( + "https://api.github.com/user", + headers={ + "Authorization": f"Bearer {access_token}", + "Accept": "application/vnd.github+json", + }, + ) + user_data = user_response.json() + + username = user_data.get("login", "") + if not username: + return RedirectResponse( + url="/comfy/login?error=Could+not+get+GitHub+username", + status_code=status.HTTP_303_SEE_OTHER, + ) + + # Check if user is allowed + allowed = [u.strip().lower() for u in GITHUB_ALLOWED_USERS if u.strip()] + if allowed and username.lower() not in allowed: + return RedirectResponse( + url="/comfy/login?error=User+not+authorized", + status_code=status.HTTP_303_SEE_OTHER, + ) + + # Create session + token = _create_session_token(username) + response = RedirectResponse(url="/comfy/", status_code=status.HTTP_303_SEE_OTHER) + response.set_cookie( + key="comfy_session", + value=token, + max_age=SESSION_MAX_AGE, + httponly=True, + samesite="lax", + ) + return response + + @router.get("/comfy/logout") async def logout() -> Response: """Clear session and redirect to login.""" @@ -244,7 +322,7 @@ def _check_auth(comfy_session: str | None, path: str = "", method: str = "GET") because modulepreload/crossorigin requests don't send cookies. OPTIONS requests (CORS preflight) are also allowed without auth. """ - if not COMFYUI_USER: + if not _is_auth_configured(): return if method == "OPTIONS": @@ -313,7 +391,7 @@ async def proxy_comfyui(request: Request, path: str, comfy_session: str | None = async def proxy_websocket(websocket: WebSocket, comfy_session: str | None = Cookie(default=None)) -> None: """Proxy WebSocket connections to ComfyUI.""" # Check auth via cookie - if COMFYUI_USER and not _verify_session_token(comfy_session): + if _is_auth_configured() and not _verify_session_token(comfy_session): await websocket.close(code=4001, reason="Unauthorized") return