From 68334dbf04fa25c1f64541710f6db03c2ba3888d Mon Sep 17 00:00:00 2001
From: dklimpel <5740567+dklimpel@users.noreply.github.com>
Date: Tue, 21 Nov 2023 12:35:18 +0100
Subject: [PATCH] Use single key in JWK if JWS does not specify `kid`

---
 authlib/jose/rfc7519/jwt.py          | 13 ++++++++++---
 tests/files/jwks_single_private.json |  5 +++++
 tests/files/jwks_single_public.json  |  5 +++++
 tests/jose/test_jwt.py               | 12 ++++++++++++
 4 files changed, 32 insertions(+), 3 deletions(-)
 create mode 100644 tests/files/jwks_single_private.json
 create mode 100644 tests/files/jwks_single_public.json

diff --git a/authlib/jose/rfc7519/jwt.py b/authlib/jose/rfc7519/jwt.py
index 3737d303..e0bba87d 100644
--- a/authlib/jose/rfc7519/jwt.py
+++ b/authlib/jose/rfc7519/jwt.py
@@ -167,9 +167,16 @@ def load_key(header, payload):
         if isinstance(key, dict) and 'keys' in key:
             keys = key['keys']
             kid = header.get('kid')
-            for k in keys:
-                if k.get('kid') == kid:
-                    return k
+
+            if kid is not None:
+                # look for the requested key
+                for k in keys:
+                    if k.get('kid') == kid:
+                        return k
+            else:
+                # use the only key
+                if len(keys) == 1:
+                    return keys[0]
             raise ValueError('Invalid JSON Web Key Set')
         return key
 
diff --git a/tests/files/jwks_single_private.json b/tests/files/jwks_single_private.json
new file mode 100644
index 00000000..8a0b33b7
--- /dev/null
+++ b/tests/files/jwks_single_private.json
@@ -0,0 +1,5 @@
+{
+  "keys": [
+    {"kty": "RSA", "n": "pF1JaMSN8TEsh4N4O_5SpEAVLivJyLH-Cgl3OQBPGgJkt8cg49oasl-5iJS-VdrILxWM9_JCJyURpUuslX4Eb4eUBtQ0x5BaPa8-S2NLdGTaL7nBOO8o8n0C5FEUU-qlEip79KE8aqOj-OC44VsIquSmOvWIQD26n3fCVlgwoRBD1gzzsDOeaSyzpKrZR851Kh6rEmF2qjJ8jt6EkxMsRNACmBomzgA4M1TTsisSUO87444pe35Z4_n5c735o2fZMrGgMwiJNh7rT8SYxtIkxngioiGnwkxGQxQ4NzPAHg-XSY0J04pNm7KqTkgtxyrqOANJLIjXlR-U9SQ90NjHVQ", "e": "AQAB", "d": "G4E84ppZwm3fLMI0YZ26iJ_sq3BKcRpQD6_r0o8ZrZmO7y4Uc-ywoP7h1lhFzaox66cokuloZpKOdGHIfK-84EkI3WeveWHPqBjmTMlN_ClQVcI48mUbLhD7Zeenhi9y9ipD2fkNWi8OJny8k4GfXrGqm50w8schrsPksnxJjvocGMT6KZNfDURKF2HlM5X1uY8VCofokXOjBEeHIfYM8e7IcmPpyXwXKonDmVVbMbefo-u-TttgeyOYaO6s3flSy6Y0CnpWi43JQ_VEARxQl6Brj1oizr8UnQQ0nNCOWwDNVtOV4eSl7PZoiiT7CxYkYnhJXECMAM5YBpm4Qk9zdQ", "p": "1g4ZGrXOuo75p9_MRIepXGpBWxip4V7B9XmO9WzPCv8nMorJntWBmsYV1I01aITxadHatO4Gl2xLniNkDyrEQzJ7w38RQgsVK-CqbnC0K9N77QPbHeC1YQd9RCNyUohOimKvb7jyv798FBU1GO5QI2eNgfnnfteSVXhD2iOoTOs", "q": "xJJ-8toxJdnLa0uUsAbql6zeNXGbUBMzu3FomKlyuWuq841jS2kIalaO_TRj5hbnE45jmCjeLgTVO6Ach3Wfk4zrqajqfFJ0zUg_Wexp49lC3RWiV4icBb85Q6bzeJD9Dn9vhjpfWVkczf_NeA1fGH_pcgfkT6Dm706GFFttLL8", "dp": "Zfx3l5NR-O8QIhzuHSSp279Afl_E6P0V2phdNa_vAaVKDrmzkHrXcl-4nPnenXrh7vIuiw_xkgnmCWWBUfylYALYlu-e0GGpZ6t2aIJIRa1QmT_CEX0zzhQcae-dk5cgHK0iO0_aUOOyAXuNPeClzAiVknz4ACZDsXdIlNFyaZs", "dq": "Z9DG4xOBKXBhEoWUPXMpqnlN0gPx9tRtWe2HRDkZsfu_CWn-qvEJ1L9qPSfSKs6ls5pb1xyeWseKpjblWlUwtgiS3cOsM4SI03H4o1FMi11PBtxKJNitLgvT_nrJ0z8fpux-xfFGMjXyFImoxmKpepLzg5nPZo6f6HscLNwsSJk", "qi": "Sk20wFvilpRKHq79xxFWiDUPHi0x0pp82dYIEntGQkKUWkbSlhgf3MAi5NEQTDmXdnB-rVeWIvEi-BXfdnNgdn8eC4zSdtF4sIAhYr5VWZo0WVWDhT7u2ccvZBFymiz8lo3gN57wGUCi9pbZqzV1-ZppX6YTNDdDCE0q-KO3Cec"}
+  ]
+}
diff --git a/tests/files/jwks_single_public.json b/tests/files/jwks_single_public.json
new file mode 100644
index 00000000..c47e1dd8
--- /dev/null
+++ b/tests/files/jwks_single_public.json
@@ -0,0 +1,5 @@
+{
+  "keys": [
+    {"kty": "RSA", "kid": "abc", "n": "pF1JaMSN8TEsh4N4O_5SpEAVLivJyLH-Cgl3OQBPGgJkt8cg49oasl-5iJS-VdrILxWM9_JCJyURpUuslX4Eb4eUBtQ0x5BaPa8-S2NLdGTaL7nBOO8o8n0C5FEUU-qlEip79KE8aqOj-OC44VsIquSmOvWIQD26n3fCVlgwoRBD1gzzsDOeaSyzpKrZR851Kh6rEmF2qjJ8jt6EkxMsRNACmBomzgA4M1TTsisSUO87444pe35Z4_n5c735o2fZMrGgMwiJNh7rT8SYxtIkxngioiGnwkxGQxQ4NzPAHg-XSY0J04pNm7KqTkgtxyrqOANJLIjXlR-U9SQ90NjHVQ", "e": "AQAB"}
+  ]
+}
diff --git a/tests/jose/test_jwt.py b/tests/jose/test_jwt.py
index 6326dd5f..c6c158fc 100644
--- a/tests/jose/test_jwt.py
+++ b/tests/jose/test_jwt.py
@@ -249,6 +249,18 @@ def test_use_jwks(self):
         claims = jwt.decode(data, pub_key)
         self.assertEqual(claims['name'], 'hi')
 
+    def test_use_jwks_single_kid(self):
+        """Test that jwks can be decoded if a kid for decoding is given
+        and encoded data has no kid and only one key is set."""
+        header = {'alg': 'RS256'}
+        payload = {'name': 'hi'}
+        private_key = read_file_path('jwks_single_private.json')
+        pub_key = read_file_path('jwks_single_public.json')
+        data = jwt.encode(header, payload, private_key)
+        self.assertEqual(data.count(b'.'), 2)
+        claims = jwt.decode(data, pub_key)
+        self.assertEqual(claims['name'], 'hi')
+
     def test_with_ec(self):
         payload = {'name': 'hi'}
         private_key = read_file_path('secp521r1-private.json')