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", 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 (
|
||||
FLASK_SECRET + "partnerapitoken"
|
||||
)
|
||||
|
@ -27,7 +27,7 @@ from sqlalchemy.orm import deferred
|
||||
from sqlalchemy.sql import and_
|
||||
from sqlalchemy_utils import ArrowType
|
||||
|
||||
from app import config
|
||||
from app import config, rate_limiter
|
||||
from app import s3
|
||||
from app.db import Session
|
||||
from app.dns_utils import get_mx_domains
|
||||
@ -1563,6 +1563,15 @@ class Alias(Base, ModelMixin):
|
||||
flush = kw.pop("flush", False)
|
||||
|
||||
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"]
|
||||
# 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
|
||||
|
||||
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
|
||||
|
||||
|
||||
@ -10,12 +11,14 @@ def initialize_redis_services(app: flask.Flask, redis_url: str):
|
||||
storage = limits.storage.RedisStorage(redis_url)
|
||||
app.session_interface = RedisSessionStore(storage.storage, storage.storage, app)
|
||||
set_redis_concurrent_lock(storage)
|
||||
rate_limit_set_redis(storage)
|
||||
elif redis_url.startswith("redis+sentinel://"):
|
||||
storage = limits.storage.RedisSentinelStorage(redis_url)
|
||||
app.session_interface = RedisSessionStore(
|
||||
storage.storage, storage.storage_slave, app
|
||||
)
|
||||
set_redis_concurrent_lock(storage)
|
||||
rate_limit_set_redis(storage)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Tried to set_redis_session with an invalid redis url: ${redis_url}"
|
||||
|
Loading…
x
Reference in New Issue
Block a user