Testing, fixes

This commit is contained in:
Piero Toffanin 2023-03-09 23:07:12 -05:00
parent ab27dcbaf3
commit 6a231d974b
6 changed files with 96 additions and 49 deletions

View file

@ -21,7 +21,7 @@ from werkzeug.exceptions import HTTPException
from werkzeug.http import http_date from werkzeug.http import http_date
from flask_babel import Babel from flask_babel import Babel
from libretranslate import flood, secret, remove_translated_files, security, storage from libretranslate import scheduler, flood, secret, remove_translated_files, security, storage
from libretranslate.language import detect_languages, improve_translation_formatting from libretranslate.language import detect_languages, improve_translation_formatting
from libretranslate.locales import (_, _lazy, get_available_locales, get_available_locale_codes, gettext_escaped, from libretranslate.locales import (_, _lazy, get_available_locales, get_available_locale_codes, gettext_escaped,
gettext_html, lazy_swag, get_alternate_locale_links) gettext_html, lazy_swag, get_alternate_locale_links)
@ -204,11 +204,12 @@ def create_app(args):
limiter = Limiter() limiter = Limiter()
if args.req_flood_threshold > 0: if not "gunicorn" in os.environ.get("SERVER_SOFTWARE", ""):
flood.setup(args.req_flood_threshold) # Gunicorn starts the scheduler in the master process
if args.api_keys and args.require_api_key_secret: scheduler.setup(args)
secret.setup()
flood.setup(args)
secret.setup(args)
measure_request = None measure_request = None
gauge_request = None gauge_request = None

View file

@ -1,11 +1,5 @@
import atexit
from multiprocessing import Value
from libretranslate.storage import get_storage from libretranslate.storage import get_storage
from apscheduler.schedulers.background import BackgroundScheduler
setup_scheduler = Value('b', False)
active = False active = False
threshold = -1 threshold = -1
@ -20,30 +14,18 @@ def forgive_banned():
if banned[ip] <= 0: if banned[ip] <= 0:
clear_list.append(ip) clear_list.append(ip)
else: else:
banned[ip] = min(threshold, banned[ip]) - 1 s.set_hash_int("banned", ip, min(threshold, banned[ip]) - 1)
for ip in clear_list: for ip in clear_list:
s.del_hash("banned", ip) s.del_hash("banned", ip)
def setup(violations_threshold=100): def setup(args):
global active global active
global threshold global threshold
active = True if args.req_flood_threshold > 0:
threshold = violations_threshold active = True
threshold = args.req_flood_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.start()
# Shut down the scheduler when exiting the app
atexit.register(lambda: scheduler.shutdown())
def report(request_ip): def report(request_ip):
if active: if active:

View file

@ -0,0 +1,23 @@
import atexit
from apscheduler.schedulers.background import BackgroundScheduler
scheduler = None
def setup(args):
from libretranslate.flood import forgive_banned
from libretranslate.secret import rotate_secrets
global scheduler
if scheduler is None:
scheduler = BackgroundScheduler()
if args.req_flood_threshold > 0:
scheduler.add_job(func=forgive_banned, trigger="interval", minutes=10)
if args.api_keys and args.require_api_key_secret:
scheduler.add_job(func=rotate_secrets, trigger="interval", minutes=30)
scheduler.start()
# Shut down the scheduler when exiting the app
atexit.register(lambda: scheduler.shutdown())

View file

@ -1,7 +1,7 @@
import atexit import atexit
import random import random
import string import string
from multiprocessing import Value from multiprocessing.dummy import Value
from libretranslate.storage import get_storage from libretranslate.storage import get_storage
from apscheduler.schedulers.background import BackgroundScheduler from apscheduler.schedulers.background import BackgroundScheduler
@ -16,6 +16,9 @@ def rotate_secrets():
secret_1 = s.get_str("secret_1") secret_1 = s.get_str("secret_1")
s.set_str("secret_0", secret_1) s.set_str("secret_0", secret_1)
s.set_str("secret_1", generate_secret()) s.set_str("secret_1", generate_secret())
print(s.get_str("secret_0"))
print(s.get_str("secret_1"))
def secret_match(secret): def secret_match(secret):
s = get_storage() s = get_storage()
@ -24,19 +27,8 @@ def secret_match(secret):
def get_current_secret(): def get_current_secret():
return get_storage().get_str("secret_1") return get_storage().get_str("secret_1")
def setup(): def setup(args):
# Only setup the scheduler and secrets on one process if args.api_keys and args.require_api_key_secret:
if not setup_secrets.value:
setup_secrets.value = True
s = get_storage() s = get_storage()
s.set_str("secret_0", generate_secret()) s.set_str("secret_0", generate_secret())
s.set_str("secret_1", generate_secret()) s.set_str("secret_1", generate_secret())
scheduler = BackgroundScheduler()
scheduler.add_job(func=rotate_secrets, trigger="interval", minutes=30)
scheduler.start()
# Shut down the scheduler when exiting the app
atexit.register(lambda: scheduler.shutdown())

View file

@ -20,7 +20,7 @@ class Storage:
def get_str(self, key): def get_str(self, key):
raise Exception("not implemented") raise Exception("not implemented")
def set_hash_value(self, ns, key, value): def set_hash_int(self, ns, key, value):
raise Exception("not implemented") raise Exception("not implemented")
def get_hash_int(self, ns, key): def get_hash_int(self, ns, key):
raise Exception("not implemented") raise Exception("not implemented")
@ -56,6 +56,11 @@ class MemoryStorage(Storage):
def get_str(self, key): def get_str(self, key):
return str(self.store.get(key, "")) return str(self.store.get(key, ""))
def set_hash_int(self, ns, key, value):
if ns not in self.store:
self.store[ns] = {}
self.store[ns][key] = int(value)
def get_hash_int(self, ns, key): def get_hash_int(self, ns, key):
d = self.store.get(ns, {}) d = self.store.get(ns, {})
return int(d.get(key, 0)) return int(d.get(key, 0))
@ -79,7 +84,10 @@ class MemoryStorage(Storage):
self.store[ns][key] -= 1 self.store[ns][key] -= 1
def get_all_hash_int(self, ns): def get_all_hash_int(self, ns):
return [{str(k): int(v)} for k,v in self.store[ns].items()] if ns in self.store:
return [{str(k): int(v)} for k,v in self.store[ns].items()]
else:
return []
def del_hash(self, ns, key): def del_hash(self, ns, key):
del self.store[ns][key] del self.store[ns][key]
@ -123,6 +131,9 @@ class RedisStorage(Storage):
else: else:
return int(v) return int(v)
def set_hash_int(self, ns, key, value):
self.conn.hset(ns, key, value)
def inc_hash_int(self, ns, key): def inc_hash_int(self, ns, key):
return int(self.conn.hincrby(ns, key)) return int(self.conn.hincrby(ns, key))
@ -130,10 +141,10 @@ class RedisStorage(Storage):
return int(self.conn.hincrby(ns, key, -1)) return int(self.conn.hincrby(ns, key, -1))
def get_all_hash_int(self, ns): def get_all_hash_int(self, ns):
return [{k.decode("utf-8"): int(v)} for k,v in self.conn.hgetall(ns).items()] return {k.decode("utf-8"): int(v) for k,v in self.conn.hgetall(ns).items()}
def del_hash(self, ns, key): def del_hash(self, ns, key):
conn.hdel(ns, key) self.conn.hdel(ns, key)
def setup(storage_uri): def setup(storage_uri):
global storage global storage

View file

@ -1,4 +1,42 @@
from prometheus_client import multiprocess from prometheus_client import multiprocess
import re
import sys
def child_exit(server, worker): def child_exit(server, worker):
multiprocess.mark_process_dead(worker.pid) multiprocess.mark_process_dead(worker.pid)
def on_starting(server):
# Parse command line arguments
proc_name = server.cfg.default_proc_name
kwargs = {}
if proc_name.startswith("wsgi:app"):
str_args = re.sub('wsgi:app\s*\(\s*(.*)\s*\)', '\\1', proc_name).strip().split(",")
for a in str_args:
if "=" in a:
k,v = a.split("=")
k = k.strip()
v = v.strip()
if v.lower() in ["true", "false"]:
v = v.lower() == "true"
elif v[0] == '"':
v = v[1:-1]
kwargs[k] = v
from libretranslate.main import get_args
sys.argv = ['--wsgi']
for k in kwargs:
ck = k.replace("_", "-")
if isinstance(kwargs[k], bool) and kwargs[k]:
sys.argv.append("--" + ck)
else:
sys.argv.append("--" + ck)
sys.argv.append(kwargs[k])
args = get_args()
from libretranslate import storage, scheduler, flood, secret
storage.setup(args.shared_storage)
scheduler.setup(args)
flood.setup(args)
secret.setup(args)