feat(auth): optional audience and issuer enforcement for JWT auth (#14151)

Author: nmgarza5

backend/onyx/auth/jwt.py

@@ -6,11 +6,21 @@
 import jwt
 import requests
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
-from jwt import InvalidTokenError, PyJWTError
+from jwt import (
+    InvalidAudienceError,
+    InvalidIssuerError,
+    InvalidTokenError,
+    MissingRequiredClaimError,
+    PyJWTError,
+)
 from jwt import decode as jwt_decode
 from jwt.algorithms import RSAAlgorithm  # ty: ignore[possibly-missing-import]
 
-from onyx.configs.app_configs import JWT_PUBLIC_KEY_URL
+from onyx.configs.app_configs import (
+    JWT_EXPECTED_AUDIENCE,
+    JWT_EXPECTED_ISSUER,
+    JWT_PUBLIC_KEY_URL,
+)
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -136,12 +146,25 @@ async def verify_jwt_token(token: str) -> dict[str, Any] | None:
             return None
 
         try:
+            # Enforced only when configured: verify_aud=True with audience=None
+            # would reject every token that carries an aud claim.
             payload = jwt_decode(
                 token,
                 public_key,
                 algorithms=["RS256"],
-                options={"verify_aud": False},
+                audience=JWT_EXPECTED_AUDIENCE,
+                issuer=JWT_EXPECTED_ISSUER,
+                options={"verify_aud": JWT_EXPECTED_AUDIENCE is not None},
             )
+        except (
+            InvalidAudienceError,
+            InvalidIssuerError,
+            MissingRequiredClaimError,
+        ) as e:
+            # Definitive claim rejection: refetched keys cannot change it, and a
+            # cache clear would let bad tokens evict the signing key for everyone.
+            logger.warning("JWT rejected by aud/iss enforcement: %s", str(e))
+            return None
         except InvalidTokenError as e:
             logger.error("Invalid JWT token: %s", str(e))
             if attempt < _PUBLIC_KEY_FETCH_ATTEMPTS - 1:

backend/onyx/configs/app_configs.py

@@ -357,7 +357,11 @@ def parse_idp_claim_map(raw: str | None) -> dict[str, list[str]]:
     MOBILE_ALLOWED_REDIRECT_URIS = _DEFAULT_MOBILE_REDIRECT_URIS
 
 # JWT Public Key URL for JWT token verification
-JWT_PUBLIC_KEY_URL: str | None = os.getenv("JWT_PUBLIC_KEY_URL", None)
+JWT_PUBLIC_KEY_URL: str | None = os.getenv("JWT_PUBLIC_KEY_URL") or None
+# Optional aud/iss scoping for JWT_PUBLIC_KEY_URL auth, off when unset. Empty
+# counts as unset because compose files pass absent vars through as "".
+JWT_EXPECTED_AUDIENCE: str | None = os.getenv("JWT_EXPECTED_AUDIENCE") or None
+JWT_EXPECTED_ISSUER: str | None = os.getenv("JWT_EXPECTED_ISSUER") or None
 
 USER_AUTH_SECRET = os.environ.get("USER_AUTH_SECRET", "")
 

backend/tests/unit/onyx/auth/test_jwt_audience_issuer.py

@@ -0,0 +1,161 @@
+"""Optional aud/iss enforcement on the JWT_PUBLIC_KEY_URL flow. Unset settings
+leave the claims unrestricted, and a configured expectation must reject
+mismatched or absent claims."""
+
+from typing import Any
+
+import jwt as pyjwt
+import pytest
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+import onyx.auth.jwt as jwt_module
+from onyx.auth.jwt import verify_jwt_token
+
+_PRIVATE_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+_PUBLIC_PEM = (
+    _PRIVATE_KEY.public_key()
+    .public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+    .decode()
+)
+
+
+def _mint(claims: dict[str, Any]) -> str:
+    return pyjwt.encode(claims, _PRIVATE_KEY, algorithm="RS256")
+
+
+@pytest.fixture
+def signed_key(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(jwt_module, "get_public_key", lambda _token: _PUBLIC_PEM)
+
+
+def _configure(
+    monkeypatch: pytest.MonkeyPatch,
+    audience: str | None,
+    issuer: str | None,
+) -> None:
+    monkeypatch.setattr(jwt_module, "JWT_EXPECTED_AUDIENCE", audience)
+    monkeypatch.setattr(jwt_module, "JWT_EXPECTED_ISSUER", issuer)
+
+
+@pytest.mark.parametrize(
+    "claims",
+    [
+        {"email": "a@b.c"},
+        {"email": "a@b.c", "aud": "some-other-service", "iss": "https://idp"},
+    ],
+)
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("signed_key")
+async def test_unset_settings_accept_any_signed_token(
+    monkeypatch: pytest.MonkeyPatch, claims: dict[str, Any]
+) -> None:
+    _configure(monkeypatch, None, None)
+    payload = await verify_jwt_token(_mint(claims))
+    assert payload is not None
+    assert payload["email"] == "a@b.c"
+
+
+@pytest.mark.parametrize(
+    "claims, accepted",
+    [
+        ({"aud": "onyx"}, True),
+        ({"aud": ["onyx", "other"]}, True),
+        ({"aud": "some-other-service"}, False),
+        ({"aud": ["other-a", "other-b"]}, False),
+        ({}, False),
+    ],
+)
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("signed_key")
+async def test_audience_enforced_when_configured(
+    monkeypatch: pytest.MonkeyPatch,
+    claims: dict[str, Any],
+    accepted: bool,
+) -> None:
+    _configure(monkeypatch, "onyx", None)
+    payload = await verify_jwt_token(_mint(claims))
+    assert (payload is not None) is accepted
+
+
+@pytest.mark.parametrize(
+    "claims, accepted",
+    [
+        ({"iss": "https://idp.example.com"}, True),
+        ({"iss": "https://evil.example.com"}, False),
+        ({}, False),
+    ],
+)
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("signed_key")
+async def test_issuer_enforced_when_configured(
+    monkeypatch: pytest.MonkeyPatch,
+    claims: dict[str, Any],
+    accepted: bool,
+) -> None:
+    _configure(monkeypatch, None, "https://idp.example.com")
+    payload = await verify_jwt_token(_mint(claims))
+    assert (payload is not None) is accepted
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("signed_key")
+async def test_both_enforced_and_matching(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _configure(monkeypatch, "onyx", "https://idp.example.com")
+    payload = await verify_jwt_token(
+        _mint({"aud": "onyx", "iss": "https://idp.example.com", "email": "a@b.c"})
+    )
+    assert payload is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("signed_key")
+async def test_both_enforced_rejects_wrong_issuer(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    _configure(monkeypatch, "onyx", "https://idp.example.com")
+    payload = await verify_jwt_token(
+        _mint({"aud": "onyx", "iss": "https://evil.example.com"})
+    )
+    assert payload is None
+
+
+def test_empty_env_counts_as_unset(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Compose files pass absent vars through as "", which must not enable
+    # enforcement (audience="" would reject every aud-less token).
+    import importlib
+
+    import onyx.configs.app_configs as app_configs
+
+    monkeypatch.setenv("JWT_EXPECTED_AUDIENCE", "")
+    monkeypatch.setenv("JWT_EXPECTED_ISSUER", "")
+    monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "")
+    reloaded = importlib.reload(app_configs)
+    try:
+        assert reloaded.JWT_EXPECTED_AUDIENCE is None
+        assert reloaded.JWT_EXPECTED_ISSUER is None
+        assert reloaded.JWT_PUBLIC_KEY_URL is None
+    finally:
+        monkeypatch.undo()
+        importlib.reload(app_configs)
+
+
+@pytest.mark.asyncio
+async def test_claim_rejection_does_not_refetch_keys(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls: list[int] = []
+
+    def _counting_get(_token: str) -> str:
+        calls.append(1)
+        return _PUBLIC_PEM
+
+    monkeypatch.setattr(jwt_module, "get_public_key", _counting_get)
+    _configure(monkeypatch, "onyx", None)
+    assert await verify_jwt_token(_mint({"aud": "some-other-service"})) is None
+    assert len(calls) == 1

cli/internal/deploy/deployfiles/embedded/docker_compose/env.template

@@ -51,6 +51,10 @@ services:
       - OPENID_CONFIG_URL=${OPENID_CONFIG_URL:-}
       - TRACK_EXTERNAL_IDP_EXPIRY=${TRACK_EXTERNAL_IDP_EXPIRY:-}
       - CORS_ALLOWED_ORIGIN=${CORS_ALLOWED_ORIGIN:-}
+      # JWT auth runs in the api_server request path only
+      - JWT_PUBLIC_KEY_URL=${JWT_PUBLIC_KEY_URL:-}
+      - JWT_EXPECTED_AUDIENCE=${JWT_EXPECTED_AUDIENCE:-}
+      - JWT_EXPECTED_ISSUER=${JWT_EXPECTED_ISSUER:-}
       # Gen AI Settings
       - GEN_AI_MAX_TOKENS=${GEN_AI_MAX_TOKENS:-}
       - LLM_SOCKET_READ_TIMEOUT=${LLM_SOCKET_READ_TIMEOUT:-}

deployment/docker_compose/docker-compose.multitenant-dev.yml

@@ -307,6 +307,8 @@ LOG_ONYX_MODEL_INTERACTIONS=False
 # CORS_ALLOWED_ORIGIN=
 # INTEGRATION_TESTS_MODE=
 # JWT_PUBLIC_KEY_URL=
+# JWT_EXPECTED_AUDIENCE=
+# JWT_EXPECTED_ISSUER=
 
 ## Gen AI Settings
 # GEN_AI_MAX_TOKENS=