Skip to content

feat: fallback to trying all JWKS keys when kid does not match #893

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
liudonggalaxy wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: fallback to trying all JWKS keys when kid does not match #893

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 29 additions & 10 deletions authlib/integrations/base_client/async_openid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
from authlib.oidc.core import ImplicitIDToken
from authlib.oidc.core import UserInfo

from .sync_openid import try_all_keys

__all__ = ["AsyncOpenIDMixin"]


Expand Down Expand Up @@ -63,20 +65,26 @@ async def parse_id_token(

jwks = await self.fetch_jwk_set()
key_set = KeySet.import_key_set(jwks)
id_token_value = token["id_token"]
try:
token = jwt.decode(
token["id_token"],
key=key_set,
algorithms=alg_values,
)
token = jwt.decode(id_token_value, key=key_set, algorithms=alg_values)
except InvalidKeyIdError:
# Refresh JWKS — keys may have rotated
jwks = await self.fetch_jwk_set(force=True)
key_set = KeySet.import_key_set(jwks)
token = jwt.decode(
token["id_token"],
key=key_set,
algorithms=alg_values,
)
try:
token = jwt.decode(
id_token_value, key=key_set, algorithms=alg_values
)
except InvalidKeyIdError:
if self.should_try_all_keys():
token = try_all_keys(
id_token_value=id_token_value,
key_set=key_set,
algorithms=alg_values,
)
else:
raise

claims = claims_cls(token.claims, token.header, claims_options, claims_params)
# https://github.com/authlib/authlib/issues/259
Expand All @@ -85,6 +93,17 @@ async def parse_id_token(
claims.validate(leeway=leeway)
return UserInfo(claims)

def should_try_all_keys(self):
"""Control whether to attempt signature verification against every key
in the JWKS when the token's kid does not match any key.

The default is False, which preserves the standard behavior of raising
InvalidKeyIdError immediately. Override this method to return True in
subclasses that integrate with providers whose JWKS keys lack a kid or
use mismatched kid values.
"""
return False

async def create_logout_url(
self,
post_logout_redirect_uri=None,
Expand Down
61 changes: 49 additions & 12 deletions authlib/integrations/base_client/sync_openid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
from joserfc import jwt
from joserfc.errors import BadSignatureError
from joserfc.errors import DecodeError
from joserfc.errors import InvalidKeyIdError
from joserfc.jwk import KeySet

Expand Down Expand Up @@ -43,6 +45,8 @@ def parse_id_token(
if "id_token" not in token:
return None

id_token_value = token["id_token"]

claims_params = dict(
nonce=nonce,
client_id=self.client_id,
Expand All @@ -63,27 +67,45 @@ def parse_id_token(

key_set = KeySet.import_key_set(self.fetch_jwk_set())
try:
token = jwt.decode(
token["id_token"],
key=key_set,
algorithms=alg_values,
)
token = jwt.decode(id_token_value, key=key_set, algorithms=alg_values)
except InvalidKeyIdError:
# Refresh JWKS — keys may have rotated
key_set = KeySet.import_key_set(self.fetch_jwk_set(force=True))
token = jwt.decode(
token["id_token"],
key=key_set,
algorithms=alg_values,
)

claims = claims_cls(token.claims, token.header, claims_options, claims_params)
try:
token = jwt.decode(
id_token_value, key=key_set, algorithms=alg_values
)
except InvalidKeyIdError:
if self.should_try_all_keys():
token = try_all_keys(
id_token_value=id_token_value,
key_set=key_set,
algorithms=alg_values,
)
else:
raise

claims = claims_cls(
token.claims, token.header, claims_options, claims_params
)
# https://github.com/authlib/authlib/issues/259
if claims.get("nonce_supported") is False:
claims.params["nonce"] = None

claims.validate(leeway=leeway)
return UserInfo(claims)

def should_try_all_keys(self):
"""Control whether to attempt signature verification against every key
in the JWKS when the token's kid does not match any key.

The default is False, which preserves the standard behavior of raising
InvalidKeyIdError immediately. Override this method to return True in
subclasses that integrate with providers whose JWKS keys lack a kid or
use mismatched kid values.
"""
return False

def create_logout_url(
self,
post_logout_redirect_uri=None,
Expand Down Expand Up @@ -120,3 +142,18 @@ def create_logout_url(

url = add_params_to_uri(end_session_endpoint, params)
return {"url": url, "state": state}


def try_all_keys(id_token_value: str, key_set: KeySet, algorithms) -> jwt.Token:
"""Try each key in the JWKS individually to find one that verifies the
token signature. Handles providers whose JWKS keys lack a kid field or
use kid values that do not match the token header.

Raises InvalidKeyIdError if no key verifies the signature.
"""
for key in key_set.keys:
try:
return jwt.decode(id_token_value, key=key, algorithms=algorithms)
except (BadSignatureError, DecodeError, ValueError):
continue
raise InvalidKeyIdError("No key in JWKS verifies the token signature.")
231 changes: 231 additions & 0 deletions tests/clients/test_flask/test_parse_id_token_kid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,231 @@
import time
from unittest import mock

import pytest
from flask import Flask
from joserfc import jwt
from joserfc.errors import InvalidKeyIdError
from joserfc.jwk import OctKey

from authlib.integrations.flask_client import OAuth
from authlib.oidc.core.grants.util import create_half_hash

from ..util import get_bearer_token


def _make_id_token(signing_key, kid=None):
"""Create a signed ID token.

:param signing_key: OctKey used to sign the token.
:param kid: kid to include in the JWT header. None omits it.
"""
token = get_bearer_token()
now = int(time.time())
claims = {
"sub": "user-42",
"nonce": "n",
"iss": "https://provider.test",
"aud": "dev",
"iat": now,
"auth_time": now,
"exp": now + 3600,
"at_hash": create_half_hash(token["access_token"], "HS256").decode("utf-8"),
}
header = {"alg": "HS256"}
if kid is not None:
header["kid"] = kid
token["id_token"] = jwt.encode(header, claims, signing_key)
return token


def _fake_jwks_refresh(jwks_dict):
"""Create a mock send function that returns the given JWKS on refresh."""
def fake_send(sess, req, **kwargs):
resp = mock.MagicMock()
resp.json = lambda: jwks_dict
resp.status_code = 200
return resp
return fake_send


def _register_client(app, jwks, jwks_uri="https://provider.test/jwks"):
oauth = OAuth(app)
return oauth.register(
"dev",
client_id="dev",
client_secret="dev",
fetch_token=get_bearer_token,
jwks=jwks,
jwks_uri=jwks_uri,
issuer="https://provider.test",
id_token_signing_alg_values_supported=["HS256"],
)


def test_default_raises_on_kid_mismatch():
"""Default behavior: kid mismatch raises InvalidKeyIdError after refresh."""
signing_key = OctKey.import_key("secret-for-mismatch-test")
token = _make_id_token(signing_key, kid="token-kid-1")
app = Flask(__name__)
app.secret_key = "!"

key_dict = signing_key.as_dict()
key_dict["kid"] = "different-kid"
extra_key = OctKey.import_key("extra-key-material-xxxxx")
extra_dict = extra_key.as_dict()
extra_dict["kid"] = "extra-kid"
jwks = {"keys": [key_dict, extra_dict]}
client = _register_client(app, jwks)

with app.test_request_context():
with mock.patch("requests.sessions.Session.send", _fake_jwks_refresh(jwks)):
with pytest.raises(InvalidKeyIdError):
client.parse_id_token(token, nonce="n")


def test_no_kid_in_token_multiple_jwks_keys_needs_fallback():
"""Token without kid + multiple JWKS keys raises unless fallback enabled."""
signing_key = OctKey.import_key("secret-for-multi-key-test")
token = _make_id_token(signing_key, kid=None)
app = Flask(__name__)
app.secret_key = "!"

other_key = OctKey.import_key("other-key-material-xxxxx")
other_dict = other_key.as_dict()
other_dict["kid"] = "key-1"
correct_dict = signing_key.as_dict()
correct_dict["kid"] = "key-2"
jwks = {"keys": [other_dict, correct_dict]}

client = _register_client(app, jwks)
client.should_try_all_keys = lambda: True

with app.test_request_context():
with mock.patch("requests.sessions.Session.send", _fake_jwks_refresh(jwks)):
user = client.parse_id_token(token, nonce="n")
assert user.sub == "user-42"


def test_key_rotation_resolved_by_refresh():
"""JWKS refresh resolves kid mismatch when provider rotated keys."""
signing_key = OctKey.import_key("secret-for-rotation-test")
token = _make_id_token(signing_key, kid="new-kid")
app = Flask(__name__)
app.secret_key = "!"

# Initial JWKS has old keys that cannot verify the token
old_key = OctKey.import_key("old-key-material-xxxxx")
old_dict = old_key.as_dict()
old_dict["kid"] = "old-kid"
old_key2 = OctKey.import_key("old-key-material-yyyyy")
old_dict2 = old_key2.as_dict()
old_dict2["kid"] = "old-kid-2"

# After refresh: correct key with matching kid + another key
correct_dict = signing_key.as_dict()
correct_dict["kid"] = "new-kid"

client = _register_client(app, {"keys": [old_dict, old_dict2]})

with app.test_request_context():
with mock.patch(
"requests.sessions.Session.send",
_fake_jwks_refresh({"keys": [correct_dict, old_dict2]}),
):
user = client.parse_id_token(token, nonce="n")
assert user.sub == "user-42"


def test_try_all_keys_with_kid_mismatch():
"""With should_try_all_keys enabled, fallback finds the correct key."""
signing_key = OctKey.import_key("secret-for-try-all-test")
token = _make_id_token(signing_key, kid="nonexistent-kid")
app = Flask(__name__)
app.secret_key = "!"

wrong_key = OctKey.import_key("wrong-secret-material")
correct_dict = signing_key.as_dict()
correct_dict["kid"] = "other-kid"
jwks = {"keys": [wrong_key.as_dict(), correct_dict]}

client = _register_client(app, jwks)
client.should_try_all_keys = lambda: True

with app.test_request_context():
with mock.patch("requests.sessions.Session.send", _fake_jwks_refresh(jwks)):
user = client.parse_id_token(token, nonce="n")
assert user.sub == "user-42"


def test_try_all_keys_when_jwks_key_has_no_kid():
"""With should_try_all_keys enabled, works when JWKS keys lack kid."""
signing_key = OctKey.import_key("secret-for-jwks-no-kid")
token = _make_id_token(signing_key, kid="token-kid")
app = Flask(__name__)
app.secret_key = "!"

# JWKS keys have no kid — triggers InvalidKeyIdError for token with kid
wrong_key = OctKey.import_key("wrong-key-for-no-kid-test")
jwks = {"keys": [wrong_key.as_dict(), signing_key.as_dict()]}
client = _register_client(app, jwks)
client.should_try_all_keys = lambda: True

with app.test_request_context():
with mock.patch("requests.sessions.Session.send", _fake_jwks_refresh(jwks)):
user = client.parse_id_token(token, nonce="n")
assert user.sub == "user-42"


def test_try_all_keys_no_valid_key_raises():
"""With should_try_all_keys enabled, raises if no key verifies."""
signing_key = OctKey.import_key("secret-for-no-match-test")
token = _make_id_token(signing_key, kid="token-kid")
app = Flask(__name__)
app.secret_key = "!"

wrong_key1 = OctKey.import_key("completely-different-secret")
wrong_dict1 = wrong_key1.as_dict()
wrong_dict1["kid"] = "wrong-1"
wrong_key2 = OctKey.import_key("another-different-secret")
wrong_dict2 = wrong_key2.as_dict()
wrong_dict2["kid"] = "wrong-2"
jwks = {"keys": [wrong_dict1, wrong_dict2]}

client = _register_client(app, jwks)
client.should_try_all_keys = lambda: True

with app.test_request_context():
with mock.patch("requests.sessions.Session.send", _fake_jwks_refresh(jwks)):
with pytest.raises(InvalidKeyIdError):
client.parse_id_token(token, nonce="n")


def test_no_jwks_uri_raises_runtime_error():
"""When jwks_uri is absent, JWKS refresh raises RuntimeError."""
signing_key = OctKey.import_key("secret-for-no-uri-test")
token = _make_id_token(signing_key, kid="token-kid")
app = Flask(__name__)
app.secret_key = "!"

# JWKS has the correct key material but wrong kid, triggering refresh
key_dict = signing_key.as_dict()
key_dict["kid"] = "embedded-kid"
extra_key = OctKey.import_key("extra-no-uri-key-material")
extra_dict = extra_key.as_dict()
extra_dict["kid"] = "extra-kid"

oauth = OAuth(app)
client = oauth.register(
"dev",
client_id="dev",
client_secret="dev",
fetch_token=get_bearer_token,
jwks={"keys": [key_dict, extra_dict]},
# No jwks_uri provided
issuer="https://provider.test",
id_token_signing_alg_values_supported=["HS256"],
)

with app.test_request_context():
with pytest.raises(RuntimeError, match="jwks_uri"):
client.parse_id_token(token, nonce="n")
Loading
Loading