This commit is contained in:
parent
0f60f7cec9
commit
e9faf93878
@ -492,6 +492,31 @@ NAMESERVERS = setup_nameservers()
|
|||||||
DISABLE_CREATE_CONTACTS_FOR_FREE_USERS = os.environ.get(
|
DISABLE_CREATE_CONTACTS_FOR_FREE_USERS = os.environ.get(
|
||||||
"DISABLE_CREATE_CONTACTS_FOR_FREE_USERS", False
|
"DISABLE_CREATE_CONTACTS_FOR_FREE_USERS", False
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# Expect format hits,seconds:hits,seconds...
|
||||||
|
# Example 1,10:4,60 means 1 in the last 10 secs or 4 in the last 60 secs
|
||||||
|
def getRateLimitFromConfig(
|
||||||
|
env_var: string, default: string = ""
|
||||||
|
) -> list[tuple[int, int]]:
|
||||||
|
value = os.environ.get(env_var, default)
|
||||||
|
if not value:
|
||||||
|
return []
|
||||||
|
entries = [entry for entry in value.split(":")]
|
||||||
|
limits = []
|
||||||
|
for entry in entries:
|
||||||
|
fields = entry.split(",")
|
||||||
|
limit = (int(fields[0]), int(fields[1]))
|
||||||
|
limits.append(limit)
|
||||||
|
return limits
|
||||||
|
|
||||||
|
|
||||||
|
ALIAS_CREATE_RATE_LIMIT_FREE = getRateLimitFromConfig(
|
||||||
|
"ALIAS_CREATE_RATE_LIMIT_FREE", "10,900:50,3600"
|
||||||
|
)
|
||||||
|
ALIAS_CREATE_RATE_LIMIT_PAID = getRateLimitFromConfig(
|
||||||
|
"ALIAS_CREATE_RATE_LIMIT_PAID", "50,900:200,3600"
|
||||||
|
)
|
||||||
PARTNER_API_TOKEN_SECRET = os.environ.get("PARTNER_API_TOKEN_SECRET") or (
|
PARTNER_API_TOKEN_SECRET = os.environ.get("PARTNER_API_TOKEN_SECRET") or (
|
||||||
FLASK_SECRET + "partnerapitoken"
|
FLASK_SECRET + "partnerapitoken"
|
||||||
)
|
)
|
||||||
|
@ -27,7 +27,7 @@ from sqlalchemy.orm import deferred
|
|||||||
from sqlalchemy.sql import and_
|
from sqlalchemy.sql import and_
|
||||||
from sqlalchemy_utils import ArrowType
|
from sqlalchemy_utils import ArrowType
|
||||||
|
|
||||||
from app import config
|
from app import config, rate_limiter
|
||||||
from app import s3
|
from app import s3
|
||||||
from app.db import Session
|
from app.db import Session
|
||||||
from app.dns_utils import get_mx_domains
|
from app.dns_utils import get_mx_domains
|
||||||
@ -1563,6 +1563,15 @@ class Alias(Base, ModelMixin):
|
|||||||
flush = kw.pop("flush", False)
|
flush = kw.pop("flush", False)
|
||||||
|
|
||||||
new_alias = cls(**kw)
|
new_alias = cls(**kw)
|
||||||
|
user = User.get(new_alias.user_id)
|
||||||
|
if user.is_premium():
|
||||||
|
limits = config.ALIAS_CREATE_RATE_LIMIT_PAID
|
||||||
|
else:
|
||||||
|
limits = config.ALIAS_CREATE_RATE_LIMIT_FREE
|
||||||
|
# limits is array of (hits,days)
|
||||||
|
for limit in limits:
|
||||||
|
key = f"alias_create_{limit[1]}d:{user.id}"
|
||||||
|
rate_limiter.check_bucket_limit(key, limit[0], limit[1])
|
||||||
|
|
||||||
email = kw["email"]
|
email = kw["email"]
|
||||||
# make sure email is lowercase and doesn't have any whitespace
|
# make sure email is lowercase and doesn't have any whitespace
|
||||||
|
31
app/app/rate_limiter.py
Normal file
31
app/app/rate_limiter.py
Normal file
@ -0,0 +1,31 @@
|
|||||||
|
from datetime import datetime
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
import redis.exceptions
|
||||||
|
import werkzeug.exceptions
|
||||||
|
from limits.storage import RedisStorage
|
||||||
|
|
||||||
|
from app.log import log
|
||||||
|
|
||||||
|
lock_redis: Optional[RedisStorage] = None
|
||||||
|
|
||||||
|
|
||||||
|
def set_redis_concurrent_lock(redis: RedisStorage):
|
||||||
|
global lock_redis
|
||||||
|
lock_redis = redis
|
||||||
|
|
||||||
|
|
||||||
|
def check_bucket_limit(
|
||||||
|
lock_name: Optional[str] = None,
|
||||||
|
max_hits: int = 5,
|
||||||
|
bucket_seconds: int = 3600,
|
||||||
|
):
|
||||||
|
# Calculate current bucket time
|
||||||
|
bucket_id = int(datetime.utcnow().timestamp()) % bucket_seconds
|
||||||
|
bucket_lock_name = f"bl:{lock_name}:{bucket_id}"
|
||||||
|
try:
|
||||||
|
value = lock_redis.incr(bucket_lock_name, bucket_seconds)
|
||||||
|
if value > max_hits:
|
||||||
|
raise werkzeug.exceptions.TooManyRequests()
|
||||||
|
except redis.exceptions.RedisError:
|
||||||
|
log.e("Cannot connect to redis")
|
@ -2,6 +2,7 @@ import flask
|
|||||||
import limits.storage
|
import limits.storage
|
||||||
|
|
||||||
from app.parallel_limiter import set_redis_concurrent_lock
|
from app.parallel_limiter import set_redis_concurrent_lock
|
||||||
|
from app.rate_limiter import set_redis_concurrent_lock as rate_limit_set_redis
|
||||||
from app.session import RedisSessionStore
|
from app.session import RedisSessionStore
|
||||||
|
|
||||||
|
|
||||||
@ -10,12 +11,14 @@ def initialize_redis_services(app: flask.Flask, redis_url: str):
|
|||||||
storage = limits.storage.RedisStorage(redis_url)
|
storage = limits.storage.RedisStorage(redis_url)
|
||||||
app.session_interface = RedisSessionStore(storage.storage, storage.storage, app)
|
app.session_interface = RedisSessionStore(storage.storage, storage.storage, app)
|
||||||
set_redis_concurrent_lock(storage)
|
set_redis_concurrent_lock(storage)
|
||||||
|
rate_limit_set_redis(storage)
|
||||||
elif redis_url.startswith("redis+sentinel://"):
|
elif redis_url.startswith("redis+sentinel://"):
|
||||||
storage = limits.storage.RedisSentinelStorage(redis_url)
|
storage = limits.storage.RedisSentinelStorage(redis_url)
|
||||||
app.session_interface = RedisSessionStore(
|
app.session_interface = RedisSessionStore(
|
||||||
storage.storage, storage.storage_slave, app
|
storage.storage, storage.storage_slave, app
|
||||||
)
|
)
|
||||||
set_redis_concurrent_lock(storage)
|
set_redis_concurrent_lock(storage)
|
||||||
|
rate_limit_set_redis(storage)
|
||||||
else:
|
else:
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
f"Tried to set_redis_session with an invalid redis url: ${redis_url}"
|
f"Tried to set_redis_session with an invalid redis url: ${redis_url}"
|
||||||
|
Loading…
x
Reference in New Issue
Block a user