Add env FXA_TOKEN_AUTH_VERSION

Setting FXA_TOKEN_AUTH_VERSION to 2024 (the default) will use the
existing authentication for FxA bearer tokens. Setting it to 2025 will
use the new authentication method. When the new authentication is
proven, the 2024 version can be deleted.

api/authentication.py

@@ -14,10 +14,17 @@
 from allauth.socialaccount.models import SocialAccount
 from codetiming import Timer
 from markus.utils import generate_tag
-from rest_framework.authentication import TokenAuthentication
+from rest_framework.authentication import (
+    BaseAuthentication,
+    TokenAuthentication,
+    get_authorization_header,
+)
 from rest_framework.exceptions import (
     APIException,
     AuthenticationFailed,
+    NotFound,
+    ParseError,
+    PermissionDenied,
 )
 from rest_framework.request import Request
 
@@ -28,6 +35,28 @@
     settings.SOCIALACCOUNT_PROVIDERS["fxa"]["OAUTH_ENDPOINT"]
 )
 
+# Specify the version strings in FXA_TOKEN_AUTH_VERSION
+#
+# The older version ("2024") works, but has a few issues.
+# The cache key changes between Python instances, so little or no cache hits are used.
+# Fetching a profile takes a few seconds, in which time another process can create a
+# SocialAccount, leading to IntegrityError. Some of these are tracked in MPP-3505.
+#
+# The newer version ("2025") addresses these issues, works more like a standard DRF
+# authentication class, expands the logged data, and tracks the time to call Accounts
+# introspection and profile APIs. However, it is unproven, so we're using an
+# environment variable to be able to try it in stage before production, and to
+# revert with a config change only.
+#
+# The names are designed to be annoying so they will be removed. The old code has
+# the suffix _2024 and the new code _2025 (when needed). When the new code is
+# proven, the old code can be removed with minimal name changes.
+#
+# ruff thinks the strings "2024" and "2025" are passwords (check S105 / S106).
+# These constants allow telling ruff to ignore them once.
+FXA_TOKEN_AUTH_OLD_AND_PROVEN = "2024"  # noqa: S105
+FXA_TOKEN_AUTH_NEW_AND_BUSTED = "2025"  # noqa: S105
+
 
 class CachedFxaIntrospectResponse(TypedDict, total=False):
     """The data stored in the cache to avoid multiple introspection requests."""
@@ -62,6 +91,11 @@ class FxaIntrospectCompleteData(TypedDict):
     exp: NotRequired[int]
 
 
+def get_cache_key_2024(token):
+    """note: hash() returns different results in different Python processes."""
+    return hash(token)
+
+
 def get_cache_key(token: str) -> str:
     return f"introspect_result:v1:{sha256(token.encode()).hexdigest()}"
 
@@ -312,6 +346,90 @@ def load_introspection_result_from_cache(
     return response
 
 
+def introspect_token_2024(token: str) -> dict[str, Any]:
+    try:
+        fxa_resp = requests.post(
+            INTROSPECT_TOKEN_URL,
+            json={"token": token},
+            timeout=settings.FXA_REQUESTS_TIMEOUT_SECONDS,
+        )
+    except Exception as exc:
+        logger.error(
+            "Could not introspect token with FXA.",
+            extra={"error_cls": type(exc), "error": shlex.quote(str(exc))},
+        )
+        raise AuthenticationFailed("Could not introspect token with FXA.")
+
+    fxa_resp_data = {"status_code": fxa_resp.status_code, "json": {}}
+    try:
+        fxa_resp_data["json"] = fxa_resp.json()
+    except requests.exceptions.JSONDecodeError:
+        logger.error(
+            "JSONDecodeError from FXA introspect response.",
+            extra={"fxa_response": shlex.quote(fxa_resp.text)},
+        )
+        raise AuthenticationFailed("JSONDecodeError from FXA introspect response")
+    return fxa_resp_data
+
+
+def get_fxa_uid_from_oauth_token_2024(token: str, use_cache: bool = True) -> str:
+    # set a default cache_timeout, but this will be overridden to match
+    # the 'exp' time in the JWT returned by FxA
+    cache_timeout = 60
+    cache_key = get_cache_key_2024(token)
+
+    if not use_cache:
+        fxa_resp_data = introspect_token_2024(token)
+    else:
+        # set a default fxa_resp_data, so any error during introspection
+        # will still cache for at least cache_timeout to prevent an outage
+        # from causing useless run-away repetitive introspection requests
+        fxa_resp_data = {"status_code": None, "json": {}}
+        try:
+            cached_fxa_resp_data = cache.get(cache_key)
+
+            if cached_fxa_resp_data:
+                fxa_resp_data = cached_fxa_resp_data
+            else:
+                # no cached data, get new
+                fxa_resp_data = introspect_token_2024(token)
+        except AuthenticationFailed:
+            raise
+        finally:
+            # Store potential valid response, errors, inactive users, etc. from FxA
+            # for at least 60 seconds. Valid access_token cache extended after checking.
+            cache.set(cache_key, fxa_resp_data, cache_timeout)
+
+    if fxa_resp_data["status_code"] is None:
+        raise APIException("Previous FXA call failed, wait to retry.")
+
+    if not fxa_resp_data["status_code"] == 200:
+        raise APIException("Did not receive a 200 response from FXA.")
+
+    if not fxa_resp_data["json"].get("active"):
+        raise AuthenticationFailed("FXA returned active: False for token.")
+
+    # FxA user is active, check for the associated Relay account
+    if (raw_fxa_uid := fxa_resp_data.get("json", {}).get("sub")) is None:
+        raise NotFound("FXA did not return an FXA UID.")
+    fxa_uid = str(raw_fxa_uid)
+
+    # cache valid access_token and fxa_resp_data until access_token expiration
+    # TODO: revisit this since the token can expire before its time
+    if isinstance(fxa_resp_data.get("json", {}).get("exp"), int):
+        # Note: FXA iat and exp are timestamps in *milliseconds*
+        fxa_token_exp_time = int(fxa_resp_data["json"]["exp"] / 1000)
+        now_time = int(datetime.now(UTC).timestamp())
+        fxa_token_exp_cache_timeout = fxa_token_exp_time - now_time
+        if fxa_token_exp_cache_timeout > cache_timeout:
+            # cache until access_token expires (matched Relay user)
+            # this handles cases where the token already expired
+            cache_timeout = fxa_token_exp_cache_timeout
+    cache.set(cache_key, fxa_resp_data, cache_timeout)
+
+    return fxa_uid
+
+
 def introspect_token(token: str) -> IntrospectionResponse | IntrospectionError:
     """
     Validate an Accounts OAuth token with the introspect API.
@@ -431,7 +549,76 @@ def introspect_and_cache_token(
     return fxa_resp
 
 
-class FxaTokenAuthentication(TokenAuthentication):
+class FxaTokenAuthentication(BaseAuthentication):
+    """Pick 2024 or 2025 version based on settings"""
+
+    _impl: FxaTokenAuthentication2024 | FxaTokenAuthentication2025
+
+    def __init__(self) -> None:
+        if settings.FXA_TOKEN_AUTH_VERSION == FXA_TOKEN_AUTH_NEW_AND_BUSTED:
+            self._impl = FxaTokenAuthentication2025()
+        else:
+            self._impl = FxaTokenAuthentication2024()
+
+    def authenticate_header(self, request: Request) -> Any | str | None:
+        return self._impl.authenticate_header(request)
+
+    def authenticate(
+        self, request: Request
+    ) -> None | tuple[User | AnonymousUser, IntrospectionResponse]:
+        return self._impl.authenticate(request)
+
+
+class FxaTokenAuthentication2024(BaseAuthentication):
+    def authenticate_header(self, request):
+        # Note: we need to implement this function to make DRF return a 401 status code
+        # when we raise AuthenticationFailed, rather than a 403. See:
+        # https://www.django-rest-framework.org/api-guide/authentication/#custom-authentication
+        return "Bearer"
+
+    def authenticate(self, request):
+        authorization = get_authorization_header(request).decode()
+        if not authorization or not authorization.startswith("Bearer "):
+            # If the request has no Bearer token, return None to attempt the next
+            # auth scheme in the REST_FRAMEWORK AUTHENTICATION_CLASSES list
+            return None
+
+        token = authorization.split(" ")[1]
+        if token == "":
+            raise ParseError("Missing FXA Token after 'Bearer'.")
+
+        use_cache = True
+        method = request.method
+        if method in ["POST", "DELETE", "PUT"]:
+            use_cache = False
+            if method == "POST" and request.path == "/api/v1/relayaddresses/":
+                use_cache = True
+        fxa_uid = get_fxa_uid_from_oauth_token_2024(token, use_cache)
+        try:
+            # MPP-3021: select_related user object to save DB query
+            sa = SocialAccount.objects.filter(
+                uid=fxa_uid, provider="fxa"
+            ).select_related("user")[0]
+        except IndexError:
+            raise PermissionDenied(
+                "Authenticated user does not have a Relay account."
+                " Have they accepted the terms?"
+            )
+        user = sa.user
+
+        if not user.is_active:
+            raise PermissionDenied(
+                "Authenticated user does not have an active Relay account."
+                " Have they been deactivated?"
+            )
+
+        if user:
+            return (user, token)
+        else:
+            raise NotFound()
+
+
+class FxaTokenAuthentication2025(TokenAuthentication):
     """
     Implement authentication with a Mozilla Account bearer token.