fix(oauth): add locking to prevent race conditions in grant exchange (#85570)

Pulled from #82052 to break out these changes into smaller PRs.

Use locks to ensure an `ApiGrant` cannot be used twice during a race
condition that would result in multiple access/refresh token pairs.

Author: mdtro

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
@@ -276,17 +278,33 @@ def handle_async_deletion(
 
     @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/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
@@ -35,12 +36,22 @@ class GrantExchanger:
     def run(self):
         with transaction.atomic(using=router.db_for_write(ApiToken)):
             try:
-                self._validate()
-                token = self._create_token()
+                lock = locks.get(
+                    ApiGrant.get_lock_key(self.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():
+                    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()
+                    # 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",

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)