Feat: Add Audience + Scopes Param to PKCE and DeviceCode Auth Flows (flyteorg#1876)

This PR adds the audience parameter into the PKCE/AuthorizationClient and DeviceCodeAuthenticator auth flows. This param is derived from auth_helper.RemoteClientConfigStore.get_client_config()

Also, the scopes configuration - derived from config.yaml by way of configuration.PlatformConfig is now exposed to the aforementioned auth flows.

Changes are also made to the way the scopes parameter is parsed in the token_client module to bring it in line with the AuthorizationClient class flow.

This now makes it so that Auth0 can use the PKCE and DeviceCode Auth flows without errors.

Signed-off-by: tnam <[email protected]>

Author: PudgyPigeon

flytekit/clients/auth/auth_client.py

@@ -179,6 +179,7 @@ def __init__(
         endpoint: str,
         auth_endpoint: str,
         token_endpoint: str,
+        audience: typing.Optional[str] = None,
         scopes: typing.Optional[typing.List[str]] = None,
         client_id: typing.Optional[str] = None,
         redirect_uri: typing.Optional[str] = None,
@@ -196,6 +197,7 @@ def __init__(
         :param endpoint: str endpoint to connect to
         :param auth_endpoint: str endpoint where auth metadata can be found
         :param token_endpoint: str endpoint to retrieve token from
+        :param audience: (optional) Audience parameter for Auth0
         :param scopes: list[str] oauth2 scopes
         :param client_id: oauth2 client id
         :param redirect_uri: oauth2 redirect uri
@@ -227,6 +229,7 @@ def __init__(
             self._remote = endpoint_metadata
         self._token_endpoint = token_endpoint
         self._client_id = client_id
+        self._audience = audience
         self._scopes = scopes or []
         self._redirect_uri = redirect_uri
         state = _generate_state_parameter()
@@ -246,6 +249,10 @@ def __init__(
             "state": state,
         }
 
+        # Conditionally add audience param if provided - value is not None
+        if self._audience:
+            self._request_auth_code_params["audience"] = self._audience
+
         if request_auth_code_params:
             # Allow adding additional parameters to the request_auth_code_params
             self._request_auth_code_params.update(request_auth_code_params)

flytekit/clients/auth/authenticator.py

@@ -88,12 +88,19 @@ def refresh_credentials(self):
 class PKCEAuthenticator(Authenticator):
     """
     This Authenticator encapsulates the entire PKCE flow and automatically opens a browser window for login
+
+    For Auth0 - you will need to manually configure your config.yaml to include a scopes list of the syntax:
+    admin.scopes: ["offline_access", "offline", "all", "openid"] and/or similar scopes in order to get the refresh token +
+    caching. Otherwise, it will just receive the access token alone. Your FlyteCTL Helm config however should only
+    contain ["offline", "all"] - as OIDC scopes are ungrantable in Auth0 customer APIs. They are simply requested
+    for in the POST request during the token caching process.
     """
 
     def __init__(
         self,
         endpoint: str,
         cfg_store: ClientConfigStore,
+        scopes: typing.Optional[typing.List[str]] = None,
         header_key: typing.Optional[str] = None,
         verify: typing.Optional[typing.Union[bool, str]] = None,
         session: typing.Optional[requests.Session] = None,
@@ -104,6 +111,7 @@ def __init__(
         super().__init__(endpoint, header_key, KeyringStore.retrieve(endpoint), verify=verify)
         self._cfg_store = cfg_store
         self._auth_client = None
+        self._scopes = scopes
         self._session = session or requests.Session()
 
     def _initialize_auth_client(self):
@@ -120,7 +128,11 @@ def _initialize_auth_client(self):
                 endpoint=self._endpoint,
                 redirect_uri=cfg.redirect_uri,
                 client_id=cfg.client_id,
-                scopes=cfg.scopes,
+                # Audience only needed for Auth0 - Taken from client config
+                audience=cfg.audience,
+                scopes=self._scopes or cfg.scopes,
+                # self._scopes refers to flytekit.configuration.PlatformConfig (config.yaml)
+                # cfg.scopes refers to PublicClientConfig scopes (can be defined in Helm deployments)
                 auth_endpoint=cfg.authorization_endpoint,
                 token_endpoint=cfg.token_endpoint,
                 verify=self._verify,
@@ -254,15 +266,19 @@ def __init__(
         cfg_store: ClientConfigStore,
         header_key: typing.Optional[str] = None,
         audience: typing.Optional[str] = None,
+        scopes: typing.Optional[typing.List[str]] = None,
         http_proxy_url: typing.Optional[str] = None,
         verify: typing.Optional[typing.Union[bool, str]] = None,
         session: typing.Optional[requests.Session] = None,
     ):
-        self._audience = audience
         cfg = cfg_store.get_client_config()
+        self._audience = audience or cfg.audience
         self._client_id = cfg.client_id
         self._device_auth_endpoint = cfg.device_authorization_endpoint
-        self._scope = cfg.scopes
+        # Input param: scopes refers to flytekit.configuration.PlatformConfig (config.yaml)
+        # cfg.scopes refers to PublicClientConfig scopes (can be defined in Helm deployments)
+        # Use "scope" from object instantiation if value is not None - otherwise, default to cfg.scopes
+        self._scopes = scopes or cfg.scopes
         self._token_endpoint = cfg.token_endpoint
         if self._device_auth_endpoint is None:
             raise AuthenticationError(
@@ -282,7 +298,7 @@ def refresh_credentials(self):
             self._device_auth_endpoint,
             self._client_id,
             self._audience,
-            self._scope,
+            self._scopes,
             self._http_proxy_url,
             self._verify,
             self._session,
@@ -296,6 +312,8 @@ def refresh_credentials(self):
                 resp,
                 self._token_endpoint,
                 client_id=self._client_id,
+                audience=self._audience,
+                scopes=self._scopes,
                 http_proxy_url=self._http_proxy_url,
                 verify=self._verify,
             )

flytekit/clients/auth/token_client.py

@@ -99,7 +99,7 @@ def get_token(
     if device_code:
         body["device_code"] = device_code
     if scopes is not None:
-        body["scope"] = ",".join(scopes)
+        body["scope"] = " ".join(s.strip("' ") for s in scopes).strip("[]'")
     if audience:
         body["audience"] = audience
 
@@ -135,7 +135,7 @@ def get_device_code(
     Retrieves the device Authentication code that can be done to authenticate the request using a browser on a
     separate device
     """
-    _scope = " ".join(scope) if scope is not None else ""
+    _scope = " ".join(s.strip("' ") for s in scope).strip("[]'") if scope is not None else ""
     payload = {"client_id": client_id, "scope": _scope, "audience": audience}
     proxies = {"https": http_proxy_url, "http": http_proxy_url} if http_proxy_url else None
     if not session:
@@ -150,6 +150,8 @@ def poll_token_endpoint(
     resp: DeviceCodeResponse,
     token_endpoint: str,
     client_id: str,
+    audience: typing.Optional[str] = None,
+    scopes: typing.Optional[str] = None,
     http_proxy_url: typing.Optional[str] = None,
     verify: typing.Optional[typing.Union[bool, str]] = None,
 ) -> typing.Tuple[str, int]:
@@ -162,6 +164,8 @@ def poll_token_endpoint(
                 token_endpoint,
                 grant_type=GrantType.DEVICE_CODE,
                 client_id=client_id,
+                audience=audience,
+                scopes=scopes,
                 device_code=resp.device_code,
                 http_proxy_url=http_proxy_url,
                 verify=verify,

flytekit/clients/auth_helper.py

@@ -71,7 +71,7 @@ def get_authenticator(cfg: PlatformConfig, cfg_store: ClientConfigStore) -> Auth
     session = get_session(cfg)
 
     if cfg_auth == AuthType.STANDARD or cfg_auth == AuthType.PKCE:
-        return PKCEAuthenticator(cfg.endpoint, cfg_store, verify=verify, session=session)
+        return PKCEAuthenticator(cfg.endpoint, cfg_store, scopes=cfg.scopes, verify=verify, session=session)
     elif cfg_auth == AuthType.BASIC or cfg_auth == AuthType.CLIENT_CREDENTIALS or cfg_auth == AuthType.CLIENTSECRET:
         return ClientCredentialsAuthenticator(
             endpoint=cfg.endpoint,
@@ -97,6 +97,7 @@ def get_authenticator(cfg: PlatformConfig, cfg_store: ClientConfigStore) -> Auth
             endpoint=cfg.endpoint,
             cfg_store=cfg_store,
             audience=cfg.audience,
+            scopes=cfg.scopes,
             http_proxy_url=cfg.http_proxy_url,
             verify=verify,
             session=session,

flytekit/configuration/internal.py

@@ -110,12 +110,16 @@ class Credentials(object):
     """
 
     SCOPES = ConfigEntry(LegacyConfigEntry(SECTION, "scopes", list), YamlConfigEntry("admin.scopes", list))
+    """
+    This setting can be used to manually pass in scopes into authenticator flows - eg.) for Auth0 compatibility
+    """
 
     AUTH_MODE = ConfigEntry(LegacyConfigEntry(SECTION, "auth_mode"), YamlConfigEntry("admin.authType"))
     """
     The auth mode defines the behavior used to request and refresh credentials. The currently supported modes include:
     - 'standard' or 'Pkce': This uses the pkce-enhanced authorization code flow by opening a browser window to initiate
             credentials access.
+    - "DeviceFlow": This uses the Device Authorization Flow
     - 'basic', 'client_credentials' or 'clientSecret': This uses symmetric key auth in which the end user enters a
             client id and a client secret and public key encryption is used to facilitate authentication.
     - None: No auth will be attempted.