ref(auth): Use shared viewercontext library

pyproject.toml

@@ -98,6 +98,7 @@ dependencies = [
   "sentry-options>=1.2.15",
   "sentry-protos>=0.91.0",
   "sentry-redis-tools>=0.5.0",
+  "sentry-viewercontext>=0.1.0",
   "sentry-relay>=0.10.1",
   "sentry-scm>=1.9.0",
   "sentry-sdk[http2]>=2.72.0",

src/sentry/api/authentication.py

@@ -49,6 +49,7 @@
 from sentry.sentry_apps.models.sentry_app_installation import SentryAppInstallation
 from sentry.sentry_apps.token_exchange.util import GrantTypes
 from sentry.silo.base import SiloLimit, SiloMode
+from sentry.types.superuser import from_viewer_context_superuser
 from sentry.users.models.user import User
 from sentry.users.services.user import RpcUser
 from sentry.users.services.user.service import user_service
@@ -984,7 +985,7 @@ def authenticate(self, request: Request) -> tuple[Any, Any] | None:
         vc = viewer_context_from_header(header)
 
         user_id = get_acting_user_id(vc)
-        superuser = vc.actor.superuser if vc is not None else None
+        superuser = from_viewer_context_superuser(vc.actor.superuser) if vc is not None else None
 
         if vc is None or user_id is None:
             # TODO(jstanley): Temporary logging for debugging non-public prod 401s

src/sentry/auth/access.py

@@ -45,12 +45,11 @@
 from sentry.seer.agent_token import is_agent_auth
 from sentry.sentry_apps.models.sentry_app import SentryApp
 from sentry.types.request import _HttpRequestWithUser
-from sentry.types.superuser import SuperuserAccess
 from sentry.users.models.user import User
 from sentry.users.services.user import RpcUser
 from sentry.users.services.user.service import user_service
 from sentry.utils import metrics
-from sentry.viewer_context import set_viewer_context_superuser
+from sentry.viewer_context import SuperuserAccess, set_viewer_context_superuser
 
 __all__ = (
     "from_user",

src/sentry/middleware/viewer_context.py

@@ -9,6 +9,7 @@
 
 from sentry.auth.services.auth import AuthenticatedToken
 from sentry.seer.agent_token import is_agent_auth
+from sentry.types.superuser import to_viewer_context_superuser
 from sentry.utils.auth import is_user_from_viewer_context
 from sentry.viewer_context import (
     Actor,
@@ -98,7 +99,11 @@ def _viewer_context_from_request(request: HttpRequest) -> ViewerContext:
 
     # An agent token is a non-user actor: the request user is anonymous, so read the
     # delegating user it acts on behalf of from the credential.
-    superuser = auth.superuser if isinstance(auth, AuthenticatedToken) else None
+    superuser = (
+        to_viewer_context_superuser(auth.superuser)
+        if isinstance(auth, AuthenticatedToken)
+        else None
+    )
     if auth is not None and is_agent_auth(auth):
         actor = Actor(type=ActorType.AGENT, delegated_user_id=auth.user_id, superuser=superuser)
     else:

src/sentry/seer/endpoints/organization_agent_token.py

@@ -28,6 +28,7 @@
 from sentry.models.organization import Organization
 from sentry.seer import agent_token
 from sentry.seer.endpoints.agent_request import AgentTokenRequestData, AgentTokenRequestSerializer
+from sentry.types.superuser import from_viewer_context_superuser
 from sentry.viewer_context import get_acting_user_id, get_viewer_context
 
 
@@ -107,7 +108,7 @@ def post(self, request: Request, organization: Organization) -> Response:
 
         viewer = get_viewer_context()
         superuser = (
-            viewer.actor.superuser
+            from_viewer_context_superuser(viewer.actor.superuser)
             if viewer is not None
             and get_acting_user_id(viewer) == user_id
             and viewer.organization_id == organization.id

src/sentry/types/superuser.py

@@ -1,6 +1,7 @@
 from datetime import timedelta
 
 from pydantic import BaseModel, StrictInt
+from sentry_viewercontext import SuperuserAccess as ViewerContextSuperuserAccess
 
 SUPERUSER_ACCESS_TTL = timedelta(minutes=5)
 
@@ -12,3 +13,21 @@ class SuperuserAccess(BaseModel):
 
     class Config:
         frozen = True
+
+
+def from_viewer_context_superuser(
+    superuser: ViewerContextSuperuserAccess | None,
+) -> SuperuserAccess | None:
+    """Convert the shared library's elevation into the Pydantic model RPC and agent tokens use."""
+    if superuser is None:
+        return None
+    return SuperuserAccess(expires_at=superuser.expires_at)
+
+
+def to_viewer_context_superuser(
+    superuser: SuperuserAccess | None,
+) -> ViewerContextSuperuserAccess | None:
+    """Convert Pydantic elevation into the shared library's type for a ViewerContext."""
+    if superuser is None:
+        return None
+    return ViewerContextSuperuserAccess(expires_at=superuser.expires_at)

src/sentry/viewer_context.py

@@ -1,32 +1,47 @@
 from __future__ import annotations
 
-import contextlib
-import contextvars
 import dataclasses
-import enum
-import hashlib
 import logging
-import time
-from collections.abc import Generator, Mapping
+from collections.abc import Mapping
 from random import random
-from typing import Any, NamedTuple
+from typing import Any
 
-import jwt as pyjwt
 import sentry_sdk
 from django.conf import settings
+from sentry_viewercontext import (
+    Actor,
+    ActorId,
+    ActorType,
+    SuperuserAccess,
+    ViewerContext,
+    decode_jwt,
+    encode_jwt,
+    get_viewer_context,
+    is_jwt,
+    key_id,
+    viewer_context_scope,
+    viewer_context_var,
+)
 
 from sentry.silo.base import SiloMode
-from sentry.types.superuser import SUPERUSER_ACCESS_TTL, SuperuserAccess
+
+__all__ = [
+    "Actor",
+    "ActorId",
+    "ActorType",
+    "SuperuserAccess",
+    "ViewerContext",
+    "get_viewer_context",
+    "viewer_context_scope",
+]
 
 logger = logging.getLogger(__name__)
 
 # Sentinel for `observe_viewer_context_propagation(ctx=...)`: distinguishes
 # "caller passed None" from "caller didn't pass ctx".
 _OBSERVE_USE_CONTEXTVAR: Any = object()
 
-_viewer_context_var: contextvars.ContextVar[ViewerContext | None] = contextvars.ContextVar(
-    "viewer_context", default=None
-)
+_viewer_context_var = viewer_context_var
 
 """
 ViewerContext is an object propagated across codebase (and crossing service boundary) to deliver
@@ -39,106 +54,6 @@
 """
 
 
-class ActorType(enum.StrEnum):
-    USER = "user"
-    SERVICE_ACCOUNT = "service_account"
-    SYSTEM = "system"
-    INTEGRATION = "integration"
-    AGENT = "agent"
-    UNKNOWN = "unknown"
-
-
-class ActorId(NamedTuple):
-    type: ActorType
-    id: int | None
-
-
-@dataclasses.dataclass(frozen=True)
-class Actor:
-    """Typed caller identity, separate from any human delegating its work."""
-
-    type: ActorType
-    actor_id: int | None = None
-    delegated_user_id: int | None = None
-    superuser: SuperuserAccess | None = dataclasses.field(default=None, repr=False)
-
-    @property
-    def id(self) -> ActorId:
-        return ActorId(self.type, self.actor_id)
-
-    def serialize(self) -> dict[str, Any]:
-        result: dict[str, Any] = {"type": self.type.value}
-        if self.actor_id is not None:
-            result["id"] = self.actor_id
-        if self.delegated_user_id is not None:
-            result["delegated_user_id"] = self.delegated_user_id
-        if self.superuser is not None:
-            result["superuser"] = self.superuser.dict()
-        return result
-
-    @classmethod
-    def deserialize(cls, data: Any) -> Actor:
-        if not isinstance(data, dict) or "type" not in data:
-            raise ValueError("Actor must be an object with a type")
-        superuser = data.get("superuser")
-        return cls(
-            type=ActorType(data["type"]),
-            actor_id=data.get("id"),
-            delegated_user_id=data.get("delegated_user_id"),
-            superuser=SuperuserAccess.parse_obj(superuser) if superuser is not None else None,
-        )
-
-
-@dataclasses.dataclass(frozen=True)
-class ViewerContext:
-    """Caller identity and organization/project scope for the current unit of work."""
-
-    actor: Actor
-    organization_id: int | None = None
-    project_id: int | None = None
-
-    def serialize(self) -> dict[str, Any]:
-        result: dict[str, Any] = {"version": 2, "actor": self.actor.serialize()}
-        if self.organization_id is not None:
-            result["organization_id"] = self.organization_id
-        if self.project_id is not None:
-            result["project_id"] = self.project_id
-        return result
-
-    @classmethod
-    def deserialize(cls, data: dict[str, Any]) -> ViewerContext:
-        if not isinstance(data, dict):
-            raise ValueError("Viewer context must be an object")
-        if type(data.get("version")) is not int or data["version"] != 2:
-            raise ValueError("Expected a V2 viewer context")
-        if {"user_id", "actor_type", "superuser", "token"}.intersection(data):
-            raise ValueError("V2 contexts cannot carry V1 identity or credentials")
-        return cls(
-            actor=Actor.deserialize(data.get("actor")),
-            organization_id=data.get("organization_id"),
-            project_id=data.get("project_id"),
-        )
-
-
-@contextlib.contextmanager
-def viewer_context_scope(ctx: ViewerContext) -> Generator[None]:
-    """Enter a viewer context for the duration of a unit of work.
-
-    Always use this instead of raw ``_viewer_context_var.set()`` —
-    it guarantees cleanup via ``reset(token)`` even on exceptions.
-    """
-    tok = _viewer_context_var.set(ctx)
-    try:
-        yield
-    finally:
-        _viewer_context_var.reset(tok)
-
-
-def get_viewer_context() -> ViewerContext | None:
-    """Return the current ``ViewerContext``, or ``None`` if not set."""
-    return _viewer_context_var.get()
-
-
 def get_acting_user_id(ctx: ViewerContext | None) -> int | None:
     """Return the human user *ctx* acts as, or ``None``.
 
@@ -271,19 +186,6 @@ def set_viewer_context_project(project_id: int) -> None:
 # JWT encoding / decoding for cross-service propagation
 # ---------------------------------------------------------------------------
 
-_JWT_STANDARD_CLAIMS = frozenset({"iat", "exp", "iss", "aud", "nbf", "jti", "sub"})
-# JWT header field identifying which key was used for signing (RFC 7515 §4.1.4).
-_JWT_KEY_ID_HEADER = "kid"
-
-
-def _key_id(key: str) -> str:
-    """Short stable identifier for a key (first 8 hex chars of SHA-256).
-
-    Embedded as ``kid`` in the JWT header so the receiver can look up the
-    correct verification key without trying every key it knows about.
-    """
-    return hashlib.sha256(key.encode("utf-8")).hexdigest()[:8]
-
 
 def _get_signing_key(key: str | None = None) -> str:
     """Return the key to use for JWT signing.
@@ -312,7 +214,7 @@ def _get_verification_keys() -> dict[str, str]:
 
     seer_secret = getattr(settings, "SEER_API_SHARED_SECRET", "")
     if seer_secret:
-        keys[_key_id(seer_secret)] = seer_secret
+        keys[key_id(seer_secret)] = seer_secret
 
     return keys
 
@@ -324,34 +226,13 @@ def encode_viewer_context(
     ttl: int | None = None,
 ) -> str:
     """Encode a :class:`ViewerContext` as a signed HS256 JWT."""
-    secret = _get_signing_key(key)
-
     if ttl is None:
         ttl = getattr(settings, "VIEWER_CONTEXT_JWT_TTL", 900)
-
-    now = time.time()
-    expires_at = now + ttl
-    superuser = viewer_context.actor.superuser
-    if superuser is not None:
-        expires_at = min(
-            expires_at,
-            now + SUPERUSER_ACCESS_TTL.total_seconds(),
-            superuser.expires_at,
-        )
-    payload: dict[str, Any] = {
-        **viewer_context.serialize(),
-        "iat": now,
-        "exp": expires_at,
-        "iss": "sentry",
-    }
     if viewer_context.organization_id is not None and _organization_is_early_adopter(
         viewer_context.organization_id
     ):
-        payload["organization_is_early_adopter"] = True
-
-    return pyjwt.encode(
-        payload, secret, algorithm="HS256", headers={_JWT_KEY_ID_HEADER: _key_id(secret)}
-    )
+        viewer_context = dataclasses.replace(viewer_context, organization_is_early_adopter=True)
+    return encode_jwt(viewer_context, key=_get_signing_key(key), ttl=ttl)
 
 
 def _organization_is_early_adopter(organization_id: int) -> bool:
@@ -381,31 +262,13 @@ def decode_viewer_context(
 ) -> ViewerContext:
     """Decode and verify an HS256 JWT into a :class:`ViewerContext`.
 
-    When *key* is provided it is used directly.  Otherwise all keys
-    from ``_get_verification_keys()`` are tried, kid-matched key first.
+    When *key* is provided it is used directly.  Otherwise the key matching the
+    token's ``kid`` among ``_get_verification_keys()`` is used.
     """
-    if key is not None:
-        secret = key
-    else:
-        keys_by_kid = _get_verification_keys()
-        if not keys_by_kid:
-            raise ValueError("No verification keys available.")
-
-        kid = pyjwt.get_unverified_header(token).get(_JWT_KEY_ID_HEADER)
-        secret = keys_by_kid.get(kid, "") if kid else ""
-        if not secret:
-            raise pyjwt.exceptions.InvalidKeyError(f"No verification key matches kid={kid!r}")
-
-    claims = pyjwt.decode(
-        token,
-        secret,
-        algorithms=["HS256"],
-        options={"require": ["iat", "exp", "iss"]},
-        issuer="sentry",
-        leeway=leeway,
-    )
-    vc_data = {ck: cv for ck, cv in claims.items() if ck not in _JWT_STANDARD_CLAIMS}
-    return ViewerContext.deserialize(vc_data)
+    keys = [key] if key is not None else list(_get_verification_keys().values())
+    if not keys:
+        raise ValueError("No verification keys available.")
+    return decode_jwt(token, keys=keys, leeway=leeway)
 
 
 def viewer_context_from_header(
@@ -423,13 +286,4 @@ def viewer_context_from_header(
 
 
 def is_jwt_viewer_context(header_value: str) -> bool:
-    """Check whether the header value is a JWT by attempting to read its header.
-
-    Uses PyJWT's own parser — raises ``DecodeError`` on anything that
-    isn't a valid JWT structure (raw JSON, empty string, etc.).
-    """
-    try:
-        pyjwt.get_unverified_header(header_value)
-        return True
-    except pyjwt.exceptions.DecodeError:
-        return False
+    return is_jwt(header_value)