[feature] Added WebSocket endpoint which broadcasts all location updates #1157

This websocket endpoint is needed for the implementation of dynamic mobile maps.

Closes #1157

---------

Co-authored-by: Gagan Deep <pandafy.dev@gmail.com>

Author: 2 people

openwisp_controller/geo/apps.py

@@ -1,7 +1,12 @@
+import json
+
+import channels.layers
 import swapper
+from asgiref.sync import async_to_sync
 from django.conf import settings
 from django.db import transaction
 from django.db.models import Case, Count, Sum, When
+from django.db.models.signals import post_save
 from django.utils.translation import gettext_lazy as _
 from django_loci.apps import LociConfig
 from swapper import get_model_name

openwisp_controller/geo/channels/consumers.py

@@ -1,5 +1,6 @@
 import swapper
-from django_loci.channels.base import BaseLocationBroadcast
+from asgiref.sync import async_to_sync
+from django_loci.channels.base import BaseCommonLocationBroadcast, BaseLocationBroadcast
 
 Location = swapper.load_model("geo", "Location")
 
@@ -17,3 +18,25 @@ def is_authorized(self, user, location):
         ):
             return False
         return result
+
+
+class CommonLocationBroadcast(BaseCommonLocationBroadcast):
+    model = Location
+
+    def join_groups(self, user):
+        """
+        Subscribe user to all organizations they manage or bypass if superuser.
+        """
+        if user.is_superuser:
+            super().join_groups(user)
+            self.group_name = [self.group_name]
+            return
+        self.group_name = []
+        for org in user.organizations_managed:
+            group = f"loci.mobile-location.organization.{org}"
+            self.group_name.append(group)
+            async_to_sync(self.channel_layer.group_add)(group, self.channel_name)
+
+    def disconnect(self, close_code):
+        for group in getattr(self, "group_name", []):
+            async_to_sync(self.channel_layer.group_discard)(group, self.channel_name)

openwisp_controller/geo/channels/routing.py

@@ -2,19 +2,27 @@
 from channels.routing import ProtocolTypeRouter, URLRouter
 from channels.security.websocket import AllowedHostsOriginValidator
 from django.urls import path
-from django_loci.channels.base import location_broadcast_path
+from django_loci.channels.base import (
+    common_location_broadcast_path,
+    location_broadcast_path,
+)
 from openwisp_notifications.websockets.routing import (
     get_routes as get_notification_routes,
 )
 
-from .consumers import LocationBroadcast
+from .consumers import CommonLocationBroadcast, LocationBroadcast
 
 
 def get_routes():
     return [
         path(
             location_broadcast_path, LocationBroadcast.as_asgi(), name="LocationChannel"
-        )
+        ),
+        path(
+            common_location_broadcast_path,
+            CommonLocationBroadcast.as_asgi(),
+            name="CommonLocationChannel",
+        ),
     ]
 
 

openwisp_controller/geo/tests/pytest.py

@@ -1,70 +1,44 @@
-import importlib
+import asyncio
 import os
+from contextlib import suppress
 from unittest import skipIf
 
 import pytest
 from channels.db import database_sync_to_async
 from channels.routing import ProtocolTypeRouter
-from channels.testing import WebsocketCommunicator
 from django.conf import settings
-from django.contrib.auth import get_user_model, login
+from django.contrib.auth import get_permission_codename, get_user_model
 from django.contrib.auth.models import Permission
-from django.http.request import HttpRequest
 from django.utils.module_loading import import_string
+from django_loci.tests import TestChannelsMixin
 from swapper import load_model
 
+from openwisp_controller.geo.channels.consumers import (
+    CommonLocationBroadcast,
+    LocationBroadcast,
+)
+from openwisp_users.tests.utils import TestOrganizationMixin
+
 from .utils import TestGeoMixin
 
 Device = load_model("config", "Device")
 Location = load_model("geo", "Location")
 DeviceLocation = load_model("geo", "DeviceLocation")
 User = get_user_model()
 OrganizationUser = load_model("openwisp_users", "OrganizationUser")
+Group = load_model("openwisp_users", "Group")
 
 
 @skipIf(os.environ.get("SAMPLE_APP", False), "Running tests on SAMPLE_APP")
-class TestChannels(TestGeoMixin):
+class TestChannels(TestGeoMixin, TestChannelsMixin, TestOrganizationMixin):
+    location_consumer = LocationBroadcast
+    common_location_consumer = CommonLocationBroadcast
     application = import_string(getattr(settings, "ASGI_APPLICATION"))
     object_model = Device
     location_model = Location
     object_location_model = DeviceLocation
     user_model = get_user_model()
 
-    def _force_login(self, user, backend=None):
-        engine = importlib.import_module(settings.SESSION_ENGINE)
-        request = HttpRequest()
-        request.session = engine.SessionStore()
-        login(request, user, backend)
-        request.session.save()
-        return request.session
-
-    async def _get_request_dict(self, pk=None, user=None):
-        if not pk:
-            location = await database_sync_to_async(self._create_location)(
-                is_mobile=True
-            )
-            await database_sync_to_async(self._create_object_location)(
-                location=location
-            )
-            pk = location.pk
-        path = "/ws/loci/location/{0}/".format(pk)
-        session = None
-        if user:
-            session = await database_sync_to_async(self._force_login)(user)
-        return {"pk": pk, "path": path, "session": session}
-
-    def _get_communicator(self, request_vars, user=None):
-        communicator = WebsocketCommunicator(self.application, request_vars["path"])
-        if user:
-            communicator.scope.update(
-                {
-                    "user": user,
-                    "session": request_vars["session"],
-                    "url_route": {"kwargs": {"pk": request_vars["pk"]}},
-                }
-            )
-        return communicator
-
     @pytest.mark.asyncio
     @pytest.mark.django_db(transaction=True)
     async def test_consumer_staff_but_no_change_permission(self):
@@ -74,37 +48,153 @@ async def test_consumer_staff_but_no_change_permission(self):
         location = await database_sync_to_async(self._create_location)(is_mobile=True)
         await database_sync_to_async(self._create_object_location)(location=location)
         pk = location.pk
-        request_vars = await self._get_request_dict(user=user, pk=pk)
-        communicator = self._get_communicator(request_vars, user)
+        request_vars = await self._get_specific_location_request_dict(pk=pk, user=user)
+        communicator = self._get_specific_location_communicator(request_vars, user)
         connected, _ = await communicator.connect()
         assert not connected
         await communicator.disconnect()
         # add permission to change location and repeat
-        perm = await database_sync_to_async(
-            (
-                await database_sync_to_async(Permission.objects.filter)(
-                    name="Can change location"
-                )
-            ).first
-        )()
+        perm = await Permission.objects.filter(
+            codename=f"change_{self.location_model._meta.model_name}",
+            content_type__app_label=self.location_model._meta.app_label,
+        ).afirst()
         await database_sync_to_async(user.user_permissions.add)(perm)
         user = await database_sync_to_async(User.objects.get)(pk=user.pk)
-        request_vars = await self._get_request_dict(user=user, pk=pk)
-        communicator = self._get_communicator(request_vars, user)
+        request_vars = await self._get_specific_location_request_dict(pk=pk, user=user)
+        communicator = self._get_specific_location_communicator(request_vars, user)
         connected, _ = await communicator.connect()
         assert not connected
         await communicator.disconnect()
         # add user to organization
         await database_sync_to_async(OrganizationUser.objects.create)(
-            organization=location.organization, user=user, is_admin=True
+            organization=location.organization,
+            user=user,
+            is_admin=True,
         )
         await database_sync_to_async(location.organization.save)()
         user = await database_sync_to_async(User.objects.get)(pk=user.pk)
-        request_vars = await self._get_request_dict(user=user, pk=pk)
-        communicator = self._get_communicator(request_vars, user)
+        request_vars = await self._get_specific_location_request_dict(pk=pk, user=user)
+        communicator = self._get_specific_location_communicator(request_vars, user)
         connected, _ = await communicator.connect()
         assert connected
         await communicator.disconnect()
 
+    @pytest.mark.asyncio
+    @pytest.mark.django_db(transaction=True)
+    async def test_common_location_consumer_staff_but_no_change_permission(self):
+        user = await database_sync_to_async(self._create_user)(is_staff=True)
+        location = await database_sync_to_async(self._create_location)(is_mobile=True)
+        await database_sync_to_async(self._create_object_location)(location=location)
+        pk = location.pk
+        request_vars = await self._get_common_location_request_dict(pk=pk, user=user)
+        communicator = self._get_common_location_communicator(request_vars, user)
+        connected, _ = await communicator.connect()
+        assert not connected
+        await communicator.disconnect()
+        # After granting change permission, the user can connect to the common
+        # location endpoint, but must receive updates only for locations
+        # belonging to their organization.
+        perm = await Permission.objects.filter(
+            codename=f"change_{self.location_model._meta.model_name}",
+            content_type__app_label=self.location_model._meta.app_label,
+        ).afirst()
+        await database_sync_to_async(user.user_permissions.add)(perm)
+        user = await database_sync_to_async(User.objects.get)(pk=user.pk)
+        request_vars = await self._get_common_location_request_dict(pk=pk, user=user)
+        communicator = self._get_common_location_communicator(request_vars, user)
+        connected, _ = await communicator.connect()
+        assert connected
+        await communicator.disconnect()
+
+    @pytest.mark.asyncio
+    @pytest.mark.django_db(transaction=True)
+    async def test_common_location_org_isolation(self):
+        administrator = await Group.objects.acreate(name="Administrator")
+        perm = await Permission.objects.filter(
+            codename=get_permission_codename("change", self.location_model._meta),
+        ).afirst()
+        await administrator.permissions.aadd(perm)
+        org1 = await database_sync_to_async(self._get_org)(org_name="test1")
+        org2 = await database_sync_to_async(self._get_org)(org_name="test2")
+        org1_location = await database_sync_to_async(self._create_location)(
+            is_mobile=True, organization=org1
+        )
+        org2_location = await database_sync_to_async(self._create_location)(
+            is_mobile=True, organization=org2
+        )
+        org1_user = await database_sync_to_async(self._create_administrator)(
+            organizations=[org1],
+            username="user1",
+            password="password",
+            email="user1@test.org",
+        )
+        org2_user = await database_sync_to_async(self._create_administrator)(
+            organizations=[org2],
+            username="user2",
+            password="password",
+            email="user2@test.org",
+        )
+        admin = await database_sync_to_async(self._get_admin)()
+        org1_communicator = self._get_common_location_communicator(
+            await self._get_common_location_request_dict(
+                pk=org1_location.pk, user=org1_user
+            ),
+            org1_user,
+        )
+        org2_communicator = self._get_common_location_communicator(
+            await self._get_common_location_request_dict(
+                pk=org2_location.pk, user=org2_user
+            ),
+            org2_user,
+        )
+        admin_communicator = self._get_common_location_communicator(
+            await self._get_common_location_request_dict(
+                pk=org1_location.pk, user=admin
+            ),
+            admin,
+        )
+        connected, _ = await org1_communicator.connect()
+        assert connected
+        connected, _ = await org2_communicator.connect()
+        assert connected
+        connected, _ = await admin_communicator.connect()
+        assert connected
+
+        # Updating co-ordinates for org1_location should notify org1_user and admin,
+        await self._save_location(str(org1_location.pk))
+        org1_response = await org1_communicator.receive_json_from(timeout=1)
+        assert org1_response["id"] == str(org1_location.pk)
+        admin_response = await admin_communicator.receive_json_from(timeout=1)
+        assert admin_response["id"] == str(org1_location.pk)
+        with pytest.raises(asyncio.TimeoutError):
+            await org2_communicator.receive_json_from(timeout=1)
+
+        with suppress(asyncio.CancelledError):
+            await org2_communicator.disconnect()
+
+        org2_communicator = self._get_common_location_communicator(
+            await self._get_common_location_request_dict(
+                pk=org2_location.pk, user=org2_user
+            ),
+            org2_user,
+        )
+        connected, _ = await org2_communicator.connect()
+        assert connected
+
+        # Updating co-ordinates for org2_location should notify org2_user and admin,
+        await self._save_location(str(org2_location.pk))
+        org2_response = await org2_communicator.receive_json_from(timeout=1)
+        assert org2_response["id"] == str(org2_location.pk)
+        admin_response = await admin_communicator.receive_json_from(timeout=1)
+        assert admin_response["id"] == str(org2_location.pk)
+        with pytest.raises(asyncio.TimeoutError):
+            await org1_communicator.receive_json_from(timeout=1)
+
+        # The task is been cancelled if not completed in the given timeout
+        with suppress(asyncio.CancelledError):
+            await org1_communicator.disconnect()
+            await org2_communicator.disconnect()
+        await admin_communicator.disconnect()
+
     def test_asgi_application_router(self):
         assert isinstance(self.application, ProtocolTypeRouter)