Introduce userinfo policy

Signed-off-by: Kostis Triantafyllakis <kostastriantaf@gmail.com>

Author: ctriant

doc/server/contents/conf.rst

@@ -408,7 +408,11 @@ An example::
               "normal",
               "aggregated",
               "distributed"
-            ]
+            ],
+            "policy": {
+              "function": "/path/to/callable",
+              "kwargs": {}
+            }
           }
         },
         "revocation": {
@@ -747,6 +751,10 @@ the following::
     "userinfo": {
         "class": "oidc_provider.users.UserInfo",
         "kwargs": {
+            "policy": {
+              "function": "/path/to/callable",
+              "kwargs": {}
+            },
             "claims_map": {
                 "phone_number": "telephone",
                 "family_name": "last_name",
@@ -760,6 +768,17 @@ the following::
         }
     }
 
+The policy for userinfo endpoint is optional and can also be configured in a client's metadata, for example::
+
+    "userinfo": {
+        "kwargs": {
+            "policy": {
+              "function": "/path/to/callable",
+              "kwargs": {}
+            }
+        }
+    }
+
 ================================
 Special Configuration directives
 ================================

src/idpyoidc/server/oidc/userinfo.py

@@ -10,11 +10,13 @@
 from cryptojwt.jwt import utc_time_sans_frac
 
 from idpyoidc import claims
+from idpyoidc.util import importer
 from idpyoidc.message import Message
 from idpyoidc.message import oidc
 from idpyoidc.message.oauth2 import ResponseMessage
 from idpyoidc.server.endpoint import Endpoint
 from idpyoidc.server.exception import ClientAuthenticationError
+from idpyoidc.exception import ImproperlyConfigured
 from idpyoidc.server.util import OAUTH2_NOCACHE_HEADERS
 
 logger = logging.getLogger(__name__)
@@ -46,18 +48,28 @@ def __init__(self, upstream_get: Callable, add_claims_by_scope: Optional[bool] =
         # Add the issuer ID as an allowed JWT target
         self.allowed_targets.append("")
 
-    def get_client_id_from_token(self, context, token, request=None):
-        _info = context.session_manager.get_session_info_by_token(
+        if kwargs is None:
+            self.config = {
+                "policy": {
+                    "function": "/path/to/callable",
+                    "kwargs": {}
+                },
+            }
+        else:
+            self.config = kwargs
+
+    def get_client_id_from_token(self, endpoint_context, token, request=None):
+        _info = endpoint_context.session_manager.get_session_info_by_token(
             token, handler_key="access_token"
         )
         return _info["client_id"]
 
     def do_response(
-        self,
-        response_args: Optional[Union[Message, dict]] = None,
-        request: Optional[Union[Message, dict]] = None,
-        client_id: Optional[str] = "",
-        **kwargs
+            self,
+            response_args: Optional[Union[Message, dict]] = None,
+            request: Optional[Union[Message, dict]] = None,
+            client_id: Optional[str] = "",
+            **kwargs
     ) -> dict:
 
         if "error" in kwargs and kwargs["error"]:
@@ -157,6 +169,12 @@ def process_request(self, request=None, **kwargs):
             info["sub"] = _grant.sub
             if _grant.add_acr_value("userinfo"):
                 info["acr"] = _grant.authentication_event["authn_info"]
+
+            if "userinfo" in _cntxt.cdb[request["client_id"]]:
+                self.config["policy"] = _cntxt.cdb[request["client_id"]]["userinfo"]["policy"]
+
+            if "policy" in self.config:
+                info = self._enforce_policy(request, info, token, self.config)
         else:
             info = {
                 "error": "invalid_request",
@@ -190,3 +208,26 @@ def parse_request(self, request, http_info=None, **kwargs):
             request["access_token"] = auth_info["token"]
 
         return request
+
+    def _enforce_policy(self, request, response_info, token, config):
+        policy = config["policy"]
+        callable = policy["function"]
+        kwargs = policy.get("kwargs", {})
+
+        if isinstance(callable, str):
+            try:
+                fn = importer(callable)
+            except Exception:
+                raise ImproperlyConfigured(f"Error importing {callable} policy callable")
+        else:
+            fn = callable
+
+        try:
+            return fn(request, token, response_info, **kwargs)
+        except Exception as e:
+            logger.error(f"Error while executing the {fn} policy callable: {e}")
+            return self.error_cls(error="server_error", error_description="Internal server error")
+
+
+def validate_userinfo_policy(request, token, response_info, **kwargs):
+    return response_info

tests/test_server_26_oidc_userinfo_endpoint.py

@@ -21,10 +21,12 @@
 from idpyoidc.server.scopes import SCOPE2CLAIMS
 from idpyoidc.server.user_authn.authn_context import INTERNETPROTOCOLPASSWORD
 from idpyoidc.server.user_info import UserInfo
+from idpyoidc.server.oidc.userinfo import validate_userinfo_policy
 from idpyoidc.time_util import utc_time_sans_frac
 from tests import CRYPT_CONFIG
 from tests import SESSION_PARAMS
 
+
 KEYDEFS = [
     {"type": "RSA", "key": "", "use": ["sig"]},
     {"type": "EC", "crv": "P-256", "use": ["sig"]},
@@ -637,3 +639,52 @@ def test_process_request_absent_userinfo_conf(self):
 
         with pytest.raises(ImproperlyConfigured):
             code = self._mint_code(grant, session_id)
+
+    def test_userinfo_policy(self):
+        _auth_req = AUTH_REQ.copy()
+
+        session_id = self._create_session(_auth_req)
+        grant = self.session_manager[session_id]
+        access_token = self._mint_token("access_token", grant, session_id)
+
+        http_info = {"headers": {"authorization": "Bearer {}".format(access_token.value)}}
+
+        def _custom_validate_userinfo_policy(request, token, response_info, **kwargs):
+            return {"custom": "policy"}
+
+        self.endpoint.config["policy"] = {}
+        self.endpoint.config["policy"]["function"] = _custom_validate_userinfo_policy
+
+        _req = self.endpoint.parse_request({}, http_info=http_info)
+        args = self.endpoint.process_request(_req)
+        assert args
+        res = self.endpoint.do_response(request=_req, **args)
+        _response = json.loads(res["response"])
+        assert "custom" in _response
+
+    def test_userinfo_policy_per_client(self):
+        _auth_req = AUTH_REQ.copy()
+
+        session_id = self._create_session(_auth_req)
+        grant = self.session_manager[session_id]
+        access_token = self._mint_token("access_token", grant, session_id)
+
+        http_info = {"headers": {"authorization": "Bearer {}".format(access_token.value)}}
+
+        def _custom_validate_userinfo_policy(request, token, response_info, **kwargs):
+            return {"custom": "policy"}
+
+        self.context.cdb["client_1"]["userinfo"] = {
+            "policy": {
+                "function": _custom_validate_userinfo_policy,
+                "kwargs": {}
+            }
+        }
+
+        _req = self.endpoint.parse_request({}, http_info=http_info)
+        args = self.endpoint.process_request(_req)
+        assert args
+        res = self.endpoint.do_response(request=_req, **args)
+        _response = json.loads(res["response"])
+        assert "custom" in _response
+