Fix DecideBackendByTargetIdP and introduce DecideBackendByDiscoIdP

Signed-off-by: Ivan Kanakarakis <[email protected]>

Author: c00kiemon5ter

src/satosa/micro_services/custom_routing.py

@@ -22,22 +22,55 @@ class CustomRoutingError(SATOSAError):
 
 class DecideBackendByTargetIdP(RequestMicroService):
     """
-    Select which backend should be used based on who is the SAML IDP
+    Select target backend based on the target issuer.
     """
 
     def __init__(self, config:dict, *args, **kwargs):
         """
         Constructor.
+
         :param config: microservice configuration loaded from yaml file
         :type config: Dict[str, Dict[str, str]]
         """
         super().__init__(*args, **kwargs)
+
         self.target_mapping = config['target_mapping']
-        self.endpoint_paths = config['endpoint_paths']
         self.default_backend = config['default_backend']
 
-        if not isinstance(self.endpoint_paths, list):
-            raise SATOSAConfigurationError()
+    def process(self, context:Context, data:InternalData):
+        """
+        Set context.target_backend based on the target issuer (context.target_entity_id)
+
+        :param context: request context
+        :param data: the internal request
+        """
+        target_issuer = context.get_decoration(Context.KEY_TARGET_ENTITYID)
+        if not target_issuer:
+            return super().process(context, data)
+
+        target_backend = (
+            self.target_mapping.get(target_issuer)
+            or self.default_backend
+        )
+
+        report = {
+            'msg': 'decided target backend by target issuer',
+            'target_issuer': target_issuer,
+            'target_backend': target_backend,
+        }
+        logger.info(report)
+
+        context.target_backend = target_backend
+        return super().process(context, data)
+
+
+class DecideBackendByDiscoIdP(DecideBackendByTargetIdP):
+    def __init__(self, config:dict, *args, **kwargs):
+        super().__init__(config, *args, **kwargs)
+
+        self.disco_endpoints = config['disco_endpoints']
+        if not isinstance(self.disco_endpoints, list):
+            raise CustomRoutingError('disco_endpoints must be a list of str')
 
     def register_endpoints(self):
         """
@@ -54,69 +87,20 @@ def register_endpoints(self):
                  [(regexp, Callable[[satosa.context.Context], satosa.response.Response]), ...]
         """
 
-        # this intercepts disco response
         return [
-            (path , self.backend_by_entityid)
-            for path in self.endpoint_paths
+            (path , self._handle_disco_response)
+            for path in self.disco_endpoints
         ]
 
-    def _get_request_entity_id(self, context):
-        return (
-            context.get_decoration(Context.KEY_TARGET_ENTITYID) or
-            context.request.get('entityID')
-        )
-
-    def _get_backend(self, context:Context, entity_id:str) -> str:
-        """
-        returns the Target Backend to use
-        """
-        return (
-            self.target_mapping.get(entity_id) or
-            self.default_backend
-        )
-
-    def process(self, context:Context, data:dict):
-        """
-        Will modify the context.target_backend attribute based on the target entityid.
-        :param context: request context
-        :param data: the internal request
-        """
-        entity_id = self._get_request_entity_id(context)
-        if entity_id:
-            self._rewrite_context(entity_id, context)
-        return super().process(context, data)
-
-    def _rewrite_context(self, entity_id:str, context:Context) -> None:
-        tr_backend = self._get_backend(context, entity_id)
-        context.decorate(Context.KEY_TARGET_ENTITYID, entity_id)
-        context.target_frontend = context.target_frontend or context.state.get('ROUTER')
-        native_backend = context.target_backend
-        msg = (f'Found DecideBackendByTarget ({self.name} microservice) '
-               f'redirecting {entity_id} from {native_backend} '
-               f'backend to {tr_backend}')
-        logger.info(msg)
-        context.target_backend = tr_backend
-
-    def backend_by_entityid(self, context:Context):
-        entity_id = self._get_request_entity_id(context)
-
-        if entity_id:
-            self._rewrite_context(entity_id, context)
-        else:
-            raise CustomRoutingError(
-                f"{self.__class__.__name__} "
-                "can't find any valid entity_id in the context."
-            )
-
-        if not context.state.get('ROUTER'):
-            raise SATOSAStateError(
-                f"{self.__class__.__name__} "
-                "can't find any valid state in the context."
-            )
+    def _handle_disco_response(self, context:Context):
+        target_issuer_from_disco = context.request.get('entityID')
+        if not target_issuer_from_disco:
+            raise CustomRoutingError('no valid entity_id in the disco response')
 
-        data_serialized = context.state.get(self.name, {}).get("internal", {})
+        context.decorate(Context.KEY_TARGET_ENTITYID, target_issuer_from_disco)
+        data_serialized = context.state.get(self.name, {}).get('internal', {})
         data = InternalData.from_dict(data_serialized)
-        return super().process(context, data)
+        return self.process(context, data)
 
 
 class DecideBackendByRequester(RequestMicroService):

tests/satosa/micro_services/test_custom_routing.py

@@ -1,14 +1,18 @@
 from base64 import urlsafe_b64encode
+from unittest import TestCase
 
 import pytest
 
 from satosa.context import Context
+from satosa.state import State
 from satosa.exception import SATOSAError, SATOSAConfigurationError, SATOSAStateError
 from satosa.internal import InternalData
 from satosa.micro_services.custom_routing import DecideIfRequesterIsAllowed
+from satosa.micro_services.custom_routing import DecideBackendByDiscoIdP
 from satosa.micro_services.custom_routing import DecideBackendByTargetIdP
 from satosa.micro_services.custom_routing import CustomRoutingError
 
+
 TARGET_ENTITY = "entity1"
 
 
@@ -160,61 +164,92 @@ def test_missing_target_entity_id_from_context(self, context):
             decide_service.process(context, req)
 
 
-class TestDecideBackendByTargetIdP:
-    rules = {
-        'default_backend': 'Saml2',
-        'endpoint_paths': ['.*/disco'],
-        'target_mapping': {'http://idpspid.testunical.it:8088': 'spidSaml2'}
-    }
-
-    def create_decide_service(self, rules):
-        decide_service = DecideBackendByTargetIdP(
-                config=rules,
-                name="test_decide_service",
-                base_url="https://satosa.example.com"
-        )
-        decide_service.next = lambda ctx, data: data
-        return decide_service
+class TestDecideBackendByTargetIdP(TestCase):
+    def setUp(self):
+        context = Context()
+        context.state = State()
 
-
-    def test_missing_state(self, target_context):
-        decide_service = self.create_decide_service(self.rules)
-        target_context.request = {
-            'entityID': 'http://idpspid.testunical.it:8088',
+        config = {
+            'default_backend': 'default_backend',
+            'target_mapping': {
+                'mapped_idp.example.org': 'mapped_backend',
+            },
+            'disco_endpoints': [
+                '.*/disco',
+            ],
         }
-        req = InternalData(requester="test_requester")
-        req.requester = "somebody else"
-        assert decide_service.process(target_context, req)
-
-        with pytest.raises(SATOSAStateError):
-            decide_service.backend_by_entityid(target_context)
 
-
-    def test_unmatching_target(self, target_context):
-        """
-            It would rely on the default backend
-        """
-        decide_service = self.create_decide_service(self.rules)
-        target_context.request = {
-            'entityID': 'unknow-entity-id',
+        plugin = DecideBackendByTargetIdP(
+            config=config,
+            name='test_decide_service',
+            base_url='https://satosa.example.org',
+        )
+        plugin.next = lambda ctx, data: (ctx, data)
+
+        self.config = config
+        self.context = context
+        self.plugin = plugin
+
+    def test_when_target_is_not_set_do_skip(self):
+        data = InternalData(requester='test_requester')
+        newctx, newdata = self.plugin.process(self.context, data)
+        assert not newctx.target_backend
+
+    def test_when_target_is_not_mapped_choose_default_backend(self):
+        self.context.decorate(Context.KEY_TARGET_ENTITYID, 'idp.example.org')
+        data = InternalData(requester='test_requester')
+        newctx, newdata = self.plugin.process(self.context, data)
+        assert newctx.target_backend == 'default_backend'
+
+    def test_when_target_is_mapped_choose_mapping_backend(self):
+        self.context.decorate(Context.KEY_TARGET_ENTITYID, 'mapped_idp.example.org')
+        data = InternalData(requester='test_requester')
+        data.requester = 'somebody else'
+        newctx, newdata = self.plugin.process(self.context, data)
+        assert newctx.target_backend == 'mapped_backend'
+
+
+class TestDecideBackendByDiscoIdP(TestCase):
+    def setUp(self):
+        context = Context()
+        context.state = State()
+
+        config = {
+            'default_backend': 'default_backend',
+            'target_mapping': {
+                'mapped_idp.example.org': 'mapped_backend',
+            },
+            'disco_endpoints': [
+                '.*/disco',
+            ],
         }
-        target_context.state['ROUTER'] = 'Saml2'
 
-        req = InternalData(requester="test_requester")
-        assert decide_service.process(target_context, req)
+        plugin = DecideBackendByDiscoIdP(
+            config=config,
+            name='test_decide_service',
+            base_url='https://satosa.example.org',
+        )
+        plugin.next = lambda ctx, data: (ctx, data)
 
-        res = decide_service.backend_by_entityid(target_context)
-        assert isinstance(res, InternalData)
+        self.config = config
+        self.context = context
+        self.plugin = plugin
 
-    def test_matching_target(self, target_context):
-        decide_service = self.create_decide_service(self.rules)
-        target_context.request = {
-            'entityID': 'http://idpspid.testunical.it:8088-entity-id'
+    def test_when_target_is_not_set_raise_error(self):
+        self.context.request = {}
+        with pytest.raises(CustomRoutingError):
+            self.plugin._handle_disco_response(self.context)
+
+    def test_when_target_is_not_mapped_choose_default_backend(self):
+        self.context.request = {
+            'entityID': 'idp.example.org',
         }
-        target_context.state['ROUTER'] = 'Saml2'
+        newctx, newdata = self.plugin._handle_disco_response(self.context)
+        assert newctx.target_backend == 'default_backend'
 
-        req = InternalData(requester="test_requester")
-        req.requester = "somebody else"
-        assert decide_service.process(target_context, req)
-        res = decide_service.backend_by_entityid(target_context)
-        assert isinstance(res, InternalData)
+    def test_when_target_is_mapped_choose_mapping_backend(self):
+        self.context.request = {
+            'entityID': 'mapped_idp.example.org',
+        }
+        newctx, newdata = self.plugin._handle_disco_response(self.context)
+        assert newctx.target_backend == 'mapped_backend'