From e9faf938783a0018cca7abe3242a5abbd629edc6 Mon Sep 17 00:00:00 2001
From: MrMeeb <charles.macdonald@me.com>
Date: Sat, 3 Feb 2024 16:55:23 +0000
Subject: [PATCH] 4.38.0

---
 app/app/config.py         | 25 +++++++++++++++++++++++++
 app/app/models.py         | 11 ++++++++++-
 app/app/rate_limiter.py   | 31 +++++++++++++++++++++++++++++++
 app/app/redis_services.py |  3 +++
 4 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 app/app/rate_limiter.py

diff --git a/app/app/config.py b/app/app/config.py
index 30b0a28..e83de69 100644
--- a/app/app/config.py
+++ b/app/app/config.py
@@ -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"
 )
diff --git a/app/app/models.py b/app/app/models.py
index 7db2236..2eb786e 100644
--- a/app/app/models.py
+++ b/app/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 = 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
diff --git a/app/app/rate_limiter.py b/app/app/rate_limiter.py
new file mode 100644
index 0000000..dac3da6
--- /dev/null
+++ b/app/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/app/redis_services.py b/app/app/redis_services.py
index 9ee98a9..288c481 100644
--- a/app/app/redis_services.py
+++ b/app/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}"