diff --git a/src/ui/gunicorn.conf.py b/src/ui/gunicorn.conf.py index f14ea0d45..b7d625280 100644 --- a/src/ui/gunicorn.conf.py +++ b/src/ui/gunicorn.conf.py @@ -1,8 +1,11 @@ from contextlib import suppress from hashlib import sha256 +from json import JSONDecodeError, dumps, loads from os import cpu_count, getenv, getpid, sep, urandom from os.path import join from pathlib import Path +from signal import SIGINT, SIGTERM, signal +from threading import Lock from regex import compile as re_compile from sys import path as sys_path from time import sleep @@ -36,7 +39,7 @@ workers = MAX_WORKERS worker_class = "gthread" threads = int(getenv("MAX_THREADS", MAX_WORKERS * 2)) max_requests_jitter = min(8, MAX_WORKERS) -graceful_timeout = 5 +graceful_timeout = 30 DEBUG = getenv("DEBUG", False) @@ -46,12 +49,16 @@ if DEBUG: reload = True reload_extra_files = [file.as_posix() for file in Path(sep, "usr", "share", "bunkerweb", "ui", "templates").iterdir()] +LOCK = Lock() + def on_starting(server): + TMP_DIR.mkdir(parents=True, exist_ok=True) if not getenv("FLASK_SECRET") and not TMP_DIR.joinpath(".flask_secret").is_file(): - TMP_DIR.mkdir(parents=True, exist_ok=True) TMP_DIR.joinpath(".flask_secret").write_text(sha256(urandom(32)).hexdigest(), encoding="utf-8") + TMP_DIR.joinpath(".ui.json").write_text("{}", encoding="utf-8") + LOGGER = setup_logger("UI") db = Database(LOGGER, ui=True) @@ -125,6 +132,29 @@ def on_starting(server): LOGGER.info("UI is ready") +def handle_stop(signum=None, frame=None): + if not TMP_DIR.joinpath(".ui.json").is_file(): + return + + ui_data = "Error" + while ui_data == "Error": + with suppress(JSONDecodeError): + ui_data = loads(TMP_DIR.joinpath(".ui.json").read_text(encoding="utf-8")) + + ui_data["SERVER_STOPPING"] = True + + with LOCK: + TMP_DIR.joinpath(".ui.json").write_text(dumps(ui_data), encoding="utf-8") + + +signal(SIGINT, handle_stop) +signal(SIGTERM, handle_stop) + + +def on_reload(server): + handle_stop() + + def when_ready(server): RUN_DIR.mkdir(parents=True, exist_ok=True) RUN_DIR.joinpath("ui.pid").write_text(str(getpid()), encoding="utf-8") @@ -135,3 +165,4 @@ def on_exit(server): RUN_DIR.joinpath("ui.pid").unlink(missing_ok=True) TMP_DIR.joinpath("ui.healthy").unlink(missing_ok=True) TMP_DIR.joinpath(".flask_secret").unlink(missing_ok=True) + TMP_DIR.joinpath(".ui.json").unlink(missing_ok=True) diff --git a/src/ui/main.py b/src/ui/main.py index de02a31fa..52ed5d2a9 100755 --- a/src/ui/main.py +++ b/src/ui/main.py @@ -21,7 +21,7 @@ from datetime import datetime, timedelta, timezone from dateutil.parser import parse as dateutil_parse from docker import DockerClient from docker.errors import NotFound as docker_NotFound, APIError as docker_APIError, DockerException -from flask import Flask, Response, flash, jsonify, redirect, render_template, request, send_file, session, url_for +from flask import Flask, Response, flash, jsonify, make_response, redirect, render_template, request, send_file, session, url_for from flask_login import current_user, LoginManager, login_required, login_user, logout_user from flask_wtf.csrf import CSRFProtect, CSRFError from hashlib import sha256 @@ -458,11 +458,16 @@ def handle_csrf_error(_): @app.before_request def before_request(): + ui_data = get_ui_data() + + if ui_data.get("SERVER_STOPPING", False): + response = make_response(jsonify({"message": "Server is shutting down, try again later."}), 503) + response.headers["Retry-After"] = 30 # Clients should retry after 30 seconds # type: ignore + return response + app.config["SCRIPT_NONCE"] = sha256(urandom(32)).hexdigest() if not request.path.startswith(("/css", "/images", "/js", "/json", "/webfonts")): - ui_data = get_ui_data() - if ( app.config["DB"].database_uri and app.config["DB"].readonly @@ -645,10 +650,12 @@ def setup(): random_url=f"/{''.join(choice(ascii_letters + digits) for _ in range(10))}", ) + @app.route("/setup/loading", methods=["GET"]) def setup_loading(): return render_template("setup_loading.html") + @app.route("/totp", methods=["GET", "POST"]) @login_required def totp():