From 896360af19a37c9a2a4634ec88021c4f69bdb141 Mon Sep 17 00:00:00 2001 From: lokeshrangineni <19699092+lokeshrangineni@users.noreply.github.com> Date: Sat, 24 Aug 2024 05:16:46 -0400 Subject: [PATCH] feat: Refactoring code to get oidc end points from discovery URL. (#4429) * refactoring the permissions side server side code to get the OIDC end points from the discovery URL. Also removing the auth_server_url config from oidc auth config. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> * refactoring the permissions side server side code to get the OIDC end points from the discovery URL. Also removing the auth_server_url config from oidc auth config. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> * refactoring the permissions side server side code to get the OIDC end points from the discovery URL. Also removing the auth_server_url config from oidc auth config. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> * refactoring the permissions side server side code to get the OIDC end points from the discovery URL. Also removing the auth_server_url config from oidc auth config. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> * Fixing the issue with pre-commit hook template. Accidentally this was reverted in previous rebase and reverting it now. Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> --------- Signed-off-by: Lokesh Rangineni <19699092+lokeshrangineni@users.noreply.github.com> --- .../components/authz_manager.md | 3 +- .../permissions/auth/oidc_token_parser.py | 15 ++++--- sdk/python/feast/permissions/auth_model.py | 1 - .../oidc_authentication_client_manager.py | 21 ++-------- sdk/python/feast/permissions/oidc_service.py | 40 +++++++++++++++++++ sdk/python/tests/conftest.py | 1 - .../feature_repos/repo_configuration.py | 1 - .../universal/data_sources/file.py | 1 - .../infra/scaffolding/test_repo_config.py | 4 -- .../tests/unit/permissions/auth/conftest.py | 3 +- .../permissions/auth/server/mock_utils.py | 9 +++++ .../permissions/auth/test_token_parser.py | 9 ++++- 12 files changed, 73 insertions(+), 35 deletions(-) create mode 100644 sdk/python/feast/permissions/oidc_service.py diff --git a/docs/getting-started/components/authz_manager.md b/docs/getting-started/components/authz_manager.md index 09ca4d1366..876dd84f2e 100644 --- a/docs/getting-started/components/authz_manager.md +++ b/docs/getting-started/components/authz_manager.md @@ -68,8 +68,7 @@ auth: type: oidc client_id: _CLIENT_ID__ client_secret: _CLIENT_SECRET__ - realm: _REALM__ - auth_server_url: _OIDC_SERVER_URL_ + realm: _REALM__ auth_discovery_url: _OIDC_SERVER_URL_/realms/master/.well-known/openid-configuration ... ``` diff --git a/sdk/python/feast/permissions/auth/oidc_token_parser.py b/sdk/python/feast/permissions/auth/oidc_token_parser.py index 921a585bc2..fce9fdcbb2 100644 --- a/sdk/python/feast/permissions/auth/oidc_token_parser.py +++ b/sdk/python/feast/permissions/auth/oidc_token_parser.py @@ -11,6 +11,7 @@ from feast.permissions.auth.token_parser import TokenParser from feast.permissions.auth_model import OidcAuthConfig +from feast.permissions.oidc_service import OIDCDiscoveryService from feast.permissions.user import User logger = logging.getLogger(__name__) @@ -27,6 +28,9 @@ class OidcTokenParser(TokenParser): def __init__(self, auth_config: OidcAuthConfig): self._auth_config = auth_config + self.oidc_discovery_service = OIDCDiscoveryService( + self._auth_config.auth_discovery_url + ) async def _validate_token(self, access_token: str): """ @@ -38,9 +42,9 @@ async def _validate_token(self, access_token: str): request.headers = {"Authorization": f"Bearer {access_token}"} oauth_2_scheme = OAuth2AuthorizationCodeBearer( - tokenUrl=f"{self._auth_config.auth_server_url}/realms/{self._auth_config.realm}/protocol/openid-connect/token", - authorizationUrl=f"{self._auth_config.auth_server_url}/realms/{self._auth_config.realm}/protocol/openid-connect/auth", - refreshUrl=f"{self._auth_config.auth_server_url}/realms/{self._auth_config.realm}/protocol/openid-connect/token", + tokenUrl=self.oidc_discovery_service.get_token_url(), + authorizationUrl=self.oidc_discovery_service.get_authorization_url(), + refreshUrl=self.oidc_discovery_service.get_refresh_url(), ) await oauth_2_scheme(request=request) @@ -62,9 +66,10 @@ async def user_details_from_access_token(self, access_token: str) -> User: except Exception as e: raise AuthenticationError(f"Invalid token: {e}") - url = f"{self._auth_config.auth_server_url}/realms/{self._auth_config.realm}/protocol/openid-connect/certs" optional_custom_headers = {"User-agent": "custom-user-agent"} - jwks_client = PyJWKClient(url, headers=optional_custom_headers) + jwks_client = PyJWKClient( + self.oidc_discovery_service.get_jwks_url(), headers=optional_custom_headers + ) try: signing_key = jwks_client.get_signing_key_from_jwt(access_token) diff --git a/sdk/python/feast/permissions/auth_model.py b/sdk/python/feast/permissions/auth_model.py index afb0a22bc9..28eeb951a7 100644 --- a/sdk/python/feast/permissions/auth_model.py +++ b/sdk/python/feast/permissions/auth_model.py @@ -8,7 +8,6 @@ class AuthConfig(FeastConfigBaseModel): class OidcAuthConfig(AuthConfig): - auth_server_url: Optional[str] = None auth_discovery_url: str client_id: str client_secret: Optional[str] = None diff --git a/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py b/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py index 544764aae0..6744a1d2ad 100644 --- a/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py +++ b/sdk/python/feast/permissions/client/oidc_authentication_client_manager.py @@ -4,6 +4,7 @@ from feast.permissions.auth_model import OidcAuthConfig from feast.permissions.client.auth_client_manager import AuthenticationClientManager +from feast.permissions.oidc_service import OIDCDiscoveryService logger = logging.getLogger(__name__) @@ -12,25 +13,11 @@ class OidcAuthClientManager(AuthenticationClientManager): def __init__(self, auth_config: OidcAuthConfig): self.auth_config = auth_config - def _get_token_endpoint(self): - response = requests.get(self.auth_config.auth_discovery_url) - if response.status_code == 200: - oidc_config = response.json() - if not oidc_config["token_endpoint"]: - raise RuntimeError( - " OIDC token_endpoint is not available from discovery url response." - ) - return oidc_config["token_endpoint"].replace( - "master", self.auth_config.realm - ) - else: - raise RuntimeError( - f"Error fetching OIDC token endpoint configuration: {response.status_code} - {response.text}" - ) - def get_token(self): # Fetch the token endpoint from the discovery URL - token_endpoint = self._get_token_endpoint() + token_endpoint = OIDCDiscoveryService( + self.auth_config.auth_discovery_url + ).get_token_url() token_request_body = { "grant_type": "password", diff --git a/sdk/python/feast/permissions/oidc_service.py b/sdk/python/feast/permissions/oidc_service.py new file mode 100644 index 0000000000..73d0ec8f1b --- /dev/null +++ b/sdk/python/feast/permissions/oidc_service.py @@ -0,0 +1,40 @@ +import requests + + +class OIDCDiscoveryService: + def __init__(self, discovery_url: str): + self.discovery_url = discovery_url + self._discovery_data = None # Initialize it lazily. + + @property + def discovery_data(self): + """Lazily fetches and caches the OIDC discovery data.""" + if self._discovery_data is None: + self._discovery_data = self._fetch_discovery_data() + return self._discovery_data + + def _fetch_discovery_data(self) -> dict: + try: + response = requests.get(self.discovery_url) + response.raise_for_status() + return response.json() + except requests.RequestException as e: + raise RuntimeError( + f"Error fetching OIDC discovery response, discovery url - {self.discovery_url}, exception - {e} " + ) + + def get_authorization_url(self) -> str: + """Returns the authorization endpoint URL.""" + return self.discovery_data.get("authorization_endpoint") + + def get_token_url(self) -> str: + """Returns the token endpoint URL.""" + return self.discovery_data.get("token_endpoint") + + def get_jwks_url(self) -> str: + """Returns the jwks endpoint URL.""" + return self.discovery_data.get("jwks_uri") + + def get_refresh_url(self) -> str: + """Returns the refresh token URL (usually same as token URL).""" + return self.get_token_url() diff --git a/sdk/python/tests/conftest.py b/sdk/python/tests/conftest.py index 74aa68e984..d40f699b6b 100644 --- a/sdk/python/tests/conftest.py +++ b/sdk/python/tests/conftest.py @@ -463,7 +463,6 @@ def is_integration_test(all_markers_from_module): username: reader_writer password: password realm: master - auth_server_url: KEYCLOAK_URL_PLACE_HOLDER auth_discovery_url: KEYCLOAK_URL_PLACE_HOLDER/realms/master/.well-known/openid-configuration """), ], diff --git a/sdk/python/tests/integration/feature_repos/repo_configuration.py b/sdk/python/tests/integration/feature_repos/repo_configuration.py index 235c909d5f..0bf737f616 100644 --- a/sdk/python/tests/integration/feature_repos/repo_configuration.py +++ b/sdk/python/tests/integration/feature_repos/repo_configuration.py @@ -464,7 +464,6 @@ def setup(self): password="password", realm="master", type="oidc", - auth_server_url=keycloak_url, auth_discovery_url=f"{keycloak_url}/realms/master/.well-known" f"/openid-configuration", ) diff --git a/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py b/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py index b600699f81..adbb248a20 100644 --- a/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py +++ b/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py @@ -449,7 +449,6 @@ def __init__(self, project_name: str, *args, **kwargs): username: reader_writer password: password realm: master - auth_server_url: {keycloak_url} auth_discovery_url: {keycloak_url}/realms/master/.well-known/openid-configuration """ self.auth_config = auth_config_template.format(keycloak_url=self.keycloak_url) diff --git a/sdk/python/tests/unit/infra/scaffolding/test_repo_config.py b/sdk/python/tests/unit/infra/scaffolding/test_repo_config.py index 0725d6d261..5331d350e2 100644 --- a/sdk/python/tests/unit/infra/scaffolding/test_repo_config.py +++ b/sdk/python/tests/unit/infra/scaffolding/test_repo_config.py @@ -214,7 +214,6 @@ def test_auth_config(): username: test_user_name password: test_password realm: master - auth_server_url: http://localhost:8712 auth_discovery_url: http://localhost:8080/realms/master/.well-known/openid-configuration registry: "registry.db" provider: local @@ -237,7 +236,6 @@ def test_auth_config(): username: test_user_name password: test_password realm: master - auth_server_url: http://localhost:8712 auth_discovery_url: http://localhost:8080/realms/master/.well-known/openid-configuration registry: "registry.db" provider: local @@ -260,7 +258,6 @@ def test_auth_config(): username: test_user_name password: test_password realm: master - auth_server_url: http://localhost:8080 auth_discovery_url: http://localhost:8080/realms/master/.well-known/openid-configuration registry: "registry.db" provider: local @@ -278,7 +275,6 @@ def test_auth_config(): assert oidc_repo_config.auth_config.username == "test_user_name" assert oidc_repo_config.auth_config.password == "test_password" assert oidc_repo_config.auth_config.realm == "master" - assert oidc_repo_config.auth_config.auth_server_url == "http://localhost:8080" assert ( oidc_repo_config.auth_config.auth_discovery_url == "http://localhost:8080/realms/master/.well-known/openid-configuration" diff --git a/sdk/python/tests/unit/permissions/auth/conftest.py b/sdk/python/tests/unit/permissions/auth/conftest.py index dc71aba23b..0d6acd7fb2 100644 --- a/sdk/python/tests/unit/permissions/auth/conftest.py +++ b/sdk/python/tests/unit/permissions/auth/conftest.py @@ -73,8 +73,7 @@ def clusterrolebindings(sa_name, namespace) -> dict: @pytest.fixture def oidc_config() -> OidcAuthConfig: return OidcAuthConfig( - auth_server_url="", - auth_discovery_url="", + auth_discovery_url="https://localhost:8080/realms/master/.well-known/openid-configuration", client_id=_CLIENT_ID, client_secret="", username="", diff --git a/sdk/python/tests/unit/permissions/auth/server/mock_utils.py b/sdk/python/tests/unit/permissions/auth/server/mock_utils.py index 8f598774ee..12f7785b05 100644 --- a/sdk/python/tests/unit/permissions/auth/server/mock_utils.py +++ b/sdk/python/tests/unit/permissions/auth/server/mock_utils.py @@ -42,6 +42,15 @@ async def mock_oath2(self, request): lambda url, data, headers: token_response, ) + monkeypatch.setattr( + "feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data", + lambda self, *args, **kwargs: { + "authorization_endpoint": "https://localhost:8080/realms/master/protocol/openid-connect/auth", + "token_endpoint": "https://localhost:8080/realms/master/protocol/openid-connect/token", + "jwks_uri": "https://localhost:8080/realms/master/protocol/openid-connect/certs", + }, + ) + def mock_kubernetes(request, monkeypatch): sa_name = request.getfixturevalue("sa_name") diff --git a/sdk/python/tests/unit/permissions/auth/test_token_parser.py b/sdk/python/tests/unit/permissions/auth/test_token_parser.py index 6ae9094f81..cb153a17c9 100644 --- a/sdk/python/tests/unit/permissions/auth/test_token_parser.py +++ b/sdk/python/tests/unit/permissions/auth/test_token_parser.py @@ -21,13 +21,20 @@ ) @patch("feast.permissions.auth.oidc_token_parser.PyJWKClient.get_signing_key_from_jwt") @patch("feast.permissions.auth.oidc_token_parser.jwt.decode") +@patch("feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data") def test_oidc_token_validation_success( - mock_jwt, mock_signing_key, mock_oauth2, oidc_config + mock_discovery_data, mock_jwt, mock_signing_key, mock_oauth2, oidc_config ): signing_key = MagicMock() signing_key.key = "a-key" mock_signing_key.return_value = signing_key + mock_discovery_data.return_value = { + "authorization_endpoint": "https://localhost:8080/realms/master/protocol/openid-connect/auth", + "token_endpoint": "https://localhost:8080/realms/master/protocol/openid-connect/token", + "jwks_uri": "https://localhost:8080/realms/master/protocol/openid-connect/certs", + } + user_data = { "preferred_username": "my-name", "resource_access": {_CLIENT_ID: {"roles": ["reader", "writer"]}},