CVE-2024-2032: Race Condition within a Thread (CWE-366)
value: A race condition vulnerability exists in zenml-io/zenml versions up to and including 0.55.3, which allows for the creation of multiple users with the same username when requests are sent in parallel. This issue was fixed in version 0.55.5. The vulnerability arises due to insufficient handling of concurrent user creation requests, leading to data inconsistencies and potential authentication problems. Specifically, concurrent processes may overwrite or corrupt user data, complicating user identification and posing security risks. This issue is particularly concerning for APIs that rely on usernames as input parameters, such as PUT /api/v1/users/test_race, where it could lead to further complications.
CWE description: If two threads of execution use a resource simultaneously, there exists the possibility that resources may be used while invalid, in turn making the state of execution undefined.
Published: 2024-06-06T19:15Z
Repository: https://github.com/zenml-io/zenml
Language: Python
Files changed: src/zenml/zen_stores/migrations/versions/72675226b2de_unique_users.py, src/zenml/zen_stores/schemas/user_schemas.py, src/zenml/zen_stores/sql_zen_store.py, tests/integration/functional/zen_stores/test_zen_store.py
Commit date: 2024-03-04 10:05:43 +0100
Commit message: Uniquely constrained users table (#2483)

* use db constrains for users

* lint a bit

* coderabbitai

* set rate limits for login api

* Revert "set rate limits for login api"

This reverts commit 78ce40e738e73ce7d9398c73559f3d29a502c24f.

* bring checks back

* Auto-update of Starter template

* Auto-update of NLP template

* review suggestions

* resolve branching

* Auto-update of Starter template

* Auto-update of NLP template

* lint

---------

Co-authored-by: GitHub Actions <actions@github.com>
Diff: diff --git a/src/zenml/zen_stores/migrations/versions/72675226b2de_unique_users.py b/src/zenml/zen_stores/migrations/versions/72675226b2de_unique_users.py
new file mode 100644
index 000000000..9c3532cbe
--- /dev/null
+++ b/src/zenml/zen_stores/migrations/versions/72675226b2de_unique_users.py
@@ -0,0 +1,31 @@
+"""unique users [72675226b2de].
+
+Revision ID: 72675226b2de
+Revises: 0.55.4
+Create Date: 2024-02-29 14:58:25.584731
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "72675226b2de"
+down_revision = "0.55.4"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "uq_user_name_is_service_account", ["name", "is_service_account"]
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "uq_user_name_is_service_account", type_="unique"
+        )
diff --git a/src/zenml/zen_stores/schemas/user_schemas.py b/src/zenml/zen_stores/schemas/user_schemas.py
index 56129fd0b..744301b7c 100644
--- a/src/zenml/zen_stores/schemas/user_schemas.py
+++ b/src/zenml/zen_stores/schemas/user_schemas.py
@@ -15,11 +15,11 @@
 
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.models import (
     ServiceAccountRequest,
     ServiceAccountResponse,
@@ -63,10 +63,11 @@ if TYPE_CHECKING:
 
 class UserSchema(NamedSchema, table=True):
     """SQL Model for users."""
 
     __tablename__ = "user"
+    __table_args__ = (UniqueConstraint("name", "is_service_account"),)
 
     is_service_account: bool = Field(default=False)
     full_name: str
     description: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
     email: Optional[str] = Field(nullable=True)
diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py
index d28fa8a19..c19dee22b 100644
--- a/src/zenml/zen_stores/sql_zen_store.py
+++ b/src/zenml/zen_stores/sql_zen_store.py
@@ -5144,32 +5144,33 @@ class SqlZenStore(BaseZenStore):
         Raises:
             EntityExistsError: If a user or service account with the given name
                 already exists.
         """
         with Session(self.engine) as session:
+            # Check if a service account with the given name already
+            # exists
+            err_msg = (
+                f"Unable to create service account with name "
+                f"'{service_account.name}': Found existing service "
+                "account with this name."
+            )
+            try:
+                self._get_account_schema(
+                    service_account.name, session=session, service_account=True
+                )
+                raise EntityExistsError(err_msg)
+            except KeyError:
+                pass
+
             # Create the service account
             new_account = UserSchema.from_service_account_request(
                 service_account
             )
             session.add(new_account)
+            # on commit an IntegrityError may arise we let it bubble up
+            session.commit()
 
-            # Check if a service account with the given name already
-            # exists
-            service_accounts = session.execute(
-                select(UserSchema).where(
-                    UserSchema.name == service_account.name,
-                    UserSchema.is_service_account.is_(True),  # type: ignore[attr-defined]
-                )
-            ).fetchall()
-            if len(service_accounts) == 1:
-                session.commit()
-            else:
-                raise EntityExistsError(
-                    f"Unable to create service account with name "
-                    f"'{service_account.name}': Found existing service "
-                    "account with this name."
-                )
             return new_account.to_service_account_model(include_metadata=True)
 
     def get_service_account(
         self,
         service_account_name_or_id: Union[str, UUID],
@@ -7355,28 +7356,31 @@ class SqlZenStore(BaseZenStore):
         Raises:
             EntityExistsError: If a user or service account with the given name
                 already exists.
         """
         with Session(self.engine) as session:
+            # Check if a user account with the given name already exists
+            err_msg = (
+                f"Unable to create user with name '{user.name}': "
+                f"Found an existing user account with this name."
+            )
+            try:
+                self._get_account_schema(
+                    user.name,
+                    session=session,
+                    # Filter out service accounts
+                    service_account=False,
+                )
+                raise EntityExistsError(err_msg)
+            except KeyError:
+                pass
+
             # Create the user
             new_user = UserSchema.from_user_request(user)
             session.add(new_user)
-
-            # Check if a user account with the given name already exists
-            users = session.execute(
-                select(UserSchema).where(
-                    UserSchema.name == user.name,
-                    UserSchema.is_service_account.is_(False),  # type: ignore[attr-defined]
-                )
-            ).fetchall()
-            if len(users) == 1:
-                session.commit()
-            else:
-                raise EntityExistsError(
-                    f"Unable to create user with name '{user.name}': "
-                    f"Found an existing user account with this name."
-                )
+            # on commit an IntegrityError may arise we let it bubble up
+            session.commit()
             return new_user.to_model(include_metadata=True)
 
     def get_user(
         self,
         user_name_or_id: Optional[Union[str, UUID]] = None,
diff --git a/tests/integration/functional/zen_stores/test_zen_store.py b/tests/integration/functional/zen_stores/test_zen_store.py
index ed6c72529..1c0261b54 100644
--- a/tests/integration/functional/zen_stores/test_zen_store.py
+++ b/tests/integration/functional/zen_stores/test_zen_store.py
@@ -20,10 +20,11 @@ from threading import Thread
 from typing import Dict, List, Optional, Tuple
 from uuid import UUID, uuid4
 
 import pytest
 from pydantic import SecretStr
+from sqlalchemy.exc import IntegrityError
 
 from tests.integration.functional.utils import sample_name
 from tests.integration.functional.zen_stores.utils import (
     CodeRepositoryContext,
     ComponentContext,
@@ -420,11 +421,11 @@ def test_creating_users_in_parallel_do_not_duplicate_fails(
         if a duplicate user exists. This is used to simulate race
         conditions in parallel user creation.
         """
         try:
             clean_client.zen_store.create_user(user_request)
-        except EntityExistsError:
+        except (EntityExistsError, IntegrityError):
             pass
 
     user_name = "test_user"
     password = "P@ssw0rd"
     count = 100
@@ -461,11 +462,11 @@ def test_creating_service_accounts_in_parallel_do_not_duplicate_fails(
         """
         try:
             clean_client.zen_store.create_service_account(
                 service_account_request
             )
-        except EntityExistsError:
+        except (EntityExistsError, IntegrityError):
             pass
 
     user_name = "test_user"
     count = 100
 

Vulnerable code: from sqlalchemy import TEXT, Column
            # Check if a service account with the given name already
            # exists
            service_accounts = session.execute(
                select(UserSchema).where(
                    UserSchema.name == service_account.name,
                    UserSchema.is_service_account.is_(True),  # type: ignore[attr-defined]
                )
            ).fetchall()
            if len(service_accounts) == 1:
                session.commit()
            else:
                raise EntityExistsError(
                    f"Unable to create service account with name "
                    f"'{service_account.name}': Found existing service "
                    "account with this name."
                )

            # Check if a user account with the given name already exists
            users = session.execute(
                select(UserSchema).where(
                    UserSchema.name == user.name,
                    UserSchema.is_service_account.is_(False),  # type: ignore[attr-defined]
                )
            ).fetchall()
            if len(users) == 1:
                session.commit()
            else:
                raise EntityExistsError(
                    f"Unable to create user with name '{user.name}': "
                    f"Found an existing user account with this name."
                )
        except EntityExistsError:
        except EntityExistsError:
Fixed code: """unique users [72675226b2de].

Revision ID: 72675226b2de
Revises: 0.55.4
Create Date: 2024-02-29 14:58:25.584731

"""

from alembic import op

# revision identifiers, used by Alembic.
revision = "72675226b2de"
down_revision = "0.55.4"
branch_labels = None
depends_on = None


def upgrade() -> None:
    """Upgrade database schema and/or data, creating a new revision."""
    with op.batch_alter_table("user", schema=None) as batch_op:
        batch_op.create_unique_constraint(
            "uq_user_name_is_service_account", ["name", "is_service_account"]
        )


def downgrade() -> None:
    """Downgrade database schema and/or data back to the previous revision."""
    with op.batch_alter_table("user", schema=None) as batch_op:
        batch_op.drop_constraint(
            "uq_user_name_is_service_account", type_="unique"
        )
from sqlalchemy import TEXT, Column, UniqueConstraint
    __table_args__ = (UniqueConstraint("name", "is_service_account"),)
            # Check if a service account with the given name already
            # exists
            err_msg = (
                f"Unable to create service account with name "
                f"'{service_account.name}': Found existing service "
                "account with this name."
            )
            try:
                self._get_account_schema(
                    service_account.name, session=session, service_account=True
                )
                raise EntityExistsError(err_msg)
            except KeyError:
                pass

            # on commit an IntegrityError may arise we let it bubble up
            session.commit()
            # Check if a user account with the given name already exists
            err_msg = (
                f"Unable to create user with name '{user.name}': "
                f"Found an existing user account with this name."
            )
            try:
                self._get_account_schema(
                    user.name,
                    session=session,
                    # Filter out service accounts
                    service_account=False,
                )
                raise EntityExistsError(err_msg)
            except KeyError:
                pass

            # on commit an IntegrityError may arise we let it bubble up
            session.commit()
from sqlalchemy.exc import IntegrityError
        except (EntityExistsError, IntegrityError):
        except (EntityExistsError, IntegrityError):