Shared storage implementation

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,10 +266,12 @@ 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")
                     if args.get_api_key_link:
@@ -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())

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,

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,9 +32,12 @@ def setup(violations_threshold=100):
     active = True
     threshold = violations_threshold
 
+    # Only setup the scheduler and secrets on one process
+    if not setup_scheduler.value:
+        setup_scheduler.value = True
+
         scheduler = BackgroundScheduler()
         scheduler.add_job(func=forgive_banned, trigger="interval", minutes=30)
-    scheduler.add_job(func=rotate_secrets, trigger="interval", minutes=30)
         
         scheduler.start()
 
@@ -57,19 +47,19 @@ def setup(violations_threshold=100):
 
 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

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="<Storage URI>",
+        help="Shared storage URI to use for multi-process data sharing (e.g. via gunicorn)",
+    )
     parser.add_argument(
         "--load-only",
         type=operator.methodcaller("split", ","),

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 %}