controller: notify socket implementation in python

diff --git a/python/knot_resolver/controller/errors.py b/python/knot_resolver/controller/errors.py
new file mode 100644 (file)
index 0000000..b49d151
--- /dev/null
+++ b/python/knot_resolver/controller/errors.py
@@ -0,0 +1,20 @@
+from knot_resolver.errors import BaseKresError
+
+
+class ControllerError(BaseKresError):
+    """Class exception for all errors used in controller submodules."""
+
+    def __init__(self, msg: str) -> None:
+        super().__init__()
+        self._msg = msg
+
+    def __str__(self) -> str:
+        return self._msg
+
+
+class ControllerNotifySocketError(ControllerError):
+    """Exception class for notify socket errors."""
+
+    def __init__(self, msg: str, error_pointer: str = "") -> None:
+        msg = f"notify socket error: {msg}"
+        super().__init__(msg, error_pointer)

diff --git a/python/knot_resolver/controller/notify/notify_socket.py b/python/knot_resolver/controller/notify/notify_socket.py
new file mode 100644 (file)
index 0000000..2396f19
--- /dev/null
+++ b/python/knot_resolver/controller/notify/notify_socket.py
@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import os
+import socket
+import struct
+from pathlib import Path
+
+from knot_resolver.controller.errors import ControllerNotifySocketError
+from knot_resolver.logging import get_logger
+
+NOTIFY_SOCKET_NAME = "supervisor-notify-socket"
+NOTIFY_SOCKET_ENV_VAR = "NOTIFY_SOCKET"
+
+CREDENTIALS_FORMAT = "3i"
+CREDENTIALS_SIZE = struct.calcsize(CREDENTIALS_FORMAT)
+RECEIVE_BUFFER_SIZE = 2048
+
+logger = get_logger(__name__)
+
+
+def init_notify_socket() -> socket.socket:
+    notify_socket_path = Path.cwd() / NOTIFY_SOCKET_NAME
+    notify_socket_path.unlink(missing_ok=True)
+
+    notify_socket = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM | socket.SOCK_NONBLOCK)
+    try:
+        notify_socket.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
+        notify_socket.bind(str(notify_socket_path))
+    except OSError as e:
+        notify_socket.close()
+        notify_socket_path.unlink(missing_ok=True)
+        msg = f"failed to initialize notify socket at '{notify_socket_path}'"
+        raise ControllerNotifySocketError(msg) from e
+
+    os.environ[NOTIFY_SOCKET_ENV_VAR] = str(notify_socket_path)
+    return notify_socket
+
+
+def read_notify_socket(sock: socket.socket) -> tuple[int, bytes] | None:
+    try:
+        data, ancdata, flags, _ = sock.recvmsg(
+            RECEIVE_BUFFER_SIZE,
+            socket.CMSG_SPACE(CREDENTIALS_SIZE),
+        )
+    except BlockingIOError:
+        return None
+
+    if flags & socket.MSG_TRUNC:
+        raise ControllerNotifySocketError("received datagram was truncated")
+
+    pid = next(
+        (
+            struct.unpack(CREDENTIALS_FORMAT, cdata[:CREDENTIALS_SIZE])[0]
+            for level, ctype, cdata in ancdata
+            if level == socket.SOL_SOCKET and ctype == socket.SCM_CREDENTIALS
+        ),
+        None,
+    )
+
+    if pid is None:
+        logger.warning("ignoring received data without credentials: %s", data)
+        return None
+
+    return pid, data
+
+
+def send_notify_message(**values: str) -> None:
+    notify_socket_path = os.getenv(NOTIFY_SOCKET_ENV_VAR)
+
+    if notify_socket_path is None:
+        logger.warning("failed to get $NOTIFY_SOCKET environment variable")
+        return
+
+    if notify_socket_path.startswith("@"):
+        notify_socket_path = notify_socket_path.replace("@", "\0", 1)
+
+    try:
+        notify_socket = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+        notify_socket.connect(notify_socket_path)
+    except OSError:
+        logger.exception("failed to connect to $NOTIFY_SOCKET at '%s'", notify_socket_path)
+        return
+
+    payload = "\n".join((f"{key}={value}" for key, value in values.items()))
+    try:
+        notify_socket.send(payload.encode("utf8"))
+    except OSError:
+        logger.exception("failed to send notify message to $NOTIFY_SOCKET at '%s'", notify_socket_path)
+
+    notify_socket.close()

diff --git a/tests/python/knot_resolver/controller/notify/test_notify_socket.py b/tests/python/knot_resolver/controller/notify/test_notify_socket.py
new file mode 100644 (file)
index 0000000..190049b
--- /dev/null
+++ b/tests/python/knot_resolver/controller/notify/test_notify_socket.py
@@ -0,0 +1,24 @@
+import os
+
+import pytest
+
+from knot_resolver.controller.notify.notify_socket import init_notify_socket, read_notify_socket, send_notify_message
+
+
+@pytest.mark.parametrize(
+    "msg",
+    [
+        {"READY": 1},
+        {"RELOADING": 1},
+        {"STOPPING": 1},
+    ],
+)
+def test_notify_socket(msg: dict[str, int]) -> None:
+    key, value = next(iter(msg.items()))
+
+    notify_socket = init_notify_socket()
+    send_notify_message(**msg)
+    pid, message = read_notify_socket(notify_socket)
+
+    assert pid == os.getpid()
+    assert message.decode() == f"{key}={value}"