Allow health check messages with fixed encryption keys

to check if listener is online

docs/CHANGELOG.rst

@@ -18,6 +18,10 @@ Changed
 - Authenticate worker with the listener (and vice versa) and encrypt the 
   communication between them using ``CurveZMQ`` protocol
 
+Added
+-----
+- Add fixed pair of encryption keys for ``listener`` service health checks
+
 
 ===================
 39.0.0 - 2024-05-09

resolwe/flow/managers/listener/listener.py

@@ -80,6 +80,10 @@
 LISTENER_PUBLIC_KEY = env_public_key.encode()
 LISTENER_PRIVATE_KEY = env_private_key.encode()
 
+# This is a special key that is used to check if listener is running.
+LIVENESS_CHECK_PUBLIC_KEY = b"pielqA({EHts^?MtURnndo0$)ocr46=?Xiv>-Sn5"
+LIVENESS_CHECK_PRIVATE_KEY = b"5>r36/f^OjoVNMMY[fxr=ep!UO#uL?JPg2ci(td4"
+
 if not settings.DEBUG:
     assert (
         LISTENER_PRIVATE_KEY != DEFAULT_LISTENER_PRIVATE_KEY
@@ -105,6 +109,12 @@ async def callback(self, domain, key):
         """
         try:
             assert domain == "*", "Only domain '*' is supported."
+
+            # Allow the message with liveness keys to proceed. Make sure to check later
+            #  that it does not have any permission.
+            if key == LIVENESS_CHECK_PUBLIC_KEY:
+                return True
+
             status, data_id = await database_sync_to_async(
                 Worker.objects.filter(public_key=key)
                 .values_list("status", "data_id")
@@ -693,12 +703,6 @@ def get_data():
                         "Exception updating unresponsive peer status."
                     )
 
-    async def handle_liveness_probe(
-        self, message: Message, peer_identity: PeerIdentity
-    ) -> Response[bool]:
-        """Respond to the liveness probe."""
-        return message.respond_ok(True)
-
     def _handle_lock_message_error(
         self, lock_status: RedisLockStatus, received_message: Message
     ) -> Response:
@@ -748,6 +752,11 @@ async def default_command_handler(
         response: Optional[Response] = None
         extend_lock_task: Optional[asyncio.Task] = None
 
+        # Check if this is a health check and respond OK regardless the content of the
+        # message.
+        if received_message.client_id == LIVENESS_CHECK_PUBLIC_KEY:
+            return received_message.respond_ok("OK")
+
         try:
             data_id = abs(int(peer_identity))
         except Exception:

resolwe/flow/tests/test_executors.py

@@ -26,7 +26,11 @@
 )
 from resolwe.flow.executors.zeromq_utils import ZMQCommunicator
 from resolwe.flow.managers.dispatcher import Manager
-from resolwe.flow.managers.listener.listener import LISTENER_PUBLIC_KEY
+from resolwe.flow.managers.listener.listener import (
+    LISTENER_PUBLIC_KEY,
+    LIVENESS_CHECK_PRIVATE_KEY,
+    LIVENESS_CHECK_PUBLIC_KEY,
+)
 from resolwe.flow.managers.listener.redis_cache import redis_cache
 from resolwe.flow.models import Data, DataDependency, Process, Worker
 from resolwe.flow.models.annotations import (
@@ -382,7 +386,6 @@ async def send_single_message(
                 zmq_socket, "init_container <-> listener", logger
             )
             async with communicator:
-                print("Sending message")
                 future = asyncio.ensure_future(
                     communicator.send_command(Message.command("update_status", "PP"))
                 )
@@ -435,6 +438,18 @@ async def send_single_message(
                 )
             )
 
+        # Peer with id 'liveness_probe' must be able to get a response even with fake
+        # keys.
+        response = asyncio.new_event_loop().run_until_complete(
+            send_single_message(
+                b"liveness_probe",
+                LIVENESS_CHECK_PUBLIC_KEY,
+                LIVENESS_CHECK_PRIVATE_KEY,
+                LISTENER_PUBLIC_KEY,
+            )
+        )
+        self.assertEqual(response.status, ResponseStatus.OK)
+
         # Non-matching keys and data id message must be rejected.
         response = asyncio.new_event_loop().run_until_complete(
             send_single_message(