feat: Refactoring code to get oidc end points from discovery URL. (fe…

…ast-dev#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>
lokeshrangineni · Aug 24, 2024 · 896360a · 896360a
1 parent dda0088
commit 896360a
Show file tree

Hide file tree

Showing 12 changed files with 73 additions and 35 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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"]}},