JwtIdProvider

jwt.py

class JwtIdProvider(IdentityProvider):
    def __init__(self, jwk_uri: str, aud: str, context: ServicerContext) -> None:
        self.jwks_client = PyJWKClient(jwk_uri)
        self.aud = aud
        self.context = context

    async def get_user(self) -> User | None:
        for header, value in self.context.invocation_metadata():
            if header != "authorization":
                continue

            if not value.startswith("Bearer "):
                continue

            token = value.removeprefix("Bearer").strip()
            try:
                signing_key = await self.jwks_client.get_signing_key_from_jwt(token)
                decoded = jwt.decode(
                    token,
                    signing_key.key,
                    audience=self.aud,
                    algorithms=["RS256"],
                )
            except jwt.PyJWTError as e:
                logger.warning("Failed to decode JWT: %s", e)
                raise ApiError(status=StatusCode.UNAUTHENTICATED, details="Invalid JWT") from e
            roles = decoded.get("resource_access", {}).get(self.aud, {}).get("roles", [])
            return User(
                id=decoded.get("email") or decoded.get("preferred_username"),
                roles=roles,
            )
        return None