diff --git a/inference/core/env.py b/inference/core/env.py index 7c8be62b08..764b583f93 100644 --- a/inference/core/env.py +++ b/inference/core/env.py @@ -157,6 +157,35 @@ for entry in KAFKA_WORKFLOWS_SINKS_WHITELISTED_BOOTSTRAP_SERVERS.split(",") if entry.strip() ] +# Operator policy for the MQTT Reader / MQTT Writer Workflow blocks, which open +# an outbound connection to the `host`/`port` taken from the workflow definition +# or its inputs. When True (default, preserves behaviour) that value is +# honoured, subject to the allowlist below. When False the workflow-provided +# value is ignored and the blocks connect to the first entry of +# MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS instead; with no entry configured +# there, the MQTT blocks are disabled (every run reports an error). +MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST = str2bool( + os.getenv("MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", True) +) +# Optional comma-separated list of `host[:port]` MQTT brokers the MQTT Workflow +# blocks may connect to. When set and a user-provided host is allowed, the +# workflow's `host:port` must match an entry or the run reports an error; an +# entry without a port matches that host on any port. When a user-provided host +# is not allowed, the first entry is what the blocks connect to (port 1883 when +# the entry has none). Hosts are compared after trimming whitespace, lowercasing +# and stripping IPv6 brackets; there is no DNS resolution. Empty entries are +# ignored. A variable that is set but holds no entries is an empty allowlist: +# nothing is permitted. Default None (no allowlist). +MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS = os.getenv( + "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS" +) +if MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS is not None: + # a list, not a set: the first entry is the operator's broker + MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS = [ + entry.strip() + for entry in MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS.split(",") + if entry.strip() + ] # List of allowed origins ALLOW_ORIGINS = os.getenv("ALLOW_ORIGINS", "*") diff --git a/inference/core/interfaces/workflows_configuration.py b/inference/core/interfaces/workflows_configuration.py index 2e5a924855..0bb5716c60 100644 --- a/inference/core/interfaces/workflows_configuration.py +++ b/inference/core/interfaces/workflows_configuration.py @@ -79,6 +79,14 @@ def build_configuration_from_env() -> WorkflowsConfiguration: if env.KAFKA_WORKFLOWS_SINKS_WHITELISTED_BOOTSTRAP_SERVERS is None else tuple(env.KAFKA_WORKFLOWS_SINKS_WHITELISTED_BOOTSTRAP_SERVERS) ), + allow_mqtt_blocks_user_provided_host=env.MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST, + # NOT sorted: the first entry is the operator's broker when + # user-provided hosts are not allowed. + mqtt_blocks_whitelisted_hosts=( + None + if env.MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS is None + else tuple(env.MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS) + ), ), tensor=TensorConfiguration( representation_enabled=env.ENABLE_TENSOR_DATA_REPRESENTATION, diff --git a/tests/inference/unit_tests/core/interfaces/test_workflows_configuration.py b/tests/inference/unit_tests/core/interfaces/test_workflows_configuration.py index 0022f2e08c..7fd87dde9e 100644 --- a/tests/inference/unit_tests/core/interfaces/test_workflows_configuration.py +++ b/tests/inference/unit_tests/core/interfaces/test_workflows_configuration.py @@ -99,6 +99,18 @@ else list(c.engine.kafka_sinks_whitelisted_bootstrap_servers) ), ), + ( + "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", + lambda c: c.engine.allow_mqtt_blocks_user_provided_host, + ), + ( + "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", + lambda c: ( + None + if c.engine.mqtt_blocks_whitelisted_hosts is None + else list(c.engine.mqtt_blocks_whitelisted_hosts) + ), + ), ("ENABLE_TENSOR_DATA_REPRESENTATION", lambda c: c.tensor.representation_enabled), ("WORKFLOWS_IMAGE_TENSOR_DEVICE", lambda c: c.tensor.image_tensor_device), ( @@ -226,7 +238,7 @@ def test_the_field_table_matches_the_facade_exports() -> None: "missing_from_table": sorted(exported - tabled), "missing_from_facade": sorted(tabled - exported), } - assert len(tabled) == 78, len(tabled) + assert len(tabled) == 80, len(tabled) def test_every_name_workflows_imports_from_the_facade_is_exported() -> None: @@ -291,6 +303,11 @@ def test_server_configuration_equals_env_field_by_field(name, reader) -> None: "KAFKA_WORKFLOWS_SINKS_WHITELISTED_BOOTSTRAP_SERVERS", ["kafka-2:9092", "kafka-1:9092"], ), + ("MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False), + ("MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", None), + ("MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", []), + # order preserved, not sorted + ("MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker-2:1883", "broker-1"]), ("LAMBDA", True), ], ) diff --git a/tests/workflows/integration_tests/execution/conftest.py b/tests/workflows/integration_tests/execution/conftest.py index 1b06e12d94..9f57e0021e 100644 --- a/tests/workflows/integration_tests/execution/conftest.py +++ b/tests/workflows/integration_tests/execution/conftest.py @@ -1,8 +1,13 @@ import os import os.path import socket +import ssl import tempfile -from typing import Generator +import threading +from collections import deque +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from typing import Dict, Generator, Optional, Tuple import cv2 import numpy as np @@ -148,8 +153,142 @@ def face_image() -> np.ndarray: # Below taken from https://github.com/eclipse-paho/paho.mqtt.python/blob/d45de3737879cfe7a6acc361631fa5cb1ef584bb/tests/testsupport/broker.py +# and extended with SUBSCRIBE/SUBACK, PINGREQ/PINGRESP, retained messages and +# broker-initiated PUBLISH so subscriber blocks can be tested as well. +MQTT_CONNECT = 0x10 +MQTT_PUBLISH = 0x30 +MQTT_PUBACK = 0x40 +MQTT_PUBREC = 0x50 +MQTT_PUBREL = 0x60 +MQTT_PUBCOMP = 0x70 +MQTT_SUBSCRIBE = 0x80 +MQTT_SUBACK = 0x90 +MQTT_PINGREQ = 0xC0 +MQTT_PINGRESP = 0xD0 +MQTT_DISCONNECT = 0xE0 +MQTT_SUBACK_FAILURE = 0x80 + + +def _encode_remaining_length(length: int) -> bytes: + encoded = b"" + while True: + byte, length = length % 128, length // 128 + if length > 0: + byte |= 0x80 + encoded += bytes([byte]) + if length == 0: + return encoded + + +def _split_mqtt_packets(buffer: bytes): + """Split a byte buffer into complete MQTT control packets; returns (packets, rest).""" + packets = [] + while buffer: + length, multiplier, index = 0, 1, 1 + while True: + if index >= len(buffer): + return packets, buffer + byte = buffer[index] + length += (byte & 0x7F) * multiplier + multiplier *= 128 + index += 1 + if not byte & 0x80: + break + end = index + length + if end > len(buffer): + return packets, buffer + packets.append(buffer[:end]) + buffer = buffer[end:] + return packets, buffer + + +def mqtt_topic_matches(topic_filter: str, topic: str) -> bool: + """MQTT topic-filter matching with `+` (one level) and `#` (rest) wildcards.""" + filter_levels = topic_filter.split("/") + topic_levels = topic.split("/") + for index, level in enumerate(filter_levels): + if level == "#": + return True + if index >= len(topic_levels): + return False + if level != "+" and level != topic_levels[index]: + return False + return len(filter_levels) == len(topic_levels) + + +class _Session: + """What a broker keeps for one client id: subscriptions, an offline queue and + the connection currently bound to it (None while the client is away).""" + + def __init__(self, client_id: str, clean: bool): + self.client_id = client_id + self.clean = clean + self.subscriptions: list = [] + self.queue: deque = deque() + self.connection: Optional["_Connection"] = None + + def matches(self, topic: str) -> Optional[int]: + """Highest granted QoS among the subscriptions matching `topic`, or None.""" + granted = [ + qos + for topic_filter, qos in self.subscriptions + if mqtt_topic_matches(topic_filter, topic) + ] + return max(granted) if granted else None + + +class _Connection: + """Per-connection state: the socket, its receive buffer and packet counter.""" + + def __init__(self, sock): + self.sock = sock + self.buffer = b"" + self.connack_sent = False + self.stop = False + self.next_packet_id = 1 + self.session: Optional[_Session] = None + self.send_lock = threading.Lock() + + def close(self) -> None: + self.stop = True + sock, self.sock = self.sock, None + if sock is None: + return + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + sock.close() + + class FakeMQTTBroker: - def __init__(self, connack_reason_code: int = 0, listening: bool = True): + """A minimal MQTT 3.1.1 broker for tests. + + By default it serves one connection per `start()` call, answers CONNECT, + SUBSCRIBE and PINGREQ, re-sends retained messages on subscription and + records everything else (including a client's PUBLISH, which is never + acknowledged) in `messages`. + + `session_aware=True` adds what the persistent-session tests need: CONNECT is + parsed (client id, clean-session flag), sessions are kept per id in + `sessions`, a returning id gets CONNACK "session present" and its queued + messages, a second connection with a live id evicts the first, a client's + PUBLISH is acknowledged and routed to subscribers (queued for an offline + persistent session when published at QoS 1 or 2) and `serve()` accepts any + number of concurrent connections. Delivery is capped at the granted QoS, and + `suback_reason_code` is that cap (0 by default), so a persistent-session test + passes `suback_reason_code=2` for the broker to queue anything. + """ + + def __init__( + self, + connack_reason_code: int = 0, + listening: bool = True, + suback_reason_code: int = 0, + keep_serving: bool = False, + tls_context: Optional[ssl.SSLContext] = None, + session_aware: bool = False, + ): # Bind to "localhost" for maximum performance, as described in: # http://docs.python.org/howto/sockets.html#ipc sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) @@ -160,55 +299,463 @@ def __init__(self, connack_reason_code: int = 0, listening: bool = True): self.messages = [] self.messages_count_to_wait_for = 2 self.connack_reason_code = connack_reason_code + # granted QoS returned in SUBACK; 0x80 means the subscription is refused + self.suback_reason_code = suback_reason_code + # keep_serving: ignore recv timeouts and the message count, serve until + # finish() / drop_connection(); needed for subscriber tests + self.keep_serving = keep_serving + # tls_context: a server-side SSLContext; the accepted connection is + # wrapped with it before any MQTT packet is read + self.tls_context = tls_context + self.session_aware = session_aware + self.handshake_failures = 0 + self.subscriptions = [] + self.retained = {} + self.connections_accepted = 0 + # session_aware only: (client_id, clean_session) per CONNECT, in order + self.connects = [] + # session_aware only: live and remembered sessions by client id + self.sessions: Dict[str, _Session] = {} + self._anonymous_sessions = 0 + self._connections: list = [] + self._closed = False + self._lock = threading.RLock() sock.settimeout(5) self._sock = sock - self._conn = None # bound but not listening: connection attempts are refused until listen() if listening: self.listen() def listen(self): - self._sock.listen(1) + self._sock.listen(5) + + # -- accepting ----------------------------------------------------------- + + def _accept(self) -> Optional[_Connection]: + sock = self._sock + if sock is None: + raise ValueError("Socket is not open") + try: + conn, address = sock.accept() + except OSError: + # finish() closed the listening socket while no client had + # connected (a test whose block fails before connecting), or the + # accept timed out + return None + conn.settimeout(1) + if self.tls_context is not None: + try: + conn = self.tls_context.wrap_socket(conn, server_side=True) + except (ssl.SSLError, OSError): + # the client refused our certificate or never spoke TLS: + # treat it like a closed connection + self.handshake_failures += 1 + conn.close() + return None + connection = _Connection(conn) + with self._lock: + self._connections.append(connection) + self.connections_accepted += 1 + return connection def start(self): + """Accept one connection and serve it in the calling thread.""" if self._sock is None: raise ValueError("Socket is not open") - if self._conn is not None: + if any(c.sock is not None for c in self._connections): raise ValueError("Connection is already open") - - conn, address = self._sock.accept() - conn.settimeout(1) - self._conn = conn - connack_sent = False - while len(self.messages) < self.messages_count_to_wait_for or not connack_sent: - packet = self.receive_packet(1000) - print(f"Received {packet}") - if not packet: - break - if packet.startswith(b"\x10"): - print("sending CONNACK") - self._conn.send(b"\x20\x02\x00" + bytes([self.connack_reason_code])) - connack_sent = True + connection = self._accept() + if connection is None: + return + self._serve_connection(connection) + + def serve(self): + """Accept connections until finish(), each served on its own thread.""" + while not self._closed: + connection = self._accept() + if connection is None: continue - self.messages.append(packet) + threading.Thread( + target=self._serve_connection, args=(connection,), daemon=True + ).start() + + def _serve_connection(self, connection: _Connection) -> None: + try: + while ( + not self._closed + and not connection.stop + and ( + self.keep_serving + or len(self.messages) < self.messages_count_to_wait_for + or not connection.connack_sent + ) + ): + try: + chunk = self._receive(connection, 1000) + except socket.timeout: + if self.keep_serving: + continue + raise + except (OSError, ValueError): + # connection closed, possibly from the test thread + break + print(f"Received {chunk}") + if not chunk: + break + # exact byte: CONNECT's flags nibble is 0; a TLS ClientHello starts 0x16 + if ( + not connection.connack_sent + and not connection.buffer + and chunk[0] != MQTT_CONNECT + ): + # not an MQTT CONNECT (e.g. a TLS ClientHello against this plain + # broker): a real broker drops the connection on the protocol + # error, which is what lets the client fail fast + break + connection.buffer += chunk + packets, connection.buffer = _split_mqtt_packets(connection.buffer) + try: + for packet in packets: + if self._handle_packet(connection, packet): + connection.connack_sent = True + except (OSError, ValueError): + break + finally: + self._release(connection) + + def _release(self, connection: _Connection) -> None: + with self._lock: + connection.close() + if connection in self._connections: + self._connections.remove(connection) + session = connection.session + if session is not None and session.connection is connection: + # a clean session ends with its connection; a persistent one + # stays, with its subscriptions and queue, until the id returns + session.connection = None + if session.clean: + self.sessions.pop(session.client_id, None) + + # -- packets ------------------------------------------------------------- + + def _handle_packet(self, connection: _Connection, packet: bytes) -> bool: + """Handle one control packet; returns True when a CONNACK was sent.""" + packet_type = packet[0] & 0xF0 + if packet_type == MQTT_CONNECT: + self._handle_connect(connection, packet) + return True + if packet_type == MQTT_SUBSCRIBE: + self._handle_subscribe(connection, packet) + return False + if packet_type == MQTT_PINGREQ: + self._send(connection, bytes([MQTT_PINGRESP, 0])) + return False + if packet_type in (MQTT_PUBACK, MQTT_PUBREC, MQTT_PUBREL, MQTT_PUBCOMP): + # QoS 1/2 handshakes, not messages: a broker-initiated QoS 2 PUBLISH + # is answered with PUBREC and completed with PUBREL; a client's + # PUBREL (its own QoS 2 publish) is completed with PUBCOMP + if packet_type == MQTT_PUBREC: + self._send(connection, bytes([MQTT_PUBREL | 0x02, 2]) + packet[2:4]) + if packet_type == MQTT_PUBREL: + self._send(connection, bytes([MQTT_PUBCOMP, 2]) + packet[2:4]) + return False + if packet_type == MQTT_DISCONNECT: + return False + self.messages.append(packet) + if packet_type == MQTT_PUBLISH and self.session_aware: + self._handle_publish(connection, packet) + return False + + @staticmethod + def _body(packet: bytes) -> bytes: + """Strip the fixed header (type byte + remaining-length varint).""" + index = 1 + while packet[index] & 0x80: + index += 1 + return packet[index + 1 :] + + @staticmethod + def _read_string(body: bytes, index: int) -> Tuple[str, int]: + length = int.from_bytes(body[index : index + 2], "big") + index += 2 + return body[index : index + length].decode("utf-8"), index + length + + def _handle_connect(self, connection: _Connection, packet: bytes) -> None: + print("sending CONNACK") + if not self.session_aware: + self._send(connection, b"\x20\x02\x00" + bytes([self.connack_reason_code])) + return + body = self._body(packet) + _, index = self._read_string(body, 0) # protocol name + index += 1 # protocol level + clean_session = bool(body[index] & 0x02) + index += 3 # connect flags + keepalive + client_id, _ = self._read_string(body, index) + self.connects.append((client_id, clean_session)) + if self.connack_reason_code != 0: + self._send(connection, b"\x20\x02\x00" + bytes([self.connack_reason_code])) + return + with self._lock: + if not client_id: + # a broker assigns an id to an anonymous (clean) client + self._anonymous_sessions += 1 + client_id = f"anonymous-{self._anonymous_sessions}" + session = self.sessions.get(client_id) + if session is not None and session.connection is not None: + # a second connection with a live id: the broker drops the first + session.connection.close() + session.connection = None + session_present = session is not None and not clean_session + if not session_present: + session = _Session(client_id, clean=clean_session) + self.sessions[client_id] = session + session.connection = connection + connection.session = session + backlog = list(session.queue) + session.queue.clear() + self._send(connection, b"\x20\x02" + bytes([int(session_present), 0])) + # a real broker delivers the queued messages right after CONNACK, before + # the client's SUBSCRIBE arrives + for topic, payload, qos in backlog: + self._publish_to(connection, topic, payload, qos=qos) + + def _handle_subscribe(self, connection: _Connection, packet: bytes) -> None: + body = self._body(packet) + packet_id = body[0:2] + index = 2 + granted = [] + requested = [] + while index < len(body): + topic_filter, index = self._read_string(body, index) + requested_qos = body[index] + index += 1 + requested.append((topic_filter, requested_qos)) + self.subscriptions.append((topic_filter, requested_qos)) + if self.suback_reason_code == MQTT_SUBACK_FAILURE: + granted.append(MQTT_SUBACK_FAILURE) + else: + granted.append(min(requested_qos, self.suback_reason_code)) + suback = ( + bytes([MQTT_SUBACK]) + + _encode_remaining_length(2 + len(granted)) + + packet_id + + bytes(granted) + ) + self._send(connection, suback) + if self.suback_reason_code == MQTT_SUBACK_FAILURE: + return + if connection.session is not None: + with self._lock: + for (topic_filter, _), granted_qos in zip(requested, granted): + connection.session.subscriptions = [ + s + for s in connection.session.subscriptions + if s[0] != topic_filter + ] + [(topic_filter, granted_qos)] + for topic_filter, _ in requested: + for topic, payload in list(self.retained.items()): + if mqtt_topic_matches(topic_filter, topic): + self._publish_to(connection, topic, payload, retain=True) + + def _handle_publish(self, connection: _Connection, packet: bytes) -> None: + """Acknowledge a client's PUBLISH and route it to the subscribers.""" + flags = packet[0] & 0x0F + qos = (flags >> 1) & 0x03 + retain = bool(flags & 0x01) + body = self._body(packet) + topic, index = self._read_string(body, 0) + packet_id = b"" + if qos > 0: + packet_id = body[index : index + 2] + index += 2 + payload = body[index:] + if qos == 1: + self._send(connection, bytes([MQTT_PUBACK, 2]) + packet_id) + elif qos == 2: + # completed with PUBCOMP when the client's PUBREL arrives + self._send(connection, bytes([MQTT_PUBREC, 2]) + packet_id) + if retain: + self.retained[topic] = payload + self._route(topic, payload, qos) + + # -- publishing ---------------------------------------------------------- + + def publish( + self, topic: str, payload: bytes, retain: bool = False, qos: int = 0 + ) -> None: + """Publish from the broker side: to the connected client, or, when + session aware, to every matching subscriber (queued for an offline + persistent session when `qos` is 1 or 2).""" + if isinstance(payload, str): + payload = payload.encode("utf-8") + if self.session_aware: + if retain: + self.retained[topic] = payload + self._route(topic, payload, qos, retain=retain) + return + connection = self._current_connection() + if connection is None: + raise ValueError("Connection is not open") + self._publish_to(connection, topic, payload, retain=retain, qos=qos) + + def _route( + self, topic: str, payload: bytes, qos: int, retain: bool = False + ) -> None: + with self._lock: + targets = [] + for session in list(self.sessions.values()): + granted = session.matches(topic) + if granted is None: + continue + effective_qos = min(qos, granted) + if session.connection is not None: + targets.append((session.connection, effective_qos)) + elif not session.clean and effective_qos > 0: + # the broker queues only QoS 1 and 2 for an offline session + session.queue.append((topic, payload, effective_qos)) + for connection, effective_qos in targets: + try: + self._publish_to( + connection, topic, payload, retain=retain, qos=effective_qos + ) + except (OSError, ValueError): + pass + + def _publish_to( + self, + connection: _Connection, + topic: str, + payload: bytes, + retain: bool = False, + qos: int = 0, + ) -> None: + """Send a PUBLISH packet from the broker to one client.""" + if isinstance(payload, str): + payload = payload.encode("utf-8") + topic_bytes = topic.encode("utf-8") + variable_header = len(topic_bytes).to_bytes(2, "big") + topic_bytes + if qos > 0: + with connection.send_lock: + packet_id = connection.next_packet_id + connection.next_packet_id = packet_id % 65535 + 1 + variable_header += packet_id.to_bytes(2, "big") + body = variable_header + payload + header = bytes([MQTT_PUBLISH | (qos << 1) | int(retain)]) + self._send(connection, header + _encode_remaining_length(len(body)) + body) + + # -- lifecycle ----------------------------------------------------------- + + def _current_connection(self) -> Optional[_Connection]: + with self._lock: + live = [c for c in self._connections if c.sock is not None] + return live[-1] if live else None + + def drop_connection(self) -> None: + """Close the (most recent) client connection so a later start() or the + serving loop can accept a new one.""" + connection = self._current_connection() + if connection is not None: + connection.close() def finish(self): - if self._conn is not None: - self._conn.close() - self._conn = None + self._closed = True + with self._lock: + connections = list(self._connections) + for connection in connections: + connection.close() if self._sock is not None: self._sock.close() self._sock = None def receive_packet(self, num_bytes): - if self._conn is None: + connection = self._current_connection() + if connection is None: + raise ValueError("Connection is not open") + return self._receive(connection, num_bytes) + + @staticmethod + def _receive(connection: _Connection, num_bytes: int) -> bytes: + # capture locally: drop_connection()/finish() clear the attribute + # from another thread + sock = connection.sock + if sock is None: raise ValueError("Connection is not open") + return sock.recv(num_bytes) - packet_in = self._conn.recv(num_bytes) - return packet_in + @staticmethod + def _send(connection: _Connection, data: bytes) -> None: + with connection.send_lock: + sock = connection.sock + if sock is None: + raise ValueError("Connection is not open") + sock.sendall(data) + + +def _write_test_certificate(tmp_path, name, issuer=None, dns_names=()): + """Self-signed CA (no issuer) or a leaf signed by `issuer`, written as PEM. + + Mirrors `certificate()` in tests/inference/unit_tests/core/interfaces/http/ + test_mtls_enforcement.py; the leaf carries subjectAltName entries so paho's + hostname verification accepts `localhost`. + """ + from cryptography import x509 + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import ec + from cryptography.x509.oid import NameOID + + key = ec.generate_private_key(ec.SECP256R1()) + subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, name)]) + now = datetime.now(timezone.utc) + builder = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(issuer[0].subject if issuer else subject) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - timedelta(minutes=1)) + .not_valid_after(now + timedelta(days=1)) + .add_extension( + x509.BasicConstraints(ca=issuer is None, path_length=None), critical=True + ) + ) + if dns_names: + builder = builder.add_extension( + x509.SubjectAlternativeName( + [x509.DNSName(dns_name) for dns_name in dns_names] + ), + critical=False, + ) + cert = builder.sign(issuer[1] if issuer else key, hashes.SHA256()) + cert_path, key_path = tmp_path / (name + ".pem"), tmp_path / (name + ".key") + cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + key_path.write_bytes( + key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + ) + return cert, key, cert_path, key_path + + +@pytest.fixture(scope="function") +def mqtt_test_certificates(tmp_path): + """A throwaway CA and a `localhost` server certificate for TLS broker tests. + + Yields `ca_path` (PEM the blocks are pointed at through `ca_certificate_path`) + and `server_context` (what `FakeMQTTBroker(tls_context=...)` wraps with). + """ + pytest.importorskip("cryptography") + ca = _write_test_certificate(tmp_path, "mqtt-test-ca") + server = _write_test_certificate( + tmp_path, "mqtt-test-server", issuer=ca, dns_names=("localhost",) + ) + server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_context.load_cert_chain(certfile=str(server[2]), keyfile=str(server[3])) + return SimpleNamespace(ca_path=str(ca[2]), server_context=server_context) @pytest.fixture(scope="function") diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py new file mode 100644 index 0000000000..85128b6ea8 --- /dev/null +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -0,0 +1,838 @@ +import os +import threading +import time +from typing import Callable + +import pytest +from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import MQTTReaderBlockV1 + +from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.conftest import FakeMQTTBroker + +RUNNING = b'{"state": "RUNNING"}' +PAUSED = b'{"state": "PAUSED"}' + + +def wait_until(predicate: Callable[[], bool], timeout: float = 5.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return False + + +@pytest.fixture +def broker(): + broker = FakeMQTTBroker(keep_serving=True) + thread = threading.Thread(target=broker.start, daemon=True) + thread.start() + yield broker + broker.finish() + thread.join(timeout=2) + + +def reader_kwargs(broker: FakeMQTTBroker, **overrides) -> dict: + kwargs = { + "host": broker.host, + "port": broker.port, + "topic": "plc/#", + "timeout": 2.0, + } + kwargs.update(overrides) + return kwargs + + +@pytest.mark.timeout(15) +def test_retained_message_returned_on_first_run(broker): + # given + broker.retained["plc/state"] = RUNNING + block = MQTTReaderBlockV1() + + try: + # when + first = block.run(**reader_kwargs(broker)) + second = block.run(**reader_kwargs(broker)) + + # then + assert first == { + "value": '{"state": "RUNNING"}', + "payload": {"state": "RUNNING"}, + "topic": "plc/state", + "is_new": True, + "error_status": False, + "error_message": None, + } + assert second == {**first, "is_new": False} + assert broker.subscriptions == [("plc/#", 0)] + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_empty_first_run_then_latest_message_delivered(broker): + # given + block = MQTTReaderBlockV1() + + try: + # when + empty = block.run(**reader_kwargs(broker)) + for payload in (b"A", b"B", PAUSED): + broker.publish("plc/state", payload) + # the one-slot buffer holds a message as soon as A lands; wait for the + # newest one to have replaced it + assert wait_until( + lambda: bool(block._state.messages) + and block._state.messages[-1][1] == PAUSED + ) + result = block.run(**reader_kwargs(broker)) + + # then + assert empty["value"] is None + assert empty["is_new"] is False + assert empty["error_status"] is False + assert result["value"] == '{"state": "PAUSED"}' + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_sequential_mode_delivers_backlog_one_per_run(broker): + # given + block = MQTTReaderBlockV1() + kwargs = reader_kwargs(broker, read_mode="sequential") + + try: + # when + block.run(**kwargs) + broker.publish("plc/a", b"A") + broker.publish("plc/b", b"B") + broker.publish("plc/c", b"C") + assert wait_until(lambda: len(block._state.messages) == 3) + results = [block.run(**kwargs) for _ in range(4)] + + # then + assert [r["value"] for r in results] == ["A", "B", "C", "C"] + assert [r["topic"] for r in results] == ["plc/a", "plc/b", "plc/c", "plc/c"] + assert [r["is_new"] for r in results] == [True, True, True, False] + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_qos2_message_delivered_once(broker): + # given + block = MQTTReaderBlockV1() + kwargs = reader_kwargs(broker, qos=2, read_mode="sequential") + + try: + # when + block.run(**kwargs) + broker.publish("plc/state", RUNNING, qos=2) + assert wait_until(lambda: len(block._state.messages) == 1) + time.sleep(0.2) + first = block.run(**kwargs) + second = block.run(**kwargs) + + # then + assert broker.subscriptions == [("plc/#", 2)] + assert first["payload"] == {"state": "RUNNING"} + assert first["is_new"] is True + assert second["is_new"] is False + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_refused_subscription_reported_and_connection_kept(): + # given + broker = FakeMQTTBroker(keep_serving=True, suback_reason_code=0x80) + thread = threading.Thread(target=broker.start, daemon=True) + thread.start() + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**reader_kwargs(broker)) + + # then + assert result["error_status"] is True + assert "refused" in result["error_message"].lower() + assert block._state.connected.is_set() + assert block._client is not None + finally: + block.close() + broker.finish() + thread.join(timeout=2) + + +@pytest.mark.timeout(20) +def test_subscription_restored_after_reconnect(broker): + # given - a retained state the broker re-sends on every subscription + broker.retained["plc/state"] = RUNNING + block = MQTTReaderBlockV1() + + try: + first = block.run(**reader_kwargs(broker)) + assert first["payload"] == {"state": "RUNNING"} + assert first["is_new"] is True + + # when - the broker drops the connection and accepts a new one + broker.drop_connection() + assert wait_until(lambda: not block._state.connected.is_set()) + second_thread = threading.Thread(target=broker.start, daemon=True) + second_thread.start() + assert wait_until(lambda: len(broker.subscriptions) == 2, timeout=10) + assert wait_until(lambda: len(block._state.messages) == 1) + redelivered = block.run(**reader_kwargs(broker)) + broker.publish("plc/state", PAUSED) + assert wait_until( + lambda: bool(block._state.messages) + and block._state.messages[-1][1] == PAUSED + ) + updated = block.run(**reader_kwargs(broker)) + + # then - the subscription is live again, the retained copy is not new + assert broker.connections_accepted == 2 + assert redelivered == {**first, "is_new": False} + assert updated["payload"] == {"state": "PAUSED"} + assert updated["is_new"] is True + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_outage_repeats_last_message_without_waiting_then_recovers(broker, caplog): + # given - a warm reader holding RUNNING + broker.retained["plc/state"] = RUNNING + block = MQTTReaderBlockV1() + kwargs = reader_kwargs(broker, timeout=2.0) + second_thread = None + + try: + first = block.run(**kwargs) + assert first["payload"] == {"state": "RUNNING"} + + # when - the connection drops and three frames arrive while it is down + broker.drop_connection() + assert wait_until(lambda: not block._state.connected.is_set()) + with caplog.at_level("INFO", logger="inference"): + during = [] + for _ in range(3): + started = time.monotonic() + during.append(block.run(**kwargs)) + assert time.monotonic() - started < 0.5 + + # and the broker comes back: the retained copy is re-sent on the + # resubscribe (the buffer holds one message in latest mode), then a + # new message is published + second_thread = threading.Thread(target=broker.start, daemon=True) + second_thread.start() + assert wait_until(lambda: len(broker.subscriptions) == 2, timeout=10) + assert wait_until(lambda: len(block._state.messages) == 1, timeout=10) + after_reconnect = block.run(**kwargs) + broker.publish("plc/state", PAUSED) + assert wait_until( + lambda: bool(block._state.messages) + and block._state.messages[-1][1] == PAUSED + ) + updated = block.run(**kwargs) + + # then - every run during the outage repeated RUNNING with the error set + for result in during: + assert result == { + **first, + "is_new": False, + "error_status": True, + "error_message": result["error_message"], + } + assert "not connected" in result["error_message"] + errors = [ + r + for r in caplog.records + if r.levelname == "ERROR" and r.getMessage().startswith("MQTT Reader ") + ] + assert len(errors) == 1 + # the retained copy re-sent on resubscribe is not new, but it is live: + # it closes the outage; the next run returns the new message + assert after_reconnect == {**first, "is_new": False} + assert updated["payload"] == {"state": "PAUSED"} + assert updated["is_new"] is True + assert updated["error_status"] is False + recoveries = [ + r.getMessage() for r in caplog.records if "restored after" in r.getMessage() + ] + assert len(recoveries) == 1 + assert "3 runs returned the last message" in recoveries[0] + assert broker.connections_accepted == 2 + finally: + block.close() + if second_thread is not None: + second_thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_unreachable_broker_leaves_nothing_behind_and_later_run_recovers(): + # given - bound but not listening, so the TCP connect is refused + broker = FakeMQTTBroker(listening=False, keep_serving=True) + block = MQTTReaderBlockV1() + thread = None + + try: + # when + first = block.run(**reader_kwargs(broker, timeout=0.3)) + assert block._client is None + + broker.listen() + broker.retained["plc/state"] = RUNNING + thread = threading.Thread(target=broker.start, daemon=True) + thread.start() + second = block.run(**reader_kwargs(broker)) + + # then + assert first["error_status"] is True + assert second["error_status"] is False + assert second["payload"] == {"state": "RUNNING"} + finally: + block.close() + broker.finish() + if thread is not None: + thread.join(timeout=2) + + +def refusing_broker(reason_code: int): + """A broker that answers every CONNECT with `reason_code`, accepting any + number of connections so retries, if any, are counted.""" + broker = FakeMQTTBroker(connack_reason_code=reason_code, keep_serving=True) + thread = threading.Thread(target=broker.serve, daemon=True) + thread.start() + return broker, thread + + +@pytest.mark.timeout(15) +def test_refused_connection_is_reported_once_and_not_retried(): + # given - a broker that refuses the credentials + broker, thread = refusing_broker(5) + block = MQTTReaderBlockV1() + + try: + # when - the first run hits the refusal, then the block lives on past + # several reconnect intervals (max delay is 2 * timeout) + started = time.monotonic() + first = block.run(**reader_kwargs(broker, timeout=0.5)) + first_elapsed = time.monotonic() - started + time.sleep(2.5) + started = time.monotonic() + second = block.run(**reader_kwargs(broker, timeout=0.5)) + second_elapsed = time.monotonic() - started + + # then - the refusal is named at once, the loop stopped, no retry happened + assert first["error_status"] is True + assert "not authorised" in first["error_message"] + assert "code 5" in first["error_message"] + assert "client_id" in first["error_message"] + assert "Raise 'timeout'" not in first["error_message"] + assert first_elapsed < 0.4 + assert second == first + assert second_elapsed < 0.1 + assert broker.connections_accepted == 1 + assert broker.subscriptions == [] + finally: + block.close() + broker.finish() + thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_unavailable_broker_keeps_retrying_and_is_named(): + # given - a broker that answers "unavailable", a passing condition + broker, thread = refusing_broker(3) + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**reader_kwargs(broker, timeout=0.5)) + time.sleep(2.5) + + # then - the reason is reported and the client kept reconnecting + assert result["error_status"] is True + assert "broker unavailable" in result["error_message"] + assert "retrying in the background" in result["error_message"] + assert broker.connections_accepted >= 2 + finally: + block.close() + broker.finish() + thread.join(timeout=2) + + +@pytest.fixture +def tls_broker(mqtt_test_certificates): + broker = FakeMQTTBroker( + keep_serving=True, tls_context=mqtt_test_certificates.server_context + ) + thread = threading.Thread(target=broker.start, daemon=True) + thread.start() + yield broker + broker.finish() + thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_retained_message_over_tls_with_ca_certificate( + tls_broker, mqtt_test_certificates +): + # given + tls_broker.retained["plc/state"] = RUNNING + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + try: + # when + result = block.run( + **reader_kwargs( + tls_broker, + encryption="tls", + ca_certificate_path=mqtt_test_certificates.ca_path, + # the TLS handshake, CONNACK, SUBACK and the retained message + # each get this budget; CI runners are slower than a laptop + timeout=5.0, + ) + ) + + # then + assert result["error_status"] is False, result["error_message"] + assert result["payload"] == {"state": "RUNNING"} + assert result["is_new"] is True + assert tls_broker.connections_accepted == 1 + assert tls_broker.handshake_failures == 0 + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_tls_without_ca_certificate_is_rejected_by_verification(tls_broker): + # given - the test CA is not in the system trust store + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**reader_kwargs(tls_broker, encryption="tls")) + + # then - the client refused the certificate; nothing is kept + assert result["error_status"] is True + assert "not connected" in result["error_message"].lower() + assert "TLS is enabled" in result["error_message"] + assert block._client is None + assert wait_until(lambda: tls_broker.handshake_failures == 1) + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_tls_with_unreadable_ca_bundle_is_reported_by_paho(tls_broker, tmp_path): + # given - a real file that is not a PEM bundle: this exercises paho's + # loader, not a mock + bad_bundle = tmp_path / "not-a-bundle.pem" + bad_bundle.write_text("this is not a certificate") + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + try: + # when + result = block.run( + **reader_kwargs( + tls_broker, encryption="tls", ca_certificate_path=str(bad_bundle) + ) + ) + + # then + assert result["error_status"] is True + assert "could not load CA bundle" in result["error_message"] + assert block._client is None + assert tls_broker.connections_accepted == 0 + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_tls_against_plain_broker_reports_not_connected(broker): + # given - the broker never speaks TLS; like a real broker it drops a + # connection whose first bytes are not an MQTT CONNECT, so the client's + # handshake fails at once + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**reader_kwargs(broker, encryption="tls", timeout=1.0)) + + # then + assert result["error_status"] is True + assert "not connected" in result["error_message"].lower() + assert block._client is None + finally: + block.close() + + +# --- persistent sessions (client_id) ------------------------------------------ + + +@pytest.fixture +def session_broker(): + # granted QoS caps delivery, so the broker must grant up to QoS 2 to queue + # anything for an offline persistent session + broker = FakeMQTTBroker(keep_serving=True, session_aware=True, suback_reason_code=2) + thread = threading.Thread(target=broker.serve, daemon=True) + thread.start() + yield broker + broker.finish() + thread.join(timeout=2) + + +def persistent_kwargs(broker: FakeMQTTBroker, **overrides) -> dict: + kwargs = reader_kwargs(broker, client_id="line1", qos=1, read_mode="sequential") + kwargs.update(overrides) + return kwargs + + +def session_is_away(broker: FakeMQTTBroker, client_id: str = "line1") -> bool: + session = broker.sessions.get(client_id) + return session is not None and session.connection is None + + +def subscribe_and_leave(broker: FakeMQTTBroker, **overrides) -> None: + """A block instance with a persistent session subscribes, then goes away.""" + block = MQTTReaderBlockV1() + try: + result = block.run(**persistent_kwargs(broker, **overrides)) + assert result["error_status"] is False, result["error_message"] + finally: + block.close() + assert wait_until(lambda: session_is_away(broker)) + + +@pytest.mark.timeout(20) +def test_persistent_session_delivers_backlog_after_restart(session_broker, caplog): + # given - a subscriber with a fixed id went away and three QoS 1 messages + # were published meanwhile + with caplog.at_level("INFO", logger="inference"): + subscribe_and_leave(session_broker) + assert session_broker.connects == [("line1", False)] + assert "session present: False" in caplog.text + caplog.clear() + for payload in (b"A", b"B", b"C"): + session_broker.publish("plc/state", payload, qos=1) + block = MQTTReaderBlockV1() + + try: + # when - a new instance (a restarted pipeline) resumes the session + with caplog.at_level("INFO", logger="inference"): + results = [block.run(**persistent_kwargs(session_broker)) for _ in range(3)] + + # then - the broker resumed the session and delivered the backlog in order + assert session_broker.connects[-1] == ("line1", False) + assert "session present: True" in caplog.text + assert [r["value"] for r in results] == ["A", "B", "C"] + assert [r["is_new"] for r in results] == [True, True, True] + assert [r["error_status"] for r in results] == [False, False, False] + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_persistent_session_in_latest_mode_keeps_newest_of_backlog(session_broker): + # given + subscribe_and_leave(session_broker, read_mode="latest") + for payload in (b"A", b"B", b"C"): + session_broker.publish("plc/state", payload, qos=1) + block = MQTTReaderBlockV1() + + try: + # when + first = block.run(**persistent_kwargs(session_broker, read_mode="latest")) + second = block.run(**persistent_kwargs(session_broker, read_mode="latest")) + + # then - only the newest message of the burst survives + assert first["value"] == "C" + assert first["is_new"] is True + assert second["is_new"] is False + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_persistent_session_does_not_queue_qos0_messages(session_broker): + # given - the publisher used QoS 0, which a broker never queues + subscribe_and_leave(session_broker) + for payload in (b"A", b"B", b"C"): + session_broker.publish("plc/state", payload, qos=0) + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**persistent_kwargs(session_broker)) + + # then - the session resumed, but there was nothing to deliver + assert session_broker.connects[-1] == ("line1", False) + assert result["error_status"] is False + assert result["value"] is None + assert result["is_new"] is False + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_persistent_session_survives_reconnect_of_the_same_instance(session_broker): + # given - a live persistent subscription + block = MQTTReaderBlockV1() + + try: + first = block.run(**persistent_kwargs(session_broker)) + assert first["error_status"] is False + + # when - the connection drops, a message is published while the block is + # away, and the background loop reconnects + session_broker.drop_connection() + assert wait_until(lambda: not block._state.connected.is_set()) + assert wait_until(lambda: session_is_away(session_broker)) + session_broker.publish("plc/state", PAUSED, qos=1) + assert wait_until(lambda: len(block._state.messages) == 1, timeout=10) + result = block.run(**persistent_kwargs(session_broker)) + + # then - the queued message reached the same instance after the reconnect + assert session_broker.connections_accepted == 2 + assert session_broker.connects == [("line1", False), ("line1", False)] + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_persistent_session_outage_repeats_last_then_delivers_queued(session_broker): + # given - a warm persistent subscriber holding RUNNING + block = MQTTReaderBlockV1() + kwargs = persistent_kwargs(session_broker, timeout=2.0) + + try: + block.run(**kwargs) + session_broker.publish("plc/state", RUNNING, qos=1) + assert wait_until(lambda: len(block._state.messages) == 1) + first = block.run(**kwargs) + assert first["payload"] == {"state": "RUNNING"} + assert first["is_new"] is True + + # when - the connection drops, a QoS 1 message is queued while it is away, + # and a frame arrives during the outage + session_broker.drop_connection() + assert wait_until(lambda: not block._state.connected.is_set()) + assert wait_until(lambda: session_is_away(session_broker)) + session_broker.publish("plc/state", PAUSED, qos=1) + started = time.monotonic() + during = block.run(**kwargs) + assert time.monotonic() - started < 0.5 + + # and paho reconnects on its own + assert wait_until(lambda: len(block._state.messages) == 1, timeout=10) + after = block.run(**kwargs) + + # then - the outage frame repeated RUNNING with the error set, and the + # queued message came through afterwards + assert during == { + **first, + "is_new": False, + "error_status": True, + "error_message": during["error_message"], + } + assert "not connected" in during["error_message"] + assert after["payload"] == {"state": "PAUSED"} + assert after["is_new"] is True + assert after["error_status"] is False + assert session_broker.connections_accepted == 2 + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_second_connection_with_the_same_client_id_evicts_the_first(session_broker): + # given - two block instances sharing one client id on one broker; a retained + # state makes each first run return at once instead of waiting `timeout` for + # a retained message, so the evicted first cannot reconnect and evict the + # second while the second's first run is still in progress + session_broker.retained["plc/state"] = RUNNING + first = MQTTReaderBlockV1() + second = MQTTReaderBlockV1() + + try: + assert first.run(**persistent_kwargs(session_broker))["error_status"] is False + + # when - the second connects with the same id + result = second.run(**persistent_kwargs(session_broker)) + + # then - the broker drops the first; its reconnect would evict the second + # in turn, so the first is closed here (only the first eviction is tested) + assert wait_until(lambda: not first._state.connected.is_set()) + first.close() + assert result["error_status"] is False + assert session_broker.connects[:2] == [("line1", False), ("line1", False)] + assert wait_until( + lambda: second._state.connected.is_set() + and session_broker.sessions["line1"].connection is not None, + timeout=10, + ) + session_broker.publish("plc/state", PAUSED, qos=1) + assert wait_until(lambda: len(second._state.messages) == 1, timeout=10) + updated = second.run(**persistent_kwargs(session_broker)) + assert updated["payload"] == {"state": "PAUSED"} + assert updated["is_new"] is True + finally: + first.close() + second.close() + + +@pytest.mark.timeout(20) +def test_writer_without_client_id_delivers_to_an_away_persistent_subscriber( + session_broker, +): + # given - a persistent subscriber that went away, and the unchanged MQTT Writer + from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v1 import ( + MQTTWriterSinkBlockV1, + ) + + subscribe_and_leave(session_broker) + writer = MQTTWriterSinkBlockV1() + reader = MQTTReaderBlockV1() + + try: + # when - the writer publishes at QoS 1 with a broker-generated id + published = writer.run( + host=session_broker.host, + port=session_broker.port, + topic="plc/state", + message=PAUSED.decode("utf-8"), + qos=1, + timeout=2.0, + ) + writer.close() + result = reader.run(**persistent_kwargs(session_broker)) + + # then - queued delivery depends only on the publish QoS, not on a writer id + assert published["error_status"] is False, published["message"] + assert published["message"] == "Message published successfully" + assert session_broker.connects[1] == ("", True) + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + finally: + writer.close() + reader.close() + + +GATED_WORKFLOW = { + "version": "1.0", + "inputs": [ + {"type": "WorkflowParameter", "name": "host"}, + {"type": "WorkflowParameter", "name": "port"}, + ], + "steps": [ + { + "type": "roboflow_enterprise/mqtt_reader@v1", + "name": "reader", + "host": "$inputs.host", + "port": "$inputs.port", + "topic": "plc/state", + "timeout": 2.0, + }, + { + "type": "roboflow_core/continue_if@v1", + "name": "gate", + "condition_statement": { + "type": "StatementGroup", + "statements": [ + { + "type": "UnaryStatement", + "operand": { + "type": "DynamicOperand", + "operand_name": "raw", + "operations": [ + { + "type": "StringMatches", + "regex": '"state"\\s*:\\s*"RUNNING"', + } + ], + }, + "operator": {"type": "(Boolean) is True"}, + } + ], + }, + "evaluation_parameters": {"raw": "$steps.reader.value"}, + "next_steps": ["$steps.shout"], + }, + { + "type": "roboflow_core/property_definition@v1", + "name": "shout", + "data": "$steps.reader.value", + "operations": [{"type": "StringToUpperCase"}], + }, + ], + "outputs": [ + {"type": "JsonField", "name": "state", "selector": "$steps.reader.payload"}, + {"type": "JsonField", "name": "is_new", "selector": "$steps.reader.is_new"}, + {"type": "JsonField", "name": "shouted", "selector": "$steps.shout.output"}, + ], +} + + +@pytest.fixture +def enterprise_blocks_enabled(monkeypatch): + from inference.core.env import ENTERPRISE_BLOCKS_PLUGIN + from inference.core.workflows.execution_engine.introspection import blocks_loader + from inference.core.workflows.execution_engine.v1.compiler.core import ( + COMPILATION_CACHE, + ) + + plugins = [p for p in os.getenv("WORKFLOWS_PLUGINS", "").split(",") if p] + if ENTERPRISE_BLOCKS_PLUGIN not in plugins: # env.py may have expanded it already; + plugins.append(ENTERPRISE_BLOCKS_PLUGIN) # a duplicate trips the clash check + monkeypatch.setenv("WORKFLOWS_PLUGINS", ",".join(plugins)) + blocks_loader.clear_caches() + yield + blocks_loader.clear_caches() + # drop graphs compiled with enterprise blocks from the process-wide cache + with COMPILATION_CACHE._cache_lock: + COMPILATION_CACHE._cache.clear() + COMPILATION_CACHE._keys_buffer.clear() + + +@pytest.mark.timeout(20) +def test_workflow_gates_downstream_step_on_received_state( + broker, enterprise_blocks_enabled +): + # given - a retained RUNNING state, then a PAUSED update between frames + broker.retained["plc/state"] = RUNNING + execution_engine = ExecutionEngine.init( + workflow_definition=GATED_WORKFLOW, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + runtime_parameters = {"host": broker.host, "port": broker.port} + reader_block = execution_engine._engine._compiled_workflow.steps["reader"].step + + try: + # when + running_frame = execution_engine.run(runtime_parameters=runtime_parameters) + broker.publish("plc/state", PAUSED) + assert wait_until(lambda: len(reader_block._state.messages) == 1) + paused_frame = execution_engine.run(runtime_parameters=runtime_parameters) + repeated_frame = execution_engine.run(runtime_parameters=runtime_parameters) + + # then - one block instance served every frame + assert running_frame[0]["state"] == {"state": "RUNNING"} + assert running_frame[0]["is_new"] is True + assert running_frame[0]["shouted"] == '{"STATE": "RUNNING"}' + assert paused_frame[0]["state"] == {"state": "PAUSED"} + assert paused_frame[0]["is_new"] is True + assert paused_frame[0].get("shouted") is None + assert repeated_frame[0]["state"] == {"state": "PAUSED"} + assert repeated_frame[0]["is_new"] is False + assert broker.connections_accepted == 1 + finally: + reader_block.close() diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py index cf313ce021..6590fe4b22 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py @@ -88,6 +88,8 @@ def test_rejected_connack_is_not_treated_as_connected(): # then assert result["error_status"] is True + assert "not authorised" in result["message"] + assert "code 5" in result["message"] assert not block._connected.is_set() assert broker.messages == [] finally: @@ -95,6 +97,48 @@ def test_rejected_connack_is_not_treated_as_connected(): broker.finish() +@pytest.mark.timeout(15) +def test_refused_connection_is_reported_once_and_not_retried(): + # given - a broker that refuses the credentials and would accept any retry + broker = FakeMQTTBroker(connack_reason_code=5, keep_serving=True) + broker_thread = threading.Thread(target=broker.serve, daemon=True) + broker_thread.start() + block = MQTTWriterSinkBlockV1() + kwargs = dict( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="should not be delivered", + timeout=0.5, + ) + + try: + # when - the first run hits the refusal, then the block lives on past + # several reconnect intervals (max delay is 2 * timeout) + started = time.monotonic() + first = block.run(**kwargs) + first_elapsed = time.monotonic() - started + time.sleep(2.5) + started = time.monotonic() + second = block.run(**kwargs) + second_elapsed = time.monotonic() - started + + # then - the refusal is named at once, the loop stopped, nothing published + assert first["error_status"] is True + assert "not authorised" in first["message"] + assert "Check username and password" in first["message"] + assert "Raise 'timeout'" not in first["message"] + assert first_elapsed < 0.4 + assert second == first + assert second_elapsed < 0.1 + assert broker.connections_accepted == 1 + assert broker.messages == [] + finally: + block.close() + broker.finish() + broker_thread.join(timeout=2) + + @pytest.mark.timeout(15) def test_block_recovers_when_broker_becomes_available_after_first_run(): # given - the port is bound but refuses connections, so the first TCP @@ -248,6 +292,38 @@ def test_second_broker_is_rejected_instead_of_publishing_to_first(fake_mqtt_brok other_broker.finish() +@pytest.mark.timeout(15) +def test_publish_over_tls_with_ca_certificate(mqtt_test_certificates): + # given + broker = FakeMQTTBroker(tls_context=mqtt_test_certificates.server_context) + broker.messages_count_to_wait_for = 1 + broker_thread = threading.Thread(target=broker.start) + broker_thread.start() + block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) + + try: + # when + result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="encrypted payload", + timeout=5.0, + encryption="tls", + ca_certificate_path=mqtt_test_certificates.ca_path, + ) + broker_thread.join(timeout=2) + + # then + assert result["error_status"] is False, result["message"] + assert result["message"] == "Message published successfully" + assert b"encrypted payload" in broker.messages[-1] + assert broker.handshake_failures == 0 + finally: + block.close() + broker.finish() + + MQTT_SINK_WORKFLOW = { "version": "1.0", "inputs": [ diff --git a/tests/workflows/unit_tests/core_steps/sinks/test_mqtt_env.py b/tests/workflows/unit_tests/core_steps/sinks/test_mqtt_env.py new file mode 100644 index 0000000000..77bda17220 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/sinks/test_mqtt_env.py @@ -0,0 +1,62 @@ +"""Server-side parity: `inference.core.env` parsing of the MQTT blocks policy. + +Retained here (not in `workflows/tests`) because it imports the server package. +""" + +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict + +ALLOW_USER_HOST = "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST" +WHITELISTED_HOSTS = "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS" +_REPO_ROOT = str(Path(__file__).resolve().parents[5]) + + +def _import_env(extra_env: Dict[str, str]) -> Dict[str, Any]: + # env.py computes the values at import time: read them in a fresh interpreter + env = {**os.environ} + env.pop(ALLOW_USER_HOST, None) + env.pop(WHITELISTED_HOSTS, None) + env.update(extra_env) + env["PYTHONPATH"] = _REPO_ROOT + os.pathsep + env.get("PYTHONPATH", "") + code = ( + "import json; from inference.core import env; " + f"print('RESULT' + json.dumps([env.{ALLOW_USER_HOST}, " + f"env.{WHITELISTED_HOSTS}]))" + ) + result = subprocess.run( + [sys.executable, "-c", code], + env=env, + cwd=_REPO_ROOT, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + line = [l for l in result.stdout.splitlines() if l.startswith("RESULT")][-1] + allow_user_provided, allowlist = json.loads(line[len("RESULT") :]) + return {"allow_user_provided": allow_user_provided, "allowlist": allowlist} + + +def test_env_defaults_allow_user_hosts_without_allowlist() -> None: + values = _import_env({}) + + assert values == {"allow_user_provided": True, "allowlist": None} + + +def test_env_parses_the_allowlist_in_order_without_empty_entries() -> None: + values = _import_env( + { + ALLOW_USER_HOST: "False", + WHITELISTED_HOSTS: " broker-2:1883, ,broker-1 ,broker-2:1883,", + } + ) + + # trimmed, blank entries dropped, otherwise exactly what the operator wrote + assert values == { + "allow_user_provided": False, + "allowlist": ["broker-2:1883", "broker-1", "broker-2:1883"], + } diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index e07f3acdec..3cd5590e8e 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -16,7 +16,18 @@ for contributor and maintainer responsibilities. ## Unreleased -### Execution engine +### Fixed + +- MQTT Writer: a broker that refuses the connection (bad user name or password, not authorised, unacceptable protocol version) is now reported in the outputs with the broker's reason, and the client stops reconnecting instead of retrying the same credentials about once a second, which tripped brokers' authentication rate limiting. The refusal is reported as soon as the broker answers rather than after `timeout`, and every later run repeats it without touching the broker; fix the configuration and restart the pipeline. A broker answering "unavailable" keeps retrying in the background. + +### Added + +- MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect, and a 15 s keepalive reports a silently lost broker within about 25 s. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. A broker that refuses the connection (bad user name or password, not authorised, rejected client id) is reported with the broker's reason as soon as it answers, and the client stops reconnecting; a broker answering "unavailable" keeps retrying. After the first subscription a run never waits for the broker: while the connection is down, runs return the last message with `error_status` set instead of an empty message, buffered messages first; the outage is logged once at its start and once at recovery. +- Persistent sessions for the MQTT Reader: an optional `client_id` (unique per broker, for example the pipeline name from a workflow input) connects with that id and a non-clean session, so the broker keeps the subscription and queues QoS 1/2 messages while the block is away and delivers the backlog when the same id reconnects, after a pipeline or process restart included. Requires `qos` 1 or 2 (a run with `client_id` and `qos` 0 reports an error); the publisher must also publish at QoS 1 or higher. Meant for `InferencePipeline`s; `sequential` works through the backlog one message per run. Left empty, the behaviour is unchanged. +- TLS for the MQTT Reader and MQTT Writer: an `encryption` dropdown (`none`, default, or `tls`) encrypts the connection and verifies the broker's certificate against the system trust store; `ca_certificate_path` (shown for TLS only) points at a PEM bundle for a private CA and requires `ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`. The port is not switched automatically and verification cannot be disabled. +- Operator policy for the MQTT Reader and MQTT Writer broker address: `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (comma-separated `host[:port]` allowlist, no DNS) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (default `True`; when `False` the workflow's host and port are ignored and the first allowlist entry is used). Defaults preserve existing behaviour. + +### Execution Engine Change - Reject cyclic saved inner-workflow references during resolution with a composition error instead of `RecursionError`. Enforce nesting depth and total inner-workflow diff --git a/workflows/roboflow_workflows/configuration.py b/workflows/roboflow_workflows/configuration.py index a57be85f91..451ca8cee7 100644 --- a/workflows/roboflow_workflows/configuration.py +++ b/workflows/roboflow_workflows/configuration.py @@ -59,6 +59,11 @@ class EngineConfiguration: # None-vs-list distinction; the facade rebuilds the ORDERED list. allow_kafka_sinks_user_provided_bootstrap_servers: bool = True kafka_sinks_whitelisted_bootstrap_servers: Optional[Tuple[str, ...]] = None + # MQTT blocks broker-address policy, same shape as the Kafka pair: the + # facade rebuilds the ORDERED list and the first entry is the operator's + # broker when user-provided hosts are not allowed. + allow_mqtt_blocks_user_provided_host: bool = True + mqtt_blocks_whitelisted_hosts: Optional[Tuple[str, ...]] = None @dataclass(frozen=True) diff --git a/workflows/roboflow_workflows/enterprise_blocks/loader.py b/workflows/roboflow_workflows/enterprise_blocks/loader.py index 9981b8d31c..a9142f0a45 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/loader.py +++ b/workflows/roboflow_workflows/enterprise_blocks/loader.py @@ -12,6 +12,7 @@ from roboflow_workflows.enterprise_blocks.sinks.microsoft_sql_server.v1 import ( MicrosoftSQLServerSinkBlockV1, ) +from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import MQTTReaderBlockV1 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v1 import ( MQTTWriterSinkBlockV1, ) @@ -43,6 +44,7 @@ def load_enterprise_blocks() -> List[Type[WorkflowBlock]]: return [ OPCWriterSinkBlockV1, MQTTWriterSinkBlockV1, + MQTTReaderBlockV1, PLCBlockV1, PLCReaderBlockV1, PLCWriterBlockV1, diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py new file mode 100644 index 0000000000..ded0da20c0 --- /dev/null +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py @@ -0,0 +1,263 @@ +"""Shared pieces of the MQTT Reader and MQTT Writer blocks. + +Both blocks open an outbound connection to the broker named by the workflow. +`resolve_broker_address()` applies the operator's policy from +`MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` and +`MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (see `roboflow_workflows.environment`) +before either block builds a client, and `configure_tls()` enables +server-verified TLS on the client, gating a workflow-chosen CA path behind the +engine's file-system permission. `normalise_client_id()` turns the reader's +optional `client_id` into either a usable id or "unset", and +`connection_refused_message()` with `PERMANENT_CONNACK_CODES` turn a broker's +CONNACK refusal into the outputs' error text and the decision whether paho may +keep reconnecting. +""" + +import logging +import os +from typing import Any, List, Optional, Tuple + +import paho.mqtt.client as mqtt +from roboflow_workflows.environment import ( + MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST, + MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS, +) + +logger = logging.getLogger("inference") + +DEFAULT_MQTT_PORT = 1883 +# paho defaults to 60 s; a silent link loss (no TCP close) is only noticed +# through the keepalive, within about 1.5 times this value +MQTT_KEEPALIVE_SECONDS = 15 + + +class ConfigurationError(ValueError): + """The operator's policy or the deployment's permissions forbid the connection.""" + + +def configure_tls( + client: Any, + *, + use_tls: bool, + ca_certificate_path: Optional[str], + allow_access_to_file_system: bool, +) -> None: + """Enable server-verified TLS on a paho client before it connects. + + Args: + client: The paho ``Client`` to configure; left untouched when ``use_tls`` + is False or when an error is raised. + use_tls: Whether the connection must be wrapped in TLS. + ca_certificate_path: Optional PEM bundle for a broker whose issuer is not in + the system trust store. Ignored when ``use_tls`` is False. + allow_access_to_file_system: The engine's file-system permission; a CA path + is a server-side path chosen by the workflow and is refused without it. + + Raises: + ConfigurationError: When the CA path is set but file-system access is + disabled, or when the CA bundle cannot be loaded. + """ + if not use_tls: + return + if ca_certificate_path and not allow_access_to_file_system: + # the path is chosen by the workflow and read by the server process + raise ConfigurationError( + "ca_certificate_path needs access to the local file system, which is " + "disabled on this deployment (ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=False " + "or the engine's allow_access_to_file_system init parameter)." + ) + if ca_certificate_path and not os.path.isfile(ca_certificate_path): + # OpenSSL's loader blocks forever on a FIFO and errors on a directory; + # refuse anything that is not a regular file before it is opened + raise ConfigurationError( + f"TLS could not be configured: ca_certificate_path {ca_certificate_path!r} " + "is not a readable file." + ) + try: + if ca_certificate_path: + client.tls_set(ca_certs=str(ca_certificate_path)) + else: + # system trust store; hostname verification stays on + client.tls_set() + except Exception as e: + raise ConfigurationError( + f"TLS could not be configured: could not load CA bundle " + f"{ca_certificate_path!r} ({e})." + if ca_certificate_path + else f"TLS could not be configured: {e}." + ) from e + + +def normalise_client_id(value: Any) -> Optional[str]: + """Normalise the reader's optional ``client_id`` input. + + A persistent MQTT session exists only for a fixed client id, so the presence + of the id is the switch. An empty or whitespace-only value is "unset", so a + stray space never creates a persistent session. + + Args: + value: The workflow-provided ``client_id``: ``None``, a literal or a + selector-resolved value. + + Returns: + The id with surrounding whitespace removed, or ``None`` when unset. + + Raises: + ConfigurationError: When the value is not a string (a wiring error). + """ + if value is None: + return None + if not isinstance(value, str): + raise ConfigurationError( + f"Invalid client_id: {value!r}. Must be a string or left empty." + ) + value = value.strip() + return value or None + + +# CONNACK refusals a retry cannot change: unacceptable protocol version (1), +# identifier rejected (2), bad user name or password (4), not authorised (5). +# 3 (broker unavailable) is transient and keeps paho's reconnect. +PERMANENT_CONNACK_CODES = frozenset({1, 2, 4, 5}) +TRANSIENT_CONNACK_CODE = 3 + + +def connection_refused_message(reason_code: int, *, block_inputs: str) -> str: + """Return the outputs' error text for a broker that refused the CONNECT. + + Args: + reason_code: The CONNACK return code, 1 or higher. + block_inputs: The block's inputs the user should check, for example + ``"username and password"``. + + Returns: + A message naming the broker's reason. For a permanent refusal it says the + client does not retry; for "broker unavailable" it says the client keeps + retrying in the background. + """ + reason = mqtt.connack_string(reason_code) + if reason_code == TRANSIENT_CONNACK_CODE: + return ( + f"MQTT broker refused the connection: {reason} (code {reason_code}); " + "the client keeps retrying in the background." + ) + return ( + f"MQTT broker refused the connection: {reason} (code {reason_code}). " + f"Check {block_inputs}; the client does not retry until the block is closed " + "or the pipeline restarted." + ) + + +def _normalise_host(host: str) -> str: + host = host.strip().lower() + if host.startswith("[") and host.endswith("]"): + host = host[1:-1] + return host + + +def split_host_port(entry: str) -> Tuple[str, Optional[int]]: + """Split an allowlist entry into a normalised host and an optional port. + + Accepts ``host``, ``host:port``, ``[ipv6]``, ``[ipv6]:port`` and a bare + IPv6 address. There is no DNS resolution. + + Args: + entry: One operator-written allowlist entry. + + Returns: + ``(host, port)`` with the host lower-cased, whitespace trimmed and IPv6 + brackets stripped; ``port`` is ``None`` when the entry names none. + """ + entry = entry.strip() + if entry.startswith("["): + closing = entry.find("]") + if closing == -1: + return _normalise_host(entry), None + rest = entry[closing + 1 :] + port = int(rest[1:]) if rest.startswith(":") and rest[1:].isdigit() else None + return _normalise_host(entry[: closing + 1]), port + if entry.count(":") == 1: + host, port_text = entry.split(":") + if port_text.strip().isdigit(): + return _normalise_host(host), int(port_text) + return _normalise_host(entry), None + # no colon: a plain host; several colons: a bare IPv6 address + return _normalise_host(entry), None + + +def normalise_broker_address(host: str, port: int) -> str: + """Return the ``host:port`` form used when comparing addresses. + + Args: + host: Broker host as written in the workflow. + port: Broker port. + + Returns: + ``host:port`` with the host lower-cased and IPv6 brackets stripped. + """ + return f"{_normalise_host(str(host))}:{int(port)}" + + +def _operator_entries(allowlist: Optional[List[str]]) -> List[str]: + if allowlist is None: + return [] + return [str(entry).strip() for entry in allowlist if str(entry).strip()] + + +def resolve_broker_address( + host: str, port: int, log_override: bool = True +) -> Tuple[str, int]: + """Apply the operator's broker policy to a workflow-provided address. + + Args: + host: Broker host taken from the workflow. + port: Broker port taken from the workflow, already validated as an int. + log_override: Emit the warning when the workflow value is replaced by + the operator's broker; blocks pass ``False`` after their first run + so the warning is logged once per instance instead of once per frame. + + Returns: + The ``(host, port)`` the block must connect to: the workflow value when + permitted, or the operator's first allowlist entry when workflow-provided + hosts are not allowed. + + Raises: + ConfigurationError: When the policy forbids the connection. The message + never echoes the allowlist. + """ + # module attributes are read on every call so tests can patch them + user_provided_allowed = MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST + allowlist = MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS + operator_entries = _operator_entries(allowlist) + if not user_provided_allowed: + if not operator_entries: + raise ConfigurationError( + "MQTT blocks are disabled on this deployment: a workflow-provided " + "broker host is not allowed and the operator configured none." + ) + operator_host, operator_port = split_host_port(operator_entries[0]) + if operator_port is None: + operator_port = DEFAULT_MQTT_PORT + if log_override: + logger.warning( + "MQTT blocks: the workflow-provided host/port was replaced by the " + "operator-configured broker (%s:%s), because " + "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST is False.", + operator_host, + operator_port, + ) + return operator_host, operator_port + if allowlist is None: + return host, port + + wanted_host = _normalise_host(str(host)) + wanted_port = int(port) + for entry in operator_entries: + entry_host, entry_port = split_host_port(entry) + if entry_host == wanted_host and entry_port in (None, wanted_port): + return host, port + # the allowlist holds the operator's internal hostnames: never echo it + raise ConfigurationError( + f"Broker {host!r} port {wanted_port} is not permitted on this deployment: " + "the operator restricts the MQTT brokers Workflow blocks may connect to." + ) diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/__init__.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py new file mode 100644 index 0000000000..d24a1a4ae7 --- /dev/null +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -0,0 +1,975 @@ +import json +import logging +import math +import threading +import time +from collections import deque +from typing import Any, Deque, Dict, List, Literal, Optional, Tuple, Type, Union + +import paho.mqtt.client as mqtt +from pydantic import ConfigDict, Field +from typing_extensions import Annotated + +# Bind by literal name to the shared "inference" logger tree so server +# handlers / filters / propagation still apply when installed, without +# importing the server logger module. +logger = logging.getLogger("inference") +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + MQTT_KEEPALIVE_SECONDS, + PERMANENT_CONNACK_CODES, + TRANSIENT_CONNACK_CODE, + ConfigurationError, + configure_tls, + connection_refused_message, + normalise_client_id, + resolve_broker_address, +) +from roboflow_workflows.environment import GCP_SERVERLESS, LAMBDA +from roboflow_workflows.execution_engine.entities.base import OutputDefinition +from roboflow_workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + DICTIONARY_KIND, + FLOAT_KIND, + INTEGER_KIND, + SECRET_KIND, + STRING_KIND, + Selector, +) +from roboflow_workflows.prototypes.block import ( + BlockResult, + Runtime, + RuntimeInputMode, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) + +LONG_DESCRIPTION = """ +MQTT Reader block for subscribing to an MQTT topic and exposing one received message +per workflow run as a raw string (`value`) and a parsed JSON object (`payload`). + +The first run connects synchronously: the TCP connect, the broker's session +acknowledgement and the subscription acknowledgement are each bounded by `timeout`, +and the first run then waits up to `timeout` for a retained message the broker +delivers on subscription, so a cold start may take up to four times `timeout` - +raise it for remote brokers. Afterwards a background network loop maintains the +connection, re-establishes the subscription after every reconnect and buffers +incoming messages between runs. + +`read_mode` selects what a run returns from that buffer: + +- `latest` (default): the newest message received since the previous run; anything + older is skipped. +- `sequential`: the next unread message, one per run, in arrival order. A backlog is + worked through one per run; when more than 1000 messages queue up between two runs + the oldest are dropped. + +When no new message arrived, the outputs of the previous run are repeated with +`is_new` set to False. Retained messages delivered by the broker on subscription are +treated like any other message, except that a retained copy the broker re-sends after +a reconnect is not reported as new when its topic and payload match the message +returned last. + +The subscription and the buffer live in the block instance only. An InferencePipeline +keeps them for its lifetime, a restart begins with an empty buffer, and over the HTTP +API every request builds a fresh instance, so only a retained message can be returned +there. By default the session is not persistent either: messages published while the +block was not subscribed are not delivered. + +Set `client_id` to a stable name that is unique on the broker (for example the camera +or pipeline name from a workflow input) to keep a persistent session instead. The +presence of the id is the switch: the block then connects with that id and asks the +broker to keep the session, so the broker remembers the subscription and queues QoS 1 +and 2 messages while the block is away and delivers the backlog when a block with the +same id reconnects, after a pipeline or process restart included. `qos` must then be 1 +or 2 and the publisher must also publish at QoS 1 or higher, because a broker queues +nothing for QoS 0. Use `sequential` to work through the backlog one message per run; +`latest` keeps only the newest of it. QoS 1 may redeliver a message whose +acknowledgement was lost, and `sequential` surfaces that duplicate; QoS 2 avoids it. +Two connections with the same id disconnect each other, so never share an id between +pipelines. Persistent sessions are meant for InferencePipelines: over the HTTP API each +request would resume the session, consume the whole backlog and return one message. + +After the first subscription a run never waits for the broker. While the connection is +down the outputs repeat the last message received with `error_status` True and `is_new` +False, so a downstream step can gate on `error_status`; messages buffered before the drop +are still returned first, and with `client_id` the messages the broker queued meanwhile +follow after the reconnect. A connection refused on a reconnect is reported the same way, +with the broker's reason, until the pipeline is restarted. The log carries one ERROR line +when an outage starts and one INFO line with its duration when a live message arrives +again. + +One block instance subscribes on a single broker connection: changing host, port, +credentials, timeout, topic, QoS, read mode, client id or the TLS settings between runs +is rejected as a configuration error. While the broker is unreachable every run waits up to `timeout` +for the background reconnect before reporting the failure. + +The block is not available on the Roboflow hosted platform (`GCP_SERVERLESS` or +`LAMBDA`): every run fails there without opening a connection. It is intended for +self-hosted inference servers and InferencePipelines that can reach the broker. + +The server operator may restrict which brokers the MQTT blocks connect to with +`MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (an allowlist of `host[:port]` entries) and +`MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (when False, the workflow's host and +port are ignored and the first allowlist entry is used); a run the policy forbids is +reported in the outputs. The connection uses a 15 s keepalive, so a broker that +disappears without closing the connection is noticed within about 25 s. + +A broker that refuses the connection (bad user name or password, not authorised, a +rejected client id or an unacceptable protocol version) is reported in the outputs with +the broker's reason, and the block stops reconnecting: retrying with the same credentials +cannot succeed and would only trip the broker's authentication rate limiting. Fix the +configuration and restart the pipeline. A broker answering "unavailable" is a passing +condition, so the client keeps retrying in the background. + +Set `encryption` to `tls` to encrypt the connection and verify the broker's certificate +against the system trust store; the port is not switched automatically, so set it to the +broker's TLS port (usually 8883). A broker signed by a private CA needs +`ca_certificate_path`, a PEM bundle on the machine running inference, which requires +local file system access for Workflow blocks +(`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`). Certificate verification cannot be +disabled. + +Outputs: + - value (str): Raw message payload decoded as UTF-8, or None when nothing was + received yet. + - payload (dict): The payload parsed as a JSON object, or {} when the payload is + not a JSON object. + - topic (str): The topic the message arrived on (useful with wildcard filters). + - is_new (bool): True if this run surfaced a message the previous run did not. + - error_status (bool): True if the run failed (invalid configuration, refused + subscription, undecodable payload) or the broker connection + is down or was refused after a healthy start; in the latter + cases value, payload and topic repeat the last message + received. + - error_message (str): Details of the failure, or None. + +Failures are returned in the outputs and logged; the workflow keeps running. +""" + +SUBSCRIPTION_REFUSED_QOS = 0x80 +READER_CONNECTION_INPUTS = "username, password and client_id" +NOT_CONNECTED_WITHIN_TIMEOUT = ( + "MQTT broker not connected (connection was not established within timeout); " + "the client keeps retrying in the background. Raise 'timeout' if the " + "broker needs longer to connect." +) +NOT_CONNECTED_REPEATING_LAST = ( + "MQTT broker not connected; the client keeps reconnecting in the background and " + "the outputs repeat the last message received." +) +NOT_RESUBSCRIBED_REPEATING_LAST = ( + "MQTT subscription not re-established yet after a reconnect; the outputs repeat " + "the last message received." +) +TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} +LATEST_BUFFER_SIZE = 1 +SEQUENTIAL_BUFFER_SIZE = 1000 + +ReceivedMessage = Tuple[str, bytes, bool] + + +class MQTTReaderState: + """State shared between paho's network thread and the engine thread. + + Passed to paho as ``userdata`` so the callbacks never hold a reference to the + block itself (a bound-method callback would form a reference cycle and delay + the block's cleanup). + + Attributes: + connected: Set once the broker accepted the session; cleared on disconnect. + connack: Set on every CONNACK, accepted or refused, so the first run can + stop waiting as soon as the broker answered; cleared on a transport + disconnect, kept while a refusal is recorded. + refused_code: The CONNACK return code of the last refusal, or None after an + accepted CONNACK. A permanent code (see ``PERMANENT_CONNACK_CODES``) means + the network loop was stopped and every run reports the refusal. + subscribed: Set once the broker granted the subscription; cleared on disconnect. + subscribe_failed: True when the broker refused the subscription. + topic: Topic filter to subscribe to on every (re)connect. + qos: Subscription QoS requested from the broker. + messages: Buffer of ``(topic, payload, retain)`` tuples appended by the + network thread and consumed by ``run()``; ``maxlen`` drops the oldest. + message_received: Set once the first message was buffered; lets the first + run wait for a retained message the broker sends right after SUBACK. + """ + + def __init__(self, topic: str, qos: int, buffer_size: int): + self.connected = threading.Event() + self.connack = threading.Event() + self.refused_code: Optional[int] = None + self.subscribed = threading.Event() + self.message_received = threading.Event() + self.subscribe_failed = False + self.topic = topic + self.qos = qos + self.messages: Deque[ReceivedMessage] = deque(maxlen=buffer_size) + + +def _not_connected_message(error: BaseException, encryption: str) -> str: + if encryption == "tls": + return ( + f"MQTT broker not connected ({error}). TLS is enabled: check that port is " + "the broker's TLS port and that ca_certificate_path matches the broker's " + "certificate authority." + ) + return ( + f"MQTT broker not connected ({error}). Raise 'timeout' if the broker needs " + "longer to connect." + ) + + +def mqtt_on_connect( + client, state: MQTTReaderState, flags, reason_code, properties=None +): + # paho invokes on_connect for accepted and rejected CONNACK alike; + # only reason_code 0 means an established MQTT session. The outcome is + # recorded before `connack` wakes a waiting run, so the run never reads a + # half-updated state + if reason_code != 0: + state.refused_code = reason_code + state.connected.clear() + state.connack.set() + logger.error( + "MQTT connection refused: %s (code %s)", + mqtt.connack_string(reason_code), + reason_code, + ) + if reason_code in PERMANENT_CONNACK_CODES: + # paho would otherwise reconnect with the same credentials forever; + # disconnect() puts it in the disconnecting state, so its loop ends + # after this CONNACK instead of retrying + client.disconnect() + return + + state.refused_code = None + logger.info( + "MQTT client connected (session present: %s)", + bool(flags.get("session present")), + ) + state.connected.set() + state.connack.set() + # a clean-session broker forgets subscriptions on every reconnect, so the + # subscription is (re)established here rather than once after connect(); on + # a persistent session the repeated SUBSCRIBE is idempotent + state.subscribe_failed = False + state.subscribed.clear() + try: + client.subscribe(state.topic, qos=state.qos) + except Exception as e: + # paho rejects e.g. an empty topic locally; report it through the + # refused path instead of leaving the run waiting for a SUBACK + logger.error("MQTT subscribe to topic %r failed: %s", state.topic, e) + state.subscribe_failed = True + + +def mqtt_on_connect_fail(client, state: MQTTReaderState): + # paho 1.6.1 invokes this callback with exactly (client, userdata) + logger.error("MQTT client failed to establish connection with broker") + state.connected.clear() + + +def mqtt_on_subscribe(client, state: MQTTReaderState, mid, granted_qos): + if any(qos == SUBSCRIPTION_REFUSED_QOS for qos in granted_qos): + logger.error("MQTT broker refused subscription to topic %s", state.topic) + state.subscribe_failed = True + state.subscribed.clear() + return + + logger.info("MQTT client subscribed to topic %s", state.topic) + state.subscribe_failed = False + state.subscribed.set() + + +def mqtt_on_message(client, state: MQTTReaderState, message): + # the only work done on the network thread: stash the raw message; + # decoding and parsing happen in run(), where a failure is reported + # instead of ending the network loop + state.messages.append((message.topic, message.payload, bool(message.retain))) + state.message_received.set() + + +def mqtt_on_disconnect(client, state: MQTTReaderState, reason_code, properties=None): + logger.info("MQTT client disconnected, result code %s", reason_code) + state.connected.clear() + state.subscribed.clear() + if state.refused_code is None: + # a transport drop: the next run waits for the reconnect's CONNACK; after + # a refusal the answer stays visible so run() reports it without waiting + state.connack.clear() + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "MQTT Reader", + "version": "v1", + "short_description": "Reads a message from an MQTT topic on every run: the newest one, or the next unread one.", + "long_description": LONG_DESCRIPTION, + "license": "Roboflow Enterprise License", + "block_type": "transformation", + "ui_manifest": { + "section": "industrial", + "icon": "fal fa-network-wired", + "blockPriority": 11, + "enterprise_only": True, + "local_only": True, + }, + } + ) + type: Literal["roboflow_enterprise/mqtt_reader@v1"] + host: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="Host of the MQTT broker.", + examples=["localhost", "$inputs.mqtt_host"], + ) + port: Union[ + Annotated[int, Field(ge=1, le=65535)], + Selector(kind=[INTEGER_KIND]), + ] = Field( + description="Port of the MQTT broker (1-65535).", + examples=[1883, "$inputs.mqtt_port"], + ) + topic: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="MQTT topic filter to subscribe to. Single-level (`+`) and " + "multi-level (`#`) wildcards are allowed; the `topic` output tells which " + "topic each message arrived on.", + examples=["plc/state", "sensors/+/temperature", "$inputs.mqtt_topic"], + ) + read_mode: Literal["latest", "sequential"] = Field( + default="latest", + description="What a run returns from the messages received since the previous " + "run: `latest` returns the newest and skips older ones; `sequential` returns " + "the next unread message, one per run, in arrival order.", + examples=["latest", "sequential"], + json_schema_extra={ + "values_metadata": { + "latest": { + "name": "Latest", + "description": "Return the newest message; skip anything older.", + }, + "sequential": { + "name": "Sequential", + "description": "Return the next unread message, one per run, in arrival order.", + }, + }, + }, + ) + qos: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=0, + description="Quality of Service level requested for the subscription (0, 1 or 2). " + "When client_id is set, qos must be 1 or 2: the broker only queues QoS 1 and 2 " + "messages for a persistent session, and a run with client_id and qos 0 reports " + "an error. The publisher must also publish at QoS 1 or higher for its messages " + "to be queued.", + examples=[0, 1, 2], + ) + client_id: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Leave empty for a fresh session on every connection (a " + "broker-generated id). Set a stable name that is unique on the broker, for " + "example the camera or pipeline name from a workflow input, to keep a " + "persistent session: the broker then remembers the subscription and queues " + "QoS 1/2 messages while this block is away and delivers them when it " + "reconnects. Two connections with the same id disconnect each other.", + examples=["$inputs.pipeline_name", "line1-camera3"], + ) + username: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Username for MQTT broker authentication.", + examples=["$inputs.mqtt_username"], + ) + password: Optional[Union[Selector(kind=[SECRET_KIND]), str]] = Field( + default=None, + description="Password for MQTT broker authentication. Requires username.", + examples=["$inputs.mqtt_password"], + json_schema_extra={"private": True}, + ) + timeout: Union[ + Annotated[float, Field(gt=0, allow_inf_nan=False)], + Selector(kind=[FLOAT_KIND]), + ] = Field( + default=0.5, + description="Timeout in seconds for the first connection: the TCP connect, the " + "broker's connection acknowledgement (accepted or refused) and the subscription " + "acknowledgement. Later runs never wait: while the connection is down they return " + "the last message with error_status set, and the client reconnects in the " + "background. Must be a finite number greater than 0.", + examples=[0.5], + ) + encryption: Literal["none", "tls"] = Field( + default="none", + description="Transport security for the broker connection. `tls` encrypts the " + "connection and verifies the broker's certificate; set port to the broker's " + "TLS port (usually 8883).", + examples=["none", "tls"], + json_schema_extra={ + "values_metadata": { + "none": {"name": "None", "description": "Plain TCP (default)."}, + "tls": { + "name": "TLS", + "description": "Encrypted, broker certificate verified.", + }, + }, + }, + ) + ca_certificate_path: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Path to a PEM CA bundle on the machine running inference, for " + "brokers whose certificate issuer is not in the system trust store. Requires " + "local file system access for Workflow blocks " + "(`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`); otherwise the run " + "reports an error. Used only when encryption is `tls`.", + examples=["/etc/ssl/certs/factory-ca.pem"], + json_schema_extra={"relevant_for": TLS_RELEVANT}, + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="value", kind=[STRING_KIND]), + OutputDefinition(name="payload", kind=[DICTIONARY_KIND]), + OutputDefinition(name="topic", kind=[STRING_KIND]), + OutputDefinition(name="is_new", kind=[BOOLEAN_KIND]), + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="error_message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.4.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + RuntimeRestriction( + severity=Severity.HARD, + note=( + "On the Roboflow hosted platform every run returns " + "`error_status=true` with no message, and no MQTT connection is " + "opened from the hosted platform." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + ), + RuntimeRestriction( + severity=Severity.SOFT, + note=( + "The subscription and message buffer live in this block " + "instance. Over HTTP every request builds a fresh instance, so " + "only a retained message can be returned; with client_id set, " + "each request resumes the persistent session, consumes its whole " + "backlog and returns one message, and concurrent requests with " + "the same id disconnect each other. Use an InferencePipeline " + "for a live subscription or a persistent session." + ), + applies_to_runtimes=[ + Runtime.SELF_HOSTED_CPU, + Runtime.SELF_HOSTED_GPU, + Runtime.DEDICATED_DEPLOYMENT, + ], + applies_to_input_modes=[RuntimeInputMode.IMAGE], + ), + ] + + +class MQTTReaderBlockV1(WorkflowBlock): + def __init__(self, allow_access_to_file_system: bool = False): + # gates `ca_certificate_path`, a server-side path chosen by the workflow; + # False by default so a hand-constructed block is safe + self._allow_access_to_file_system = allow_access_to_file_system + self._client: Optional[mqtt.Client] = None + self._state: Optional[MQTTReaderState] = None + self._connection_identity: Optional[Tuple] = None + # outputs of the last surfaced message, repeated with is_new=False + self._last: Optional[Dict[str, Any]] = None + # the first read after subscribing waits for a retained message, which + # the broker sends right after SUBACK as a separate packet + self._awaiting_retained = False + # readiness waits are a cold-start concern: once the first run reached + # SUBACK, later runs never block on the broker + self._subscribed_once = False + # an outage is bracketed in the log: one ERROR line when a run first + # observes it, one INFO line with duration and run count on recovery + self._outage_started_at: Optional[float] = None + self._outage_runs = 0 + # the resolved broker address, kept for the log lines + self._host: Optional[str] = None + self._port: Optional[int] = None + self._lifecycle_lock = threading.Lock() + + def close(self) -> None: + with self._lifecycle_lock: + client = self._client + state = self._state + if client is None: + return + self._client = None + self._state = None + self._connection_identity = None + self._last = None + self._awaiting_retained = False + self._subscribed_once = False + self._outage_started_at = None + self._outage_runs = 0 + self._host = None + self._port = None + try: + client.disconnect() + except Exception as e: + logger.error("Failed to disconnect MQTT client: %s", e) + finally: + try: + # loop_stop() joins the network thread without a timeout; + # accepted so the thread never outlives the block + client.loop_stop() + finally: + # only after the join can no callback re-set the events + if state is not None: + state.connected.clear() + state.subscribed.clear() + + def __del__(self): + try: + self.close() + except Exception as e: + logger.error("Failed to close MQTT client: %s", e) + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["allow_access_to_file_system"] + + def run( + self, + host: str, + port: int, + topic: str, + read_mode: str = "latest", + qos: int = 0, + username: Optional[str] = None, + password: Optional[str] = None, + timeout: float = 0.5, + encryption: str = "none", + ca_certificate_path: Optional[str] = None, + client_id: Optional[str] = None, + ) -> BlockResult: + if GCP_SERVERLESS or LAMBDA: + # never open an outbound broker connection from hosted workers + return self._handle_failure( + "MQTT Reader is not available on the Roboflow hosted platform" + ) + # selector-resolved values bypass manifest constraints, so validate here + try: + timeout_seconds = float(timeout) + except (TypeError, ValueError, OverflowError): + timeout_seconds = math.nan + if ( + not math.isfinite(timeout_seconds) + or not 0 < timeout_seconds <= threading.TIMEOUT_MAX + ): + return self._handle_failure( + f"Invalid timeout: {timeout!r}. Timeout must be a positive finite " + "number of seconds within the platform limit." + ) + timeout = timeout_seconds + # selector-supplied ports arrive uncoerced (numeric strings, floats + # from JSON), so coerce like the manifest would before range-checking; + # paho validates only port <= 0 and above 65535 the socket call raises + # OverflowError inside the network loop, silently killing it + try: + if isinstance(port, bool) or ( + isinstance(port, float) and not port.is_integer() + ): + raise ValueError + port_number = int(port) + except (TypeError, ValueError): + port_number = 0 + if not 1 <= port_number <= 65535: + return self._handle_failure( + f"Invalid port: {port!r}. Port must be an integer between 1 and 65535." + ) + port = port_number + try: + if isinstance(qos, bool) or ( + isinstance(qos, float) and not qos.is_integer() + ): + raise ValueError + qos_number = int(qos) + except (TypeError, ValueError): + qos_number = -1 + if qos_number not in (0, 1, 2): + return self._handle_failure(f"Invalid qos: {qos!r}. QoS must be 0, 1 or 2.") + qos = qos_number + if read_mode not in ("latest", "sequential"): + return self._handle_failure( + f"Invalid read_mode: {read_mode!r}. Must be 'latest' or 'sequential'." + ) + if not isinstance(topic, str) or not topic: + return self._handle_failure( + f"Invalid topic: {topic!r}. Topic must be a non-empty string." + ) + if password is not None and username is None: + return self._handle_failure( + "Password provided without username. Set username to enable MQTT authentication." + ) + if encryption not in ("none", "tls"): + return self._handle_failure( + f"Invalid encryption: {encryption!r}. Must be 'none' or 'tls'." + ) + # an empty editor field is "unset"; anything but a string is a wiring error + if ca_certificate_path is not None and not isinstance(ca_certificate_path, str): + return self._handle_failure( + f"Invalid ca_certificate_path: {ca_certificate_path!r}. Must be a path." + ) + ca_certificate_path = ca_certificate_path or None + try: + # blank means "no persistent session"; the presence of the id is the switch + client_id = normalise_client_id(client_id) + except ConfigurationError as e: + return self._handle_failure(str(e)) + if client_id is not None and qos == 0: + # the broker queues only QoS 1 and 2 for an offline session: a + # persistent session at QoS 0 would look right and receive nothing + return self._handle_failure( + "client_id needs qos 1 or 2: the broker only queues QoS 1 and 2 " + "messages for a persistent session." + ) + try: + # the operator's broker policy; the RESOLVED address is what the + # block connects to and what the connection identity is built from + host, port = resolve_broker_address( + host, port, log_override=self._client is None + ) + except ConfigurationError as e: + return self._handle_failure(str(e)) + with self._lifecycle_lock: + return self._connect_and_read( + host=host, + port=port, + topic=topic, + read_mode=read_mode, + qos=qos, + username=username, + password=password, + timeout=timeout, + encryption=encryption, + ca_certificate_path=ca_certificate_path, + client_id=client_id, + ) + + def _connect_and_read( + self, + host: str, + port: int, + topic: str, + read_mode: str, + qos: int, + username: Optional[str], + password: Optional[str], + timeout: float, + encryption: str, + ca_certificate_path: Optional[str], + client_id: Optional[str], + ) -> BlockResult: + connection_identity = ( + host, + port, + username, + password, + timeout, + topic, + qos, + read_mode, + encryption, + ca_certificate_path, + client_id, + ) + if self._client is None: + buffer_size = ( + LATEST_BUFFER_SIZE if read_mode == "latest" else SEQUENTIAL_BUFFER_SIZE + ) + state = MQTTReaderState(topic=topic, qos=qos, buffer_size=buffer_size) + client = None + try: + if client_id is None: + # broker-generated id, clean session: nothing outlives the connection + client = mqtt.Client(userdata=state) + else: + # a persistent session: the broker keeps the subscription and + # queues QoS 1/2 messages while this id is away + client = mqtt.Client( + client_id=client_id, clean_session=False, userdata=state + ) + # a raising callback would end the network thread; log instead + client.suppress_exceptions = True + if username is not None: + client.username_pw_set(username, password) + # TLS must be configured before connect(); the CA path is gated + # by the engine's file-system permission + configure_tls( + client, + use_tls=encryption == "tls", + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=self._allow_access_to_file_system, + ) + client.on_connect = mqtt_on_connect + client.on_connect_fail = mqtt_on_connect_fail + client.on_subscribe = mqtt_on_subscribe + client.on_message = mqtt_on_message + client.on_disconnect = mqtt_on_disconnect + # min_delay stays below the readiness wait (= timeout) so the + # first reconnect attempt after a connection drop can finish + # within a single run's wait instead of starting as it expires + client.reconnect_delay_set(min_delay=timeout / 2, max_delay=2 * timeout) + # paho 1.6.1 has no public setter for its synchronous connect + # timeout; without this the TCP phase runs under paho's 5s + # default instead of the block's timeout + client._connect_timeout = timeout + # the TCP connect happens here, bounded by _connect_timeout + # (DNS resolution is not - it runs under the OS resolver + # timeout); the CONNACK wait below covers the handshake rest + client.connect(host, port, keepalive=MQTT_KEEPALIVE_SECONDS) + client.loop_start() + except ConfigurationError as e: + # TLS setup refused before any socket was opened: nothing kept + return self._handle_failure(str(e)) + except OSError as e: + # broker unreachable: nothing is kept and no loop was started, + # so a failed one-shot run leaves no background thread behind; + # the next run on this instance retries with a fresh client + return self._handle_failure( + _not_connected_message(e, encryption=encryption) + ) + except Exception as e: + if client is not None: + try: + client.loop_stop() + except Exception: + pass + return self._handle_failure(f"Failed to initialize MQTT client: {e}") + self._client = client + self._state = state + self._connection_identity = connection_identity + self._awaiting_retained = True + self._host, self._port = host, port + elif connection_identity != self._connection_identity: + return self._handle_failure( + "MQTT connection parameters (host, port, credentials, timeout, topic, " + "qos, read_mode, encryption, ca_certificate_path or client_id) changed " + "between runs; this block subscribes only with the configuration of its " + "first run." + ) + state = self._state + if not self._subscribed_once: + # cold start: bounded waits for the broker's answers, so the first + # run of a pipeline gets a message or a precise failure + failure = self._wait_until_subscribed(state, timeout) + if failure is not None and not state.messages: + return failure + if failure is None: + self._subscribed_once = True + if self._awaiting_retained: + # one-off, bounded: a retained message (or nothing) arrives + # right after SUBACK; without this wait the first run would + # race it + self._awaiting_retained = False + state.message_received.wait(timeout=timeout) + # never wait once subscribed; always drain first, because messages + # received before a drop or a refusal are real whatever happened to the + # connection afterwards (a persistent session's backlog even lands + # before the SUBACK that may refuse the subscription) + try: + received_topic, payload, retained = state.messages.popleft() + except IndexError: + if state.subscribe_failed: + # a known refusal is permanent for this configuration + return self._handle_failure( + f"MQTT broker refused subscription to topic {state.topic!r}." + ) + if state.refused_code in PERMANENT_CONNACK_CODES: + # a reconnect refused with 1/2/4/5 stopped the loop: no fresh + # data can arrive until the pipeline is restarted + return self._repeat_last_with_error( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) + if not state.connected.is_set(): + return self._repeat_last_with_error( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + if state.refused_code == TRANSIENT_CONNACK_CODE + else NOT_CONNECTED_REPEATING_LAST + ) + if not state.subscribed.is_set(): + return self._repeat_last_with_error(NOT_RESUBSCRIBED_REPEATING_LAST) + # nothing new since the previous run + return self._repeat_last() + # a live message ends an outage, a retained duplicate included + self._log_recovery_if_needed() + try: + value = payload.decode("utf-8") + except UnicodeDecodeError: + return self._handle_failure( + f"Message received on topic {received_topic!r} is not valid UTF-8." + ) + parsed_payload: Dict[str, Any] = {} + try: + parsed = json.loads(value) + if isinstance(parsed, dict): + parsed_payload = parsed + except ValueError: + pass + if ( + retained + and self._last is not None + and self._last["topic"] == received_topic + and self._last["value"] == value + ): + # the broker re-sends the retained message on every resubscribe + # (reconnect); an unchanged copy is not a new message + return self._repeat_last() + self._last = { + "value": value, + "payload": parsed_payload, + "topic": received_topic, + } + return { + **self._last, + "is_new": True, + "error_status": False, + "error_message": None, + } + + def _wait_until_subscribed( + self, state: MQTTReaderState, timeout: float + ) -> Optional[BlockResult]: + """The first run's readiness waits; returns the failure result or None.""" + # a permanent refusal stopped the network loop: report it on every run + # without waiting and without touching the broker again + if state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) + # the background loop owns (re)connecting and (re)subscribing; runs + # only wait for readiness. The wait ends on any CONNACK, so a refusal is + # reported as soon as the broker answers instead of after the timeout + if not state.connack.wait(timeout=timeout): + return self._handle_failure(NOT_CONNECTED_WITHIN_TIMEOUT) + if state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) + if not state.connected.wait(timeout=timeout): + if state.refused_code == TRANSIENT_CONNACK_CODE: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) + return self._handle_failure(NOT_CONNECTED_WITHIN_TIMEOUT) + # a known refusal is permanent for this configuration: report it + # without paying the acknowledgement wait on every run + if state.subscribe_failed or ( + not state.subscribed.wait(timeout=timeout) and state.subscribe_failed + ): + return self._handle_failure( + f"MQTT broker refused subscription to topic {state.topic!r}." + ) + if not state.subscribed.is_set(): + return self._handle_failure( + "MQTT subscription was not acknowledged within timeout. Raise " + "'timeout' if the broker needs longer to respond." + ) + return None + + def _repeat_last(self) -> BlockResult: + if self._last is None: + return self._empty_result() + return { + **self._last, + "is_new": False, + "error_status": False, + "error_message": None, + } + + def _repeat_last_with_error(self, message: str) -> BlockResult: + """The outputs while no fresh data can arrive: the last message with the + error fields set, or the empty failure result when nothing was received yet. + + The first call of an outage logs one ERROR line; later calls log nothing + (a video pipeline would otherwise write one line per frame). + """ + if self._outage_started_at is None: + self._outage_started_at = time.monotonic() + self._outage_runs = 0 + # host, port and topic identify the block instance in a shared log + logger.error( + "MQTT Reader %s:%s %r: %s", + self._host, + self._port, + self._state.topic if self._state is not None else None, + message, + ) + self._outage_runs += 1 + if self._last is None: + return { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": True, + "error_message": message, + } + return { + **self._last, + "is_new": False, + "error_status": True, + "error_message": message, + } + + def _log_recovery_if_needed(self) -> None: + """Close the outage bracket in the log once a live message arrived again.""" + if self._outage_started_at is None: + return + logger.info( + "MQTT Reader %s:%s %r: connection restored after %.1f s; %d runs returned " + "the last message received", + self._host, + self._port, + self._state.topic if self._state is not None else None, + time.monotonic() - self._outage_started_at, + self._outage_runs, + ) + self._outage_started_at = None + self._outage_runs = 0 + + @staticmethod + def _empty_result() -> BlockResult: + return { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": False, + "error_message": None, + } + + def _handle_failure(self, message: str) -> BlockResult: + logger.error("MQTT Reader failure: %s", message) + return { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": True, + "error_message": message, + } diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py index e313a3704e..b7a69afec3 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py @@ -12,6 +12,14 @@ # importing the server logger module. logger = logging.getLogger("inference") from roboflow_workflows.core_steps.sinks.noop import disabled_sink_response +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + PERMANENT_CONNACK_CODES, + TRANSIENT_CONNACK_CODE, + ConfigurationError, + configure_tls, + connection_refused_message, + resolve_broker_address, +) from roboflow_workflows.execution_engine.entities.base import OutputDefinition from roboflow_workflows.execution_engine.entities.types import ( BOOLEAN_KIND, @@ -41,8 +49,33 @@ background loop. Over the HTTP API every request builds a fresh block instance, so each request pays the bounded connect and a failed request is final for that request. One block instance publishes to a single broker -connection: changing host, port, credentials or timeout between runs is -rejected as a configuration error. +connection: changing host, port, credentials, timeout or the TLS settings +between runs is rejected as a configuration error. + +The server operator may restrict which brokers the MQTT blocks connect to +with `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (an allowlist of `host[:port]` +entries) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (when False, the +workflow's host and port are ignored and the first allowlist entry is used); +a run the policy forbids is reported like any other failure. + +Set `encryption` to `tls` to encrypt the connection and verify the broker's +certificate against the system trust store; the port is not switched +automatically, so set it to the broker's TLS port (usually 8883). A broker +signed by a private CA needs `ca_certificate_path`, a PEM bundle on the +machine running inference, which requires local file system access for +Workflow blocks (`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`). +Certificate verification cannot be disabled. A peer that accepts the TCP +connection but never completes the TLS handshake is given up on after the +MQTT keepalive interval (60 s), not after `timeout`; real brokers close such +a connection immediately. + +A broker that refuses the connection (bad user name or password, not +authorised, an unacceptable protocol version) is reported in the outputs with +the broker's reason, and the block stops reconnecting: retrying with the same +credentials cannot succeed and would only trip the broker's authentication +rate limiting. Fix the configuration and restart the pipeline. A broker +answering "unavailable" is a passing condition, so the client keeps retrying +in the background. Outputs: - error_status (bool): Indicates if an error occurred during the MQTT publishing process. @@ -60,26 +93,98 @@ """ -def mqtt_on_connect(client, userdata, flags, reason_code, properties=None): +TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} + + +def _not_connected_message(error: BaseException, encryption: str) -> str: + if encryption == "tls": + return ( + f"MQTT broker not connected ({error}). TLS is enabled: check that port is " + "the broker's TLS port and that ca_certificate_path matches the broker's " + "certificate authority." + ) + return ( + f"MQTT broker not connected ({error}). Raise 'timeout' if the broker needs " + "longer to connect." + ) + + +WRITER_CONNECTION_INPUTS = "username and password" +NOT_CONNECTED_WITHIN_TIMEOUT = ( + "MQTT broker not connected (connection was not established within timeout); " + "the client keeps retrying in the background. Raise 'timeout' if the " + "broker needs longer to connect." +) + + +class MQTTWriterState: + """Connection state shared between paho's network thread and the engine thread. + + Passed to paho as ``userdata`` so the callbacks never hold a reference to the + block itself. + + Attributes: + connected: Set once the broker accepted the session; cleared on disconnect. + connack: Set on every CONNACK, accepted or refused, so a run can stop + waiting as soon as the broker answered; cleared on a transport + disconnect, kept while a refusal is recorded. + refused_code: The CONNACK return code of the last refusal, or None after an + accepted CONNACK. A permanent code (see ``PERMANENT_CONNACK_CODES``) means + the network loop was stopped and every run reports the refusal. + """ + + def __init__(self): + self.connected = threading.Event() + self.connack = threading.Event() + self.refused_code: Optional[int] = None + + def reset(self) -> None: + self.connected.clear() + self.connack.clear() + self.refused_code = None + + +def mqtt_on_connect( + client, state: MQTTWriterState, flags, reason_code, properties=None +): # paho invokes on_connect for accepted and rejected CONNACK alike; - # only reason_code 0 means an established MQTT session + # only reason_code 0 means an established MQTT session. The outcome is + # recorded before `connack` wakes a waiting run, so the run never reads a + # half-updated state if reason_code == 0: + state.refused_code = None + state.connected.set() + state.connack.set() logger.info("MQTT client connected") - userdata.set() - else: - logger.error("MQTT connection refused, result code %s", reason_code) - userdata.clear() + return + state.refused_code = reason_code + state.connected.clear() + state.connack.set() + logger.error( + "MQTT connection refused: %s (code %s)", + mqtt.connack_string(reason_code), + reason_code, + ) + if reason_code in PERMANENT_CONNACK_CODES: + # paho would otherwise reconnect with the same credentials forever; + # disconnect() puts it in the disconnecting state, so its loop ends + # after this CONNACK instead of retrying + client.disconnect() -def mqtt_on_connect_fail(client, userdata): +def mqtt_on_connect_fail(client, state: MQTTWriterState): # paho 1.6.1 invokes this callback with exactly (client, userdata) logger.error("MQTT client failed to establish connection with broker") - userdata.clear() + state.connected.clear() -def mqtt_on_disconnect(client, userdata, reason_code, properties=None): +def mqtt_on_disconnect(client, state: MQTTWriterState, reason_code, properties=None): logger.info("MQTT client disconnected, result code %s", reason_code) - userdata.clear() + state.connected.clear() + if state.refused_code is None: + # a transport drop: the next run waits for the reconnect's CONNACK; after + # a refusal the answer stays visible so run() reports it without waiting + state.connack.clear() class BlockManifest(WorkflowBlockManifest): @@ -156,6 +261,32 @@ class BlockManifest(WorkflowBlockManifest): "requests; in streaming pipelines it stops processing.", examples=[False], ) + encryption: Literal["none", "tls"] = Field( + default="none", + description="Transport security for the broker connection. `tls` encrypts the " + "connection and verifies the broker's certificate; set port to the broker's " + "TLS port (usually 8883).", + examples=["none", "tls"], + json_schema_extra={ + "values_metadata": { + "none": {"name": "None", "description": "Plain TCP (default)."}, + "tls": { + "name": "TLS", + "description": "Encrypted, broker certificate verified.", + }, + }, + }, + ) + ca_certificate_path: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Path to a PEM CA bundle on the machine running inference, for " + "brokers whose certificate issuer is not in the system trust store. Requires " + "local file system access for Workflow blocks " + "(`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`); otherwise the run " + "reports an error. Used only when encryption is `tls`.", + examples=["/etc/ssl/certs/factory-ca.pem"], + json_schema_extra={"relevant_for": TLS_RELEVANT}, + ) @classmethod def describe_outputs(cls) -> List[OutputDefinition]: @@ -170,9 +301,16 @@ def get_execution_engine_compatibility(cls) -> Optional[str]: class MQTTWriterSinkBlockV1(WorkflowBlock): - def __init__(self, disable_sinks: bool = False): + def __init__( + self, disable_sinks: bool = False, allow_access_to_file_system: bool = False + ): + # gates `ca_certificate_path`, a server-side path chosen by the workflow; + # False by default so a hand-constructed block is safe + self._allow_access_to_file_system = allow_access_to_file_system self.mqtt_client: Optional[mqtt.Client] = None - self._connected = threading.Event() + self._connection = MQTTWriterState() + # the same Event object as the state's, kept under its historical name + self._connected = self._connection.connected self._connection_identity: Optional[Tuple] = None self._lifecycle_lock = threading.Lock() self._disable_sinks = disable_sinks @@ -194,8 +332,9 @@ def close(self) -> None: # accepted so the thread never outlives the block client.loop_stop() finally: - # only after the join can no callback re-set the event - self._connected.clear() + # only after the join can no callback re-set the events; a + # client rebuilt by a later run starts without a stale refusal + self._connection.reset() def __del__(self): try: @@ -209,7 +348,7 @@ def get_manifest(cls) -> Type[WorkflowBlockManifest]: @classmethod def get_init_parameters(cls) -> List[str]: - return ["disable_sinks"] + return ["disable_sinks", "allow_access_to_file_system"] def run( self, @@ -223,6 +362,8 @@ def run( retain: bool = False, timeout: float = 0.5, fail_fast: bool = False, + encryption: str = "none", + ca_certificate_path: Optional[str] = None, ) -> BlockResult: if self._disable_sinks: return disabled_sink_response() @@ -278,6 +419,26 @@ def run( "Password provided without username. Set username to enable MQTT authentication.", fail_fast=fail_fast, ) + if encryption not in ("none", "tls"): + return self._handle_failure( + f"Invalid encryption: {encryption!r}. Must be 'none' or 'tls'.", + fail_fast=fail_fast, + ) + # an empty editor field is "unset"; anything but a string is a wiring error + if ca_certificate_path is not None and not isinstance(ca_certificate_path, str): + return self._handle_failure( + f"Invalid ca_certificate_path: {ca_certificate_path!r}. Must be a path.", + fail_fast=fail_fast, + ) + ca_certificate_path = ca_certificate_path or None + try: + # the operator's broker policy; the RESOLVED address is what the + # block connects to and what the connection identity is built from + host, port = resolve_broker_address( + host, port, log_override=self.mqtt_client is None + ) + except ConfigurationError as e: + return self._handle_failure(str(e), fail_fast=fail_fast) with self._lifecycle_lock: return self._connect_and_publish( host=host, @@ -290,6 +451,8 @@ def run( retain=retain, timeout=timeout, fail_fast=fail_fast, + encryption=encryption, + ca_certificate_path=ca_certificate_path, ) def _connect_and_publish( @@ -304,14 +467,32 @@ def _connect_and_publish( retain: bool, timeout: float, fail_fast: bool, + encryption: str = "none", + ca_certificate_path: Optional[str] = None, ) -> BlockResult: - connection_identity = (host, port, username, password, timeout) + connection_identity = ( + host, + port, + username, + password, + timeout, + encryption, + ca_certificate_path, + ) if self.mqtt_client is None: client = None try: - client = mqtt.Client(userdata=self._connected) + client = mqtt.Client(userdata=self._connection) if username is not None: client.username_pw_set(username, password) + # TLS must be configured before connect(); the CA path is gated + # by the engine's file-system permission + configure_tls( + client, + use_tls=encryption == "tls", + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=self._allow_access_to_file_system, + ) client.on_connect = mqtt_on_connect client.on_connect_fail = mqtt_on_connect_fail client.on_disconnect = mqtt_on_disconnect @@ -328,13 +509,15 @@ def _connect_and_publish( # timeout); the CONNACK wait below covers the handshake rest client.connect(host, port) client.loop_start() + except ConfigurationError as e: + # TLS setup refused before any socket was opened: nothing kept + return self._handle_failure(str(e), fail_fast=fail_fast) except OSError as e: # broker unreachable: nothing is kept and no loop was started, # so a failed one-shot run leaves no background thread behind; # the next run on this instance retries with a fresh client return self._handle_failure( - f"MQTT broker not connected ({e}). Raise 'timeout' if the " - "broker needs longer to connect.", + _not_connected_message(e, encryption=encryption), fail_fast=fail_fast, ) except Exception as e: @@ -350,19 +533,47 @@ def _connect_and_publish( self._connection_identity = connection_identity elif connection_identity != self._connection_identity: return self._handle_failure( - "MQTT connection parameters (host, port, credentials or timeout) changed " - "between runs; this block publishes only to the connection configured " - "on its first run.", + "MQTT connection parameters (host, port, credentials, timeout, encryption " + "or ca_certificate_path) changed between runs; this block publishes only " + "to the connection configured on its first run.", fail_fast=fail_fast, ) - # the background loop owns (re)connecting; runs only wait for readiness - if not self._connected.wait(timeout=timeout): + state = self._connection + # a permanent refusal stopped the network loop: report it on every run + # without waiting and without touching the broker again + if state.refused_code in PERMANENT_CONNACK_CODES: return self._handle_failure( - "MQTT broker not connected (connection was not established within timeout); " - "the client keeps retrying in the background. Raise 'timeout' if the " - "broker needs longer to connect.", + connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ), fail_fast=fail_fast, ) + # the background loop owns (re)connecting; runs only wait for readiness. + # The wait ends on any CONNACK, so a refusal is reported as soon as the + # broker answers instead of after the timeout + if not state.connected.is_set(): + if not state.connack.wait(timeout=timeout): + return self._handle_failure( + NOT_CONNECTED_WITHIN_TIMEOUT, fail_fast=fail_fast + ) + if state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ), + fail_fast=fail_fast, + ) + if not state.connected.wait(timeout=timeout): + if state.refused_code == TRANSIENT_CONNACK_CODE: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ), + fail_fast=fail_fast, + ) + return self._handle_failure( + NOT_CONNECTED_WITHIN_TIMEOUT, fail_fast=fail_fast + ) try: res: mqtt.MQTTMessageInfo = self.mqtt_client.publish( topic, message, qos=qos, retain=retain diff --git a/workflows/roboflow_workflows/environment.py b/workflows/roboflow_workflows/environment.py index 4ac92fbe4d..51ffc8c440 100644 --- a/workflows/roboflow_workflows/environment.py +++ b/workflows/roboflow_workflows/environment.py @@ -68,6 +68,16 @@ if _CONFIGURATION.engine.kafka_sinks_whitelisted_bootstrap_servers is None else list(_CONFIGURATION.engine.kafka_sinks_whitelisted_bootstrap_servers) ) +MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST = ( + _CONFIGURATION.engine.allow_mqtt_blocks_user_provided_host +) +# A list, not a set: the first entry is the broker the MQTT blocks connect to +# when user-provided hosts are not allowed. +MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS = ( + None + if _CONFIGURATION.engine.mqtt_blocks_whitelisted_hosts is None + else list(_CONFIGURATION.engine.mqtt_blocks_whitelisted_hosts) +) # --- tensor representation --- ENABLE_TENSOR_DATA_REPRESENTATION = _CONFIGURATION.tensor.representation_enabled diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py new file mode 100644 index 0000000000..f11a59fd3e --- /dev/null +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py @@ -0,0 +1,329 @@ +import logging +import os +from unittest.mock import MagicMock + +import pytest +from roboflow_workflows.enterprise_blocks.sinks import mqtt_common +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + DEFAULT_MQTT_PORT, + MQTT_KEEPALIVE_SECONDS, + PERMANENT_CONNACK_CODES, + ConfigurationError, + configure_tls, + connection_refused_message, + normalise_broker_address, + normalise_client_id, + resolve_broker_address, + split_host_port, +) + +ALLOW = "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST" +ALLOWLIST = "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS" + + +@pytest.fixture +def policy(monkeypatch): + def configure(allow: bool = True, allowlist=None): + monkeypatch.setattr(mqtt_common, ALLOW, allow) + monkeypatch.setattr(mqtt_common, ALLOWLIST, allowlist) + + configure() + return configure + + +@pytest.mark.parametrize( + "entry, expected", + [ + ("broker.local", ("broker.local", None)), + (" Broker.LOCAL ", ("broker.local", None)), + ("broker.local:1883", ("broker.local", 1883)), + ("Broker.LOCAL:8883 ", ("broker.local", 8883)), + ("10.0.0.5", ("10.0.0.5", None)), + ("[::1]", ("::1", None)), + ("[::1]:8883", ("::1", 8883)), + ("fe80::1", ("fe80::1", None)), + ("broker.local:notaport", ("broker.local:notaport", None)), + ], +) +def test_split_host_port(entry, expected): + assert split_host_port(entry) == expected + + +@pytest.mark.parametrize( + "host, port, expected", + [ + ("Broker.LOCAL", 1883, "broker.local:1883"), + ("[::1]", 8883, "::1:8883"), + (" 10.0.0.5 ", "1883", "10.0.0.5:1883"), + ], +) +def test_normalise_broker_address(host, port, expected): + assert normalise_broker_address(host, port) == expected + + +@pytest.mark.parametrize( + "value, expected", + [ + (None, None), + ("", None), + (" ", None), + ("\t\n", None), + ("line1", "line1"), + (" line1-camera3 ", "line1-camera3"), + ], +) +def test_normalise_client_id_treats_blank_as_unset_and_strips(value, expected): + assert normalise_client_id(value) == expected + + +@pytest.mark.parametrize("value", [5, 1.5, True, ["line1"], {"id": "line1"}]) +def test_normalise_client_id_rejects_non_strings_naming_the_field(value): + with pytest.raises(ConfigurationError, match="client_id"): + normalise_client_id(value) + + +def test_permanent_connack_codes_exclude_success_and_broker_unavailable(): + assert PERMANENT_CONNACK_CODES == {1, 2, 4, 5} + assert 0 not in PERMANENT_CONNACK_CODES + assert 3 not in PERMANENT_CONNACK_CODES + + +@pytest.mark.parametrize( + "code, reason", + [ + (1, "unacceptable protocol version"), + (2, "identifier rejected"), + (4, "bad user name or password"), + (5, "not authorised"), + ], +) +def test_permanent_refusal_message_names_reason_inputs_and_no_retry(code, reason): + message = connection_refused_message(code, block_inputs="username and password") + + assert reason in message + assert f"code {code}" in message + assert "Check username and password" in message + assert "does not retry" in message + assert "Raise 'timeout'" not in message + + +def test_broker_unavailable_message_says_retrying(): + message = connection_refused_message(3, block_inputs="username and password") + + assert "broker unavailable" in message + assert "code 3" in message + assert "retrying in the background" in message + assert "does not retry" not in message + + +def test_unknown_refusal_code_uses_permanent_wording(): + message = connection_refused_message(7, block_inputs="username and password") + + assert "unknown reason" in message + assert "code 7" in message + assert "does not retry" in message + + +def test_keepalive_is_shorter_than_paho_default(): + assert 0 < MQTT_KEEPALIVE_SECONDS < 60 + + +def test_default_policy_passes_user_value_through_unchanged(policy): + assert resolve_broker_address("Broker.LOCAL", 1883) == ("Broker.LOCAL", 1883) + + +@pytest.mark.parametrize( + "host, port", + [ + ("broker.local", 1883), + ("Broker.LOCAL", 1883), + ("[::1]", 8883), + ("10.0.0.5", 1883), + ("10.0.0.5", 8883), + ], +) +def test_allowlist_accepts_listed_host_and_port_or_portless_host(policy, host, port): + policy(allowlist=["broker.local:1883", "[::1]:8883", "10.0.0.5"]) + + assert resolve_broker_address(host, port) == (host, port) + + +@pytest.mark.parametrize( + "host, port", + [ + ("other.local", 1883), + ("broker.local", 8883), + ("[::1]", 1883), + ], +) +def test_allowlist_rejects_unlisted_address_without_revealing_the_allowlist( + policy, host, port +): + policy(allowlist=["broker.local:1883", "[::1]:8883", "10.0.0.5"]) + + with pytest.raises(ConfigurationError) as error: + resolve_broker_address(host, port) + + message = str(error.value) + assert "not permitted" in message + assert "10.0.0.5" not in message + assert "broker.local" not in message or host == "broker.local" + + +def test_empty_allowlist_permits_nothing(policy): + policy(allowlist=[]) + + with pytest.raises(ConfigurationError): + resolve_broker_address("broker.local", 1883) + + +def test_allowlist_ignores_blank_entries(policy): + policy(allowlist=[" ", "broker.local", ""]) + + assert resolve_broker_address("broker.local", 1883) == ("broker.local", 1883) + + +def test_user_host_not_allowed_uses_first_operator_entry(policy, caplog): + policy(allow=False, allowlist=["Operator.Broker:8883", "other:1883"]) + + with caplog.at_level(logging.WARNING, logger="inference"): + resolved = resolve_broker_address("workflow.host", 1883) + + assert resolved == ("operator.broker", 8883) + assert "replaced by the operator-configured broker" in caplog.text + + +def test_user_host_not_allowed_defaults_port(policy): + policy(allow=False, allowlist=["operator.broker"]) + + assert resolve_broker_address("workflow.host", 1) == ( + "operator.broker", + DEFAULT_MQTT_PORT, + ) + + +def test_override_warning_can_be_silenced(policy, caplog): + policy(allow=False, allowlist=["operator.broker"]) + + with caplog.at_level(logging.WARNING, logger="inference"): + resolve_broker_address("workflow.host", 1883, log_override=False) + + assert caplog.text == "" + + +def test_user_host_not_allowed_without_operator_broker_is_disabled(policy): + policy(allow=False, allowlist=None) + + with pytest.raises(ConfigurationError, match="disabled"): + resolve_broker_address("workflow.host", 1883) + + +class TestConfigureTLS: + @pytest.mark.parametrize("ca_certificate_path", [None, "", "/ca.pem"]) + @pytest.mark.parametrize("allowed", [True, False]) + def test_disabled_tls_touches_nothing(self, ca_certificate_path, allowed): + client = MagicMock() + + configure_tls( + client, + use_tls=False, + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=allowed, + ) + + assert client.method_calls == [] + + @pytest.mark.parametrize("ca_certificate_path", [None, ""]) + def test_tls_without_ca_path_uses_system_trust_store(self, ca_certificate_path): + client = MagicMock() + + configure_tls( + client, + use_tls=True, + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=False, + ) + + client.tls_set.assert_called_once_with() + client.tls_insecure_set.assert_not_called() + + def test_ca_path_used_with_file_system_access(self, tmp_path): + client = MagicMock() + ca_path = tmp_path / "factory-ca.pem" + ca_path.write_text("content is irrelevant: paho is faked here") + + configure_tls( + client, + use_tls=True, + ca_certificate_path=str(ca_path), + allow_access_to_file_system=True, + ) + + client.tls_set.assert_called_once_with(ca_certs=str(ca_path)) + client.tls_insecure_set.assert_not_called() + + def test_ca_path_refused_without_file_system_access(self): + client = MagicMock() + + with pytest.raises(ConfigurationError) as error: + configure_tls( + client, + use_tls=True, + ca_certificate_path="/etc/passwd", + allow_access_to_file_system=False, + ) + + assert "ca_certificate_path" in str(error.value) + assert "ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE" in str(error.value) + assert client.method_calls == [] + + @pytest.mark.parametrize("kind", ["missing", "directory", "fifo"]) + def test_ca_path_must_be_a_regular_file(self, tmp_path, kind): + # OpenSSL's loader blocks forever on a FIFO; the guard runs before it + client = MagicMock() + if kind == "missing": + path = tmp_path / "missing.pem" + elif kind == "directory": + path = tmp_path + else: + path = tmp_path / "pipe" + os.mkfifo(path) + + with pytest.raises(ConfigurationError, match="is not a readable file"): + configure_tls( + client, + use_tls=True, + ca_certificate_path=str(path), + allow_access_to_file_system=True, + ) + + client.tls_set.assert_not_called() + + def test_unloadable_ca_bundle_reported(self, tmp_path): + client = MagicMock() + client.tls_set.side_effect = ValueError("no certificate or crl found") + ca_path = tmp_path / "not-a-pem.pem" + ca_path.write_text("garbage") + + with pytest.raises(ConfigurationError, match="could not load CA bundle") as e: + configure_tls( + client, + use_tls=True, + ca_certificate_path=str(ca_path), + allow_access_to_file_system=True, + ) + + assert str(ca_path) in str(e.value) + assert "no certificate or crl found" in str(e.value) + + def test_system_store_failure_reported(self): + client = MagicMock() + client.tls_set.side_effect = ValueError("bad context") + + with pytest.raises(ConfigurationError, match="TLS could not be configured"): + configure_tls( + client, + use_tls=True, + ca_certificate_path=None, + allow_access_to_file_system=False, + ) diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py new file mode 100644 index 0000000000..8417fc8dda --- /dev/null +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -0,0 +1,1655 @@ +import threading +import time +from types import SimpleNamespace +from typing import List, get_args +from unittest.mock import patch + +import pytest +from pydantic import ValidationError +from roboflow_workflows.enterprise_blocks.sinks import mqtt_common +from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader import v1 +from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import ( + LATEST_BUFFER_SIZE, + NOT_CONNECTED_REPEATING_LAST, + NOT_CONNECTED_WITHIN_TIMEOUT, + NOT_RESUBSCRIBED_REPEATING_LAST, + SEQUENTIAL_BUFFER_SIZE, + SUBSCRIPTION_REFUSED_QOS, + BlockManifest, + MQTTReaderBlockV1, + MQTTReaderState, + mqtt_on_connect, + mqtt_on_connect_fail, + mqtt_on_disconnect, + mqtt_on_message, + mqtt_on_subscribe, +) + +CLIENT_CLASS_PATH = ( + "roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1.mqtt.Client" +) + + +def manifest_payload(**overrides) -> dict: + payload = { + "type": "roboflow_enterprise/mqtt_reader@v1", + "name": "reader", + "host": "localhost", + "port": 1883, + "topic": "plc/state", + } + payload.update(overrides) + return payload + + +def run_kwargs(**overrides) -> dict: + kwargs = { + "host": "localhost", + "port": 1883, + "topic": "plc/state", + "timeout": 0.05, + } + kwargs.update(overrides) + return kwargs + + +class FakeClient: + """Stand-in for paho's Client that drives the block's real callbacks. + + connect() succeeds unless `connect_error` is set, loop_start() fires + on_connect with `connack_reason_code`, subscribe() answers with + `granted_qos` unless `ack_subscriptions` is False, and deliver() + pushes a message through on_message. + """ + + connect_error = None + connack_reason_code = 0 + granted_qos = 0 + ack_subscriptions = True + fire_on_connect = True + # CONNACK flags handed to on_connect; a test sets {"session present": 1} + connect_flags: dict = {} + + def __init__(self, client_id="", clean_session=None, userdata=None): + # paho 1.6.1's constructor arguments the block chooses + self.client_id = client_id + self.clean_session = clean_session + self.userdata = userdata + # ordered record of the calls that matter for TLS: must precede connect() + self.calls = [] + self.on_connect = None + self.on_connect_fail = None + self.on_subscribe = None + self.on_message = None + self.on_disconnect = None + self.suppress_exceptions = False + self.credentials = None + self.connected_to = None + self.subscriptions: List[tuple] = [] + self.loop_started = False + self.loop_stopped = False + self.disconnected = False + self.reconnect_delays = None + + def username_pw_set(self, username, password=None): + self.credentials = (username, password) + + def reconnect_delay_set(self, min_delay, max_delay): + self.reconnect_delays = (min_delay, max_delay) + + def tls_set(self, ca_certs=None, **kwargs): + self.calls.append(("tls_set", ca_certs)) + + def tls_insecure_set(self, value): + self.calls.append(("tls_insecure_set", value)) + + def connect(self, host, port, keepalive=60): + self.calls.append(("connect", host, port)) + if self.connect_error is not None: + raise self.connect_error + self.connected_to = (host, port) + self.keepalive = keepalive + + def loop_start(self): + self.loop_started = True + if self.fire_on_connect: + self.on_connect( + self, self.userdata, dict(self.connect_flags), self.connack_reason_code + ) + + def loop_stop(self): + self.loop_stopped = True + + def disconnect(self): + self.disconnected = True + + def subscribe(self, topic, qos=0): + self.subscriptions.append((topic, qos)) + if self.ack_subscriptions: + self.on_subscribe(self, self.userdata, 1, [self.granted_qos]) + + def deliver(self, payload, topic="plc/state", retain=False): + message = SimpleNamespace(topic=topic, payload=payload, retain=retain) + self.on_message(self, self.userdata, message) + + +@pytest.fixture +def clients() -> List[FakeClient]: + created: List[FakeClient] = [] + + def factory(client_id="", clean_session=None, userdata=None): + client = FakeClient( + client_id=client_id, clean_session=clean_session, userdata=userdata + ) + created.append(client) + return client + + with patch(CLIENT_CLASS_PATH, side_effect=factory): + yield created + + +@pytest.fixture +def block() -> MQTTReaderBlockV1: + return MQTTReaderBlockV1() + + +class TestManifest: + def test_primary_identifier_is_namespaced(self): + identifiers = get_args(BlockManifest.model_fields["type"].annotation) + + assert identifiers == ("roboflow_enterprise/mqtt_reader@v1",) + + def test_defaults(self): + manifest = BlockManifest.model_validate(manifest_payload()) + + assert manifest.read_mode == "latest" + assert manifest.qos == 0 + assert manifest.timeout == 0.5 + assert manifest.username is None + assert manifest.password is None + assert manifest.client_id is None + + @pytest.mark.parametrize("client_id", ["line1-camera3", "$inputs.pipeline_name"]) + def test_client_id_accepts_literal_and_selector(self, client_id): + manifest = BlockManifest.model_validate(manifest_payload(client_id=client_id)) + + assert manifest.client_id == client_id + + def test_timeout_description_scopes_the_wait_to_the_first_connection(self): + schema = BlockManifest.model_json_schema() + + description = schema["properties"]["timeout"]["description"] + assert "first connection" in description + assert "Later runs never wait" in description + + def test_outputs_unchanged_by_outage_handling(self): + names = [output.name for output in BlockManifest.describe_outputs()] + + assert "error_status" in names + assert "error_message" in names + assert len(names) == 6 + + def test_qos_description_states_the_persistent_session_rule(self): + schema = BlockManifest.model_json_schema() + + assert "client_id" in schema["properties"]["qos"]["description"] + assert "1 or 2" in schema["properties"]["qos"]["description"] + + @pytest.mark.parametrize("read_mode", ["everything", "next", "", None]) + def test_read_mode_must_be_latest_or_sequential(self, read_mode): + with pytest.raises(ValidationError): + BlockManifest.model_validate(manifest_payload(read_mode=read_mode)) + + @pytest.mark.parametrize("read_mode", ["latest", "sequential"]) + def test_valid_read_modes_accepted(self, read_mode): + manifest = BlockManifest.model_validate(manifest_payload(read_mode=read_mode)) + + assert manifest.read_mode == read_mode + + @pytest.mark.parametrize("port", [0, -1, 65536]) + def test_literal_port_must_be_within_tcp_range(self, port): + with pytest.raises(ValidationError): + BlockManifest.model_validate(manifest_payload(port=port)) + + @pytest.mark.parametrize( + "timeout", [0, -1, float("nan"), float("inf"), float("-inf")] + ) + def test_literal_timeout_must_be_finite_and_positive(self, timeout): + with pytest.raises(ValidationError): + BlockManifest.model_validate(manifest_payload(timeout=timeout)) + + def test_selectors_accepted_for_connection_fields(self): + manifest = BlockManifest.model_validate( + manifest_payload( + host="$inputs.host", + port="$inputs.port", + topic="$inputs.topic", + username="$inputs.user", + password="$inputs.secret", + timeout="$inputs.timeout", + qos="$inputs.qos", + ) + ) + + assert manifest.password == "$inputs.secret" + + def test_password_is_marked_private(self): + schema = BlockManifest.model_json_schema() + + assert schema["properties"]["password"]["private"] is True + + def test_describe_outputs_order(self): + names = [output.name for output in BlockManifest.describe_outputs()] + + assert names == [ + "value", + "payload", + "topic", + "is_new", + "error_status", + "error_message", + ] + + def test_restrictions_hard_for_hosted_serverless_only(self): + restrictions = BlockManifest.get_restrictions() + + assert [r.severity.value for r in restrictions] == ["hard", "soft"] + assert restrictions[0].applies_to_runtimes == [v1.Runtime.HOSTED_SERVERLESS] + assert restrictions[1].applies_to_runtimes == [ + v1.Runtime.SELF_HOSTED_CPU, + v1.Runtime.SELF_HOSTED_GPU, + v1.Runtime.DEDICATED_DEPLOYMENT, + ] + assert restrictions[1].applies_to_input_modes == [v1.RuntimeInputMode.IMAGE] + # the HTTP note must warn about persistent sessions, not only retained messages + assert "client_id" in restrictions[1].note + assert "InferencePipeline" in restrictions[1].note + + def test_block_declares_file_system_init_parameter_only(self): + assert MQTTReaderBlockV1.get_init_parameters() == [ + "allow_access_to_file_system" + ] + + +class TestCallbacks: + def test_state_buffer_sizes_follow_read_mode_constants(self): + assert LATEST_BUFFER_SIZE == 1 + assert SEQUENTIAL_BUFFER_SIZE == 1000 + assert MQTTReaderState("t", 0, LATEST_BUFFER_SIZE).messages.maxlen == 1 + + def test_on_connect_subscribes_and_sets_event_for_accepted_connack(self): + state = MQTTReaderState(topic="plc/#", qos=1, buffer_size=1) + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + mqtt_on_connect(client, state, {}, 0) + + assert client.subscriptions == [("plc/#", 1)] + assert state.connected.is_set() + + @pytest.mark.parametrize( + "flags, expected", + [ + ({}, "False"), + ({"session present": 0}, "False"), + ({"session present": 1}, "True"), + ], + ) + def test_on_connect_logs_session_present_flag(self, caplog, flags, expected): + state = MQTTReaderState(topic="plc/#", qos=1, buffer_size=1) + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + with caplog.at_level("INFO", logger="inference"): + mqtt_on_connect(client, state, flags, 0) + + assert f"session present: {expected}" in caplog.text + + def test_on_connect_resubscribes_on_every_connect(self): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + mqtt_on_connect(client, state, {}, 0) + mqtt_on_disconnect(client, state, 1) + mqtt_on_connect(client, state, {}, 0) + + assert client.subscriptions == [("plc/#", 0), ("plc/#", 0)] + assert state.subscribed.is_set() + + @pytest.mark.parametrize("reason_code", [1, 2, 3, 4, 5]) + def test_on_connect_does_not_subscribe_for_rejected_connack(self, reason_code): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + state.connected.set() + + mqtt_on_connect(client, state, {}, reason_code) + + assert client.subscriptions == [] + assert not state.connected.is_set() + assert state.connack.is_set() + assert state.refused_code == reason_code + + @pytest.mark.parametrize("reason_code", [1, 2, 4, 5]) + def test_on_connect_stops_the_loop_for_a_permanent_refusal(self, reason_code): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + + mqtt_on_connect(client, state, {}, reason_code) + + assert client.disconnected is True + + def test_on_connect_keeps_the_loop_for_broker_unavailable(self): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + + mqtt_on_connect(client, state, {}, 3) + + assert client.disconnected is False + assert state.refused_code == 3 + + @pytest.mark.parametrize("reason_code", [0, 3, 5]) + def test_on_connect_completes_the_state_before_waking_a_waiter(self, reason_code): + # a run wakes on `connack`; what it then reads must already be final + state = MQTTReaderState(topic="plc/#", qos=1, buffer_size=1) + seen = {} + + class RecordingEvent(threading.Event): + def set(self): + seen["refused_code"] = state.refused_code + seen["connected"] = state.connected.is_set() + super().set() + + state.connack = RecordingEvent() + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + mqtt_on_connect(client, state, {}, reason_code) + + assert seen == { + "refused_code": reason_code or None, + "connected": reason_code == 0, + } + + def test_on_connect_logs_the_refusal_reason(self, caplog): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + + with caplog.at_level("ERROR", logger="inference"): + mqtt_on_connect(FakeClient(userdata=state), state, {}, 5) + + assert "not authorised" in caplog.text + assert "code 5" in caplog.text + + def test_accepted_connack_after_broker_unavailable_resets_refusal(self): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + mqtt_on_connect(client, state, {}, 3) + mqtt_on_connect(client, state, {}, 0) + + assert state.refused_code is None + assert state.connected.is_set() + assert state.connack.is_set() + + def test_on_connect_fail_matches_paho_two_argument_signature(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + state.connected.set() + + mqtt_on_connect_fail(FakeClient(), state) + + assert not state.connected.is_set() + + def test_on_subscribe_flags_refused_subscription(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + + mqtt_on_subscribe(FakeClient(), state, 1, [SUBSCRIPTION_REFUSED_QOS]) + + assert state.subscribe_failed is True + assert not state.subscribed.is_set() + + @pytest.mark.parametrize("granted", [0, 1, 2]) + def test_on_subscribe_sets_event_for_granted_qos(self, granted): + state = MQTTReaderState(topic="t", qos=2, buffer_size=1) + + mqtt_on_subscribe(FakeClient(), state, 1, [granted]) + + assert state.subscribed.is_set() + assert state.subscribe_failed is False + + def test_on_message_appends_raw_tuple_only(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=10) + message = SimpleNamespace(topic="plc/a", payload=b"\xff\xfe", retain=1) + + mqtt_on_message(FakeClient(), state, message) + + assert list(state.messages) == [("plc/a", b"\xff\xfe", True)] + + def test_latest_buffer_keeps_only_newest(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=LATEST_BUFFER_SIZE) + for payload in (b"A", b"B", b"C"): + mqtt_on_message( + FakeClient(), + state, + SimpleNamespace(topic="t", payload=payload, retain=0), + ) + + assert list(state.messages) == [("t", b"C", False)] + + def test_on_disconnect_clears_both_events(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + state.connected.set() + state.subscribed.set() + + mqtt_on_disconnect(FakeClient(), state, 1) + + assert not state.connected.is_set() + assert not state.subscribed.is_set() + + def test_on_disconnect_after_transport_drop_clears_connack(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + state.connack.set() + + mqtt_on_disconnect(FakeClient(), state, 1) + + assert not state.connack.is_set() + + def test_on_disconnect_after_refusal_keeps_connack(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + mqtt_on_connect(client, state, {}, 5) + + mqtt_on_disconnect(client, state, 0) + + assert state.connack.is_set() + assert state.refused_code == 5 + + def test_on_disconnect_accepts_mqtt5_properties_argument(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + state.connected.set() + + mqtt_on_disconnect(FakeClient(), state, 0, None) + + assert not state.connected.is_set() + + +class TestInputValidation: + @pytest.mark.parametrize( + "timeout", [0, -1, float("nan"), float("inf"), "abc", None] + ) + def test_invalid_timeout_rejected_before_client_construction( + self, clients, block, timeout + ): + result = block.run(**run_kwargs(timeout=timeout)) + + assert result["error_status"] is True + assert "timeout" in result["error_message"].lower() + assert clients == [] + + @pytest.mark.parametrize("port", [0, 65536, "abc", 1.5, True, None]) + def test_invalid_port_rejected_before_client_construction( + self, clients, block, port + ): + result = block.run(**run_kwargs(port=port)) + + assert result["error_status"] is True + assert "port" in result["error_message"].lower() + assert clients == [] + + def test_selector_supplied_port_coerced_like_manifest(self, clients, block): + result = block.run(**run_kwargs(port="1883")) + + assert result["error_status"] is False + assert clients[0].connected_to == ("localhost", 1883) + + @pytest.mark.parametrize("qos", [-1, 3, "x", 1.5, True, None]) + def test_invalid_qos_rejected_before_client_construction(self, clients, block, qos): + result = block.run(**run_kwargs(qos=qos)) + + assert result["error_status"] is True + assert "qos" in result["error_message"].lower() + assert clients == [] + + def test_selector_supplied_qos_used_for_subscription(self, clients, block): + block.run(**run_kwargs(qos="2")) + + assert clients[0].subscriptions == [("plc/state", 2)] + + def test_invalid_read_mode_rejected(self, clients, block): + result = block.run(**run_kwargs(read_mode="everything")) + + assert result["error_status"] is True + assert "read_mode" in result["error_message"] + assert clients == [] + + def test_password_without_username_rejected(self, clients, block): + result = block.run(**run_kwargs(password="secret")) + + assert result["error_status"] is True + assert "username" in result["error_message"].lower() + assert clients == [] + + def test_username_and_password_configure_authentication(self, clients, block): + block.run(**run_kwargs(username="user", password="secret")) + + assert clients[0].credentials == ("user", "secret") + + @pytest.mark.parametrize("flag", ["GCP_SERVERLESS", "LAMBDA"]) + def test_hosted_platform_rejected_without_client(self, clients, block, flag): + with patch.object(v1, flag, True): + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "hosted" in result["error_message"].lower() + assert clients == [] + + @pytest.mark.parametrize("topic", ["", None, 5]) + def test_invalid_topic_rejected_before_client_construction( + self, clients, block, topic + ): + result = block.run(**run_kwargs(topic=topic)) + + assert result["error_status"] is True + assert "topic" in result["error_message"].lower() + assert clients == [] + + @pytest.mark.parametrize("client_id", [None, "", " "]) + def test_blank_client_id_builds_a_clean_session_client( + self, clients, block, client_id + ): + result = block.run(**run_kwargs(client_id=client_id)) + + assert result["error_status"] is False + assert clients[0].client_id == "" + assert clients[0].clean_session is None + assert clients[0].userdata is block._state + + def test_client_id_builds_a_persistent_session_client(self, clients, block): + result = block.run(**run_kwargs(client_id=" line1 ", qos=1)) + + assert result["error_status"] is False + assert clients[0].client_id == "line1" + assert clients[0].clean_session is False + assert clients[0].userdata is block._state + assert clients[0].subscriptions == [("plc/state", 1)] + + def test_client_id_with_qos_zero_rejected_before_client_construction( + self, clients, block + ): + result = block.run(**run_kwargs(client_id="line1")) + + assert result["error_status"] is True + assert "client_id" in result["error_message"] + assert "qos 1 or 2" in result["error_message"] + assert clients == [] + + @pytest.mark.parametrize("client_id", [5, 1.5, True, ["line1"]]) + def test_non_string_client_id_rejected_before_client_construction( + self, clients, block, client_id + ): + result = block.run(**run_kwargs(client_id=client_id, qos=1)) + + assert result["error_status"] is True + assert "client_id" in result["error_message"] + assert clients == [] + + def test_error_outputs_carry_every_declared_key(self, clients, block): + result = block.run(**run_kwargs(port=0)) + + assert result == { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": True, + "error_message": result["error_message"], + } + + +class TestConnectionLifecycle: + def test_first_run_connects_subscribes_and_starts_loop(self, clients, block): + result = block.run(**run_kwargs(username="u", password="p", timeout=0.25)) + + client = clients[0] + assert result["error_status"] is False + assert client.connected_to == ("localhost", 1883) + assert client.loop_started is True + assert client.suppress_exceptions is True + assert client.subscriptions == [("plc/state", 0)] + assert client.on_connect is mqtt_on_connect + assert client.on_connect_fail is mqtt_on_connect_fail + assert client.on_subscribe is mqtt_on_subscribe + assert client.on_message is mqtt_on_message + assert client.on_disconnect is mqtt_on_disconnect + assert client.reconnect_delays == (0.125, 0.5) + assert client._connect_timeout == 0.25 + assert client.keepalive == 15 + + def test_client_reused_across_runs(self, clients, block): + block.run(**run_kwargs()) + block.run(**run_kwargs()) + + assert len(clients) == 1 + + def test_unreachable_broker_on_first_run_leaves_nothing_behind( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connect_error", OSError("connection refused")) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not connected" in result["error_message"].lower() + assert block._client is None + assert clients[0].loop_started is False + + def test_next_run_after_failed_connect_retries_with_fresh_client( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connect_error", OSError("connection refused")) + block.run(**run_kwargs()) + + monkeypatch.setattr(FakeClient, "connect_error", None) + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert len(clients) == 2 + + def test_setup_failure_stops_loop_and_resets_client(self, clients, block): + with patch.object( + FakeClient, "reconnect_delay_set", side_effect=RuntimeError("boom") + ): + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "initialize" in result["error_message"].lower() + assert block._client is None + assert clients[0].loop_stopped is True + + def test_connection_not_established_within_timeout_reported( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "fire_on_connect", False) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not connected" in result["error_message"].lower() + assert block._client is clients[0] + + def test_rejected_connack_reported_as_a_refusal(self, clients, block, monkeypatch): + monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not authorised" in result["error_message"] + assert "code 5" in result["error_message"] + assert "client_id" in result["error_message"] + assert "Raise 'timeout'" not in result["error_message"] + assert clients[0].subscriptions == [] + assert clients[0].disconnected is True + + @pytest.mark.parametrize( + "reason_code, reason", + [ + (1, "unacceptable protocol version"), + (2, "identifier rejected"), + (4, "bad user name or password"), + ], + ) + def test_other_permanent_refusals_name_their_reason( + self, clients, block, monkeypatch, reason_code, reason + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", reason_code) + result = block.run(**run_kwargs()) + + assert reason in result["error_message"] + assert clients[0].disconnected is True + + def test_refusal_is_reported_without_waiting_for_timeout( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert result["error_status"] is True + assert time.monotonic() - started < 0.5 + + def test_next_run_after_refusal_repeats_it_without_a_new_connection( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + first = block.run(**run_kwargs(timeout=2.0)) + + started = time.monotonic() + second = block.run(**run_kwargs(timeout=2.0)) + + assert second == first + assert time.monotonic() - started < 0.5 + assert len(clients) == 1 + assert [call for call in clients[0].calls if call[0] == "connect"] == [ + ("connect", "localhost", 1883) + ] + + def test_broker_unavailable_keeps_retrying_and_names_the_reason( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", 3) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "broker unavailable" in result["error_message"] + assert "retrying in the background" in result["error_message"] + assert clients[0].disconnected is False + + def test_run_succeeds_once_an_unavailable_broker_accepts( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", 3) + block.run(**run_kwargs()) + client = clients[0] + client.on_connect(client, client.userdata, {}, 0) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert block._state.refused_code is None + + def test_no_connack_at_all_reports_not_connected(self, clients, block, monkeypatch): + monkeypatch.setattr(FakeClient, "fire_on_connect", False) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert result["error_message"] == NOT_CONNECTED_WITHIN_TIMEOUT + + def test_close_after_refusal_completes(self, clients, block, monkeypatch): + monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + block.run(**run_kwargs()) + + block.close() + + assert block._client is None + assert clients[0].loop_stopped is True + + def test_refused_subscription_reported_and_client_kept( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "granted_qos", SUBSCRIPTION_REFUSED_QOS) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "refused" in result["error_message"].lower() + assert block._client is clients[0] + assert block._state.connected.is_set() + + def test_refused_subscription_does_not_wait_for_timeout( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "granted_qos", SUBSCRIPTION_REFUSED_QOS) + block.run(**run_kwargs(timeout=1.0)) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=1.0)) + + assert result["error_status"] is True + assert "refused" in result["error_message"].lower() + assert time.monotonic() - started < 0.5 + + def test_subscribe_error_reported_through_refused_path( + self, clients, block, monkeypatch + ): + def failing_subscribe(client, topic, qos=0): + raise ValueError("Invalid topic.") + + monkeypatch.setattr(FakeClient, "subscribe", failing_subscribe) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "refused" in result["error_message"].lower() + assert block._state.connected.is_set() + + def test_unacknowledged_subscription_reported(self, clients, block, monkeypatch): + monkeypatch.setattr(FakeClient, "ack_subscriptions", False) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "acknowledged" in result["error_message"].lower() + + @pytest.mark.parametrize( + "change", + [ + {"host": "other"}, + {"port": 8883}, + {"username": "u"}, + {"username": "u", "password": "p"}, + {"timeout": 0.1}, + {"topic": "plc/other"}, + {"qos": 1}, + {"read_mode": "sequential"}, + {"qos": 1, "client_id": "line1"}, + ], + ) + def test_changed_parameters_rejected_and_connection_kept( + self, clients, block, change + ): + block.run(**run_kwargs()) + + result = block.run(**run_kwargs(**change)) + + assert result["error_status"] is True + assert "changed between runs" in result["error_message"] + assert len(clients) == 1 + assert clients[0].disconnected is False + + @pytest.mark.parametrize("later_client_id", [None, "", "line2"]) + def test_dropped_or_changed_client_id_rejected_and_connection_kept( + self, clients, block, later_client_id + ): + block.run(**run_kwargs(client_id="line1", qos=1)) + + result = block.run(**run_kwargs(client_id=later_client_id, qos=1)) + + assert result["error_status"] is True + assert "client_id" in result["error_message"] + assert "changed between runs" in result["error_message"] + assert len(clients) == 1 + assert clients[0].disconnected is False + + def test_same_client_id_with_different_whitespace_is_not_a_change( + self, clients, block + ): + block.run(**run_kwargs(client_id="line1", qos=1)) + + result = block.run(**run_kwargs(client_id=" line1 ", qos=1)) + + assert result["error_status"] is False + assert len(clients) == 1 + + def test_close_disconnects_stops_loop_and_clears_state(self, clients, block): + block.run(**run_kwargs()) + state = block._state + + block.close() + + assert clients[0].disconnected is True + assert clients[0].loop_stopped is True + assert block._client is None + assert block._state is None + assert not state.connected.is_set() + assert not state.subscribed.is_set() + + def test_close_resets_last_outputs(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b"x") + block.run(**run_kwargs()) + + block.close() + result = block.run(**run_kwargs()) + + assert len(clients) == 2 + assert result["value"] is None + assert result["is_new"] is False + + def test_close_without_client_is_noop(self, block): + block.close() + + assert block._client is None + + def test_close_when_disconnect_raises_still_stops_loop(self, clients, block): + block.run(**run_kwargs()) + with patch.object(FakeClient, "disconnect", side_effect=RuntimeError("boom")): + block.close() + + assert clients[0].loop_stopped is True + assert block._client is None + + +class TestReading: + def test_empty_first_run(self, clients, block): + result = block.run(**run_kwargs()) + + assert result == { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": False, + "error_message": None, + } + + def test_json_message_surfaced_with_parsed_payload(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}') + + result = block.run(**run_kwargs()) + + assert result == { + "value": '{"state": "RUNNING"}', + "payload": {"state": "RUNNING"}, + "topic": "plc/state", + "is_new": True, + "error_status": False, + "error_message": None, + } + + def test_retained_message_delivered_before_first_run_is_surfaced( + self, clients, block + ): + original_loop_start = FakeClient.loop_start + + def loop_start_with_retained(client): + original_loop_start(client) + client.deliver(b'{"state": "RUNNING"}', retain=True) + + with patch.object(FakeClient, "loop_start", loop_start_with_retained): + result = block.run(**run_kwargs()) + + assert result["value"] == '{"state": "RUNNING"}' + assert result["is_new"] is True + + def test_retained_redelivery_of_same_message_is_not_new(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}', retain=True) + first = block.run(**run_kwargs()) + # a reconnect resubscribes and the broker re-sends the retained copy + clients[0].deliver(b'{"state": "RUNNING"}', retain=True) + + second = block.run(**run_kwargs()) + + assert first["is_new"] is True + assert second == {**first, "is_new": False} + + def test_retained_message_with_new_payload_is_new(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}', retain=True) + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "PAUSED"}', retain=True) + + result = block.run(**run_kwargs()) + + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + + def test_non_retained_duplicate_payload_is_new(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b"tick") + block.run(**run_kwargs()) + clients[0].deliver(b"tick") + + result = block.run(**run_kwargs()) + + assert result["is_new"] is True + + def test_nothing_new_repeats_previous_outputs_with_is_new_false( + self, clients, block + ): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}', topic="plc/line1") + first = block.run(**run_kwargs()) + + second = block.run(**run_kwargs()) + + assert first["is_new"] is True + assert second == {**first, "is_new": False} + + def test_latest_mode_returns_newest_and_skips_older(self, clients, block): + block.run(**run_kwargs()) + for payload in (b"A", b"B", b"C"): + clients[0].deliver(payload) + + result = block.run(**run_kwargs()) + repeat = block.run(**run_kwargs()) + + assert result["value"] == "C" + assert result["is_new"] is True + assert repeat["value"] == "C" + assert repeat["is_new"] is False + + def test_sequential_mode_returns_messages_in_order_one_per_run( + self, clients, block + ): + block.run(**run_kwargs(read_mode="sequential")) + for payload in (b"A", b"B", b"C"): + clients[0].deliver(payload) + + results = [block.run(**run_kwargs(read_mode="sequential")) for _ in range(4)] + + assert [r["value"] for r in results] == ["A", "B", "C", "C"] + assert [r["is_new"] for r in results] == [True, True, True, False] + + def test_sequential_mode_drops_oldest_beyond_buffer(self, clients, block): + block.run(**run_kwargs(read_mode="sequential")) + for index in range(SEQUENTIAL_BUFFER_SIZE + 1): + clients[0].deliver(str(index).encode()) + + result = block.run(**run_kwargs(read_mode="sequential")) + + assert result["value"] == "1" + assert len(block._state.messages) == SEQUENTIAL_BUFFER_SIZE - 1 + + def test_topic_output_reports_message_topic_for_wildcard_filters( + self, clients, block + ): + block.run(**run_kwargs(topic="plc/#")) + clients[0].deliver(b"1", topic="plc/line2/state") + + result = block.run(**run_kwargs(topic="plc/#")) + + assert result["topic"] == "plc/line2/state" + + def test_non_json_payload_gives_text_and_empty_payload(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b"hello") + + result = block.run(**run_kwargs()) + + assert result["value"] == "hello" + assert result["payload"] == {} + assert result["error_status"] is False + + def test_json_non_object_gives_empty_payload(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b"[1, 2]") + + result = block.run(**run_kwargs()) + + assert result["value"] == "[1, 2]" + assert result["payload"] == {} + + def test_non_utf8_payload_reported_as_error_without_previous_outputs( + self, clients, block + ): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}') + block.run(**run_kwargs()) + clients[0].deliver(b"\xff\xfe") + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "utf-8" in result["error_message"].lower() + assert result["value"] is None + assert result["payload"] == {} + assert result["is_new"] is False + + def test_message_received_on_network_thread_visible_to_run(self, clients, block): + block.run(**run_kwargs()) + delivered = threading.Event() + + def deliver(): + clients[0].deliver(b"from-thread") + delivered.set() + + threading.Thread(target=deliver).start() + assert delivered.wait(timeout=1) + + result = block.run(**run_kwargs()) + + assert result["value"] == "from-thread" + + def test_no_raw_prints_in_any_code_path(self, clients, block): + with patch("builtins.print") as mock_print: + block.run(**run_kwargs()) + clients[0].deliver(b"x") + block.run(**run_kwargs()) + block.run(**run_kwargs(port=0)) + block.close() + + mock_print.assert_not_called() + + +class TestOutageBookkeeping: + """The two helpers behind the outage outputs, driven directly.""" + + def warm_block_with_last(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}') + block.run(**run_kwargs()) + return block + + def test_repeat_with_error_returns_last_message_with_error_fields( + self, clients, block + ): + self.warm_block_with_last(clients, block) + + result = block._repeat_last_with_error("boom") + + assert result == { + "value": '{"state": "RUNNING"}', + "payload": {"state": "RUNNING"}, + "topic": "plc/state", + "is_new": False, + "error_status": True, + "error_message": "boom", + } + + def test_repeat_with_error_without_last_message_is_the_empty_failure( + self, clients, block + ): + block.run(**run_kwargs()) + + result = block._repeat_last_with_error("boom") + + assert result == { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": True, + "error_message": "boom", + } + + def test_outage_logs_one_error_line_naming_the_instance( + self, clients, block, caplog + ): + self.warm_block_with_last(clients, block) + + with caplog.at_level("ERROR", logger="inference"): + block._repeat_last_with_error("boom") + block._repeat_last_with_error("boom") + + errors = [r for r in caplog.records if r.levelname == "ERROR"] + assert len(errors) == 1 + assert "MQTT Reader localhost:1883 'plc/state': boom" in errors[0].getMessage() + assert block._outage_runs == 2 + + def test_recovery_logs_one_info_line_with_duration_and_count_and_resets( + self, clients, block, caplog + ): + self.warm_block_with_last(clients, block) + block._repeat_last_with_error("boom") + block._repeat_last_with_error("boom") + + with caplog.at_level("INFO", logger="inference"): + block._log_recovery_if_needed() + block._log_recovery_if_needed() + + infos = [ + r + for r in caplog.records + if r.levelname == "INFO" and "restored after" in r.getMessage() + ] + assert len(infos) == 1 + assert "2 runs returned the last message" in infos[0].getMessage() + assert "localhost:1883 'plc/state'" in infos[0].getMessage() + assert block._outage_started_at is None + assert block._outage_runs == 0 + + def test_recovery_while_healthy_logs_nothing(self, clients, block, caplog): + self.warm_block_with_last(clients, block) + + with caplog.at_level("INFO", logger="inference"): + block._log_recovery_if_needed() + + assert "restored after" not in caplog.text + + def test_first_run_keeps_the_resolved_address_for_the_log(self, clients, block): + block.run(**run_kwargs()) + + assert (block._host, block._port) == ("localhost", 1883) + + def test_close_resets_the_bookkeeping(self, clients, block): + self.warm_block_with_last(clients, block) + block._repeat_last_with_error("boom") + block._subscribed_once = True + + block.close() + + assert block._subscribed_once is False + assert block._outage_started_at is None + assert block._outage_runs == 0 + assert (block._host, block._port) == (None, None) + + +def drop(client: FakeClient) -> None: + client.on_disconnect(client, client.userdata, 1) + + +def reconnect(client: FakeClient, reason_code: int = 0) -> None: + client.on_connect(client, client.userdata, {}, reason_code) + + +RUNNING = b'{"state": "RUNNING"}' +PAUSED = b'{"state": "PAUSED"}' + + +class TestOutage: + """A warm instance (first run reached SUBACK) while the broker is away.""" + + def warm(self, clients, block, **kwargs): + block.run(**run_kwargs(**kwargs)) + clients[0].deliver(RUNNING) + first = block.run(**run_kwargs(**kwargs)) + assert first["is_new"] is True + return clients[0], first + + def test_drop_repeats_last_message_with_error_without_waiting(self, clients, block): + client, first = self.warm(clients, block, timeout=2.0) + drop(client) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert time.monotonic() - started < 0.5 + assert result == { + **first, + "is_new": False, + "error_status": True, + "error_message": NOT_CONNECTED_REPEATING_LAST, + } + assert "not connected" in result["error_message"] + + def test_repeated_runs_during_outage_log_one_error_line( + self, clients, block, caplog + ): + client, _ = self.warm(clients, block) + drop(client) + + with caplog.at_level("ERROR", logger="inference"): + results = [block.run(**run_kwargs()) for _ in range(4)] + + assert len({str(r) for r in results}) == 1 + errors = [r for r in caplog.records if r.levelname == "ERROR"] + assert len(errors) == 1 + assert "localhost:1883 'plc/state'" in errors[0].getMessage() + + def test_buffered_messages_drain_before_the_outage_is_reported( + self, clients, block + ): + client, _ = self.warm(clients, block, read_mode="sequential") + client.deliver(b"A") + client.deliver(b"B") + drop(client) + + results = [block.run(**run_kwargs(read_mode="sequential")) for _ in range(3)] + + assert [r["value"] for r in results] == ["A", "B", "B"] + assert [r["is_new"] for r in results] == [True, True, False] + assert [r["error_status"] for r in results] == [False, False, True] + + def test_drop_before_any_message_returns_empty_failure_without_waiting( + self, clients, block + ): + block.run(**run_kwargs(timeout=2.0)) + drop(clients[0]) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert time.monotonic() - started < 0.5 + assert result["payload"] == {} + assert result["value"] is None + assert result["error_status"] is True + assert result["error_message"] == NOT_CONNECTED_REPEATING_LAST + + def test_recovery_returns_live_message_and_logs_once(self, clients, block, caplog): + client, _ = self.warm(clients, block) + drop(client) + for _ in range(3): + block.run(**run_kwargs()) + + reconnect(client) + client.deliver(PAUSED) + with caplog.at_level("INFO", logger="inference"): + result = block.run(**run_kwargs()) + block.run(**run_kwargs()) + + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + assert result["error_status"] is False + infos = [ + r.getMessage() for r in caplog.records if "restored after" in r.getMessage() + ] + assert len(infos) == 1 + assert "3 runs returned the last message" in infos[0] + + def test_second_outage_logs_error_again(self, clients, block, caplog): + client, _ = self.warm(clients, block) + with caplog.at_level("ERROR", logger="inference"): + drop(client) + block.run(**run_kwargs()) + reconnect(client) + client.deliver(PAUSED) + block.run(**run_kwargs()) + drop(client) + block.run(**run_kwargs()) + + errors = [r for r in caplog.records if r.levelname == "ERROR"] + assert len(errors) == 2 + + def test_reconnected_with_nothing_new_repeats_without_error_and_no_recovery_line( + self, clients, block, caplog + ): + client, first = self.warm(clients, block) + drop(client) + block.run(**run_kwargs()) + reconnect(client) + + with caplog.at_level("INFO", logger="inference"): + result = block.run(**run_kwargs()) + + assert result == {**first, "is_new": False} + assert "restored after" not in caplog.text + + def test_connected_but_not_resubscribed_repeats_with_resubscribe_message( + self, clients, block + ): + client, first = self.warm(clients, block) + drop(client) + client.ack_subscriptions = False + reconnect(client) + + result = block.run(**run_kwargs()) + + assert result == { + **first, + "is_new": False, + "error_status": True, + "error_message": NOT_RESUBSCRIBED_REPEATING_LAST, + } + + def test_broker_unavailable_on_reconnect_repeats_with_its_reason( + self, clients, block + ): + client, first = self.warm(clients, block) + drop(client) + reconnect(client, 3) + + result = block.run(**run_kwargs()) + + assert result["value"] == first["value"] + assert result["error_status"] is True + assert "broker unavailable" in result["error_message"] + assert "retrying in the background" in result["error_message"] + assert client.disconnected is False + + def test_refused_reconnect_drains_then_repeats_with_refusal( + self, clients, block, caplog + ): + client, _ = self.warm(clients, block, read_mode="sequential") + client.deliver(b"A") + drop(client) + reconnect(client, 5) + + with caplog.at_level("ERROR", logger="inference"): + results = [ + block.run(**run_kwargs(read_mode="sequential")) for _ in range(3) + ] + + assert [r["value"] for r in results] == ["A", "A", "A"] + assert [r["error_status"] for r in results] == [False, True, True] + assert "not authorised" in results[1]["error_message"] + assert "does not retry" in results[1]["error_message"] + assert results[2] == results[1] + assert client.disconnected is True + run_path_errors = [ + r + for r in caplog.records + if r.levelname == "ERROR" and r.getMessage().startswith("MQTT Reader ") + ] + assert len(run_path_errors) == 1 + assert "restored after" not in caplog.text + + def test_refused_resubscription_drains_then_returns_empty_failure( + self, clients, block + ): + client, _ = self.warm(clients, block, read_mode="sequential") + client.deliver(b"A") + drop(client) + client.granted_qos = SUBSCRIPTION_REFUSED_QOS + reconnect(client) + + results = [block.run(**run_kwargs(read_mode="sequential")) for _ in range(2)] + + assert results[0]["value"] == "A" + assert results[0]["error_status"] is False + assert results[1]["payload"] == {} + assert results[1]["error_status"] is True + assert "refused subscription" in results[1]["error_message"] + + def test_cold_refused_subscription_drains_the_backlog_first( + self, clients, block, monkeypatch + ): + # a persistent session's backlog lands right after CONNACK, before the + # SUBACK that refuses the subscription: the messages are still real + monkeypatch.setattr(FakeClient, "granted_qos", SUBSCRIPTION_REFUSED_QOS) + original_loop_start = FakeClient.loop_start + + def loop_start_with_backlog(client): + original_loop_start(client) + client.deliver(b"A") + + with patch.object(FakeClient, "loop_start", loop_start_with_backlog): + first = block.run(**run_kwargs(read_mode="sequential")) + second = block.run(**run_kwargs(read_mode="sequential")) + + assert first["value"] == "A" + assert first["error_status"] is False + assert block._subscribed_once is False + assert second["payload"] == {} + assert "refused subscription" in second["error_message"] + + def test_cold_instance_still_waits_for_the_broker( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "fire_on_connect", False) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=0.3)) + + # a hair of timer slack is allowed on the lower bound + assert time.monotonic() - started >= 0.25 + assert result["error_message"] == NOT_CONNECTED_WITHIN_TIMEOUT + assert block._subscribed_once is False + + +class TestBrokerPolicy: + def test_allowlisted_broker_is_used(self, clients, block, monkeypatch): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["LocalHost"] + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert clients[0].connected_to == ("localhost", 1883) + + def test_unlisted_broker_rejected_before_client_construction( + self, clients, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker.internal"] + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not permitted" in result["error_message"] + assert "broker.internal" not in result["error_message"] + assert clients == [] + + def test_port_mismatch_rejected(self, clients, block, monkeypatch): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["localhost:8883"] + ) + + result = block.run(**run_kwargs(port=1883)) + + assert result["error_status"] is True + assert clients == [] + + def test_operator_broker_replaces_workflow_host_when_user_host_not_allowed( + self, clients, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["operator:8883"] + ) + + first = block.run(**run_kwargs(host="workflow-a")) + # identity is the resolved address: a different workflow host is not a change + second = block.run(**run_kwargs(host="workflow-b")) + + assert first["error_status"] is False + assert second["error_status"] is False + assert len(clients) == 1 + assert clients[0].connected_to == ("operator", 8883) + + def test_user_host_not_allowed_without_operator_broker_disables_block( + self, clients, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", None + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "disabled" in result["error_message"] + assert clients == [] + + +@pytest.fixture +def ca_file(tmp_path): + # a real regular file: configure_tls() refuses anything else before paho + # (faked in these tests) would read it + path = tmp_path / "ca.pem" + path.write_text("content is irrelevant: paho is faked in these tests") + return path + + +class TestTLS: + def test_manifest_defaults_and_dropdown_values(self): + manifest = BlockManifest.model_validate(manifest_payload()) + + assert manifest.encryption == "none" + assert manifest.ca_certificate_path is None + assert ( + BlockManifest.model_validate(manifest_payload(encryption="tls")).encryption + == "tls" + ) + + @pytest.mark.parametrize("encryption", ["ssl", "TLS", True, None, ""]) + def test_manifest_rejects_other_encryption_values(self, encryption): + with pytest.raises(ValidationError): + BlockManifest.model_validate(manifest_payload(encryption=encryption)) + + def test_ca_certificate_path_is_shown_only_for_tls(self): + schema = BlockManifest.model_json_schema()["properties"] + + assert schema["ca_certificate_path"]["relevant_for"] == { + "encryption": {"values": ["tls"], "required": True} + } + + @pytest.mark.parametrize("ca_certificate_path", [None, "", "/ca.pem"]) + def test_default_encryption_never_configures_tls( + self, clients, block, ca_certificate_path + ): + result = block.run(**run_kwargs(ca_certificate_path=ca_certificate_path)) + + assert result["error_status"] is False + assert [name for name, *_ in clients[0].calls] == ["connect"] + + def test_tls_without_ca_path_uses_system_store_before_connect(self, clients, block): + result = block.run(**run_kwargs(encryption="tls")) + + assert result["error_status"] is False + assert clients[0].calls == [ + ("tls_set", None), + ("connect", "localhost", 1883), + ] + + def test_tls_with_ca_path_and_file_system_access(self, clients, ca_file): + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) + ) + + assert result["error_status"] is False + assert clients[0].calls == [ + ("tls_set", str(ca_file)), + ("connect", "localhost", 1883), + ] + + def test_tls_with_ca_path_refused_without_file_system_access(self, clients, block): + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/etc/passwd") + ) + + assert result["error_status"] is True + assert "ca_certificate_path" in result["error_message"] + assert ( + "ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE" in result["error_message"] + ) + assert block._client is None + assert clients[0].calls == [] + assert clients[0].loop_started is False + + def test_unloadable_ca_bundle_reported_and_nothing_kept( + self, clients, monkeypatch, ca_file + ): + def failing_tls_set(client, ca_certs=None, **kwargs): + raise FileNotFoundError("no such file") + + monkeypatch.setattr(FakeClient, "tls_set", failing_tls_set) + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) + ) + + assert result["error_status"] is True + assert "could not load CA bundle" in result["error_message"] + assert str(ca_file) in result["error_message"] + assert block._client is None + assert clients[0].loop_started is False + + def test_next_run_after_tls_failure_retries_with_fresh_client(self, clients, block): + block.run(**run_kwargs(encryption="tls", ca_certificate_path="/etc/passwd")) + + result = block.run(**run_kwargs(encryption="tls")) + + assert result["error_status"] is False + assert len(clients) == 2 + + @pytest.mark.parametrize( + "change", + [{"encryption": "tls"}, {"ca_certificate_path": "other.pem"}], + ) + def test_changed_tls_parameters_rejected(self, clients, change): + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + block.run(**run_kwargs()) + + result = block.run(**run_kwargs(**change)) + + assert result["error_status"] is True + assert "changed between runs" in result["error_message"] + assert len(clients) == 1 + + def test_invalid_encryption_at_run_time_rejected(self, clients, block): + result = block.run(**run_kwargs(encryption="ssl")) + + assert result["error_status"] is True + assert "encryption" in result["error_message"] + assert clients == [] + + def test_non_string_ca_path_rejected(self, clients, block): + result = block.run(**run_kwargs(encryption="tls", ca_certificate_path=5)) + + assert result["error_status"] is True + assert "ca_certificate_path" in result["error_message"] + assert clients == [] + + def test_tls_insecure_is_never_called(self, clients, ca_file): + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + block.run(**run_kwargs(encryption="tls", ca_certificate_path=str(ca_file))) + + assert all(name != "tls_insecure_set" for name, *_ in clients[0].calls) diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py index d7fd7df461..7b1034f1a7 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py @@ -1,13 +1,17 @@ import threading +import time from typing import get_args from unittest.mock import MagicMock, patch import pytest from pydantic import ValidationError +from roboflow_workflows.enterprise_blocks.sinks import mqtt_common from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v1 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v1 import ( + NOT_CONNECTED_WITHIN_TIMEOUT, BlockManifest, MQTTWriterSinkBlockV1, + MQTTWriterState, mqtt_on_connect, mqtt_on_connect_fail, mqtt_on_disconnect, @@ -127,45 +131,252 @@ def test_valid_ports_accepted(self, port): class TestCallbacks: def test_on_connect_sets_event_only_for_accepted_connack(self): - event = threading.Event() + state = MQTTWriterState() - mqtt_on_connect(MagicMock(), event, {}, 0) + mqtt_on_connect(MagicMock(), state, {}, 0) - assert event.is_set() + assert state.connected.is_set() + assert state.connack.is_set() + assert state.refused_code is None @pytest.mark.parametrize("reason_code", [1, 2, 3, 4, 5]) def test_on_connect_clears_event_for_rejected_connack(self, reason_code): - event = threading.Event() - event.set() + state = MQTTWriterState() + state.connected.set() - mqtt_on_connect(MagicMock(), event, {}, reason_code) + mqtt_on_connect(MagicMock(), state, {}, reason_code) - assert not event.is_set() + assert not state.connected.is_set() + assert state.connack.is_set() + assert state.refused_code == reason_code + + @pytest.mark.parametrize("reason_code", [1, 2, 4, 5]) + def test_on_connect_stops_the_loop_for_a_permanent_refusal(self, reason_code): + client = MagicMock() + + mqtt_on_connect(client, MQTTWriterState(), {}, reason_code) + + client.disconnect.assert_called_once_with() + + def test_on_connect_keeps_the_loop_for_broker_unavailable(self): + client = MagicMock() + state = MQTTWriterState() + + mqtt_on_connect(client, state, {}, 3) + + client.disconnect.assert_not_called() + assert state.refused_code == 3 + + @pytest.mark.parametrize("reason_code", [0, 3, 5]) + def test_on_connect_completes_the_state_before_waking_a_waiter(self, reason_code): + # a run wakes on `connack`; what it then reads must already be final + state = MQTTWriterState() + seen = {} + + class RecordingEvent(threading.Event): + def set(self): + seen["refused_code"] = state.refused_code + seen["connected"] = state.connected.is_set() + super().set() + + state.connack = RecordingEvent() + + mqtt_on_connect(MagicMock(), state, {}, reason_code) + + assert seen == { + "refused_code": reason_code or None, + "connected": reason_code == 0, + } + + def test_on_connect_logs_the_refusal_reason(self, caplog): + with caplog.at_level("ERROR", logger="inference"): + mqtt_on_connect(MagicMock(), MQTTWriterState(), {}, 4) + + assert "bad user name or password" in caplog.text + assert "code 4" in caplog.text + + def test_accepted_connack_after_broker_unavailable_resets_refusal(self): + state = MQTTWriterState() + + mqtt_on_connect(MagicMock(), state, {}, 3) + mqtt_on_connect(MagicMock(), state, {}, 0) + + assert state.refused_code is None + assert state.connected.is_set() def test_on_connect_fail_matches_paho_two_argument_signature(self): - event = threading.Event() - event.set() + state = MQTTWriterState() + state.connected.set() # paho 1.6.1 invokes on_connect_fail with exactly (client, userdata) - mqtt_on_connect_fail(MagicMock(), event) + mqtt_on_connect_fail(MagicMock(), state) - assert not event.is_set() + assert not state.connected.is_set() def test_on_disconnect_clears_event(self): - event = threading.Event() - event.set() + state = MQTTWriterState() + state.connected.set() + + mqtt_on_disconnect(MagicMock(), state, 1) + + assert not state.connected.is_set() + + def test_on_disconnect_after_transport_drop_clears_connack(self): + state = MQTTWriterState() + state.connack.set() + + mqtt_on_disconnect(MagicMock(), state, 1) + + assert not state.connack.is_set() - mqtt_on_disconnect(MagicMock(), event, 1) + def test_on_disconnect_after_refusal_keeps_connack(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 5) - assert not event.is_set() + mqtt_on_disconnect(MagicMock(), state, 0) + + assert state.connack.is_set() + assert state.refused_code == 5 def test_on_disconnect_accepts_mqtt5_properties_argument(self): - event = threading.Event() - event.set() + state = MQTTWriterState() + state.connected.set() + + mqtt_on_disconnect(MagicMock(), state, 1, properties=None) + + assert not state.connected.is_set() + + def test_reset_clears_everything(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 5) + + state.reset() + + assert not state.connected.is_set() + assert not state.connack.is_set() + assert state.refused_code is None + + +def _answer_connect_with(mock_client_cls, reason_code: int) -> None: + """Make the mocked client's loop_start deliver a CONNACK with `reason_code` + to the state the block handed paho as userdata.""" + mock_client = mock_client_cls.return_value + + def fire_connack(): + state = mock_client_cls.call_args.kwargs["userdata"] + mqtt_on_connect(mock_client, state, {}, reason_code) + + mock_client.loop_start.side_effect = fire_connack + + +class TestRefusedConnection: + def test_refusal_reported_with_reason_and_inputs(self, mock_client_cls, block): + _answer_connect_with(mock_client_cls, 5) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not authorised" in result["message"] + assert "code 5" in result["message"] + assert "Check username and password" in result["message"] + assert "Raise 'timeout'" not in result["message"] + mock_client_cls.return_value.disconnect.assert_called_once_with() + mock_client_cls.return_value.publish.assert_not_called() + + @pytest.mark.parametrize( + "reason_code, reason", + [ + (1, "unacceptable protocol version"), + (2, "identifier rejected"), + (4, "bad user name or password"), + ], + ) + def test_other_permanent_refusals_name_their_reason( + self, mock_client_cls, block, reason_code, reason + ): + _answer_connect_with(mock_client_cls, reason_code) + + result = block.run(**run_kwargs()) + + assert reason in result["message"] + mock_client_cls.return_value.disconnect.assert_called_once_with() + + def test_refusal_is_reported_without_waiting_for_timeout( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert result["error_status"] is True + assert time.monotonic() - started < 0.5 + + def test_next_run_after_refusal_repeats_it_without_a_new_connection( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + first = block.run(**run_kwargs(timeout=2.0)) - mqtt_on_disconnect(MagicMock(), event, 1, properties=None) + started = time.monotonic() + second = block.run(**run_kwargs(timeout=2.0)) - assert not event.is_set() + assert second == first + assert time.monotonic() - started < 0.5 + assert mock_client_cls.call_count == 1 + mock_client_cls.return_value.connect.assert_called_once() + + def test_refusal_with_fail_fast_raises_with_the_reason( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + + with pytest.raises(Exception, match="not authorised"): + block.run(**run_kwargs(fail_fast=True)) + + def test_broker_unavailable_keeps_retrying_and_names_the_reason( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 3) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "broker unavailable" in result["message"] + assert "retrying in the background" in result["message"] + mock_client_cls.return_value.disconnect.assert_not_called() + + def test_run_succeeds_once_an_unavailable_broker_accepts( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 3) + block.run(**run_kwargs()) + state = mock_client_cls.call_args.kwargs["userdata"] + mqtt_on_connect(mock_client_cls.return_value, state, {}, 0) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert state.refused_code is None + + def test_no_connack_at_all_reports_not_connected(self, mock_client_cls, block): + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert result["message"] == NOT_CONNECTED_WITHIN_TIMEOUT + + def test_close_after_refusal_completes_and_resets_state( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + block.run(**run_kwargs()) + + block.close() + + assert block.mqtt_client is None + assert block._connection.refused_code is None + assert not block._connection.connack.is_set() + mock_client_cls.return_value.loop_stop.assert_called_once() class TestRunValidation: @@ -270,7 +481,7 @@ def test_client_connects_synchronously_before_background_loop( result = block.run(**run_kwargs()) assert result["error_status"] is False - mock_client_cls.assert_called_once_with(userdata=block._connected) + mock_client_cls.assert_called_once_with(userdata=block._connection) mock_client.connect.assert_called_once_with("localhost", 1883) mock_client.connect_async.assert_not_called() called_methods = [call[0] for call in mock_client.method_calls] @@ -627,3 +838,238 @@ def test_close_clears_readiness_only_after_loop_thread_joined( assert event_state_at_loop_stop == [True] assert not block._connected.is_set() + + +class TestBrokerPolicy: + def test_allowlisted_broker_is_used(self, mock_client_cls, block, monkeypatch): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["localhost:1883"] + ) + block._connected.set() + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + mock_client_cls.return_value.connect.assert_called_once_with("localhost", 1883) + + def test_unlisted_broker_rejected_before_client_construction( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker.internal"] + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not permitted" in result["message"] + assert "broker.internal" not in result["message"] + mock_client_cls.assert_not_called() + + def test_unlisted_broker_with_fail_fast_raises( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker.internal"] + ) + + with pytest.raises(RuntimeError, match="not permitted"): + block.run(**run_kwargs(fail_fast=True)) + + mock_client_cls.assert_not_called() + + def test_operator_broker_replaces_workflow_host_when_user_host_not_allowed( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["operator:8883"] + ) + block._connected.set() + + first = block.run(**run_kwargs(host="workflow-a")) + second = block.run(**run_kwargs(host="workflow-b")) + + assert first["error_status"] is False + assert second["error_status"] is False + mock_client_cls.return_value.connect.assert_called_once_with("operator", 8883) + + def test_user_host_not_allowed_without_operator_broker_disables_block( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", None + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "disabled" in result["message"] + mock_client_cls.assert_not_called() + + +def _call_names(mock_client) -> list: + return [call[0] for call in mock_client.method_calls] + + +@pytest.fixture +def ca_file(tmp_path): + # a real regular file: configure_tls() refuses anything else before paho + # (faked in these tests) would read it + path = tmp_path / "ca.pem" + path.write_text("content is irrelevant: paho is faked in these tests") + return path + + +class TestTLS: + def test_manifest_defaults_and_relevant_for(self): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v1", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + } + ) + schema = BlockManifest.model_json_schema()["properties"] + + assert manifest.encryption == "none" + assert manifest.ca_certificate_path is None + assert schema["ca_certificate_path"]["relevant_for"] == { + "encryption": {"values": ["tls"], "required": True} + } + + @pytest.mark.parametrize("encryption", ["ssl", "TLS", True, None, ""]) + def test_manifest_rejects_other_encryption_values(self, encryption): + with pytest.raises(ValidationError): + BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v1", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "encryption": encryption, + } + ) + + def test_init_parameters_include_file_system_access(self): + assert MQTTWriterSinkBlockV1.get_init_parameters() == [ + "disable_sinks", + "allow_access_to_file_system", + ] + + @pytest.mark.parametrize("ca_certificate_path", [None, "", "/ca.pem"]) + def test_default_encryption_never_configures_tls( + self, mock_client_cls, block, ca_certificate_path + ): + block._connected.set() + + result = block.run(**run_kwargs(ca_certificate_path=ca_certificate_path)) + + assert result["error_status"] is False + mock_client_cls.return_value.tls_set.assert_not_called() + mock_client_cls.return_value.tls_insecure_set.assert_not_called() + + def test_tls_without_ca_path_uses_system_store_before_connect( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + + result = block.run(**run_kwargs(encryption="tls")) + + assert result["error_status"] is False + mock_client.tls_set.assert_called_once_with() + names = _call_names(mock_client) + assert names.index("tls_set") < names.index("connect") + mock_client.tls_insecure_set.assert_not_called() + + def test_tls_with_ca_path_and_file_system_access(self, mock_client_cls, ca_file): + block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) + block._connected.set() + mock_client = mock_client_cls.return_value + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) + ) + + assert result["error_status"] is False + mock_client.tls_set.assert_called_once_with(ca_certs=str(ca_file)) + names = _call_names(mock_client) + assert names.index("tls_set") < names.index("connect") + + def test_tls_with_ca_path_refused_without_file_system_access( + self, mock_client_cls, block + ): + mock_client = mock_client_cls.return_value + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/etc/passwd") + ) + + assert result["error_status"] is True + assert "ca_certificate_path" in result["message"] + assert "ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE" in result["message"] + assert block.mqtt_client is None + mock_client.tls_set.assert_not_called() + mock_client.connect.assert_not_called() + mock_client.loop_start.assert_not_called() + + def test_tls_refused_with_fail_fast_raises(self, mock_client_cls, block): + with pytest.raises(RuntimeError, match="ca_certificate_path"): + block.run( + **run_kwargs( + encryption="tls", ca_certificate_path="/etc/passwd", fail_fast=True + ) + ) + + assert block.mqtt_client is None + + def test_unloadable_ca_bundle_reported_and_nothing_kept( + self, mock_client_cls, ca_file + ): + block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) + mock_client = mock_client_cls.return_value + mock_client.tls_set.side_effect = FileNotFoundError("no such file") + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) + ) + + assert result["error_status"] is True + assert "could not load CA bundle" in result["message"] + assert str(ca_file) in result["message"] + assert block.mqtt_client is None + mock_client.connect.assert_not_called() + mock_client.loop_start.assert_not_called() + + @pytest.mark.parametrize( + "change", + [{"encryption": "tls"}, {"ca_certificate_path": "other.pem"}], + ) + def test_changed_tls_parameters_rejected(self, mock_client_cls, change): + block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) + block._connected.set() + block.run(**run_kwargs()) + + result = block.run(**run_kwargs(**change)) + + assert result["error_status"] is True + assert "changed between runs" in result["message"] + mock_client_cls.assert_called_once() + + def test_invalid_encryption_at_run_time_rejected(self, mock_client_cls, block): + result = block.run(**run_kwargs(encryption="ssl")) + + assert result["error_status"] is True + assert "encryption" in result["message"] + mock_client_cls.assert_not_called() diff --git a/workflows/tests/unit_tests/test_configuration.py b/workflows/tests/unit_tests/test_configuration.py index b22744b341..8140b8a098 100644 --- a/workflows/tests/unit_tests/test_configuration.py +++ b/workflows/tests/unit_tests/test_configuration.py @@ -66,7 +66,7 @@ def test_configuration_is_frozen_in_every_group() -> None: total_fields += len(dataclasses.fields(value)) with pytest.raises(dataclasses.FrozenInstanceError): setattr(value, dataclasses.fields(value)[0].name, "mutated") - assert total_fields == 78, total_fields + assert total_fields == 80, total_fields def test_default_configuration_matches_env_pys_empty_environment_defaults() -> None: @@ -90,6 +90,8 @@ def test_default_configuration_matches_env_pys_empty_environment_defaults() -> N configuration.engine.allow_kafka_sinks_user_provided_bootstrap_servers is True ) assert configuration.engine.kafka_sinks_whitelisted_bootstrap_servers is None + assert configuration.engine.allow_mqtt_blocks_user_provided_host is True + assert configuration.engine.mqtt_blocks_whitelisted_hosts is None assert configuration.models.vlm_segmentation_max_polygon_vertices == 500 assert configuration.tensor.representation_enabled is False assert configuration.tensor.image_tensor_device is None @@ -412,7 +414,7 @@ def test_environment_facade_exports_every_owned_symbol() -> None: for name in vars(workflows_environment) if name.isupper() and not name.startswith("_") } - assert len(exported) == 78, sorted(exported) + assert len(exported) == 80, sorted(exported) assert isinstance(workflows_environment.WORKFLOW_DISABLED_BLOCK_TYPES, list) assert isinstance(workflows_environment.WORKFLOW_DISABLED_BLOCK_PATTERNS, list) assert isinstance(workflows_environment.ENABLE_TENSOR_DATA_REPRESENTATION, bool)