Add config updates via SSE

Author: timokoessler

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -1,4 +1,4 @@
-""" This file simply exports the CloudConnectionManager class"""
+"""This file simply exports the CloudConnectionManager class"""
 
 from aikido_zen.background_process.heartbeats import send_heartbeats_every_x_secs
 from aikido_zen.background_process.routes import Routes
@@ -10,10 +10,12 @@
 from aikido_zen.storage.users import Users
 from aikido_zen.storage.hostnames import Hostnames
 from ..realtime.start_polling_for_changes import start_polling_for_changes
+from ..realtime.listen_for_config_updates import listen_for_config_updates
 from ...helpers.get_current_unixtime_ms import get_unixtime_ms
 from ...storage.ai_statistics import AIStatistics
 from ...storage.firewall_lists import FirewallLists
 from ...storage.statistics import Statistics
+from aikido_zen.helpers.env_vars.feature_flags import is_feature_enabled
 
 # Import functions :
 from .get_manager_info import get_manager_info
@@ -67,6 +69,9 @@ def start(self, event_scheduler):
         send_heartbeats_every_x_secs(self, self.heartbeat_secs, event_scheduler)
         start_polling_for_changes(self, event_scheduler)
 
+        if is_feature_enabled("sse"):
+            listen_for_config_updates(self, event_scheduler)
+
     def report_initial_stats(self):
         """
         This is run 1m after startup, and checks if we should send out

aikido_zen/background_process/realtime/__init__.py

@@ -15,7 +15,7 @@ def get_realtime_url():
         if not realtime_url.endswith("/"):
             realtime_url += "/"
         return realtime_url
-    return "https://runtime.aikido.dev/"
+    return get_api_url()
 
 
 def get_config(token):

aikido_zen/background_process/realtime/listen_for_config_updates.py

@@ -0,0 +1,59 @@
+"""
+Mainly exports `listen_for_config_updates`
+"""
+
+import json
+
+from aikido_zen.helpers.token import Token
+from aikido_zen.helpers.logging import logger
+import aikido_zen.background_process.realtime as realtime
+from .sse_client import connect_to_sse
+
+
+def listen_for_config_updates(connection_manager, event_scheduler):
+    """
+    Connects to the realtime SSE endpoint and fetches the new config whenever
+    the server signals, through a "config-updated" event, that it has changed.
+    """
+    if not isinstance(connection_manager.token, Token):
+        logger.info("No token provided, not listening for config updates")
+        return
+    if connection_manager.serverless:
+        logger.info(
+            "Running in serverless environment, not listening for config updates"
+        )
+        return
+
+    token = connection_manager.token
+    last_updated_at = connection_manager.conf.last_updated_at
+
+    def on_event(event):
+        nonlocal last_updated_at
+
+        logger.debug("SSE event received: %s", event.event)
+        if event.event != "config-updated":
+            return
+
+        try:
+            payload = json.loads(event.data)
+            config_updated_at = payload["configUpdatedAt"]
+            if config_updated_at <= last_updated_at:
+                return
+        except (ValueError, KeyError, TypeError):
+            logger.debug("SSE config-updated event has invalid payload: %s", event.data)
+            return
+
+        logger.debug("SSE config-updated event, fetching new config")
+
+        try:
+            config = realtime.get_config(token)
+            logger.debug(
+                "SSE config fetched, configUpdatedAt: %s", config.get("configUpdatedAt")
+            )
+            last_updated_at = config.get("configUpdatedAt", config_updated_at)
+            connection_manager.update_service_config({**config, "success": True})
+            connection_manager.update_firewall_lists()
+        except Exception as e:
+            logger.error("Failed to fetch config after SSE event : %s", e)
+
+    connect_to_sse(token=token, on_event=on_event)

aikido_zen/background_process/realtime/listen_for_config_updates_test.py

@@ -0,0 +1,169 @@
+import json
+import logging
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from aikido_zen.helpers.token import Token
+from .listen_for_config_updates import listen_for_config_updates
+
+
+def make_event(event="config-updated", data=None):
+    return SimpleNamespace(
+        event=event, data=json.dumps(data) if data is not None else ""
+    )
+
+
+@pytest.fixture
+def connection_manager():
+    return MagicMock(
+        token=Token("123"),
+        serverless=None,
+        conf=MagicMock(last_updated_at=0),
+    )
+
+
+def test_no_token(caplog):
+    with patch(
+        "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse"
+    ) as mock_connect:
+        listen_for_config_updates(
+            connection_manager=MagicMock(token=None, serverless=None),
+            event_scheduler=MagicMock(),
+        )
+
+    assert "No token provided, not listening for config updates" in caplog.text
+    mock_connect.assert_not_called()
+
+
+def test_serverless_environment(caplog):
+    with patch(
+        "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse"
+    ) as mock_connect:
+        listen_for_config_updates(
+            connection_manager=MagicMock(token=Token("123"), serverless=True),
+            event_scheduler=MagicMock(),
+        )
+
+    assert (
+        "Running in serverless environment, not listening for config updates"
+        in caplog.text
+    )
+    mock_connect.assert_not_called()
+
+
+def test_connects_to_sse_with_token(connection_manager):
+    with patch(
+        "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse"
+    ) as mock_connect:
+        listen_for_config_updates(
+            connection_manager=connection_manager, event_scheduler=MagicMock()
+        )
+
+    mock_connect.assert_called_once()
+    assert mock_connect.call_args.kwargs["token"] is connection_manager.token
+    assert callable(mock_connect.call_args.kwargs["on_event"])
+
+
+def get_on_event(connection_manager):
+    with patch(
+        "aikido_zen.background_process.realtime.listen_for_config_updates.connect_to_sse"
+    ) as mock_connect:
+        listen_for_config_updates(
+            connection_manager=connection_manager, event_scheduler=MagicMock()
+        )
+    return mock_connect.call_args.kwargs["on_event"]
+
+
+def test_ignores_events_that_are_not_config_updated(connection_manager):
+    on_event = get_on_event(connection_manager)
+
+    with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config:
+        on_event(make_event(event="ping"))
+
+    mock_get_config.assert_not_called()
+    connection_manager.update_service_config.assert_not_called()
+
+
+def test_ignores_config_updated_event_with_invalid_json(connection_manager, caplog):
+    caplog.set_level(logging.DEBUG, logger="Zen")
+    on_event = get_on_event(connection_manager)
+
+    with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config:
+        on_event(SimpleNamespace(event="config-updated", data="not json"))
+
+    mock_get_config.assert_not_called()
+    connection_manager.update_service_config.assert_not_called()
+    assert "SSE config-updated event has invalid payload" in caplog.text
+
+
+def test_ignores_config_updated_event_missing_config_updated_at(connection_manager):
+    on_event = get_on_event(connection_manager)
+
+    with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config:
+        on_event(make_event(data={"foo": "bar"}))
+
+    mock_get_config.assert_not_called()
+    connection_manager.update_service_config.assert_not_called()
+
+
+def test_ignores_config_updated_event_that_is_not_newer(connection_manager):
+    connection_manager.conf.last_updated_at = 100
+    on_event = get_on_event(connection_manager)
+
+    with patch("aikido_zen.background_process.realtime.get_config") as mock_get_config:
+        on_event(make_event(data={"configUpdatedAt": 100}))
+
+    mock_get_config.assert_not_called()
+    connection_manager.update_service_config.assert_not_called()
+
+
+def test_fetches_and_applies_new_config_on_newer_event(connection_manager):
+    connection_manager.conf.last_updated_at = 100
+    on_event = get_on_event(connection_manager)
+
+    new_config = {"endpoints": [], "configUpdatedAt": 200}
+    with patch(
+        "aikido_zen.background_process.realtime.get_config", return_value=new_config
+    ) as mock_get_config:
+        on_event(make_event(data={"configUpdatedAt": 200}))
+
+    mock_get_config.assert_called_once_with(connection_manager.token)
+    connection_manager.update_service_config.assert_called_once_with(
+        {**new_config, "success": True}
+    )
+    connection_manager.update_firewall_lists.assert_called_once()
+
+
+def test_updates_last_updated_at_so_a_second_stale_event_is_ignored(
+    connection_manager,
+):
+    connection_manager.conf.last_updated_at = 100
+    on_event = get_on_event(connection_manager)
+
+    new_config = {"endpoints": [], "configUpdatedAt": 200}
+    with patch(
+        "aikido_zen.background_process.realtime.get_config", return_value=new_config
+    ) as mock_get_config:
+        on_event(make_event(data={"configUpdatedAt": 200}))
+        # A second event claiming to be newer than the original lastUpdatedAt
+        # (100), but not newer than what we just fetched (200), should now be
+        # ignored without fetching again.
+        on_event(make_event(data={"configUpdatedAt": 150}))
+
+    mock_get_config.assert_called_once()
+
+
+def test_handles_get_config_failure_gracefully(connection_manager, caplog):
+    connection_manager.conf.last_updated_at = 100
+    on_event = get_on_event(connection_manager)
+
+    with patch(
+        "aikido_zen.background_process.realtime.get_config",
+        side_effect=ValueError("Request timed out"),
+    ):
+        on_event(make_event(data={"configUpdatedAt": 200}))
+
+    connection_manager.update_service_config.assert_not_called()
+    assert "Failed to fetch config after SSE event" in caplog.text

aikido_zen/background_process/realtime/sse_client/__init__.py

@@ -0,0 +1,127 @@
+"""
+SSE (Server-Sent Events) client that connects to the realtime endpoint and
+dispatches events, automatically reconnecting with exponential backoff and
+jitter on disconnects.
+"""
+
+import http.client
+import random
+import socket
+import threading
+import time
+import urllib.error
+import urllib.request
+
+import aikido_zen.background_process.realtime as realtime
+from aikido_zen.helpers.logging import logger
+from .parser import SSEParser
+
+INITIAL_RECONNECT_SECS = 5
+MAX_RECONNECT_SECS = 60
+STABLE_CONNECTION_SECS = 30
+READ_TIMEOUT_SECS = 70
+
+_CONNECTION_ERRORS = (urllib.error.URLError, OSError, http.client.HTTPException)
+
+
+def connect_to_sse(
+    token,
+    on_event,
+    initial_reconnect_secs=INITIAL_RECONNECT_SECS,
+    read_timeout_secs=READ_TIMEOUT_SECS,
+):
+    """
+    Starts a daemon thread that connects to the realtime SSE endpoint and calls
+    `on_event(event)` for every event received. Reconnects with exponential
+    backoff (and jitter) on disconnect, until the server rejects the connection
+    with a 401 or 403 status, at which point it stops.
+
+    Returns the thread that was started.
+    """
+    thread = threading.Thread(
+        target=_reconnect_loop,
+        args=(token, on_event, initial_reconnect_secs, read_timeout_secs),
+        daemon=True,
+    )
+    thread.start()
+    return thread
+
+
+def _reconnect_loop(token, on_event, initial_reconnect_secs, read_timeout_secs):
+    reconnect_secs = initial_reconnect_secs
+    try:
+        while True:
+            start = time.monotonic()
+            outcome, status_code = _connect(token, on_event, read_timeout_secs)
+
+            if outcome == "disconnected" and status_code in (401, 403):
+                logger.info(
+                    "SSE connection rejected with status %s, stopping", status_code
+                )
+                return
+
+            if time.monotonic() - start >= STABLE_CONNECTION_SECS:
+                reconnect_secs = initial_reconnect_secs
+
+            jitter = random.random() * (reconnect_secs / 2)
+            delay_secs = reconnect_secs + jitter
+
+            logger.debug("SSE scheduling reconnect in %sms", round(delay_secs * 1000))
+
+            reconnect_secs = min(reconnect_secs * 2, MAX_RECONNECT_SECS)
+
+            time.sleep(delay_secs)
+    except Exception as e:
+        logger.error("SSE loop error : %s", e)
+
+
+def _connect(token, on_event, read_timeout_secs):
+    """
+    Opens a single SSE connection and dispatches events until the connection is
+    closed, errors out, or goes idle for longer than `read_timeout_secs`.
+
+    Returns a tuple of (outcome, status_code), where outcome is either "error"
+    (the connection could not be made, or timed out) or "disconnected" (a
+    response was received, but the connection has since ended).
+    """
+    url = f"{realtime.get_realtime_url()}api/runtime/stream"
+    logger.debug("SSE connecting to %s", url)
+
+    request = urllib.request.Request(
+        url,
+        method="GET",
+        headers={
+            "Authorization": str(token),
+            "Accept": "text/event-stream",
+            "Cache-Control": "no-cache",
+        },
+    )
+
+    try:
+        with urllib.request.urlopen(request, timeout=read_timeout_secs) as response:
+            status_code = response.getcode()
+            if status_code != 200:
+                return "disconnected", status_code
+
+            logger.debug("SSE connected successfully")
+
+            parser = SSEParser(response)
+            try:
+                for event in parser.events():
+                    logger.debug("SSE received event : %s", event)
+                    on_event(event)
+            except socket.timeout:
+                logger.debug("SSE read timeout")
+                return "error", None
+            except _CONNECTION_ERRORS as e:
+                logger.debug("SSE stream error : %s", e)
+                return "disconnected", status_code
+    except urllib.error.HTTPError as e:
+        logger.debug("SSE connection rejected with status %s", e.code)
+        return "disconnected", e.code
+    except _CONNECTION_ERRORS as e:
+        logger.debug("SSE connection error : %s", e)
+        return "error", None
+
+    logger.debug("SSE connection closed by server")
+    return "disconnected", status_code