fix: prevent race in concurrent worker registration

Clients running several threads under one worker name send their first pops together, which could hit the unique name constraint or bypass the worker limits.

- Add a per-user Postgres advisory lock (skipped on SQLite) so new-worker registrations by one user are serialized.
- Re-lookup the worker after taking the lock and refresh the user's workers before checking limits, so a sibling pop's worker is reused rather than refused.
- On IntegrityError, roll back and return the concurrent worker, or raise WrongCredentials/PolymorphicNameConflict as a later pop would.
- Move the limit checks into check_worker_limits().
- Reject worker names longer than the column width with TooLongWorkerName, since a truncated name never matches its own row.
- Add integration tests for these cases.

horde/apis/v2/base.py

@@ -48,7 +48,7 @@ from horde.classes.base.style_contract import (
 from horde.classes.base.team import Team, find_team_by_id, find_team_by_name, get_all_teams
 from horde.classes.base.user import User, UserSharedKey
 from horde.classes.base.waiting_prompt import WaitingPrompt
-from horde.classes.base.worker import Worker, WorkerMessage
+from horde.classes.base.worker import MAX_WORKER_NAME_CHARACTERS, Worker, WorkerMessage
 from horde.consts import HORDE_VERSION
 from horde.countermeasures import REJECTION_WINDOW, CounterMeasures
 from horde.database import functions as database
@@ -1215,6 +1215,15 @@ class JobPopTemplate(Resource):
         if self.user.is_anon():
             raise e.AnonForbidden(rc="AnonForbiddenWorker")
         self.worker_name = sanitize_string(self.args["name"])
+        # The stored name is cut to the column width, so a longer name never matches its own row.
+        if len(self.worker_name) > MAX_WORKER_NAME_CHARACTERS:
+            raise e.TooLong(
+                self.user.get_unique_alias(),
+                len(self.worker_name),
+                MAX_WORKER_NAME_CHARACTERS,
+                "worker name",
+                rc="TooLongWorkerName",
+            )
         self.worker = database.find_worker_by_name(self.worker_name, worker_class=self.worker_class)
         if not self.worker and database.worker_name_exists(self.worker_name):
             raise e.PolymorphicNameConflict(self.worker_name)
@@ -1226,37 +1235,73 @@ class JobPopTemplate(Resource):
                 raise e.BadRequest(
                     f"To avoid unwanted attention, please do not use '{cs.group()}' in your worker names.",
                 )
-            worker_count = self.user.count_workers()
-            if settings.mode_invite_only() and worker_count >= self.user.worker_invited:
-                raise e.WorkerInviteOnly(worker_count)
-            # Untrusted users can only have 3 workers
-            if not self.user.trusted and worker_count > 3:
-                raise e.Forbidden(
-                    "To avoid abuse, untrusted users can only have up to 3 distinct workers.",
-                    rc="TooManyWorkers",
-                )
-            # Trusted users can have up to 20 workers by default unless overriden
-            if worker_count > 20 and worker_count > self.user.worker_invited:
-                raise e.Forbidden(
-                    "To avoid abuse, tou cannot onboard more than 20 workers as a trusted user. Please contact us on Discord to adjust.",
-                    rc="TooManyWorkersTrusted",
-                )
-            if self.user.exceeding_ipaddr_restrictions(self.worker_ip):
-                raise e.TooManySameIPs(self.user.username)
+            self.check_worker_limits()
             # The profanity checks run a model prediction, so they come after the cheap refusals: a user at a worker
             # limit who keeps popping with new names does not cost a prediction per pop.
             if is_profane(self.worker_name):
                 raise e.Profanity(self.user.get_unique_alias(), self.worker_name, "worker name", rc="ProfaneWorkerName")
             if is_profane(self.args.bridge_agent):
                 raise e.Profanity(self.user.get_unique_alias(), self.args.bridge_agent, "bridge agent", rc="ProfaneBridgeAgent")
-            self.worker = self.worker_class(
-                user_id=self.user.id,
-                name=self.worker_name,
-            )
-            self.worker.create()
+            self.worker = self.register_worker()
         if self.user != self.worker.user:
             raise e.WrongCredentials(self.user.get_unique_alias(), self.worker_name)
 
+    def check_worker_limits(self):
+        """Refuse a new worker when the user is at the invite, worker-count or same-IP limit."""
+        worker_count = self.user.count_workers()
+        if settings.mode_invite_only() and worker_count >= self.user.worker_invited:
+            raise e.WorkerInviteOnly(worker_count)
+        # Untrusted users can only have 3 workers
+        if not self.user.trusted and worker_count > 3:
+            raise e.Forbidden(
+                "To avoid abuse, untrusted users can only have up to 3 distinct workers.",
+                rc="TooManyWorkers",
+            )
+        # Trusted users can have up to 20 workers by default unless overriden
+        if worker_count > 20 and worker_count > self.user.worker_invited:
+            raise e.Forbidden(
+                "To avoid abuse, tou cannot onboard more than 20 workers as a trusted user. Please contact us on Discord to adjust.",
+                rc="TooManyWorkersTrusted",
+            )
+        if self.user.exceeding_ipaddr_restrictions(self.worker_ip):
+            raise e.TooManySameIPs(self.user.username)
+
+    def register_worker(self):
+        """Create the worker for this pop or return the row a concurrent pop created under the same name.
+
+        Clients that run several threads under one worker name could send their first pops together. The per-user lock
+        serializes those pops and any other new-worker registrations by the same user, so the limits checked under it
+        count every worker the user registered first. Pops by different users take different locks, so their inserts
+        can still collide on the unique worker name. The losing pop then gets the response a pop sent after the
+        winner's commit would get.
+        """
+        database.acquire_worker_registration_lock(self.user.id)
+        # Looked up again before the limits, so a pop that lost to its own sibling thread is not refused for the
+        # worker that sibling just registered.
+        worker = database.find_worker_by_name(self.worker_name, worker_class=self.worker_class)
+        if worker:
+            return worker
+        if database.worker_name_exists(self.worker_name):
+            raise e.PolymorphicNameConflict(self.worker_name)
+        # The collection was loaded before the lock and does not include workers committed since.
+        db.session.refresh(self.user, attribute_names=["workers"])
+        self.check_worker_limits()
+        worker = self.worker_class(
+            user_id=self.user.id,
+            name=self.worker_name,
+        )
+        try:
+            worker.create()
+        except IntegrityError:
+            db.session.rollback()
+            concurrent_worker = database.find_worker_by_name(self.worker_name, worker_class=self.worker_class)
+            if concurrent_worker:
+                return concurrent_worker
+            if database.worker_name_exists(self.worker_name):
+                raise e.PolymorphicNameConflict(self.worker_name) from None
+            raise
+        return worker
+
     def check_ip(self):
         ip_timeout = CounterMeasures.retrieve_timeout(self.worker_ip)
         if ip_timeout:

horde/database/functions.py

@@ -12,7 +12,7 @@ from dataclasses import dataclass
 from datetime import datetime, timedelta
 
 import logfire
-from sqlalchemy import Boolean, and_, case, func, not_, or_
+from sqlalchemy import Boolean, and_, case, func, not_, or_, select
 from sqlalchemy.orm import contains_eager, joinedload, noload, selectinload, subqueryload
 
 import horde.classes.base.stats as stats
@@ -384,6 +384,20 @@ def worker_name_exists(worker_name):
     return False
 
 
+WORKER_REGISTRATION_LOCK_NAMESPACE = 0x574B5247  # "WKRG"
+
+
+def acquire_worker_registration_lock(user_id: int) -> None:
+    """Serialize new-worker registration for one user until the current transaction ends.
+
+    The worker limits count the user's workers before the insert, so two concurrent pops that each register a new
+    name would both pass them without this lock.
+    """
+    if SQLITE_MODE:
+        return
+    db.session.execute(select(func.pg_advisory_xact_lock(WORKER_REGISTRATION_LOCK_NAMESPACE, user_id))).scalar_one()
+
+
 def find_worker_by_id(worker_id):
     try:
         worker_uuid = uuid.UUID(worker_id)

tests/integration/test_worker_ops.py

@@ -99,6 +99,133 @@ class TestPopAuth:
         assert _worker_id_by_name(client, owner.api_key, "Lonely Scribe")
 
 
+# --------------------------------------------------------------------------- #
+# Concurrent worker registration                                              #
+# --------------------------------------------------------------------------- #
+
+
+def _store_worker(app, user_id: int, name: str, worker_kind: str = "text") -> None:
+    """Commit a worker row directly, standing in for a concurrent pop that registered it."""
+    from horde.classes.kobold.worker import TextWorker
+    from horde.classes.stable.worker import ImageWorker
+    from horde.flask import db
+
+    worker_class = {"text": TextWorker, "image": ImageWorker}[worker_kind]
+    with app.app_context():
+        db.session.add(worker_class(user_id=user_id, name=name))
+        db.session.commit()
+
+
+def _count_workers_named(app, name: str) -> int:
+    from horde.classes.base.worker import WorkerTemplate
+    from horde.flask import db
+
+    with app.app_context():
+        return db.session.query(WorkerTemplate).filter(WorkerTemplate.name == name).count()
+
+
+def _hide_first_calls(monkeypatch, function_name: str, hidden_calls: int, hidden_result) -> None:
+    """Make the first lookups of a pop miss a committed row, as they do when the row commits after them."""
+    from horde.database import functions as database
+
+    real_function = getattr(database, function_name)
+    calls = {"count": 0}
+
+    def _lookup(*args, **kwargs):
+        calls["count"] += 1
+        if calls["count"] <= hidden_calls:
+            return hidden_result
+        return real_function(*args, **kwargs)
+
+    monkeypatch.setattr(database, function_name, _lookup)
+
+
+def _allow_untrusted_ip(monkeypatch) -> None:
+    from horde.apis.v2 import base as api_base
+
+    monkeypatch.setattr(api_base.CounterMeasures, "is_ip_safe", staticmethod(lambda ipaddr: True))
+
+
+def _commit_worker_before_lock(monkeypatch, name: str) -> None:
+    """Commit a worker for the popping user just before the pop takes the registration lock.
+
+    This is the order of events when another pop by the same user registers while this one waits on the lock: the
+    limits this pop checked before the lock did not count that worker.
+    """
+    from horde.classes.kobold.worker import TextWorker
+    from horde.database import functions as database
+    from horde.flask import db
+
+    real_lock = database.acquire_worker_registration_lock
+
+    def _lock_after_concurrent_registration(user_id):
+        db.session.add(TextWorker(user_id=user_id, name=name))
+        db.session.commit()
+        real_lock(user_id)
+
+    monkeypatch.setattr(database, "acquire_worker_registration_lock", _lock_after_concurrent_registration)
+
+
+class TestPopWorkerRegistration:
+    def test_sibling_registration_while_waiting_is_reused_at_the_limit(self, app, client, make_api_user, monkeypatch):
+        owner = make_api_user(trusted=False, kudos=100)
+        for index in range(3):
+            _store_worker(app, owner.id, f"Sibling Scribe {index}")
+        _allow_untrusted_ip(monkeypatch)
+        _commit_worker_before_lock(monkeypatch, "Sibling Scribe 3")
+
+        resp = _pop_text(client, owner.api_key, "Sibling Scribe 3")
+
+        assert resp.status_code == 200, resp.get_data(as_text=True)
+        assert _count_workers_named(app, "Sibling Scribe 3") == 1
+
+    def test_insert_collision_with_another_users_worker_is_wrong_credentials(self, app, client, make_api_user, monkeypatch):
+        owner = make_api_user(trusted=True, kudos=100)
+        rival = make_api_user(trusted=True, kudos=100)
+        _store_worker(app, rival.id, "Contested Scribe")
+        _hide_first_calls(monkeypatch, "find_worker_by_name", hidden_calls=2, hidden_result=None)
+        _hide_first_calls(monkeypatch, "worker_name_exists", hidden_calls=2, hidden_result=False)
+
+        resp = _pop_text(client, owner.api_key, "Contested Scribe")
+
+        assert resp.status_code == 403, resp.get_data(as_text=True)
+        assert resp.get_json()["rc"] == "WrongCredentials"
+        assert _count_workers_named(app, "Contested Scribe") == 1
+
+    def test_insert_collision_with_another_worker_type_is_polymorphic_conflict(self, app, client, make_api_user, monkeypatch):
+        owner = make_api_user(trusted=True, kudos=100)
+        _store_worker(app, owner.id, "Painter Scribe", worker_kind="image")
+        _hide_first_calls(monkeypatch, "worker_name_exists", hidden_calls=2, hidden_result=False)
+
+        resp = _pop_text(client, owner.api_key, "Painter Scribe")
+
+        assert resp.status_code == 400, resp.get_data(as_text=True)
+        assert resp.get_json()["rc"] == "PolymorphicNameConflict"
+
+    def test_worker_limit_counts_registrations_committed_while_waiting(self, app, client, make_api_user, monkeypatch):
+        owner = make_api_user(trusted=False, kudos=100)
+        for index in range(3):
+            _store_worker(app, owner.id, f"Untrusted Scribe {index}")
+        _allow_untrusted_ip(monkeypatch)
+        _commit_worker_before_lock(monkeypatch, "Untrusted Scribe 3")
+
+        resp = _pop_text(client, owner.api_key, "Untrusted Scribe 4")
+
+        assert resp.status_code == 403, resp.get_data(as_text=True)
+        assert resp.get_json()["rc"] == "TooManyWorkers"
+        assert _count_workers_named(app, "Untrusted Scribe 4") == 0
+
+    def test_name_longer_than_the_column_is_rejected_before_lookup(self, app, client, make_api_user):
+        owner = make_api_user(trusted=True, kudos=100)
+        long_name = "Verbose Scribe " + "x" * 100
+
+        resp = _pop_text(client, owner.api_key, long_name)
+
+        assert resp.status_code == 400, resp.get_data(as_text=True)
+        assert resp.get_json()["rc"] == "TooLongWorkerName"
+        assert _count_workers_named(app, long_name[:100]) == 0
+
+
 # --------------------------------------------------------------------------- #
 # Worker details + modify permission matrix                                   #
 # --------------------------------------------------------------------------- #