Skip to content

Commit 2363186

Browse files
tmihalactmihalac
committed
Intra server communication
Signed-off-by: Theodor Mihalache <tmihalac@redhat.com>
1 parent 635a01b commit 2363186

File tree

9 files changed

+417
-45
lines changed

9 files changed

+417
-45
lines changed

‎infra/charts/feast-feature-server/templates/deployment.yaml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ spec:
3636
env:
3737
- name: FEATURE_STORE_YAML_BASE64
3838
value: {{ .Values.feature_store_yaml_base64 }}
39+
- name: INTRA_COMMUNICATION_BASE64
40+
value: {{ "intra-server-communication" | b64enc }}
3941
command:
4042
{{- if eq .Values.feast_mode "offline" }}
4143
- "feast"

‎sdk/python/feast/permissions/auth/kubernetes_token_parser.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import logging
2+
import os
23

34
import jwt
45
from kubernetes import client, config
@@ -41,10 +42,14 @@ async def user_details_from_access_token(self, access_token: str) -> User:
4142
current_user = f"{sa_namespace}:{sa_name}"
4243
logging.info(f"Received request from {sa_name} in {sa_namespace}")
4344

44-
roles = self.get_roles(sa_namespace, sa_name)
45-
logging.info(f"SA roles are: {roles}")
45+
intra_communication_base64 = os.getenv("INTRA_COMMUNICATION_BASE64")
46+
if sa_name is not None and sa_name == intra_communication_base64:
47+
return User(username=sa_name, roles=[])
48+
else:
49+
roles = self.get_roles(sa_namespace, sa_name)
50+
logging.info(f"SA roles are: {roles}")
4651

47-
return User(username=current_user, roles=roles)
52+
return User(username=current_user, roles=roles)
4853

4954
def get_roles(self, namespace: str, service_account_name: str) -> list[str]:
5055
"""

‎sdk/python/feast/permissions/auth/oidc_token_parser.py‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import logging
2+
import os
3+
from typing import Optional
24
from unittest.mock import Mock
35

46
import jwt
@@ -34,7 +36,7 @@ def __init__(self, auth_config: OidcAuthConfig):
3436

3537
async def _validate_token(self, access_token: str):
3638
"""
37-
Validate the token extracted from the headrer of the user request against the OAuth2 server.
39+
Validate the token extracted from the header of the user request against the OAuth2 server.
3840
"""
3941
# FastAPI's OAuth2AuthorizationCodeBearer requires a Request type but actually uses only the headers field
4042
# https://github.com/tiangolo/fastapi/blob/eca465f4c96acc5f6a22e92fd2211675ca8a20c8/fastapi/security/oauth2.py#L380
@@ -60,6 +62,11 @@ async def user_details_from_access_token(self, access_token: str) -> User:
6062
AuthenticationError if any error happens.
6163
"""
6264

65+
# check if intra server communication
66+
user = self._get_intra_comm_user(access_token)
67+
if user:
68+
return user
69+
6370
try:
6471
await self._validate_token(access_token)
6572
logger.info("Validated token")
@@ -108,3 +115,20 @@ async def user_details_from_access_token(self, access_token: str) -> User:
108115
except jwt.exceptions.InvalidTokenError:
109116
logger.exception("Exception while parsing the token:")
110117
raise AuthenticationError("Invalid token.")
118+
119+
def _get_intra_comm_user(self, access_token: str) -> Optional[User]:
120+
intra_communication_base64 = os.getenv("INTRA_COMMUNICATION_BASE64")
121+
122+
if intra_communication_base64:
123+
decoded_token = jwt.decode(
124+
access_token, options={"verify_signature": False}
125+
)
126+
if "preferred_username" in decoded_token:
127+
preferred_username: str = decoded_token["preferred_username"]
128+
if (
129+
preferred_username is not None
130+
and preferred_username == intra_communication_base64
131+
):
132+
return User(username=preferred_username, roles=[])
133+
134+
return None

‎sdk/python/feast/permissions/client/kubernetes_auth_client_manager.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import logging
22
import os
33

4+
import jwt
5+
46
from feast.permissions.auth_model import KubernetesAuthConfig
57
from feast.permissions.client.auth_client_manager import AuthenticationClientManager
68

@@ -13,6 +15,15 @@ def __init__(self, auth_config: KubernetesAuthConfig):
1315
self.token_file_path = "/var/run/secrets/kubernetes.io/serviceaccount/token"
1416

1517
def get_token(self):
18+
intra_communication_base64 = os.getenv("INTRA_COMMUNICATION_BASE64")
19+
# If intra server communication call
20+
if intra_communication_base64:
21+
payload = {
22+
"sub": f":::{intra_communication_base64}", # Subject claim
23+
}
24+
25+
return jwt.encode(payload, "")
26+
1627
try:
1728
token = self._read_token_from_file()
1829
return token

‎sdk/python/feast/permissions/client/oidc_authentication_client_manager.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import logging
2+
import os
23

4+
import jwt
35
import requests
46

57
from feast.permissions.auth_model import OidcAuthConfig
@@ -14,6 +16,15 @@ def __init__(self, auth_config: OidcAuthConfig):
1416
self.auth_config = auth_config
1517

1618
def get_token(self):
19+
intra_communication_base64 = os.getenv("INTRA_COMMUNICATION_BASE64")
20+
# If intra server communication call
21+
if intra_communication_base64:
22+
payload = {
23+
"preferred_username": f"{intra_communication_base64}", # Subject claim
24+
}
25+
26+
return jwt.encode(payload, "")
27+
1728
# Fetch the token endpoint from the discovery URL
1829
token_endpoint = OIDCDiscoveryService(
1930
self.auth_config.auth_discovery_url

‎sdk/python/feast/permissions/security_manager.py‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import logging
2+
import os
23
from contextvars import ContextVar
34
from typing import Callable, List, Optional, Union
45

@@ -110,6 +111,10 @@ def assert_permissions_to_update(
110111
Raises:
111112
PermissionError: If the current user is not authorized to execute all the requested actions on the given resource or on the existing one.
112113
"""
114+
sm = get_security_manager()
115+
if not is_auth_necessary(sm):
116+
return resource
117+
113118
actions = [AuthzedAction.DESCRIBE, AuthzedAction.UPDATE]
114119
try:
115120
existing_resource = getter(
@@ -142,10 +147,11 @@ def assert_permissions(
142147
Raises:
143148
PermissionError: If the current user is not authorized to execute the requested actions on the given resources.
144149
"""
150+
145151
sm = get_security_manager()
146-
if sm is None:
152+
if not is_auth_necessary(sm):
147153
return resource
148-
return sm.assert_permissions(
154+
return sm.assert_permissions( # type: ignore[union-attr]
149155
resources=[resource], actions=actions, filter_only=False
150156
)[0]
151157

@@ -165,10 +171,11 @@ def permitted_resources(
165171
Returns:
166172
list[FeastObject]]: A filtered list of the permitted resources, possibly empty.
167173
"""
174+
168175
sm = get_security_manager()
169-
if sm is None:
176+
if not is_auth_necessary(sm):
170177
return resources
171-
return sm.assert_permissions(resources=resources, actions=actions, filter_only=True)
178+
return sm.assert_permissions(resources=resources, actions=actions, filter_only=True) # type: ignore[union-attr]
172179

173180

174181
"""
@@ -201,3 +208,13 @@ def no_security_manager():
201208

202209
global _sm
203210
_sm = None
211+
212+
213+
def is_auth_necessary(sm: Optional[SecurityManager]) -> bool:
214+
intra_communication_base64 = os.getenv("INTRA_COMMUNICATION_BASE64")
215+
216+
return (
217+
sm is not None
218+
and sm.current_user is not None
219+
and sm.current_user.username != intra_communication_base64
220+
)

‎sdk/python/tests/unit/permissions/auth/test_token_parser.py‎

Lines changed: 145 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
# test_token_validator.py
2-
31
import asyncio
2+
import os
3+
from unittest import mock
44
from unittest.mock import MagicMock, patch
55

66
import assertpy
@@ -70,6 +70,75 @@ def test_oidc_token_validation_failure(mock_oauth2, oidc_config):
7070
)
7171

7272

73+
@mock.patch.dict(os.environ, {"INTRA_COMMUNICATION_BASE64": "test1234"})
74+
@pytest.mark.parametrize(
75+
"intra_communication_val, is_intra_server",
76+
[
77+
("test1234", True),
78+
("my-name", False),
79+
],
80+
)
81+
def test_oidc_inter_server_comm(
82+
intra_communication_val, is_intra_server, oidc_config, monkeypatch
83+
):
84+
async def mock_oath2(self, request):
85+
return "OK"
86+
87+
monkeypatch.setattr(
88+
"feast.permissions.auth.oidc_token_parser.OAuth2AuthorizationCodeBearer.__call__",
89+
mock_oath2,
90+
)
91+
signing_key = MagicMock()
92+
signing_key.key = "a-key"
93+
monkeypatch.setattr(
94+
"feast.permissions.auth.oidc_token_parser.PyJWKClient.get_signing_key_from_jwt",
95+
lambda self, access_token: signing_key,
96+
)
97+
98+
user_data = {
99+
"preferred_username": f"{intra_communication_val}",
100+
}
101+
102+
if not is_intra_server:
103+
user_data["resource_access"] = {_CLIENT_ID: {"roles": ["reader", "writer"]}}
104+
105+
monkeypatch.setattr(
106+
"feast.permissions.oidc_service.OIDCDiscoveryService._fetch_discovery_data",
107+
lambda self, *args, **kwargs: {
108+
"authorization_endpoint": "https://localhost:8080/realms/master/protocol/openid-connect/auth",
109+
"token_endpoint": "https://localhost:8080/realms/master/protocol/openid-connect/token",
110+
"jwks_uri": "https://localhost:8080/realms/master/protocol/openid-connect/certs",
111+
},
112+
)
113+
114+
monkeypatch.setattr(
115+
"feast.permissions.auth.oidc_token_parser.jwt.decode",
116+
lambda self, *args, **kwargs: user_data,
117+
)
118+
119+
access_token = "aaa-bbb-ccc"
120+
token_parser = OidcTokenParser(auth_config=oidc_config)
121+
user = asyncio.run(
122+
token_parser.user_details_from_access_token(access_token=access_token)
123+
)
124+
125+
if is_intra_server:
126+
assertpy.assert_that(user).is_not_none()
127+
assertpy.assert_that(user.username).is_equal_to(intra_communication_val)
128+
assertpy.assert_that(user.roles).is_equal_to([])
129+
else:
130+
assertpy.assert_that(user).is_not_none()
131+
assertpy.assert_that(user).is_type_of(User)
132+
if isinstance(user, User):
133+
assertpy.assert_that(user.username).is_equal_to("my-name")
134+
assertpy.assert_that(user.roles.sort()).is_equal_to(
135+
["reader", "writer"].sort()
136+
)
137+
assertpy.assert_that(user.has_matching_role(["reader"])).is_true()
138+
assertpy.assert_that(user.has_matching_role(["writer"])).is_true()
139+
assertpy.assert_that(user.has_matching_role(["updater"])).is_false()
140+
141+
73142
# TODO RBAC: Move role bindings to a reusable fixture
74143
@patch("feast.permissions.auth.kubernetes_token_parser.config.load_incluster_config")
75144
@patch("feast.permissions.auth.kubernetes_token_parser.jwt.decode")
@@ -127,3 +196,77 @@ def test_k8s_token_validation_failure(mock_jwt, mock_config):
127196
asyncio.run(
128197
token_parser.user_details_from_access_token(access_token=access_token)
129198
)
199+
200+
201+
@mock.patch.dict(os.environ, {"INTRA_COMMUNICATION_BASE64": "test1234"})
202+
@pytest.mark.parametrize(
203+
"intra_communication_val, is_intra_server",
204+
[
205+
("test1234", True),
206+
("my-name", False),
207+
],
208+
)
209+
def test_k8s_inter_server_comm(
210+
intra_communication_val,
211+
is_intra_server,
212+
oidc_config,
213+
request,
214+
rolebindings,
215+
clusterrolebindings,
216+
monkeypatch,
217+
):
218+
if is_intra_server:
219+
subject = f":::{intra_communication_val}"
220+
else:
221+
sa_name = request.getfixturevalue("sa_name")
222+
namespace = request.getfixturevalue("namespace")
223+
subject = f"system:serviceaccount:{namespace}:{sa_name}"
224+
rolebindings = request.getfixturevalue("rolebindings")
225+
clusterrolebindings = request.getfixturevalue("clusterrolebindings")
226+
227+
monkeypatch.setattr(
228+
"feast.permissions.auth.kubernetes_token_parser.client.RbacAuthorizationV1Api.list_namespaced_role_binding",
229+
lambda *args, **kwargs: rolebindings["items"],
230+
)
231+
monkeypatch.setattr(
232+
"feast.permissions.auth.kubernetes_token_parser.client.RbacAuthorizationV1Api.list_cluster_role_binding",
233+
lambda *args, **kwargs: clusterrolebindings["items"],
234+
)
235+
monkeypatch.setattr(
236+
"feast.permissions.client.kubernetes_auth_client_manager.KubernetesAuthClientManager.get_token",
237+
lambda self: "my-token",
238+
)
239+
240+
monkeypatch.setattr(
241+
"feast.permissions.auth.kubernetes_token_parser.config.load_incluster_config",
242+
lambda: None,
243+
)
244+
245+
monkeypatch.setattr(
246+
"feast.permissions.auth.kubernetes_token_parser.jwt.decode",
247+
lambda *args, **kwargs: {"sub": subject},
248+
)
249+
250+
roles = rolebindings["roles"]
251+
croles = clusterrolebindings["roles"]
252+
253+
access_token = "aaa-bbb-ccc"
254+
token_parser = KubernetesTokenParser()
255+
user = asyncio.run(
256+
token_parser.user_details_from_access_token(access_token=access_token)
257+
)
258+
259+
if is_intra_server:
260+
assertpy.assert_that(user).is_not_none()
261+
assertpy.assert_that(user.username).is_equal_to(intra_communication_val)
262+
assertpy.assert_that(user.roles).is_equal_to([])
263+
else:
264+
assertpy.assert_that(user).is_type_of(User)
265+
if isinstance(user, User):
266+
assertpy.assert_that(user.username).is_equal_to(f"{namespace}:{sa_name}")
267+
assertpy.assert_that(user.roles.sort()).is_equal_to((roles + croles).sort())
268+
for r in roles:
269+
assertpy.assert_that(user.has_matching_role([r])).is_true()
270+
for cr in croles:
271+
assertpy.assert_that(user.has_matching_role([cr])).is_true()
272+
assertpy.assert_that(user.has_matching_role(["foo"])).is_false()

‎sdk/python/tests/unit/permissions/test_decorator.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
@pytest.mark.parametrize(
66
"username, can_read, can_write",
77
[
8-
(None, False, False),
8+
(None, True, True),
99
("r", True, False),
1010
("w", False, True),
1111
("rw", True, True),

0 commit comments

Comments
 (0)