|
| 1 | +# Copyright 2025 Bloomberg Finance L.P. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | +# |
| 4 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | +# you may not use this file except in compliance with the License. |
| 6 | +# You may obtain a copy of the License at |
| 7 | +# |
| 8 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | +# |
| 10 | +# Unless required by applicable law or agreed to in writing, software |
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | +# See the License for the specific language governing permissions and |
| 14 | +# limitations under the License. |
| 15 | + |
| 16 | +""" |
| 17 | +blazingmq.dev.it.process.rawclient |
| 18 | +
|
| 19 | +
|
| 20 | +PURPOSE: Provide a BMQ raw client. |
| 21 | +""" |
| 22 | + |
| 23 | +import socket |
| 24 | +import json |
| 25 | +from typing import Optional, Union |
| 26 | + |
| 27 | +from blazingmq.schemas import broker |
| 28 | + |
| 29 | + |
| 30 | +class RawClient: |
| 31 | + def __init__(self): |
| 32 | + self._channel: Optional[socket.socket] = None |
| 33 | + |
| 34 | + @staticmethod |
| 35 | + def _wrap_control_event(payload: Union[str, dict]) -> bytes: |
| 36 | + """ |
| 37 | + Wraps the specified 'payload' with EventHeader and adds padding to the |
| 38 | + end. Returns the raw bytes control message. |
| 39 | +
|
| 40 | + See also: bmqp::EventHeader |
| 41 | + """ |
| 42 | + if isinstance(payload, str): |
| 43 | + payload_str = payload |
| 44 | + else: |
| 45 | + payload_str = json.dumps(payload) |
| 46 | + |
| 47 | + padding_len = 4 - len(payload_str) % 4 |
| 48 | + padding = bytes([padding_len] * padding_len) |
| 49 | + |
| 50 | + event_type = broker.EventType.CONTROL |
| 51 | + type_specific = broker.TypeSpecific.ENCODING_JSON |
| 52 | + |
| 53 | + control_header_bytes = 8 |
| 54 | + control_event_size = ( |
| 55 | + control_header_bytes + len(payload_str) + padding_len |
| 56 | + ).to_bytes(4, "big") |
| 57 | + control_event_desc = bytes([0x40 + event_type, 0x02, type_specific, 0x00]) |
| 58 | + |
| 59 | + return ( |
| 60 | + control_event_size |
| 61 | + + control_event_desc |
| 62 | + + payload_str.encode("ascii") |
| 63 | + + padding |
| 64 | + ) |
| 65 | + |
| 66 | + @staticmethod |
| 67 | + def _wrap_heartbeat_res_event() -> bytes: |
| 68 | + """ |
| 69 | + Wraps the heartbeat response event with EventHeader and returns the raw bytes. |
| 70 | + Notice that heartbeat response only has header and no body. |
| 71 | +
|
| 72 | + See also: bmqp::EventHeader |
| 73 | + """ |
| 74 | + event_type = broker.EventType.HEARTBEAT_RSP |
| 75 | + type_specific = broker.TypeSpecific.EMPTY |
| 76 | + |
| 77 | + header_bytes = 8 |
| 78 | + event_size = header_bytes.to_bytes(4, "big") |
| 79 | + event_desc = bytes([0x40 + event_type, 0x02, type_specific, 0x00]) |
| 80 | + |
| 81 | + return event_size + event_desc |
| 82 | + |
| 83 | + def _send_raw(self, message: bytes) -> None: |
| 84 | + """ |
| 85 | + Send the specified raw "message" over the channel to the broker. |
| 86 | + Return the received byte response. |
| 87 | + """ |
| 88 | + assert self._channel is not None |
| 89 | + |
| 90 | + try: |
| 91 | + self._channel.send(message) |
| 92 | + except Exception as e: |
| 93 | + raise ConnectionError(f"Failed to send message: {e}") from e |
| 94 | + |
| 95 | + def _receive_event(self) -> tuple[bytes, bytes]: |
| 96 | + """ |
| 97 | + Read the channel until the next event is received. |
| 98 | +
|
| 99 | + Return the header and received event contents excluding event header and padding |
| 100 | + bytes at the end of the message. |
| 101 | + """ |
| 102 | + # Process the event header |
| 103 | + header_bytes = 8 |
| 104 | + |
| 105 | + while True: |
| 106 | + try: |
| 107 | + # The situation when the event header is not fully received |
| 108 | + # with one 'recv' call is highly improbable. |
| 109 | + header = self._channel.recv(header_bytes) |
| 110 | + except socket.timeout as exc: |
| 111 | + raise ConnectionError("Timeout while waiting for event header") from exc |
| 112 | + except Exception as exc: |
| 113 | + raise ConnectionError(f"Failed to receive event header: {exc}") from exc |
| 114 | + |
| 115 | + if len(header) != header_bytes: |
| 116 | + raise ConnectionError( |
| 117 | + f"Failed to receive event header from the broker, " |
| 118 | + f"expected {header_bytes} bytes, got {len(header)} bytes" |
| 119 | + ) |
| 120 | + |
| 121 | + event_type = header[4] & 0b00111111 |
| 122 | + if not any(event_type == et.value for et in broker.EventType): |
| 123 | + raise ValueError( |
| 124 | + f"Unknown event type: {event_type}, " |
| 125 | + "expected one of the EventType values" |
| 126 | + ) |
| 127 | + if event_type == broker.EventType.HEARTBEAT_REQ: |
| 128 | + print("Received heartbeat request.") |
| 129 | + self._send_raw(self._wrap_heartbeat_res_event()) |
| 130 | + else: |
| 131 | + print("Received event with type: ", broker.EventType(event_type).name) |
| 132 | + break |
| 133 | + |
| 134 | + # Process the event body |
| 135 | + |
| 136 | + message = bytearray() |
| 137 | + |
| 138 | + # The first 4 bytes of the message contain message full size in bytes. |
| 139 | + # See also: bmqp_protocol.h / EventHeader |
| 140 | + remaining = int.from_bytes(header[:4], "big") - header_bytes |
| 141 | + while remaining > 0: |
| 142 | + part = self._channel.recv(remaining) |
| 143 | + if not part: |
| 144 | + raise ConnectionError( |
| 145 | + "Connection closed by broker while receiving event body." |
| 146 | + ) |
| 147 | + message += part |
| 148 | + remaining -= len(part) |
| 149 | + |
| 150 | + if len(message) < 1: |
| 151 | + raise ConnectionError("Received empty message from the broker") |
| 152 | + |
| 153 | + try: |
| 154 | + padding_bytes = message[-1] |
| 155 | + except IndexError as exc: |
| 156 | + raise ConnectionError( |
| 157 | + "Received message too short to contain padding byte" |
| 158 | + ) from exc |
| 159 | + |
| 160 | + # expect correct padding byte value |
| 161 | + if not 1 <= padding_bytes <= 4: |
| 162 | + raise ValueError( |
| 163 | + f"Invalid padding bytes value: {padding_bytes}, " |
| 164 | + "expected value in range [1, 4]" |
| 165 | + ) |
| 166 | + |
| 167 | + body = message[:-padding_bytes] |
| 168 | + |
| 169 | + return header, body |
| 170 | + |
| 171 | + def open_channel(self, host: str, port: int) -> None: |
| 172 | + """ |
| 173 | + Open a new channel to the broker using the specified 'host' / 'port'. |
| 174 | + This method is used to establish a connection for sending negotiation requests. |
| 175 | + """ |
| 176 | + assert self._channel is None |
| 177 | + |
| 178 | + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| 179 | + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) |
| 180 | + sock.connect((host, port)) |
| 181 | + sock.settimeout(10.0) |
| 182 | + |
| 183 | + self._channel = sock |
| 184 | + |
| 185 | + def decode_event_bytes(self, response_header: bytes, response_body: bytes) -> dict: |
| 186 | + """ |
| 187 | + Decode the received event header into a dictionary. |
| 188 | + This is used to parse the response from the broker. |
| 189 | + """ |
| 190 | + # The response is expected to be in JSON format |
| 191 | + type_specific = response_header[6] # 7th byte is typeSpecific |
| 192 | + |
| 193 | + if type_specific == broker.TypeSpecific.ENCODING_JSON: |
| 194 | + return json.loads(response_body.decode("utf-8")) |
| 195 | + if type_specific == broker.TypeSpecific.ENCODING_BER: |
| 196 | + print("BER encoding is not supported in open source Python.") |
| 197 | + raise ValueError("Not supported encoding in response") |
| 198 | + |
| 199 | + print(f"Unknown encoding type: {type_specific}") |
| 200 | + raise ValueError("Unknown encoding in response") |
| 201 | + |
| 202 | + def send_negotiation_request(self) -> dict: |
| 203 | + """ |
| 204 | + Send a negotiation request to the broker. |
| 205 | + """ |
| 206 | + assert self._channel is not None |
| 207 | + |
| 208 | + raw_client_identity = broker.CLIENT_IDENTITY_SCHEMA |
| 209 | + raw_client_identity["clientIdentity"]["clientType"] = "E_TCPCLIENT" |
| 210 | + |
| 211 | + self._send_raw(self._wrap_control_event(raw_client_identity)) |
| 212 | + _, response_body = self._receive_event() |
| 213 | + |
| 214 | + return json.loads(response_body) |
| 215 | + |
| 216 | + def stop(self) -> None: |
| 217 | + """ |
| 218 | + Disconnect from the broker. |
| 219 | + """ |
| 220 | + if self._channel is not None: |
| 221 | + self._channel.close() |
| 222 | + self._channel = None |
0 commit comments