diff --git a/app/models.py b/app/models.py index 7db2236a..9b9e630c 100644 --- a/app/models.py +++ b/app/models.py @@ -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 = ((50, 1), (200, 7)) + else: + limits = ((10, 1), (20, 7)) + # 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] * 86400) email = kw["email"] # make sure email is lowercase and doesn't have any whitespace diff --git a/app/rate_limiter.py b/app/rate_limiter.py new file mode 100644 index 00000000..dac3da6a --- /dev/null +++ b/app/rate_limiter.py @@ -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") diff --git a/app/redis_services.py b/app/redis_services.py index 9ee98a97..288c4811 100644 --- a/app/redis_services.py +++ b/app/redis_services.py @@ -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}"