Merge pull request #2328 from CTFd/2308-total-user-limit

Rough implementation of user registration limit

CTFd/auth.py

@@ -189,6 +189,14 @@ def register():
     if current_user.authed():
         return redirect(url_for("challenges.listing"))
 
+    num_users_limit = int(get_config("num_users", default=0))
+    num_users = Users.query.filter_by(banned=False, hidden=False).count()
+    if num_users_limit and num_users >= num_users_limit:
+        abort(
+            403,
+            description=f"Reached the maximum number of users ({num_users_limit}).",
+        )
+
     if request.method == "POST":
         name = request.form.get("name", "").strip()
         email_address = request.form.get("email", "").strip().lower()
@@ -490,6 +498,15 @@ def oauth_redirect():
 
             user = Users.query.filter_by(email=user_email).first()
             if user is None:
+                # Respect the user count limit
+                num_users_limit = int(get_config("num_users", default=0))
+                num_users = Users.query.filter_by(banned=False, hidden=False).count()
+                if num_users_limit and num_users >= num_users_limit:
+                    abort(
+                        403,
+                        description=f"Reached the maximum number of users ({num_users_limit}).",
+                    )
+
                 # Check if we are allowing registration before creating users
                 if registration_visible() or mlc_registration():
                     user = Users(

CTFd/forms/config.py

@@ -48,6 +48,9 @@ class AccountSettingsForm(BaseForm):
         widget=NumberInput(min=0),
         description="Max number of teams (Teams mode only)",
     )
+    num_users = IntegerField(
+        widget=NumberInput(min=0), description="Max number of users",
+    )
     verify_emails = SelectField(
         "Verify Emails",
         description="Control whether users must confirm their email addresses before playing",

CTFd/themes/admin/templates/configs/accounts.html

@@ -46,6 +46,14 @@
 			</small>
 		</div>
 
+		<div class="form-group">
+			{{ form.num_users.label }}
+			{{ form.num_users(class="form-control", value=num_users) }}
+			<small class="form-text text-muted">
+				{{ form.num_users.description }}
+			</small>
+		</div>
+
 		<div class="form-group">
 			{{ form.team_disbanding.label }}
 			{{ form.team_disbanding(class="form-control", value=team_disbanding) }}

tests/oauth/test_users.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from CTFd.models import Users
+from CTFd.utils import set_config
+from tests.helpers import create_ctfd, destroy_ctfd, login_with_mlc, register_user
+
+
+def test_num_users_oauth_limit():
+    """Only num_users users can be created even via MLC"""
+    app = create_ctfd()
+    app.config.update(
+        {
+            "OAUTH_CLIENT_ID": "ctfd_testing_client_id",
+            "OAUTH_CLIENT_SECRET": "ctfd_testing_client_secret",
+            "OAUTH_AUTHORIZATION_ENDPOINT": "http://auth.localhost/oauth/authorize",
+            "OAUTH_TOKEN_ENDPOINT": "http://auth.localhost/oauth/token",
+            "OAUTH_API_ENDPOINT": "http://api.localhost/user",
+        }
+    )
+    with app.app_context():
+        register_user(app)
+        # There should be the admin and our registered user
+        assert Users.query.count() == 2
+        set_config("num_users", 1)
+
+        # This registration should fail and we should still have 2 users
+        login_with_mlc(
+            app,
+            name="foobarbaz",
+            email="foobarbaz@a.com",
+            oauth_id=111,
+            scope="profile",
+            raise_for_error=False,
+        )
+        assert Users.query.count() == 2
+
+        # We increment num_users to 2 and then login again
+        set_config("num_users", 2)
+        login_with_mlc(
+            app,
+            name="foobarbaz",
+            email="foobarbaz@a.com",
+            oauth_id=111,
+            scope="profile",
+        )
+        # The above login should have succeeded
+        assert Users.query.count() == 3
+    destroy_ctfd(app)

tests/users/test_users.py

@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 
 from CTFd.models import Users
+from CTFd.utils import set_config
 from tests.helpers import (
     create_ctfd,
     destroy_ctfd,
@@ -107,3 +108,44 @@ def test_hidden_user_visibility():
             response = r.get_data(as_text=True)
             assert user_name in response
     destroy_ctfd(app)
+
+
+def test_num_users_limit():
+    """Only num_users users can be created"""
+    app = create_ctfd()
+    with app.app_context():
+        set_config("num_users", 1)
+
+        register_user(app)
+        with app.test_client() as client:
+            r = client.get("/register")
+            assert r.status_code == 403
+
+            # team should be blocked from creation
+            with client.session_transaction() as sess:
+                data = {
+                    "name": "user",
+                    "email": "user@examplectf.com",
+                    "password": "password",
+                    "nonce": sess.get("nonce"),
+                }
+            r = client.post("/register", data=data)
+            resp = r.get_data(as_text=True)
+            # This number is 2 to account for the admin and the registered user
+            assert Users.query.count() == 2
+            assert "Reached the maximum number of users" in resp
+
+            # Can the team be created after the num has been bumped
+            set_config("num_users", 2)
+            with client.session_transaction() as sess:
+                data = {
+                    "name": "user1",
+                    "email": "user1@examplectf.com",
+                    "password": "password",
+                    "nonce": sess.get("nonce"),
+                }
+            r = client.post("/register", data=data)
+            resp = r.get_data(as_text=True)
+            assert r.status_code == 302
+            assert Users.query.count() == 3
+    destroy_ctfd(app)