Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Refactoring code to get oidc end points from discovery URL. #4429

Merged
3 changes: 0 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,6 @@
[![License](https://img.shields.io/badge/License-Apache%202.0-blue)](https://github.com/feast-dev/feast/blob/master/LICENSE)
[![GitHub Release](https://img.shields.io/github/v/release/feast-dev/feast.svg?style=flat&sort=semver&color=blue)](https://github.com/feast-dev/feast/releases)

## Join us on Slack!
👋👋👋 [Come say hi on Slack!](https://join.slack.com/t/feastopensource/signup)
tokoko marked this conversation as resolved.
Show resolved Hide resolved

## Overview

Feast (**Fea**ture **St**ore) is an open source feature store for machine learning. Feast is the fastest path to manage existing infrastructure to productionize analytic data for model training and online inference.
Expand Down
3 changes: 1 addition & 2 deletions docs/getting-started/components/authz_manager.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
...
```
Expand Down
15 changes: 10 additions & 5 deletions sdk/python/feast/permissions/auth/oidc_token_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand All @@ -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):
"""
Expand All @@ -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)
Expand All @@ -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)
Expand Down
1 change: 0 additions & 1 deletion sdk/python/feast/permissions/auth_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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",
Expand Down
40 changes: 40 additions & 0 deletions sdk/python/feast/permissions/oidc_service.py
Original file line number Diff line number Diff line change
@@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍


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()
1 change: 0 additions & 1 deletion sdk/python/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
"""),
],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 0 additions & 4 deletions sdk/python/tests/unit/infra/scaffolding/test_repo_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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"
Expand Down
3 changes: 1 addition & 2 deletions sdk/python/tests/unit/permissions/auth/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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="",
Expand Down
9 changes: 9 additions & 0 deletions sdk/python/tests/unit/permissions/auth/server/mock_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"]}},
Expand Down
Loading