set rate limit for creating alias endpoint

app/api/views/new_custom_alias.py

@@ -10,7 +10,7 @@ from app.api.serializer import (
     serialize_alias_info_v2,
     get_alias_info_v2,
 )
-from app.config import MAX_NB_EMAIL_FREE_PLAN
+from app.config import MAX_NB_EMAIL_FREE_PLAN, ALIAS_LIMIT
 from app.dashboard.views.custom_alias import verify_prefix_suffix, signer
 from app.extensions import db, limiter
 from app.log import LOG
@@ -28,7 +28,7 @@ from app.utils import convert_to_id
 
 
 @api_bp.route("/alias/custom/new", methods=["POST"])
-@limiter.limit("5/minute")
+@limiter.limit(ALIAS_LIMIT)
 @require_api_auth
 def new_custom_alias():
     """
@@ -99,7 +99,7 @@ def new_custom_alias():
 
 
 @api_bp.route("/v2/alias/custom/new", methods=["POST"])
-@limiter.limit("5/minute")
+@limiter.limit(ALIAS_LIMIT)
 @require_api_auth
 def new_custom_alias_v2():
     """
@@ -194,7 +194,7 @@ def new_custom_alias_v2():
 
 
 @api_bp.route("/v3/alias/custom/new", methods=["POST"])
-@limiter.limit("5/minute")
+@limiter.limit(ALIAS_LIMIT)
 @require_api_auth
 def new_custom_alias_v3():
     """

app/api/views/new_random_alias.py

@@ -6,14 +6,14 @@ from app.api.serializer import (
     get_alias_info_v2,
     serialize_alias_info_v2,
 )
-from app.config import MAX_NB_EMAIL_FREE_PLAN
+from app.config import MAX_NB_EMAIL_FREE_PLAN, ALIAS_LIMIT
 from app.extensions import db, limiter
 from app.log import LOG
 from app.models import Alias, AliasUsedOn, AliasGeneratorEnum
 
 
 @api_bp.route("/alias/random/new", methods=["POST"])
-@limiter.limit("5/minute")
+@limiter.limit(ALIAS_LIMIT)
 @require_api_auth
 def new_random_alias():
     """

app/config.py

@@ -369,3 +369,5 @@ try:
     COINBASE_YEARLY_PRICE = float(os.environ["COINBASE_YEARLY_PRICE"])
 except Exception:
     COINBASE_YEARLY_PRICE = 30.00
+
+ALIAS_LIMIT = "100/day;50/hour;5/minute"

app/extensions.py

@@ -1,4 +1,3 @@
-from flask import request
 from flask_limiter import Limiter
 from flask_limiter.util import get_remote_address
 from flask_login import LoginManager
@@ -14,9 +13,9 @@ migrate = Migrate(db=db)
 limiter = Limiter(key_func=get_remote_address)
 
 
-@limiter.request_filter
-def ip_whitelist():
-    # Uncomment line to test rate limit in dev environment
-    # return False
-    # No limit for local development
-    return request.remote_addr == "127.0.0.1"
+# @limiter.request_filter
+# def ip_whitelist():
+#     # Uncomment line to test rate limit in dev environment
+#     # return False
+#     # No limit for local development
+#     return request.remote_addr == "127.0.0.1"

server.py

@@ -17,6 +17,7 @@ from flask import (
     jsonify,
     flash,
     session,
+    g,
 )
 from flask_admin import Admin
 from flask_cors import cross_origin, CORS
@@ -499,7 +500,11 @@ def setup_error_page(app):
 
     @app.errorhandler(429)
     def rate_limited(e):
-        LOG.warning("Client hit rate limit on path %s", request.path)
+        LOG.warning(
+            "Client hit rate limit on path %s, user:%s",
+            request.path,
+            g.user or current_user,
+        )
         if request.path.startswith("/api/"):
             return jsonify(error="Rate limit exceeded"), 429
         else:

tests/api/test_new_custom_alias.py

@@ -1,4 +1,4 @@
-from flask import url_for
+from flask import url_for, g
 
 from app.alias_utils import delete_alias
 from app.config import EMAIL_DOMAIN, MAX_NB_EMAIL_FREE_PLAN
@@ -208,3 +208,52 @@ def test_success_v3(flask_client):
     new_alias: Alias = Alias.get_by(email=r.json["alias"])
     assert new_alias.note == "test note"
     assert len(new_alias.mailboxes) == 2
+
+
+def test_custom_domain_alias(flask_client):
+    user = login(flask_client)
+
+    # create a custom domain
+    CustomDomain.create(user_id=user.id, domain="ab.cd", verified=True, commit=True)
+
+    signed_suffix = signer.sign("@ab.cd").decode()
+
+    r = flask_client.post(
+        "/api/v3/alias/custom/new",
+        json={
+            "alias_prefix": "prefix",
+            "signed_suffix": signed_suffix,
+            "mailbox_ids": [user.default_mailbox_id],
+        },
+    )
+
+    assert r.status_code == 201
+    assert r.json["alias"] == f"prefix@ab.cd"
+
+
+def test_too_many_requests(flask_client):
+    user = login(flask_client)
+
+    # create a custom domain
+    CustomDomain.create(user_id=user.id, domain="ab.cd", verified=True, commit=True)
+
+    # can't create more than 5 aliases in 1 minute
+    for i in range(7):
+        signed_suffix = signer.sign("@ab.cd").decode()
+
+        r = flask_client.post(
+            "/api/v3/alias/custom/new",
+            json={
+                "alias_prefix": f"prefix{i}",
+                "signed_suffix": signed_suffix,
+                "mailbox_ids": [user.default_mailbox_id],
+            },
+        )
+
+        # to make flask-limiter work with unit test
+        # https://github.com/alisaifee/flask-limiter/issues/147#issuecomment-642683820
+        g._rate_limiting_complete = False
+    else:
+        # last request
+        assert r.status_code == 429
+        assert r.json == {"error": "Rate limit exceeded"}

tests/api/test_new_random_alias.py

@@ -1,6 +1,6 @@
 import uuid
 
-from flask import url_for
+from flask import url_for, g
 
 from app.config import EMAIL_DOMAIN, MAX_NB_EMAIL_FREE_PLAN
 from app.models import Alias
@@ -69,11 +69,28 @@ def test_out_of_quota(flask_client):
 
     assert r.status_code == 400
     assert (
-        r.json["error"]
-        == "You have reached the limitation of a free account with the maximum of 3 aliases, please upgrade your plan to create more aliases"
+        r.json["error"] == "You have reached the limitation of a free account with "
+        "the maximum of 3 aliases, please upgrade your plan to create more aliases"
     )
 
 
+def test_too_many_requests(flask_client):
+    login(flask_client)
+
+    # can't create more than 5 aliases in 1 minute
+    for _ in range(7):
+        r = flask_client.post(
+            url_for("api.new_random_alias", hostname="www.test.com", mode="uuid"),
+        )
+        # to make flask-limiter work with unit test
+        # https://github.com/alisaifee/flask-limiter/issues/147#issuecomment-642683820
+        g._rate_limiting_complete = False
+    else:
+        # last request
+        assert r.status_code == 429
+        assert r.json == {"error": "Rate limit exceeded"}
+
+
 def is_valid_uuid(val):
     try:
         uuid.UUID(str(val))