SNOW-1825621 OAuth code flow PKCE support

DESCRIPTION.md

@@ -16,6 +16,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added a feature to verify if the connection is still good enough to send queries over.
   - Added support for base64-encoded DER private key strings in the `private_key` authentication type.
   - Added support for OAuth authorization code flow.
+  - Added support for PKCE on top of OAuth authorization flow.
 
 - v3.12.4(December 3,2024)
   - Fixed a bug where multipart uploads to Azure would be missing their MD5 hashes.

src/snowflake/connector/auth/oauth_code.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import base64
+import hashlib
 import json
 import logging
 import secrets
@@ -54,6 +55,7 @@ def __init__(
         token_request_url: str,
         redirect_uri: str,
         scope: str,
+        pkce: bool = False,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -71,6 +73,10 @@ def __init__(
         logger.debug("chose oauth state: %s", self._state)
         self._oauth_token = None
         self._protocol = "http"
+        self.pkce = pkce
+        if pkce:
+            logger.debug("oauth pkce is going to be used")
+        self._verifier: str | None = None
 
     def reset_secrets(self) -> None:
         self._oauth_token = None
@@ -102,6 +108,18 @@ def construct_url(self) -> str:
         }
         if self.scope:
             params["scope"] = self.scope
+        if self.pkce:
+            self._verifier = secrets.token_urlsafe(43)
+            # calculate challenge and verifier
+            challenge = (
+                base64.urlsafe_b64encode(
+                    hashlib.sha256(self._verifier.encode("utf-8")).digest()
+                )
+                .decode("utf-8")
+                .rstrip("=")
+            )
+            params["code_challenge"] = challenge
+            params["code_challenge_method"] = "S256"
         url_params = urllib.parse.urlencode(params)
         url = f"{self.authentication_url}?{url_params}"
         return url
@@ -184,6 +202,10 @@ def prepare(
         }
         if self.client_secret:
             fields["client_secret"] = self.client_secret
+        if self.pkce:
+            assert self._verifier is not None
+            fields["code_verifier"] = self._verifier
+
         resp = urllib3.PoolManager().request_encode_body(  # TODO: use network pool to gain use of proxy settings and so on
             "POST",
             self.token_request_url,

src/snowflake/connector/connection.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import atexit
+import collections.abc
 import logging
 import os
 import pathlib
@@ -333,6 +334,11 @@ def _get_private_bytes_from_file(
         str,
         # SNOW-1825621: OAUTH implementation
     ),
+    "oauth_security_features": (
+        ("pkce",),
+        collections.abc.Iterable,  # of strings
+        # SNOW-1825621: OAUTH PKCE
+    ),
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")
@@ -1117,6 +1123,7 @@ def __open_connection(self):
                     backoff_generator=self._backoff_generator,
                 )
             elif self._authenticator == OAUTH_AUTHORIZATION_CODE:
+                pkce = "pkce" in map(lambda e: e.lower(), self._oauth_security_features)
                 if self._oauth_client_id is None:
                     Error.errorhandler_wrapper(
                         self,
@@ -1142,6 +1149,7 @@ def __open_connection(self):
                     ),
                     redirect_uri="http://127.0.0.1:{port}/",
                     scope=self._oauth_scope,
+                    pkce=pkce,
                 )
             elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
                 self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (