Repository: jupyterhub/oauthenticator
Pull request #765 (2024-09-28 15:21:32)

Add `enable_pkce` config, True by default
Fixes https://github.com/jupyterhub/oauthenticator/issues/468


Patch:
diff --git a/oauthenticator/oauth2.py b/oauthenticator/oauth2.py
index 8c62a67..bdf9f28 100644
--- a/oauthenticator/oauth2.py
+++ b/oauthenticator/oauth2.py
@@ -5,8 +5,10 @@ Founded based on work by Kyle Kelley (@rgbkrk)
 """
 
 import base64
+import hashlib
 import json
 import os
+import secrets
 import uuid
 from functools import reduce
 from inspect import isawaitable
@@ -91,6 +93,18 @@ class OAuthLoginHandler(OAuth2Mixin, BaseHandler):
             STATE_COOKIE_NAME, state_cookie_value, expires_days=1, httponly=True
         )
 
+    def _generate_pkce_params(self):
+        # https://datatracker.ietf.org/doc/html/rfc7636#section-4
+        # It is recommended that the output of the random number generator creates
+        # a 32-octet sequence which is base64url-encoded to produce a 43-octet URL
+        # safe string to use as the code verifier.
+        code_verifier = secrets.token_urlsafe(32)
+        code_challenge = hashlib.sha256(code_verifier.encode("utf-8")).digest()
+        code_challenge_base64 = (
+            base64.urlsafe_b64encode(code_challenge).decode("utf-8").rstrip("=")
+        )
+        return code_verifier, code_challenge_base64
+
     def _generate_state_id(self):
         return uuid.uuid4().hex
 
@@ -115,7 +129,15 @@ class OAuthLoginHandler(OAuth2Mixin, BaseHandler):
         state_id = self._generate_state_id()
         next_url = self._get_next_url()
 
-        cookie_state = _serialize_state({"state_id": state_id, "next_url": next_url})
+        state = {"state_id": state_id, "next_url": next_url}
+
+        if self.authenticator.enable_pkce:
+            code_verifier, code_challenge = self._generate_pkce_params()
+            state["code_verifier"] = code_verifier
+            token_params["code_challenge"] = code_challenge
+            token_params["code_challenge_method"] = "S256"
+
+        cookie_state = _serialize_state(state)
         self.set_state_cookie(cookie_state)
 
         authorize_state = _serialize_state({"state_id": state_id})
@@ -663,6 +685,34 @@ class OAuthenticator(Authenticator):
         """,
     )
 
+    enable_pkce = Bool(
+        True,
+        config=True,
+        help="""
+            Enable Proof Key for Code Exchange (PKCE) for the OAuth2 authorization code flow.
+            For more information, see `RFC 7636 <https://datatracker.ietf.org/doc/html/rfc7636>`_.
+
+            PKCE can be used even if the authorization server does not support it. According to
+            `section 3.1 of RFC 6749 <https://www.rfc-editor.org/rfc/rfc6749#section-3.1>`_:
+
+                The authorization server MUST ignore unrecognized request parameters.
+
+            Additionally, `section 5 of RFC 7636 <https://datatracker.ietf.org/doc/html/rfc7636#section-5>`_ states:
+
+                As the OAuth 2.0 [RFC6749] server responses are unchanged by this
+                specification, client implementations of this specification do not
+                need to know if the server has implemented this specification or not
+                and SHOULD send the additional parameters as defined in Section 4 to
+                all servers.
+
+            Note that S256 is the only code challenge method supported. As per `section 4.2 of RFC 6749 
+            <https://www.rfc-editor.org/rfc/rfc6749#section-3.1>`_:
+
+                If the client is capable of using "S256", it MUST use "S256", as
+                "S256" is Mandatory To Implement (MTI) on the server.
+            """,
+    )
+
     client_id_env = ""
     client_id = Unicode(
         config=True,
@@ -980,6 +1030,18 @@ class OAuthenticator(Authenticator):
             "data": data,
         }
 
+        if self.enable_pkce:
+            # https://datatracker.ietf.org/doc/html/rfc7636#section-4.5
+            cookie_state = handler.get_state_cookie()
+            if not cookie_state:
+                raise web.HTTPError(400, "OAuth state missing from cookies")
+
+            code_verifier = _deserialize_state(cookie_state).get("code_verifier")
+            if not code_verifier:
+                raise web.HTTPError(400, "Missing code_verifier")
+
+            params.update([("code_verifier", code_verifier)])
+
         # the client_id and client_secret should not be included in the access token request params
         # when basic authentication is used
         # ref: https://www.rfc-editor.org/rfc/rfc6749#section-2.3.1


Test patch:
diff --git a/oauthenticator/tests/mocks.py b/oauthenticator/tests/mocks.py
index efeac57..d917460 100644
--- a/oauthenticator/tests/mocks.py
+++ b/oauthenticator/tests/mocks.py
@@ -15,6 +15,8 @@ from tornado.httputil import HTTPServerRequest
 from tornado.log import app_log
 from tornado.simple_httpclient import SimpleAsyncHTTPClient
 
+from ..oauth2 import _serialize_state
+
 RegExpType = type(re.compile('.'))
 
 
@@ -222,6 +224,11 @@ def setup_oauth_mock(
             method="GET", uri=f"https://hub.example.com?code={code}"
         )
         handler.hub = Mock(server=Mock(base_url='/hub/'), base_url='/hub/')
+        handler.get_state_cookie = Mock(
+            return_value=_serialize_state(
+                {"state_id": "123", "next_url": "/ABC", "code_verifier": "123"}
+            )
+        )
         return handler
 
     client.handler_for_user = handler_for_user
diff --git a/oauthenticator/tests/test_oauth2.py b/oauthenticator/tests/test_oauth2.py
index 4386e92..e0c2fde 100644
--- a/oauthenticator/tests/test_oauth2.py
+++ b/oauthenticator/tests/test_oauth2.py
@@ -33,9 +33,10 @@ TEST_STATE_ID = '123'
 TEST_NEXT_URL = '/ABC'
 
 
-async def test_login_states():
+@mark.parametrize("enable_pkce", [True, False])
+async def test_login_states(enable_pkce):
     login_request_uri = f"http://myhost/login?next={TEST_NEXT_URL}"
-    authenticator = OAuthenticator()
+    authenticator = OAuthenticator(enable_pkce=enable_pkce)
     login_handler = mock_handler(
         OAuthLoginHandler,
         uri=login_request_uri,
@@ -43,30 +44,35 @@ async def test_login_states():
     )
 
     login_handler._generate_state_id = Mock(return_value=TEST_STATE_ID)
-
+    code_verifier, code_challenge = login_handler._generate_pkce_params()
+    login_handler._generate_pkce_params = Mock(
+        return_value=(code_verifier, code_challenge)
+    )
     login_handler.set_state_cookie = Mock()
     login_handler.authorize_redirect = Mock()
 
     login_handler.get()  # no await, we've mocked the authorizer_redirect to NOT be async
 
-    expected_cookie_value = _serialize_state(
-        {
-            'state_id': TEST_STATE_ID,
-            'next_url': TEST_NEXT_URL,
-        }
-    )
+    expected_state = {
+        'state_id': TEST_STATE_ID,
+        'next_url': TEST_NEXT_URL,
+    }
+    if enable_pkce:
+        expected_state['code_verifier'] = code_verifier
+    expected_cookie_value = _serialize_state(expected_state)
 
     login_handler.set_state_cookie.assert_called_once_with(expected_cookie_value)
 
-    expected_state_param_value = _serialize_state(
-        {
-            'state_id': TEST_STATE_ID,
-        }
-    )
+    expected_state_param_value = {
+        'state': _serialize_state({'state_id': TEST_STATE_ID})
+    }
+    if enable_pkce:
+        expected_state_param_value['code_challenge'] = code_challenge
+        expected_state_param_value['code_challenge_method'] = 'S256'
 
     login_handler.authorize_redirect.assert_called_once()
     assert (
-        login_handler.authorize_redirect.call_args.kwargs['extra_params']['state']
+        login_handler.authorize_redirect.call_args.kwargs['extra_params']
         == expected_state_param_value
     )