Add PKCE parameters to UM auth flow (#275)

* Add PKCE parameters to UM auth flow

* Format with black

* Refactor params to automatically set the code challenge method

Author: awolfden

tests/test_user_management.py

@@ -621,6 +621,28 @@ def test_authorization_url_has_expected_query_params_with_state(self):
             "response_type": RESPONSE_TYPE_CODE,
         }
 
+    def test_authorization_url_has_expected_query_params_with_code_challenge(self):
+        connection_id = "connection_123"
+        redirect_uri = "https://localhost/auth/callback"
+        code_challenge = json.dumps({"code_challenge": "code_challenge_for_pkce"})
+
+        authorization_url = self.user_management.get_authorization_url(
+            connection_id=connection_id,
+            code_challenge=code_challenge,
+            redirect_uri=redirect_uri,
+        )
+
+        parsed_url = urlparse(authorization_url)
+
+        assert dict(parse_qsl(parsed_url.query)) == {
+            "code_challenge": code_challenge,
+            "code_challenge_method": "S256",
+            "client_id": workos.client_id,
+            "redirect_uri": redirect_uri,
+            "connection_id": connection_id,
+            "response_type": RESPONSE_TYPE_CODE,
+        }
+
     def test_authenticate_with_password(
         self, capture_and_mock_request, mock_auth_response
     ):
@@ -653,13 +675,15 @@ def test_authenticate_with_password(
 
     def test_authenticate_with_code(self, capture_and_mock_request, mock_auth_response):
         code = "test_code"
+        code_verifier = "test_code_verifier"
         user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36"
         ip_address = "192.0.0.1"
 
         url, request = capture_and_mock_request("post", mock_auth_response, 200)
 
         response = self.user_management.authenticate_with_code(
             code=code,
+            code_verifier=code_verifier,
             user_agent=user_agent,
             ip_address=ip_address,
         )
@@ -670,6 +694,7 @@ def test_authenticate_with_code(self, capture_and_mock_request, mock_auth_respon
         assert response["access_token"] == "access_token_12345"
         assert response["refresh_token"] == "refresh_token_12345"
         assert request["json"]["code"] == code
+        assert request["json"]["code_verifier"] == code_verifier
         assert request["json"]["user_agent"] == user_agent
         assert request["json"]["ip_address"] == ip_address
         assert request["json"]["client_id"] == "client_b27needthisforssotemxo"

workos/user_management.py

@@ -423,6 +423,7 @@ def get_authorization_url(
         domain_hint=None,
         login_hint=None,
         state=None,
+        code_challenge=None,
     ):
         """Generate an OAuth 2.0 authorization URL.
 
@@ -444,6 +445,7 @@ def get_authorization_url(
                 OktaSAML, and AzureSAML connection types. (Optional)
             state (str) - An encoded string passed to WorkOS that'd be preserved through the authentication workflow, passed
                 back as a query parameter. (Optional)
+            code_challenge (str) - Code challenge is derived from the code verifier used for the PKCE flow. (Optional)
 
         Returns:
             str: URL to redirect a User to to begin the OAuth workflow with WorkOS
@@ -476,6 +478,9 @@ def get_authorization_url(
             params["login_hint"] = login_hint
         if state is not None:
             params["state"] = state
+        if code_challenge:
+            params["code_challenge"] = code_challenge
+            params["code_challenge_method"] = "S256"
 
         prepared_request = Request(
             "GET",
@@ -534,13 +539,16 @@ def authenticate_with_password(
     def authenticate_with_code(
         self,
         code,
+        code_verifier=None,
         ip_address=None,
         user_agent=None,
     ):
         """Authenticates an OAuth user or a user that is logging in through SSO.
 
         Kwargs:
             code (str): The authorization value which was passed back as a query parameter in the callback to the Redirect URI.
+            code_verifier (str): The randomly generated string used to derive the code challenge that was passed to the authorization
+                url as part of the PKCE flow. This parameter is required when the client secret is not present. (Optional)
             ip_address (str): The IP address of the request from the user who is attempting to authenticate. (Optional)
             user_agent (str): The user agent of the request from the user who is attempting to authenticate. (Optional)
 
@@ -565,6 +573,9 @@ def authenticate_with_code(
         if user_agent:
             payload["user_agent"] = user_agent
 
+        if code_verifier:
+            payload["code_verifier"] = code_verifier
+
         response = self.request_helper.request(
             USER_AUTHENTICATE_PATH,
             method=REQUEST_METHOD_POST,