use option to control rollout rate

Author: mdtro

src/sentry/api/authentication.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import hashlib
+import random
 from collections.abc import Callable, Iterable
 from typing import Any, ClassVar
 
@@ -307,6 +308,17 @@ def authenticate(self, request: Request):
         return self.transform_auth(user_id, None)
 
 
+class TokenStrLookupRequired(Exception):
+    """
+    Used in combination with `apitoken.use-and-update-hash-rate` option.
+
+    If raised, calling code should peform API token lookups based on its
+    plaintext value and not its hashed value.
+    """
+
+    pass
+
+
 @AuthenticationSiloLimit(SiloMode.REGION, SiloMode.CONTROL)
 class UserAuthTokenAuthentication(StandardAuthentication):
     token_name = b"bearer"
@@ -328,11 +340,17 @@ def _find_or_update_token_by_hash(self, token_str: str) -> ApiToken | ApiTokenRe
 
         hashed_token = hashlib.sha256(token_str.encode()).hexdigest()
 
+        rate = options.get("apitoken.use-and-update-hash-rate")
+        random_rate = random.random()
+
         if SiloMode.get_current_mode() == SiloMode.REGION:
             try:
-                # Try to find the token by its hashed value first
-                return ApiTokenReplica.objects.get(hashed_token=hashed_token)
-            except ApiTokenReplica.DoesNotExist:
+                if rate > random_rate:
+                    # Try to find the token by its hashed value first
+                    return ApiTokenReplica.objects.get(hashed_token=hashed_token)
+                else:
+                    raise TokenStrLookupRequired
+            except (ApiTokenReplica.DoesNotExist, TokenStrLookupRequired):
                 try:
                     # If we can't find it by hash, use the plaintext string
                     return ApiTokenReplica.objects.get(token=token_str)
@@ -342,10 +360,13 @@ def _find_or_update_token_by_hash(self, token_str: str) -> ApiToken | ApiTokenRe
         else:
             try:
                 # Try to find the token by its hashed value first
-                return ApiToken.objects.select_related("user", "application").get(
-                    hashed_token=hashed_token
-                )
-            except ApiToken.DoesNotExist:
+                if rate > random_rate:
+                    return ApiToken.objects.select_related("user", "application").get(
+                        hashed_token=hashed_token
+                    )
+                else:
+                    raise TokenStrLookupRequired
+            except (ApiToken.DoesNotExist, TokenStrLookupRequired):
                 try:
                     # If we can't find it by hash, use the plaintext string
                     api_token = ApiToken.objects.select_related("user", "application").get(
@@ -355,9 +376,11 @@ def _find_or_update_token_by_hash(self, token_str: str) -> ApiToken | ApiTokenRe
                     # If the token does not exist by plaintext either, it is not a valid token
                     raise AuthenticationFailed("Invalid token")
                 else:
-                    # Update it with the hashed value if found by plaintext
-                    api_token.hashed_token = hashed_token
-                    api_token.save(update_fields=["hashed_token"])
+                    if rate > random_rate:
+                        # Update it with the hashed value if found by plaintext
+                        api_token.hashed_token = hashed_token
+                        api_token.save(update_fields=["hashed_token"])
+
                     return api_token
 
     def accepts_auth(self, auth: list[bytes]) -> bool:

src/sentry/options/defaults.py

@@ -301,6 +301,14 @@
     flags=FLAG_ALLOW_EMPTY | FLAG_PRIORITIZE_DISK | FLAG_AUTOMATOR_MODIFIABLE,
 )
 
+# Controls the rate of using the hashed value of User API tokens for lookups when logging in
+# and also updates tokens which are not hashed
+register(
+    "apitoken.use-and-update-hash-rate",
+    default=0.0,
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
+
 register(
     "api.rate-limit.org-create",
     default=5,

tests/sentry/api/test_authentication.py

@@ -207,6 +207,7 @@ def test_no_match(self):
             self.auth.authenticate(request)
 
     @override_options({"apitoken.save-hash-on-create": False})
+    @override_options({"apitoken.use-and-update-hash-rate", 1.0})
     def test_token_hashed_with_option_off(self):
         # see https://github.com/getsentry/sentry/pull/65941
         # the UserAuthTokenAuthentication middleware was updated to hash tokens as
@@ -228,6 +229,25 @@ def test_token_hashed_with_option_off(self):
         api_token.refresh_from_db()
         assert api_token.hashed_token == expected_hash
 
+    @override_options({"apitoken.save-hash-on-create": False})
+    @override_options({"apitoken.use-and-update-hash-rate", 0.0})
+    def test_token_not_hashed_with_0_rate(self):
+        api_token = ApiToken.objects.create(user=self.user, token_type=AuthTokenType.USER)
+
+        # we haven't authenticated to the API endpoint yet, so this value should be empty
+        assert api_token.hashed_token is None
+
+        request = HttpRequest()
+        request.META["HTTP_AUTHORIZATION"] = f"Bearer {api_token.token}"
+
+        # trigger the authentication middleware
+        result = self.auth.authenticate(request)
+        assert result is not None
+
+        # check for the expected hash value
+        api_token.refresh_from_db()
+        assert api_token.hashed_token is None
+
 
 @no_silo_test
 class TestTokenAuthenticationReplication(TestCase):
@@ -237,6 +257,7 @@ def setUp(self):
         self.auth = UserAuthTokenAuthentication()
 
     @override_options({"apitoken.save-hash-on-create": False})
+    @override_options({"apitoken.use-and-update-hash-rate", 1.0})
     def test_hash_is_replicated(self):
         api_token = ApiToken.objects.create(user=self.user, token_type=AuthTokenType.USER)
         expected_hash = hashlib.sha256(api_token.token.encode()).hexdigest()