fix: add locking to grants to avoid race conditions

Author: mdtro

src/sentry/api/endpoints/api_authorizations.py

@@ -12,6 +12,7 @@
 from sentry.hybridcloud.models.outbox import outbox_context
 from sentry.models.apiapplication import ApiApplicationStatus
 from sentry.models.apiauthorization import ApiAuthorization
+from sentry.models.apigrant import ApiGrant
 from sentry.models.apitoken import ApiToken
 
 
@@ -44,18 +45,33 @@ def delete(self, request: Request) -> Response:
             return Response({"authorization": ""}, status=400)
 
         try:
-            auth = ApiAuthorization.objects.get(user_id=request.user.id, id=authorization)
+            api_authorization = ApiAuthorization.objects.get(
+                user_id=request.user.id, id=authorization
+            )
         except ApiAuthorization.DoesNotExist:
             return Response(status=404)
 
         with outbox_context(transaction.atomic(using=router.db_for_write(ApiToken)), flush=False):
             for token in ApiToken.objects.filter(
                 user_id=request.user.id,
-                application=auth.application_id,
-                scoping_organization_id=auth.organization_id,
+                application=api_authorization.application_id,
+                scoping_organization_id=api_authorization.organization_id,
             ):
                 token.delete()
 
-            auth.delete()
+        # remove any grants that were created from this authorization
+        # that may not have been exchanged for a token yet
+        with outbox_context(transaction.atomic(using=router.db_for_write(ApiGrant)), flush=False):
+            for grant in ApiGrant.objects.filter(
+                user_id=request.user.id,
+                application=api_authorization.application_id,
+                organization_id=api_authorization.organization_id,
+            ):
+                grant.delete()
+
+        with outbox_context(
+            transaction.atomic(using=router.db_for_write(ApiAuthorization)), flush=False
+        ):
+            api_authorization.delete()
 
         return Response(status=204)

src/sentry/models/apigrant.py

@@ -15,6 +15,14 @@
 DEFAULT_EXPIRATION = timedelta(minutes=10)
 
 
+class InvalidGrantError(Exception):
+    pass
+
+
+class ExpiredGrantError(Exception):
+    pass
+
+
 def default_expiration():
     return timezone.now() + DEFAULT_EXPIRATION
 
@@ -97,6 +105,10 @@ def is_expired(self):
     def redirect_uri_allowed(self, uri):
         return uri == self.redirect_uri
 
+    @classmethod
+    def get_lock_key(cls, grant_id):
+        return f"api_grant:{grant_id}"
+
     @classmethod
     def sanitize_relocation_json(
         cls, json: Any, sanitizer: Sanitizer, model_name: NormalizedModelName | None = None

src/sentry/models/apitoken.py

@@ -19,7 +19,9 @@
 from sentry.db.models.fields.hybrid_cloud_foreign_key import HybridCloudForeignKey
 from sentry.hybridcloud.outbox.base import ControlOutboxProducingManager, ReplicatedControlModel
 from sentry.hybridcloud.outbox.category import OutboxCategory
-from sentry.models.apigrant import ApiGrant
+from sentry.locks import locks
+from sentry.models.apiapplication import ApiApplicationStatus
+from sentry.models.apigrant import ApiGrant, ExpiredGrantError, InvalidGrantError
 from sentry.models.apiscopes import HasApiScopes
 from sentry.types.region import find_all_region_names
 from sentry.types.token import AuthTokenType
@@ -261,17 +263,33 @@ def handle_async_replication(self, region_name: str, shard_identifier: int) -> N
 
     @classmethod
     def from_grant(cls, grant: ApiGrant):
-        with transaction.atomic(router.db_for_write(cls)):
-            api_token = cls.objects.create(
-                application=grant.application,
-                user=grant.user,
-                scope_list=grant.get_scopes(),
-                scoping_organization_id=grant.organization_id,
-            )
-
-            # remove the ApiGrant from the database to prevent reuse of the same
-            # authorization code
-            grant.delete()
+        if grant.application.status != ApiApplicationStatus.active:
+            raise InvalidGrantError()
+
+        if grant.is_expired():
+            raise ExpiredGrantError()
+
+        lock = locks.get(
+            ApiGrant.get_lock_key(grant.id),
+            duration=10,
+            name="api_grant",
+        )
+
+        # we use a lock to prevent race conditions when creating the ApiToken
+        # an attacker could send two requests to create an access/refresh token pair
+        # at the same time, using the same grant, and get two different tokens
+        with lock.acquire():
+            with transaction.atomic(router.db_for_write(cls)):
+                api_token = cls.objects.create(
+                    application=grant.application,
+                    user=grant.user,
+                    scope_list=grant.get_scopes(),
+                    scoping_organization_id=grant.organization_id,
+                )
+
+                # remove the ApiGrant from the database to prevent reuse of the same
+                # authorization code
+                grant.delete()
 
             return api_token
 

src/sentry/sentry_apps/api/endpoints/sentry_app_authorizations.py

@@ -14,6 +14,7 @@
 from sentry.sentry_apps.token_exchange.grant_exchanger import GrantExchanger
 from sentry.sentry_apps.token_exchange.refresher import Refresher
 from sentry.sentry_apps.token_exchange.util import GrantTypes
+from sentry.utils.locking import UnableToAcquireLock
 from sentry.sentry_apps.utils.errors import SentryAppIntegratorError
 
 logger = logging.getLogger(__name__)
@@ -86,6 +87,8 @@ def post(self, request: Request, installation) -> Response:
                 },
             )
             raise
+        except UnableToAcquireLock:
+            return Response({"error": "invalid_grant"}, status=409)
 
         attrs = {"state": request.data.get("state"), "application": None}
 

src/sentry/sentry_apps/token_exchange/grant_exchanger.py

@@ -6,6 +6,7 @@
 from django.utils.functional import cached_property
 
 from sentry import analytics
+from sentry.locks import locks
 from sentry.models.apiapplication import ApiApplication
 from sentry.models.apigrant import ApiGrant
 from sentry.models.apitoken import ApiToken
@@ -33,23 +34,29 @@ class GrantExchanger:
     user: User
 
     def run(self):
-        with transaction.atomic(using=router.db_for_write(ApiToken)):
-            try:
-                self._validate()
-                token = self._create_token()
-
-                # Once it's exchanged it's no longer valid and should not be
-                # exchangeable, so we delete it.
-                self._delete_grant()
-            except SentryAppIntegratorError:
-                logger.info(
-                    "grant-exchanger.context",
-                    extra={
-                        "application_id": self.application.id,
-                        "grant_id": self.grant.id,
-                    },
-                )
-                raise
+        lock = locks.get(
+            ApiGrant.get_lock_key(self.grant.id),
+            duration=10,
+            name="api_grant",
+        )
+        with lock.acquire():
+            with transaction.atomic(using=router.db_for_write(ApiToken)):
+                try:
+                    self._validate()
+                    token = self._create_token()
+
+                    # Once it's exchanged it's no longer valid and should not be
+                    # exchangeable, so we delete it.
+                    self._delete_grant()
+                except SentryAppIntegratorError:
+                    logger.info(
+                        "grant-exchanger.context",
+                        extra={
+                            "application_id": self.application.id,
+                            "grant_id": self.grant.id,
+                        },
+                    )
+                    raise
         self.record_analytics()
 
         return token

src/sentry/web/frontend/oauth_token.py

@@ -18,6 +18,7 @@
 from sentry.models.apitoken import ApiToken
 from sentry.sentry_apps.token_exchange.util import GrantTypes
 from sentry.utils import json, metrics
+from sentry.utils.locking import UnableToAcquireLock
 from sentry.web.frontend.base import control_silo_view
 from sentry.web.frontend.openidtoken import OpenIDToken
 
@@ -128,7 +129,9 @@ def post(self, request: Request) -> HttpResponse:
     def get_access_tokens(self, request: Request, application: ApiApplication) -> dict:
         code = request.POST.get("code")
         try:
-            grant = ApiGrant.objects.get(application=application, code=code)
+            grant = ApiGrant.objects.get(
+                application=application, application__status=ApiApplicationStatus.active, code=code
+            )
         except ApiGrant.DoesNotExist:
             return {"error": "invalid_grant", "reason": "invalid grant"}
 
@@ -141,7 +144,12 @@ def get_access_tokens(self, request: Request, application: ApiApplication) -> di
         elif grant.redirect_uri != redirect_uri:
             return {"error": "invalid_grant", "reason": "invalid redirect URI"}
 
-        token_data = {"token": ApiToken.from_grant(grant=grant)}
+        try:
+            token_data = {"token": ApiToken.from_grant(grant=grant)}
+        except UnableToAcquireLock:
+            # TODO(mdtro): we should return a 409 status code here
+            return {"error": "invalid_grant", "reason": "invalid grant"}
+
         if grant.has_scope("openid") and options.get("codecov.signing_secret"):
             open_id_token = OpenIDToken(
                 request.POST.get("client_id"),
@@ -150,6 +158,7 @@ def get_access_tokens(self, request: Request, application: ApiApplication) -> di
                 nonce=request.POST.get("nonce"),
             )
             token_data["id_token"] = open_id_token.get_signed_id_token(grant=grant)
+
         return token_data
 
     def get_refresh_token(self, request: Request, application: ApiApplication) -> dict:

tests/sentry/sentry_apps/token_exchange/test_grant_exchanger.py

@@ -104,6 +104,22 @@ def test_deletes_grant_on_successful_exchange(self):
         self.grant_exchanger.run()
         assert not ApiGrant.objects.filter(id=grant_id)
 
+    def test_race_condition_on_grant_exchange(self):
+        from sentry.locks import locks
+        from sentry.utils.locking import UnableToAcquireLock
+
+        # simulate a race condition on the grant exchange
+        grant_id = self.orm_install.api_grant_id
+        lock = locks.get(
+            ApiGrant.get_lock_key(grant_id),
+            duration=10,
+            name="api_grant",
+        )
+        lock.acquire()
+
+        with pytest.raises(UnableToAcquireLock):
+            self.grant_exchanger.run()
+
     @patch("sentry.analytics.record")
     def test_records_analytics(self, record):
         GrantExchanger(

tests/sentry/web/frontend/test_oauth_token.py

@@ -2,6 +2,7 @@
 
 from django.utils import timezone
 
+from sentry.locks import locks
 from sentry.models.apiapplication import ApiApplication
 from sentry.models.apigrant import ApiGrant
 from sentry.models.apitoken import ApiToken
@@ -203,6 +204,28 @@ def test_one_time_use_grant(self):
         )
         assert resp.status_code == 400
 
+    def test_grant_lock(self):
+        self.login_as(self.user)
+
+        # Simulate a concurrent request by using an existing grant
+        # that has its grant lock taken out.
+        lock = locks.get(ApiGrant.get_lock_key(self.grant.id), duration=10, name="api_grant")
+        lock.acquire()
+
+        # Attempt to create a token with the same grant
+        # This should fail because the lock is held by the previous request
+        resp = self.client.post(
+            self.path,
+            {
+                "grant_type": "authorization_code",
+                "code": self.grant.code,
+                "client_id": self.application.client_id,
+                "client_secret": self.client_secret,
+            },
+        )
+        assert resp.status_code == 400
+        assert resp.json() == {"error": "invalid_grant"}
+
     def test_invalid_redirect_uri(self):
         self.login_as(self.user)