From e00f3af7dbd2ddad4399f2109f8708bb93672af0 Mon Sep 17 00:00:00 2001 From: Piero Toffanin Date: Thu, 9 Mar 2023 16:09:04 -0500 Subject: [PATCH] Shared storage implementation --- libretranslate/app.py | 13 +++-- libretranslate/default_values.py | 5 ++ libretranslate/flood.py | 60 ++++++++++-------------- libretranslate/main.py | 7 +++ libretranslate/templates/app.js.template | 4 +- 5 files changed, 49 insertions(+), 40 deletions(-) diff --git a/libretranslate/app.py b/libretranslate/app.py index d6d7ec0..c6ee4db 100644 --- a/libretranslate/app.py +++ b/libretranslate/app.py @@ -21,7 +21,7 @@ from werkzeug.exceptions import HTTPException from werkzeug.http import http_date from flask_babel import Babel -from libretranslate import flood, remove_translated_files, security +from libretranslate import flood, secret, remove_translated_files, security, storage from libretranslate.language import detect_languages, improve_translation_formatting from libretranslate.locales import (_, _lazy, get_available_locales, get_available_locale_codes, gettext_escaped, gettext_html, lazy_swag, get_alternate_locale_links) @@ -127,6 +127,8 @@ def create_app(args): bp = Blueprint('Main app', __name__) + storage.setup(args.shared_storage) + if not args.disable_files_translation: remove_translated_files.setup(get_upload_dir()) languages = load_languages() @@ -204,6 +206,9 @@ def create_app(args): if args.req_flood_threshold > 0: flood.setup(args.req_flood_threshold) + if args.api_keys and args.require_api_key_secret: + secret.setup() + measure_request = None gauge_request = None @@ -261,9 +266,11 @@ def create_app(args): if (args.require_api_key_secret and key_missing - and not flood.secret_match(get_req_secret()) + and not secret.secret_match(get_req_secret()) ): need_key = True + + # TODO: find a way to send a "refresh" error key? if need_key: description = _("Please contact the server operator to get an API key") @@ -347,7 +354,7 @@ def create_app(args): response = Response(render_template("app.js.template", url_prefix=args.url_prefix, get_api_key_link=args.get_api_key_link, - api_secret=flood.get_current_secret() if args.require_api_key_secret else ""), content_type='application/javascript; charset=utf-8') + api_secret=secret.get_current_secret() if args.require_api_key_secret else ""), content_type='application/javascript; charset=utf-8') if args.require_api_key_secret: response.headers['Last-Modified'] = http_date(datetime.now()) diff --git a/libretranslate/default_values.py b/libretranslate/default_values.py index a9d007b..3e14580 100644 --- a/libretranslate/default_values.py +++ b/libretranslate/default_values.py @@ -136,6 +136,11 @@ _default_options_objects = [ 'default_value': False, 'value_type': 'bool' }, + { + 'name': 'SHARED_STORAGE', + 'default_value': 'memory://', + 'value_type': 'str' + }, { 'name': 'LOAD_ONLY', 'default_value': None, diff --git a/libretranslate/flood.py b/libretranslate/flood.py index 9d2d711..5a881e9 100644 --- a/libretranslate/flood.py +++ b/libretranslate/flood.py @@ -1,22 +1,20 @@ import atexit -import random -import string +from multiprocessing import Value +from libretranslate.storage import get_storage from apscheduler.schedulers.background import BackgroundScheduler -def generate_secret(): - return ''.join(random.choices(string.ascii_uppercase + string.digits, k=7)) -banned = {} +setup_scheduler = Value('b', False) active = False threshold = -1 -secrets = [generate_secret(), generate_secret()] def forgive_banned(): - global banned global threshold clear_list = [] + s = get_storage() + banned = s.get_all_hash_int("banned") for ip in banned: if banned[ip] <= 0: @@ -25,18 +23,7 @@ def forgive_banned(): banned[ip] = min(threshold, banned[ip]) - 1 for ip in clear_list: - del banned[ip] - -def rotate_secrets(): - global secrets - secrets[0] = secrets[1] - secrets[1] = generate_secret() - -def secret_match(s): - return s in secrets - -def get_current_secret(): - return secrets[1] + s.del_hash("banned", ip) def setup(violations_threshold=100): global active @@ -45,31 +32,34 @@ def setup(violations_threshold=100): active = True threshold = violations_threshold - scheduler = BackgroundScheduler() - scheduler.add_job(func=forgive_banned, trigger="interval", minutes=30) - scheduler.add_job(func=rotate_secrets, trigger="interval", minutes=30) - - scheduler.start() + # Only setup the scheduler and secrets on one process + if not setup_scheduler.value: + setup_scheduler.value = True - # Shut down the scheduler when exiting the app - atexit.register(lambda: scheduler.shutdown()) + scheduler = BackgroundScheduler() + scheduler.add_job(func=forgive_banned, trigger="interval", minutes=30) + + scheduler.start() + + # Shut down the scheduler when exiting the app + atexit.register(lambda: scheduler.shutdown()) def report(request_ip): if active: - banned[request_ip] = banned.get(request_ip, 0) - banned[request_ip] += 1 - + get_storage().inc_hash_int("banned", request_ip) def decrease(request_ip): - if banned[request_ip] > 0: - banned[request_ip] -= 1 - + s = get_storage() + if s.get_hash_int("banned", request_ip) > 0: + s.dec_hash_int("banned", request_ip) def has_violation(request_ip): - return request_ip in banned and banned[request_ip] > 0 - + s = get_storage() + return s.get_hash_int("banned", request_ip) > 0 def is_banned(request_ip): + s = get_storage() + # More than X offences? - return active and banned.get(request_ip, 0) >= threshold + return active and s.get_hash_int("banned", request_ip) >= threshold diff --git a/libretranslate/main.py b/libretranslate/main.py index 5a58885..0e02145 100644 --- a/libretranslate/main.py +++ b/libretranslate/main.py @@ -126,6 +126,13 @@ def get_args(): action="store_true", help="Require use of an API key for programmatic access to the API, unless the client also sends a secret match", ) + parser.add_argument( + "--shared-storage", + type=str, + default=DEFARGS['SHARED_STORAGE'], + metavar="", + help="Shared storage URI to use for multi-process data sharing (e.g. via gunicorn)", + ) parser.add_argument( "--load-only", type=operator.methodcaller("split", ","), diff --git a/libretranslate/templates/app.js.template b/libretranslate/templates/app.js.template index 4992d5b..2c10661 100644 --- a/libretranslate/templates/app.js.template +++ b/libretranslate/templates/app.js.template @@ -244,8 +244,8 @@ document.addEventListener('DOMContentLoaded', function(){ try{ {% if api_secret != "" %} if (this.status === 403){ - window.location.reload(true); - return; + //window.location.reload(true); + //return; } {% endif %}