From c1ce35b97143f6e939e237172362de5d3e2b9b1b Mon Sep 17 00:00:00 2001
From: Alie Langston <alangsto@wellesley.edu>
Date: Wed, 15 Jan 2025 09:48:47 -0500
Subject: [PATCH] fix: ensure that LTI 1.3 launches work

---
 CHANGELOG.rst                                  |  4 ++++
 lti_consumer/__init__.py                       |  2 +-
 lti_consumer/lti_1p3/key_handlers.py           | 15 +++++----------
 lti_consumer/lti_1p3/tests/test_consumer.py    |  1 +
 .../lti_1p3/tests/test_key_handlers.py         | 18 ++++++++++++++----
 lti_consumer/plugin/views.py                   |  1 -
 6 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 9b58fc29..aa38452e 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -16,6 +16,10 @@ Please See the `releases tab <https://github.com/openedx/xblock-lti-consumer/rel
 Unreleased
 ~~~~~~~~~~
 
+9.13.1 - 2025-01-15
+-------------------
+* Fix broken LTI 1.3 launch
+
 9.13.0 - 2025-01-08
 -------------------
 * Removed pyjwkset package and replace with pyjwt package
diff --git a/lti_consumer/__init__.py b/lti_consumer/__init__.py
index 618ced98..90123b61 100644
--- a/lti_consumer/__init__.py
+++ b/lti_consumer/__init__.py
@@ -4,4 +4,4 @@
 from .apps import LTIConsumerApp
 from .lti_xblock import LtiConsumerXBlock
 
-__version__ = '9.13.0'
+__version__ = '9.13.1'
diff --git a/lti_consumer/lti_1p3/key_handlers.py b/lti_consumer/lti_1p3/key_handlers.py
index beb6eedf..6bcdc3f5 100644
--- a/lti_consumer/lti_1p3/key_handlers.py
+++ b/lti_consumer/lti_1p3/key_handlers.py
@@ -62,7 +62,7 @@ def __init__(self, public_key=None, keyset_url=None):
                 )
                 raise exceptions.InvalidRsaKey() from err
 
-    def _get_keyset(self, kid=None):
+    def _get_keyset(self):
         """
         Get keyset from available sources.
 
@@ -86,13 +86,6 @@ def _get_keyset(self, kid=None):
                 raise exceptions.NoSuitableKeys() from err
             keyset.extend(keys.keys)
 
-        if self.public_key and kid:
-            # Fill in key id of stored key.
-            # This is needed because if the JWS is signed with a
-            # key with a kid, pyjwkest doesn't match them with
-            # keys without kid (kid=None) and fails verification
-            self.public_key.kid = kid
-
         if self.public_key:
             # Add to keyset
             keyset.append(self.public_key)
@@ -185,7 +178,7 @@ def encode_and_sign(self, message, expiration=None):
 
         # The class instance that sets up the signing operation
         # An RS 256 key is required for LTI 1.3
-        return jwt.encode(_message, self.key.key, algorithm="RS256")
+        return jwt.encode(_message, self.key.key, algorithm="RS256", headers={"kid": self.key.key_id})
 
     def get_public_jwk(self):
         """
@@ -197,7 +190,9 @@ def get_public_jwk(self):
         if self.key:
             algo_obj = jwt.get_algorithm_by_name('RS256')
             public_key = algo_obj.prepare_key(self.key.key).public_key()
-            jwk['keys'].append(json.loads(algo_obj.to_jwk(public_key)))
+            public_jwk = json.loads(algo_obj.to_jwk(public_key))
+            public_jwk['kid'] = self.key.key_id
+            jwk['keys'].append(public_jwk)
         return jwk
 
     def validate_and_decode(self, token, iss=None, aud=None, exp=True):
diff --git a/lti_consumer/lti_1p3/tests/test_consumer.py b/lti_consumer/lti_1p3/tests/test_consumer.py
index 7b9bac2e..5f2eeb25 100644
--- a/lti_consumer/lti_1p3/tests/test_consumer.py
+++ b/lti_consumer/lti_1p3/tests/test_consumer.py
@@ -122,6 +122,7 @@ def _decode_token(self, token):
         keyset = PyJWKSet.from_dict(public_keyset).keys
 
         for obj in keyset:
+            self.assertEqual(obj.key_id, RSA_KEY_ID)
             message = jwt.decode(
                 token,
                 key=obj.key,
diff --git a/lti_consumer/lti_1p3/tests/test_key_handlers.py b/lti_consumer/lti_1p3/tests/test_key_handlers.py
index e7097f44..dcf6efe2 100644
--- a/lti_consumer/lti_1p3/tests/test_key_handlers.py
+++ b/lti_consumer/lti_1p3/tests/test_key_handlers.py
@@ -11,6 +11,7 @@
 import ddt
 import jwt
 from Cryptodome.PublicKey import RSA
+from cryptography.hazmat.primitives import serialization
 from django.test.testcases import TestCase
 from jwt.api_jwk import PyJWK
 
@@ -58,6 +59,11 @@ def test_encode_and_sign(self):
             message
         )
 
+        self.assertEqual(
+            jwt.get_unverified_header(signed_token)['kid'],
+            self.rsa_key_id
+        )
+
     # pylint: disable=unused-argument
     @patch('time.time', return_value=1000)
     def test_encode_and_sign_with_exp(self, mock_time):
@@ -233,16 +239,20 @@ def test_get_empty_keyset(self):
 
     def test_get_keyset_with_pub_key(self):
         """
-        Check that getting a keyset from a RSA key.
+        Check that if there is a public key, it is returned in the keyset.
         """
         self._setup_key_handler()
 
         # pylint: disable=protected-access
-        keyset = self.key_handler._get_keyset(kid=self.rsa_key_id)
+        keyset = self.key_handler._get_keyset()
         self.assertEqual(len(keyset), 1)
+        public_key = keyset[0].key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        )[:-1]
         self.assertEqual(
-            keyset[0].kid,
-            self.rsa_key_id
+            public_key,
+            self.public_key
         )
 
     def test_validate_and_decode(self):
diff --git a/lti_consumer/plugin/views.py b/lti_consumer/plugin/views.py
index a9967e4a..2f753bf3 100644
--- a/lti_consumer/plugin/views.py
+++ b/lti_consumer/plugin/views.py
@@ -479,7 +479,6 @@ def access_token_endpoint(
     except Exception:  # pylint: disable=broad-except
         exc_info = sys.exc_info()
 
-        # import pdb; pdb.set_trace()
         # Handle errors and return a proper response
         if exc_info[0] == MissingRequiredClaim:
             # Missing request attributes