From d12e776949354fc79085488cb1dacb5a84062fb8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adri=C3=A0=20Casaj=C3=BAs?=
 <acasajus@users.noreply.github.com>
Date: Tue, 30 Jan 2024 18:29:59 +0100
Subject: [PATCH] Rate limit alias creation to prevent abuse (#2021)

* Rate limit alias creation to prevent abuse

* Limit in secs

* Calculate bucket time

* fix exception

* Tune limits
---
 app/models.py         | 11 ++++++++++-
 app/rate_limiter.py   | 31 +++++++++++++++++++++++++++++++
 app/redis_services.py |  3 +++
 3 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 app/rate_limiter.py

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