Fix: prevent username collision in set-password signup path

database/databaseConfig.py

@@ -1,6 +1,7 @@
 import os
 from dotenv import find_dotenv, load_dotenv
 from pymongo import MongoClient, TEXT
+from pymongo.errors import DuplicateKeyError as MongoDuplicateKeyError
 from utils.logger import Logger
 
 logger = Logger.get_logger("databaseConfig")
@@ -101,12 +102,86 @@ def initialize_text_index():
                 logger.info("Compound index created on user_id and created_at in image collection")
 
             # Add user collection indexes
-            if 'username_1' not in existing_user_indexes:
-                user_collection.create_index([('username', 1)], name='username_1')
-                logger.info("Index created on username in user collection")
-            if 'email_1' not in existing_user_indexes:
-                user_collection.create_index([('email', 1)], name='email_1')
-                logger.info("Index created on email in user collection")
+            # The username index MUST be unique to prevent race-condition duplicates.
+            # If an old non-unique index exists, drop it first so we can recreate
+            # it with unique=True.
+            if 'username_1' in existing_user_indexes:
+                if not existing_user_indexes['username_1'].get('unique'):
+                    # Safe upgrade: create a temporary unique index FIRST to
+                    # detect any duplicate data before we drop anything.
+                    # If duplicates exist, create_index raises DuplicateKeyError
+                    # and the original index is left completely untouched.
+                    try:
+                        user_collection.create_index(
+                            [('username', 1)],
+                            name='username_1_unique_tmp',
+                            unique=True,
+                        )
+                        # No duplicates confirmed — safe to swap.
+                        user_collection.drop_index('username_1')
+                        user_collection.create_index(
+                            [('username', 1)], name='username_1', unique=True
+                        )
+                        user_collection.drop_index('username_1_unique_tmp')
+                        logger.info("Upgraded username_1 index to unique=True")
+                    except MongoDuplicateKeyError:
+                        # Temp creation failed; old non-unique index is still in place.
+                        # Clean up the temp index just in case it was partially recorded.
+                        try:
+                            user_collection.drop_index('username_1_unique_tmp')
+                        except Exception:
+                            pass
+                        logger.error(
+                            "Cannot upgrade username index to unique: duplicate usernames "
+                            "exist in the collection. Manual deduplication is required "
+                            "before uniqueness can be enforced."
+                        )
+                # else: already unique — nothing to do
+            else:
+                try:
+                    user_collection.create_index([('username', 1)], name='username_1', unique=True)
+                    logger.info("Unique index created on username in user collection")
+                except MongoDuplicateKeyError:
+                    logger.error(
+                        "Cannot create unique username index: duplicate usernames exist. "
+                        "Manual deduplication required."
+                    )
+            if 'email_1' in existing_user_indexes:
+                if not existing_user_indexes['email_1'].get('unique'):
+                    # Safe upgrade: create temp unique index first to catch duplicates
+                    # before touching the existing index.
+                    try:
+                        user_collection.create_index(
+                            [('email', 1)],
+                            name='email_1_unique_tmp',
+                            unique=True,
+                        )
+                        user_collection.drop_index('email_1')
+                        user_collection.create_index(
+                            [('email', 1)], name='email_1', unique=True
+                        )
+                        user_collection.drop_index('email_1_unique_tmp')
+                        logger.info("Upgraded email_1 index to unique=True")
+                    except MongoDuplicateKeyError:
+                        try:
+                            user_collection.drop_index('email_1_unique_tmp')
+                        except Exception:
+                            pass
+                        logger.error(
+                            "Cannot upgrade email index to unique: duplicate emails "
+                            "exist in the collection. Manual deduplication is required "
+                            "before uniqueness can be enforced."
+                        )
+                # else: already unique — nothing to do
+            else:
+                try:
+                    user_collection.create_index([('email', 1)], name='email_1', unique=True)
+                    logger.info("Unique index created on email in user collection")
+                except MongoDuplicateKeyError:
+                    logger.error(
+                        "Cannot create unique email index: duplicate emails exist. "
+                        "Manual deduplication required."
+                    )
         except Exception as ie:
             logger.error(f"Error creating collection indexes: {ie}")
     except Exception as e:

routes/auth.py

@@ -1,9 +1,11 @@
 from flask import Blueprint, request, jsonify, current_app
 from datetime import datetime, timedelta, timezone
+import re
 import secrets
 import bcrypt
 from google.oauth2 import id_token
 from google.auth.transport import requests as google_requests
+import pymongo.errors
 from pymongo import ReturnDocument
 
 from bson import ObjectId
@@ -29,6 +31,51 @@
 OTP_MAX_VERIFY_ATTEMPTS = 5
 OTP_VERIFY_LOCKOUT_SECONDS = 300  # 5 minutes
 
+# Maximum sequential suffix attempts in _unique_username before a random fallback is used.
+_MAX_USERNAME_SEQUENTIAL = 10
+# Maximum full retries (generate + insert) for auto-assigned usernames before giving up.
+_MAX_SIGNUP_RETRIES = 3
+
+
+def _unique_username(base: str) -> str:
+    """Return a username derived from *base* that does not already exist.
+
+    Strips any '@' or '$' characters so the result can never be mistaken for
+    an email address or trigger NoSQL injection checks, then appends an
+    incrementing numeric suffix until a free slot is found.
+
+    A single query using $in fetches all taken variants (base, base1 … baseN) so
+    the suffix search is done in-memory without extra round-trips. Falls back
+    to a hex suffix if every sequential slot is occupied. The DB unique index
+    is the final backstop against races.
+    """
+    # Default to "user" so an all-symbol prefix (e.g. "@") never yields an
+    # empty string, which would produce an unusable username.
+    base = base.replace("@", "").replace("$", "").strip() or "user"
+
+    # Fetch every username that matches base or base<digits>.
+    # Define the specific candidates we want to check in-memory.
+    candidates = [base] + [f"{base}{i}" for i in range(1, _MAX_USERNAME_SEQUENTIAL + 1)]
+
+    taken = {
+        doc["username"]
+        for doc in db.users.find(
+            {"username": {"$in": candidates}},
+            {"username": 1, "_id": 0},
+        )
+    }
+
+    # Try the bare base first, then base1 … base_MAX in-memory.
+    if base not in taken:
+        return base
+    for counter in range(1, _MAX_USERNAME_SEQUENTIAL + 1):
+        candidate = f"{base}{counter}"
+        if candidate not in taken:
+            return candidate
+
+    # All sequential slots occupied — hex suffix to minimise collision risk.
+    return f"{base}{secrets.token_hex(4)}"
+
 
 def _validate_otp_verification(email: str):
     """Check that a valid, unexpired OTP verification session exists for email.
@@ -250,15 +297,23 @@ def complete_signup():
     )
 
     now_utc = datetime.now(timezone.utc)
-    result = db.users.insert_one({
-        "email": email,
-        "username": username,
-        "password": hashed_password,
-        "role": role,
-        "created_at": now_utc,
-        "last_active": now_utc
-    })
-    db.email_otps.delete_many({"email": email})
+    try:
+        result = db.users.insert_one({
+            "email": email,
+            "username": username,
+            "password": hashed_password,
+            "role": role,
+            "created_at": now_utc,
+            "last_active": now_utc
+        })
+    except pymongo.errors.DuplicateKeyError as e:
+        current_app.logger.warning(
+            "DuplicateKeyError on complete-signup for email=%s username=%s — race condition",
+            email, username,
+        )
+        if e.details and e.details.get("keyPattern", {}).get("email"):
+            return jsonify({"error": "Email already registered"}), 409
+        return jsonify({"error": "Username already taken. Please choose a different one."}), 409
 
     token = create_access_token(
         user_id=str(result.inserted_id),
@@ -354,8 +409,6 @@ def set_password():
         current_app.logger.warning("SET PASSWORD VALIDATION ERROR")
         return jsonify({"error": str(e)}), 400
 
-    if purpose not in ("signup", "reset"):
-        return jsonify({"error": "Invalid purpose. Must be 'signup' or 'reset'."}), 400
 
     if len(password) < 8:
         return jsonify({"error": "Password must be at least 8 characters"}), 400
@@ -376,14 +429,29 @@ def set_password():
         role = "admin" if is_admin_email(email) else "user"
 
         now_utc = datetime.now(timezone.utc)
-        user_id = db.users.insert_one({
-            "email": email,
-            "username": email.split("@")[0],
-            "password": hashed,
-            "role": role,
-            "created_at": now_utc,
-            "last_active": now_utc
-        }).inserted_id
+        user_id = None
+        for _attempt in range(_MAX_SIGNUP_RETRIES):
+            username = _unique_username(email.split("@")[0])
+            try:
+                user_id = db.users.insert_one({
+                    "email": email,
+                    "username": username,
+                    "password": hashed,
+                    "role": role,
+                    "created_at": now_utc,
+                    "last_active": now_utc
+                }).inserted_id
+                break  # success
+            except pymongo.errors.DuplicateKeyError as e:
+                current_app.logger.warning(
+                    "DuplicateKeyError on signup attempt %d/%d for email=%s username=%s",
+                    _attempt + 1, _MAX_SIGNUP_RETRIES, email, username,
+                )
+                if e.details and e.details.get("keyPattern", {}).get("email"):
+                    return jsonify({"error": "Email already registered"}), 409
+                if _attempt == _MAX_SIGNUP_RETRIES - 1:
+                    return jsonify({"error": "Could not assign a unique username. Please try again."}), 409
+
 
         # Cleanup OTPs
         db.email_otps.delete_many({"email": email})
@@ -460,16 +528,37 @@ def google_auth():
 
         # Create a minimal Google-backed user (no local password)
         now_utc = datetime.now(timezone.utc)
-        result = db.users.insert_one({
-            "email": email,
-            "username": name or email.split("@")[0],
-            "password": None,
-            "role": role,
-            "provider": "google",
-            "google_id": sub,
-            "created_at": now_utc,
-            "last_active": now_utc
-        })
+        base_name = name or email.split("@")[0]
+        result = None
+        for _attempt in range(_MAX_SIGNUP_RETRIES):
+            google_username = _unique_username(base_name)
+            try:
+                result = db.users.insert_one({
+                    "email": email,
+                    "username": google_username,
+                    "password": None,
+                    "role": role,
+                    "provider": "google",
+                    "google_id": sub,
+                    "created_at": now_utc,
+                    "last_active": now_utc
+                })
+                break  # success
+            except pymongo.errors.DuplicateKeyError as e:
+                current_app.logger.warning(
+                    "DuplicateKeyError on Google signup attempt %d/%d for email=%s username=%s",
+                    _attempt + 1, _MAX_SIGNUP_RETRIES, email, google_username,
+                )
+                if e.details and e.details.get("keyPattern", {}).get("email"):
+                    user = db.users.find_one({"email": email})
+                    if user:
+                        return jsonify({
+                            "access_token": create_access_token(user_id=str(user["_id"]), role=user.get("role", "user")),
+                            "role": user.get("role", "user")
+                        }), 200
+                    return jsonify({"error": "Failed to authenticate"}), 500
+                if _attempt == _MAX_SIGNUP_RETRIES - 1:
+                    return jsonify({"error": "Could not assign a unique username. Please try again."}), 409
         user_id = str(result.inserted_id)
     else:
         user_id = str(user["_id"])

tests/test_set_password.py

@@ -166,3 +166,82 @@ def test_set_password_signup_duplicate_email(mock_db, client):
     assert resp.status_code == 400
     data = resp.get_json()
     assert data["error"] == "User already exists"
+
+
+# ---------------------------------------------------------------------------
+# Tests — username collision handling
+# ---------------------------------------------------------------------------
+
+
+@patch("routes.auth.is_admin_email", return_value=False)
+@patch("routes.auth.db")
+@patch("routes.auth.create_access_token", return_value="mock-jwt-token")
+def test_set_password_signup_username_collision_resolved(mock_token, mock_db, mock_admin, client):
+    """When the email prefix is already taken as a username, _unique_username
+    must generate a suffixed variant instead of silently inserting a duplicate."""
+    otp_record = {
+        "email": "alice@example.com",
+        "verified": True,
+        "verified_at": datetime.now(timezone.utc),
+    }
+
+    inserted_doc = MagicMock(inserted_id="new-id-456")
+    mock_db.users.insert_one.return_value = inserted_doc
+
+    # Simulate: first find_one (existing-user check) → None,
+    #           second find_one (_unique_username "alice" taken) → existing record,
+    #           third find_one (_unique_username "alice1" free) → None,
+    #           fourth find_one (OTP verification) → otp_record.
+    mock_db.users.find_one.side_effect = [
+        None,          # existing user by email — not found
+        {"_id": "x"}, # "alice" username already taken
+        None,          # "alice1" username is free
+    ]
+    mock_db.email_otps.find_one.return_value = otp_record
+
+    resp = _post_set_password(client, {
+        "email": "alice@example.com",
+        "password": "securepassword123",
+        "purpose": "signup",
+    })
+
+    assert resp.status_code == 200
+
+    # Extract the username that was actually persisted
+    call_args = mock_db.users.insert_one.call_args[0][0]
+    assert call_args["username"] == "alice1", (
+        f"Expected 'alice1' but got '{call_args['username']}'"
+    )
+
+
+@patch("routes.auth.db")
+def test_unique_username_no_collision(mock_db):
+    """_unique_username returns the base name unchanged when it is free."""
+    from routes.auth import _unique_username
+
+    mock_db.users.find_one.return_value = None
+    assert _unique_username("bob") == "bob"
+
+
+@patch("routes.auth.db")
+def test_unique_username_strips_at_sign(mock_db):
+    """_unique_username must strip '@' so the result cannot be parsed as email."""
+    from routes.auth import _unique_username
+
+    mock_db.users.find_one.return_value = None
+    result = _unique_username("user@domain")
+    assert "@" not in result
+
+
+@patch("routes.auth.db")
+def test_unique_username_increments_until_free(mock_db):
+    """_unique_username must keep incrementing the suffix until a slot is free."""
+    from routes.auth import _unique_username
+
+    # "carol" and "carol1" are taken; "carol2" is free
+    # "carol" and "carol1" are taken; "carol2" is free
+    mock_db.users.find.return_value = [
+        {"username": "carol"},
+        {"username": "carol1"},
+    ]
+    assert _unique_username("carol") == "carol2"