From 16242f12f3a5ff5c4b1a0c94b6239ccd0290f137 Mon Sep 17 00:00:00 2001 From: lachlangrose Date: Thu, 22 Oct 2026 18:05:54 +1030 Subject: [PATCH] feat: push the model to GeoViewer from the export panel Add an optional "Push model to GeoViewer" button in a new GeoViewer tab of the export panel. It sends the stratigraphic and fault surfaces to a running GeoViewer over its local link (protocol 1.8): a small stdlib WebSocket client, the frame codec, and a short session that says hello, sends the objects, and removes the objects of the last push that are not sent again. Co-Authored-By: Claude Sonnet 5.5 --- docs/usage/3d-modeling.md | 16 ++ .../gui/modelling/steps/export_panel.py | 77 ++++++ loopstructural/main/geoviewer/__init__.py | 1 + loopstructural/main/geoviewer/link.py | 159 +++++++++++++ loopstructural/main/geoviewer/protocol.py | 179 ++++++++++++++ loopstructural/main/geoviewer/scene.py | 122 ++++++++++ .../main/geoviewer/websocket_client.py | 222 ++++++++++++++++++ loopstructural/toolbelt/preferences.py | 3 + tests/unit/geoviewer_fake_server.py | 163 +++++++++++++ tests/unit/test_geoviewer_link.py | 94 ++++++++ tests/unit/test_geoviewer_protocol.py | 120 ++++++++++ tests/unit/test_geoviewer_scene.py | 93 ++++++++ 12 files changed, 1249 insertions(+) create mode 100644 loopstructural/main/geoviewer/__init__.py create mode 100644 loopstructural/main/geoviewer/link.py create mode 100644 loopstructural/main/geoviewer/protocol.py create mode 100644 loopstructural/main/geoviewer/scene.py create mode 100644 loopstructural/main/geoviewer/websocket_client.py create mode 100644 tests/unit/geoviewer_fake_server.py create mode 100644 tests/unit/test_geoviewer_link.py create mode 100644 tests/unit/test_geoviewer_protocol.py create mode 100644 tests/unit/test_geoviewer_scene.py diff --git a/docs/usage/3d-modeling.md b/docs/usage/3d-modeling.md index f3f21909..783328d4 100644 --- a/docs/usage/3d-modeling.md +++ b/docs/usage/3d-modeling.md @@ -25,3 +25,19 @@ In LoopStructural stratigraphic surfaces can be modelled using implicit function ### Fault modelling Faults are modelled in LoopStructural by building three implicit functions defining the fault surface, fault slip vector and the fault extent. Combined with a parametric representation of the fault displacement within these coordinates a kinematic model. + +## Push the model to GeoViewer + +GeoViewer is the advanced 3D viewer of Loop. It is a separate program. The plugin can send the +surfaces of the solved model to it. This is optional: the 3D view of the plugin does not change. + +1. Start GeoViewer. Open its link popup (in the status bar) and note the port and the token. +2. In the plugin, build the model (step 4). Open **Export** in step 5 and select the + **GeoViewer** tab. +3. Type the port and the token, choose the surfaces to send, and select + **Push model to GeoViewer**. + +The plugin saves the port and the token for the next push. The surfaces go into the group +`LoopStructural model` of the GeoViewer scene. A new push replaces the surfaces of the last push. +Other objects of the scene stay. The coordinates are in the CRS of the model, and the scene of +GeoViewer keeps its own origin. diff --git a/loopstructural/gui/modelling/steps/export_panel.py b/loopstructural/gui/modelling/steps/export_panel.py index 420f47f1..836b6c20 100644 --- a/loopstructural/gui/modelling/steps/export_panel.py +++ b/loopstructural/gui/modelling/steps/export_panel.py @@ -10,6 +10,7 @@ QDoubleSpinBox, QFormLayout, QLabel, + QLineEdit, QMessageBox, QPushButton, QSpinBox, @@ -19,6 +20,9 @@ ) from loopstructural.main import model_export +from loopstructural.main.geoviewer import scene as geoviewer_scene +from loopstructural.main.geoviewer.link import DEFAULT_PORT, GeoViewerLink, LinkError +from loopstructural.toolbelt.preferences import PlgOptionsManager from ...background_task import finish_background_task, start_background_task from ...compatibility import configure_layer_combo @@ -70,6 +74,7 @@ def __init__(self, parent=None, *, model_manager=None, data_manager=None): self.tabs.addTab(self._surfaces_tab(), "Surfaces") self.tabs.addTab(self._block_model_tab(), "Block model") self.tabs.addTab(self._cross_section_tab(), "Cross-section") + self.tabs.addTab(self._geoviewer_tab(), "GeoViewer") layout.addWidget(self.tabs) self.set_model_solved(False) self.setCollapsed(True) @@ -149,6 +154,40 @@ def _cross_section_tab(self): self._update_section_inputs() return tab + def _geoviewer_tab(self): + tab = QWidget(self) + form = QFormLayout(tab) + note = QLabel( + "Send the surfaces of the model to GeoViewer, the advanced 3D viewer. " + "Start GeoViewer first, and copy the port and the token from its link popup. " + "A new push replaces the surfaces of the last push.", + tab, + ) + note.setWordWrap(True) + self.geoviewer_port = _spin(tab, DEFAULT_PORT, minimum=1, maximum=65535) + self.geoviewer_token = QLineEdit(tab) + self.geoviewer_token.setEchoMode(QLineEdit.EchoMode.Password) + self.geoviewer_stratigraphy = QCheckBox("Stratigraphic surfaces", tab) + self.geoviewer_stratigraphy.setChecked(True) + self.geoviewer_faults = QCheckBox("Fault surfaces", tab) + self.geoviewer_faults.setChecked(True) + self.geoviewer_button = QPushButton("Push model to GeoViewer", tab) + self.geoviewer_button.clicked.connect(self.push_to_geoviewer) + self._load_geoviewer_settings() + form.addRow(note) + form.addRow("Port", self.geoviewer_port) + form.addRow("Token", self.geoviewer_token) + form.addRow(self.geoviewer_stratigraphy) + form.addRow(self.geoviewer_faults) + form.addRow(self.geoviewer_button) + return tab + + def _load_geoviewer_settings(self): + port = PlgOptionsManager.get_value_from_key('geoviewer_port', DEFAULT_PORT, int) + token = PlgOptionsManager.get_value_from_key('geoviewer_token', '', str) + self.geoviewer_port.setValue(int(port or DEFAULT_PORT)) + self.geoviewer_token.setText(token or '') + def _update_section_inputs(self, *_args): is_line = self.section_kind.currentData() == 'line' self.section_line.setVisible(is_line) @@ -203,6 +242,44 @@ def _with_extension(self, file_widget, format_combo): def _warn(self, text): QMessageBox.warning(self, "Export", text) + def push_to_geoviewer(self): + """Send the surfaces of the model to a running GeoViewer. The model does not change.""" + strat = self.geoviewer_stratigraphy.isChecked() + faults = self.geoviewer_faults.isChecked() + if not (strat or faults): + return self._warn("Select the stratigraphic surfaces, the fault surfaces, or both.") + port, token = self.geoviewer_port.value(), self.geoviewer_token.text().strip() + PlgOptionsManager.set_value_from_key('geoviewer_port', port) + PlgOptionsManager.set_value_from_key('geoviewer_token', token) + manager = self.model_manager + crs = self.data_manager.get_model_crs() if self.data_manager else None + crs_text = crs.authid() if crs is not None and crs.isValid() else '' + + def target(progress): + model = getattr(manager, 'model', None) + if model is None: + raise geoviewer_scene.PushError("Build the model in step 4 first.") + try: + result = GeoViewerLink(port, token).push( + crs_text, + geoviewer_scene.proposed_origin(model), + lambda origin: geoviewer_scene.model_messages( + manager, origin, stratigraphic=strat, faults=faults, progress=progress + ), + replace_prefix=geoviewer_scene.ID_PREFIX, + progress=progress, + ) + except (LinkError, geoviewer_scene.PushError) as error: + raise RuntimeError(str(error)) from error + if result.refused or result.other_errors: + details = "; ".join(text for _id, text in result.refused) or "; ".join( + result.other_errors + ) + return f"{result.summary()} {details}" + return result.summary() + + self._run(target, "Push model to GeoViewer") + def export_surfaces(self): folder = self.surface_folder.filePath().strip() if not folder: diff --git a/loopstructural/main/geoviewer/__init__.py b/loopstructural/main/geoviewer/__init__.py new file mode 100644 index 00000000..86f039fd --- /dev/null +++ b/loopstructural/main/geoviewer/__init__.py @@ -0,0 +1 @@ +"""The link from the plugin to the GeoViewer application: push a model to a running viewer.""" diff --git a/loopstructural/main/geoviewer/link.py b/loopstructural/main/geoviewer/link.py new file mode 100644 index 00000000..b7429f2e --- /dev/null +++ b/loopstructural/main/geoviewer/link.py @@ -0,0 +1,159 @@ +"""The connection of the plugin to a running GeoViewer. + +``GeoViewerLink.push`` makes the connection, says ``hello``, sends the objects +of the model and closes the connection. Nothing keeps the link open, so the +plugin has no state to repair when the viewer stops. The viewer keeps the scene. +""" + +from typing import Callable, Iterable, List, Optional, Sequence, Tuple + +from . import protocol +from .websocket_client import WebSocketClient, WebSocketClosed, WebSocketError + +DEFAULT_PORT = 47615 +HOST = '127.0.0.1' + +# Time to wait for the answer to the hello message +_HELLO_TIMEOUT = 5.0 +# Time to wait for the errors of the viewer after the last message +_SETTLE_TIMEOUT = 0.5 + + +class LinkError(Exception): + """The push failed. The text tells the user why.""" + + +class PushResult: + """What happened in one push.""" + + def __init__(self): + self.sent = 0 + self.removed = 0 + # (object id, text) of each error that the viewer sent about an object + self.refused: List[Tuple[str, str]] = [] + self.other_errors: List[str] = [] + + def summary(self) -> str: + text = f"{self.sent} objects sent to GeoViewer" + if self.removed: + text += f", {self.removed} old objects removed" + if self.refused or self.other_errors: + text += f". The viewer refused {len(self.refused) + len(self.other_errors)} messages" + return text + "." + + +class GeoViewerLink: + """A short session with the viewer. + + Parameters + ---------- + port, token : the port and the token of the link of the viewer. The viewer + shows both in its link popup. + """ + + def __init__(self, port: int = DEFAULT_PORT, token: str = '', host: str = HOST): + self.port = int(port) + self.token = token or '' + self.host = host + + def push( + self, + crs: str, + origin: Sequence[float], + build_messages: Callable[[Sequence[float]], Iterable[Tuple[str, bytes]]], + *, + replace_prefix: Optional[str] = None, + progress: Optional[Callable[[str], None]] = None, + ) -> PushResult: + """Send objects to the viewer. + + Parameters + ---------- + crs, origin : the CRS and the origin that the plugin proposes for the + scene. If the scene of the viewer has an origin already, the viewer + keeps it and says so in its ``hello`` answer. + build_messages : called with the origin that the viewer uses. It yields + ``(object id, frame)`` pairs. The coordinates must be relative to + this origin. + replace_prefix : after the push, remove the objects of the viewer whose + id starts with this text and that were not sent again. + progress : called with a text before each step. + """ + say = progress or (lambda _text: None) + say(f"Connecting to GeoViewer on port {self.port}...") + try: + client = WebSocketClient.connect(self.host, self.port, token=self.token) + except WebSocketError as error: + raise LinkError( + f"{error} Start GeoViewer, and check the port and the token in its link popup." + ) from error + try: + return self._session(client, crs, origin, build_messages, replace_prefix, say) + except WebSocketError as error: + raise LinkError(str(error)) from error + finally: + client.close() + + def _session(self, client, crs, origin, build_messages, replace_prefix, say) -> PushResult: + result = PushResult() + client.send_binary(protocol.hello_message(crs, origin)) + scene_origin, existing = self._read_hello_answer(client, origin) + sent_ids = set() + for object_id, frame in build_messages(scene_origin): + say(f"Sending {object_id}...") + client.send_binary(frame) + sent_ids.add(object_id) + result.sent += 1 + self._collect_errors(client, result, wait=0.0) + if replace_prefix: + for object_id in existing: + if object_id.startswith(replace_prefix) and object_id not in sent_ids: + client.send_binary(protocol.remove_message(object_id)) + result.removed += 1 + self._collect_errors(client, result, wait=_SETTLE_TIMEOUT) + return result + + @staticmethod + def _read_hello_answer(client, origin): + """Return the origin of the scene and the ids that are in the scene. + + The viewer answers ``hello`` with its own ``hello`` and then with one + ``scene.inventory``. + """ + scene_origin = [float(v) for v in origin] + ids: List[str] = [] + got_hello = got_inventory = False + while not (got_hello and got_inventory): + data = client.receive(_HELLO_TIMEOUT) + if data is None: + raise LinkError("GeoViewer did not answer the hello message.") + header = protocol.decode_header(data) + kind = header.get('type') + if kind == 'hello': + scene_origin = [float(v) for v in header.get('origin', scene_origin)] + got_hello = True + elif kind == 'scene.inventory': + ids = protocol.inventory_ids(header) + got_inventory = True + elif kind == 'error': + raise LinkError(f"GeoViewer refused the link: {header.get('message', '')}") + return scene_origin, ids + + @staticmethod + def _collect_errors(client, result: PushResult, *, wait: float): + """Read the messages that wait in the socket and keep the errors.""" + while True: + try: + data = client.receive(wait if wait > 0 else 0.001) + except WebSocketClosed: + return + if data is None: + return + header = protocol.decode_header(data) + if header.get('type') != 'error': + continue + text = header.get('message', header.get('code', 'error')) + if header.get('object'): + result.refused.append((header['object'], text)) + else: + result.other_errors.append(text) diff --git a/loopstructural/main/geoviewer/protocol.py b/loopstructural/main/geoviewer/protocol.py new file mode 100644 index 00000000..90c10e6b --- /dev/null +++ b/loopstructural/main/geoviewer/protocol.py @@ -0,0 +1,179 @@ +"""The message format of the GeoViewer link (protocol 1.x). + +No QGIS and no pyvista code, only the standard library and numpy. A message is +one binary frame: + + [u32 header_len][JSON header, UTF-8][buffer 0][buffer 1]... + +All numbers are little-endian. The header lists the buffers in order. See +``docs/protocol.md`` in the GeoViewer repository for the full description. +""" + +import json +import struct +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import numpy as np + +PROTOCOL_VERSION = {'major': 1, 'minor': 8} + +# The largest frame that the viewer accepts +MAX_FRAME_BYTES = 256 * 1024 * 1024 +# Beyond this distance from the origin an f32 coordinate is not exact to 1 mm +F32_LIMIT = 4000.0 + +APPLICATION_NAME = 'plugin_loopstructural' + + +class ProtocolError(Exception): + """A frame cannot be written or read.""" + + +def encode_frame(header: Dict[str, Any], buffers: Sequence[Tuple[str, np.ndarray]] = ()) -> bytes: + """Return the frame for a header and its named buffers. + + The function adds ``version`` and ``buffers`` to the header. A buffer is a + ``(name, array)`` pair. The array must be f32, f64, u32, i32 or u8. + """ + descriptions = [] + payload = [] + for name, array in buffers: + array = np.ascontiguousarray(array) + dtype = _DTYPE_NAMES.get(array.dtype.newbyteorder('<').str) + if dtype is None: + raise ProtocolError(f"Buffer {name!r} has the unsupported type {array.dtype}.") + descriptions.append({'name': name, 'dtype': dtype, 'shape': list(array.shape)}) + payload.append(array.astype(array.dtype.newbyteorder('<'), copy=False).tobytes()) + full = dict(header) + full['version'] = PROTOCOL_VERSION + full['buffers'] = descriptions + header_bytes = json.dumps(full, separators=(',', ':'), allow_nan=False).encode('utf-8') + frame = b''.join([struct.pack(' MAX_FRAME_BYTES: + raise ProtocolError( + f"The message is {len(frame)} bytes. The viewer accepts {MAX_FRAME_BYTES} bytes at most." + ) + return frame + + +_DTYPE_NAMES = { + np.dtype(' Dict[str, Any]: + """Return the JSON header of a frame. The buffers are ignored.""" + if len(frame) < 4: + raise ProtocolError("The frame is shorter than 4 bytes.") + (length,) = struct.unpack_from(' len(frame): + raise ProtocolError("The frame stops inside the header.") + try: + return json.loads(frame[4 : 4 + length].decode('utf-8')) + except ValueError as error: + raise ProtocolError(f"The header is not valid JSON: {error}") from error + + +def hello_message(crs: str, origin: Sequence[float], version: str = '0') -> bytes: + """Return the ``hello`` frame. ``origin`` is the origin of the scene in ``crs``.""" + return encode_frame( + { + 'type': 'hello', + 'application': {'name': APPLICATION_NAME, 'version': version}, + 'crs': crs, + 'origin': [float(v) for v in origin], + } + ) + + +def clear_message() -> bytes: + """Return the ``scene.clear`` frame.""" + return encode_frame({'type': 'scene.clear'}) + + +def remove_message(object_id: str) -> bytes: + """Return the ``scene.remove`` frame for an object or a group (with its children).""" + return encode_frame({'type': 'scene.remove', 'id': object_id}) + + +def _upsert_header( + object_id: str, + object_type: str, + name: str, + parent: Optional[str], + revision: int, + style: Optional[Dict[str, Any]], +) -> Dict[str, Any]: + return { + 'type': 'scene.upsert', + 'id': object_id, + 'object_type': object_type, + 'name': name, + 'parent': parent, + 'revision': int(revision), + 'edit': [], + 'replaces': None, + 'geometry': {}, + 'params': {}, + 'data': [], + 'property_groups': [], + 'drape': None, + 'style': dict(style or {}), + } + + +def group_message( + object_id: str, name: str, parent: Optional[str] = None, revision: int = 1 +) -> bytes: + """Return the ``scene.upsert`` frame of a Group.""" + return encode_frame(_upsert_header(object_id, 'Group', name, parent, revision, None)) + + +def surface_message( + object_id: str, + name: str, + vertices, + triangles, + *, + origin: Sequence[float] = (0.0, 0.0, 0.0), + parent: Optional[str] = None, + revision: int = 1, + style: Optional[Dict[str, Any]] = None, +) -> bytes: + """Return the ``scene.upsert`` frame of a Surface. + + Parameters + ---------- + vertices : array-like, shape (n, 3) + The vertices in the coordinates of the CRS. The function subtracts + ``origin`` in f64 and then writes f32, or f64 for a model that is more + than 4 km from the origin. + triangles : array-like, shape (m, 3) + Vertex indices. + """ + vertices = np.asarray(vertices, dtype=np.float64).reshape(-1, 3) + triangles = np.asarray(triangles, dtype=np.int64).reshape(-1, 3) + if len(vertices) == 0 or len(triangles) == 0: + raise ProtocolError(f"The surface {name!r} has no triangles.") + if not np.isfinite(vertices).all(): + raise ProtocolError(f"The surface {name!r} has a vertex that is NaN or infinite.") + if triangles.min() < 0 or triangles.max() >= len(vertices): + raise ProtocolError(f"The surface {name!r} has a triangle index outside the vertices.") + relative = vertices - np.asarray(origin, dtype=np.float64) + if np.abs(relative).max() <= F32_LIMIT: + relative = relative.astype(np.float32) + header = _upsert_header(object_id, 'Surface', name, parent, revision, style) + header['geometry'] = {'vertices': 'vertices', 'triangles': 'triangles'} + return encode_frame( + header, + [('vertices', relative), ('triangles', triangles.astype(np.uint32))], + ) + + +def inventory_ids(header: Dict[str, Any]) -> List[str]: + """Return the ids of a ``scene.inventory`` header.""" + return [entry['id'] for entry in header.get('entries', [])] diff --git a/loopstructural/main/geoviewer/scene.py b/loopstructural/main/geoviewer/scene.py new file mode 100644 index 00000000..79157ef3 --- /dev/null +++ b/loopstructural/main/geoviewer/scene.py @@ -0,0 +1,122 @@ +"""Turn the results of a solved model into the objects of the GeoViewer scene. + +The functions do not import QGIS. They read the model through the model manager +(``model`` with ``get_stratigraphic_surfaces`` and ``get_fault_surfaces``). +""" + +import time +from typing import Iterator, List, Optional, Sequence, Tuple + +import numpy as np + +from . import protocol + +# Every object of the plugin has an id that starts with this text, so that a new +# push can replace the objects of the old push. +ID_PREFIX = 'loopstructural/' +ROOT_ID = ID_PREFIX + 'model' +STRATIGRAPHY_ID = ROOT_ID + '/stratigraphy' +FAULTS_ID = ROOT_ID + '/faults' + +# A fault is drawn in a different colour than the units: dark red +FAULT_COLOUR = [200, 40, 40, 255] + + +class PushError(Exception): + """The model cannot be pushed, for example because it is not solved.""" + + +def proposed_origin(model) -> List[float]: + """Return the centre of the bounding box of the model, with z = 0. + + The viewer uses it as the origin of the scene when its scene is empty. + The values are rounded to whole metres, so the origin is easy to read. + """ + box = model.bounding_box + centre = (np.asarray(box.origin, dtype=float) + np.asarray(box.maximum, dtype=float)) / 2.0 + return [float(round(centre[0])), float(round(centre[1])), 0.0] + + +def triangles_of(mesh) -> Tuple[np.ndarray, np.ndarray]: + """Return ``(vertices, triangles)`` of a pyvista surface mesh.""" + mesh = mesh.triangulate() + vertices = np.asarray(mesh.points, dtype=np.float64) + faces = np.asarray(mesh.faces, dtype=np.int64) + if faces.size == 0: + return vertices, np.zeros((0, 3), dtype=np.int64) + return vertices, faces.reshape(-1, 4)[:, 1:] + + +def _surfaces(model, stratigraphic: bool, faults: bool): + result = [] + if stratigraphic: + result += [('stratigraphy', s) for s in model.get_stratigraphic_surfaces()] + if faults: + result += [('faults', s) for s in model.get_fault_surfaces()] + return result + + +def _unique_id(base: str, used: set) -> str: + candidate, count = base, 1 + while candidate in used: + count += 1 + candidate = f"{base}_{count}" + used.add(candidate) + return candidate + + +def model_messages( + model_manager, + origin: Sequence[float], + *, + stratigraphic: bool = True, + faults: bool = True, + revision: Optional[int] = None, + progress=None, +) -> Iterator[Tuple[str, bytes]]: + """Yield ``(object id, frame)`` for the groups and the surfaces of the model. + + The groups come first, so the viewer knows the parent of each surface. A + surface with no triangles is skipped, as in the file export. + """ + model = getattr(model_manager, 'model', None) + if model is None: + raise PushError("There is no model to push. Build the model in step 4 first.") + revision = int(time.time()) if revision is None else int(revision) + surfaces = _surfaces(model, stratigraphic, faults) + if not surfaces: + raise PushError("The model has no surfaces to push.") + + kinds = {kind for kind, _ in surfaces} + yield ROOT_ID, protocol.group_message(ROOT_ID, 'LoopStructural model', None, revision) + if 'stratigraphy' in kinds: + yield STRATIGRAPHY_ID, protocol.group_message( + STRATIGRAPHY_ID, 'Stratigraphic surfaces', ROOT_ID, revision + ) + if 'faults' in kinds: + yield FAULTS_ID, protocol.group_message(FAULTS_ID, 'Fault surfaces', ROOT_ID, revision) + + used = set() + pushed = 0 + for kind, surface in surfaces: + if progress is not None: + progress(f"Reading {surface.name}...") + vertices, triangles = triangles_of(surface.vtk()) + if len(triangles) == 0: + continue + parent = STRATIGRAPHY_ID if kind == 'stratigraphy' else FAULTS_ID + object_id = _unique_id(f"{parent}/{surface.name}", used) + style = {'colour': FAULT_COLOUR} if kind == 'faults' else None + yield object_id, protocol.surface_message( + object_id, + str(surface.name), + vertices, + triangles, + origin=origin, + parent=parent, + revision=revision, + style=style, + ) + pushed += 1 + if pushed == 0: + raise PushError("All surfaces are empty. Check the bounding box and the data.") diff --git a/loopstructural/main/geoviewer/websocket_client.py b/loopstructural/main/geoviewer/websocket_client.py new file mode 100644 index 00000000..8a494ade --- /dev/null +++ b/loopstructural/main/geoviewer/websocket_client.py @@ -0,0 +1,222 @@ +"""A small WebSocket client for the link to GeoViewer. + +It uses the standard library only, because QGIS does not ship a WebSocket +package for every platform. It does what the link needs: one connection to +``127.0.0.1`` with an ``Authorization`` header, binary messages in both +directions, and the control frames (ping, pong, close). It does not do TLS. +""" + +import base64 +import hashlib +import os +import socket +import struct +from typing import Optional, Tuple + +_GUID = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11' + +OP_CONTINUATION = 0x0 +OP_TEXT = 0x1 +OP_BINARY = 0x2 +OP_CLOSE = 0x8 +OP_PING = 0x9 +OP_PONG = 0xA + +# The viewer refuses a message of more than 256 MiB; refuse a reply of more than 16 MiB +MAX_INCOMING_BYTES = 16 * 1024 * 1024 +_HANDSHAKE_MAX_BYTES = 16 * 1024 +# Time for the rest of a frame after its first bytes arrived +_FRAME_TIMEOUT = 10.0 + + +class WebSocketError(Exception): + """The connection failed, or the server broke the protocol.""" + + +class WebSocketClosed(WebSocketError): + """The server closed the connection.""" + + +def _mask(data: bytes, key: bytes) -> bytes: + if not data: + return data + repeated = (key * (len(data) // 4 + 1))[: len(data)] + value = int.from_bytes(data, 'little') ^ int.from_bytes(repeated, 'little') + return value.to_bytes(len(data), 'little') + + +class WebSocketClient: + """One client connection. Not thread safe: use it from one thread.""" + + def __init__(self, sock: socket.socket): + self._sock = sock + self._buffer = b'' + self.closed = False + + @classmethod + def connect( + cls, host: str, port: int, *, token: str = '', timeout: float = 5.0, path: str = '/' + ) -> 'WebSocketClient': + """Open the connection and do the handshake. + + Raises ``WebSocketError`` if the server is not there, or refuses the token. + """ + try: + sock = socket.create_connection((host, port), timeout=timeout) + except OSError as error: + raise WebSocketError(f"Cannot connect to {host}:{port}: {error}") from error + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + client = cls(sock) + try: + client._handshake(host, port, token, path) + except Exception: + sock.close() + raise + return client + + def _handshake(self, host: str, port: int, token: str, path: str): + key = base64.b64encode(os.urandom(16)) + lines = [ + f'GET {path} HTTP/1.1', + f'Host: {host}:{port}', + 'Upgrade: websocket', + 'Connection: Upgrade', + f'Sec-WebSocket-Key: {key.decode()}', + 'Sec-WebSocket-Version: 13', + ] + if token: + lines.append(f'Authorization: Bearer {token}') + try: + self._sock.sendall(('\r\n'.join(lines) + '\r\n\r\n').encode('ascii')) + data = b'' + while b'\r\n\r\n' not in data: + chunk = self._sock.recv(4096) + if not chunk: + raise WebSocketError("The viewer closed the connection during the handshake.") + data += chunk + if len(data) > _HANDSHAKE_MAX_BYTES: + raise WebSocketError("The handshake answer is too long.") + except OSError as error: + raise WebSocketError(f"The handshake failed: {error}") from error + head, _, rest = data.partition(b'\r\n\r\n') + self._buffer = rest + status_line, *header_lines = head.decode('latin-1').split('\r\n') + parts = status_line.split(' ', 2) + status = parts[1] if len(parts) > 1 else '' + if status == '401': + raise WebSocketError("The viewer refused the token. Check the token of the link.") + if status != '101': + raise WebSocketError(f"The viewer refused the connection: {status_line}") + headers = {} + for line in header_lines: + name, _, value = line.partition(':') + headers[name.strip().lower()] = value.strip() + expected = base64.b64encode(hashlib.sha1(key + _GUID).digest()).decode() + if headers.get('sec-websocket-accept') != expected: + raise WebSocketError("The server did not answer as a WebSocket server.") + + # -- sending ------------------------------------------------------------ + + def _send_frame(self, opcode: int, payload: bytes): + length = len(payload) + head = bytearray([0x80 | opcode]) + if length < 126: + head.append(0x80 | length) + elif length < 1 << 16: + head.append(0x80 | 126) + head += struct.pack('>H', length) + else: + head.append(0x80 | 127) + head += struct.pack('>Q', length) + key = os.urandom(4) + try: + self._sock.sendall(bytes(head) + key + _mask(payload, key)) + except OSError as error: + self.closed = True + raise WebSocketError(f"Cannot send to the viewer: {error}") from error + + def send_binary(self, payload: bytes): + """Send one binary message.""" + self._send_frame(OP_BINARY, payload) + + # -- receiving ---------------------------------------------------------- + + def _read_exactly(self, count: int) -> bytes: + while len(self._buffer) < count: + chunk = self._sock.recv(65536) + if not chunk: + self.closed = True + raise WebSocketClosed("The viewer closed the connection.") + self._buffer += chunk + data, self._buffer = self._buffer[:count], self._buffer[count:] + return data + + def _read_frame(self, wait: Optional[float]) -> Tuple[bool, int, bytes]: + # The wait applies to the start of a frame only. A frame that started + # is read to the end, so a timeout cannot leave half a frame behind. + self._sock.settimeout(wait) + first, second = self._read_exactly(2) + self._sock.settimeout(_FRAME_TIMEOUT) + fin, opcode = bool(first & 0x80), first & 0x0F + length = second & 0x7F + if length == 126: + (length,) = struct.unpack('>H', self._read_exactly(2)) + elif length == 127: + (length,) = struct.unpack('>Q', self._read_exactly(8)) + if length > MAX_INCOMING_BYTES: + raise WebSocketError("The viewer sent a message that is too large.") + key = self._read_exactly(4) if second & 0x80 else None + payload = self._read_exactly(length) + return fin, opcode, _mask(payload, key) if key else payload + + def receive(self, timeout: Optional[float] = None) -> Optional[bytes]: + """Return the next binary message, or None if nothing arrives in ``timeout`` seconds. + + Ping frames get a pong. A text message is skipped. If the server closes + the connection, ``WebSocketClosed`` is raised. + """ + message = bytearray() + message_opcode = None + try: + while True: + fin, opcode, payload = self._read_frame(timeout if not message else _FRAME_TIMEOUT) + if opcode == OP_PING: + self._send_frame(OP_PONG, payload) + elif opcode == OP_PONG: + continue + elif opcode == OP_CLOSE: + self.closed = True + try: + self._send_frame(OP_CLOSE, payload[:2]) + except WebSocketError: + pass + raise WebSocketClosed("The viewer closed the connection.") + else: + if opcode != OP_CONTINUATION: + message_opcode = opcode + message = bytearray() + message += payload + if len(message) > MAX_INCOMING_BYTES: + raise WebSocketError("The viewer sent a message that is too large.") + if fin and message_opcode == OP_BINARY: + return bytes(message) + if fin: + message = bytearray() + except (socket.timeout, TimeoutError): + return None + except OSError as error: + self.closed = True + raise WebSocketError(f"Cannot read from the viewer: {error}") from error + + def close(self): + """Close the connection. It is safe to call it more than one time.""" + if not self.closed: + self.closed = True + try: + self._send_frame(OP_CLOSE, struct.pack('>H', 1000)) + except WebSocketError: + pass + try: + self._sock.close() + except OSError: + pass diff --git a/loopstructural/toolbelt/preferences.py b/loopstructural/toolbelt/preferences.py index 158f453a..dac31f1e 100644 --- a/loopstructural/toolbelt/preferences.py +++ b/loopstructural/toolbelt/preferences.py @@ -37,6 +37,9 @@ class PlgSettingsStructure: interpolator_cpw: float = 1.0 interpolator_npw: float = 1.0 separate_dock_widgets: bool = False + # the link to the GeoViewer application (see `main.geoviewer`) + geoviewer_port: int = 47615 + geoviewer_token: str = "" class PlgOptionsManager: diff --git a/tests/unit/geoviewer_fake_server.py b/tests/unit/geoviewer_fake_server.py new file mode 100644 index 00000000..98d2d02c --- /dev/null +++ b/tests/unit/geoviewer_fake_server.py @@ -0,0 +1,163 @@ +"""A fake GeoViewer for the tests: a WebSocket server on 127.0.0.1. + +It does what the link of the viewer does (see ``docs/protocol.md`` of the +GeoViewer repository): it checks the bearer token, answers ``hello`` with a +``hello`` and a ``scene.inventory``, and keeps the frames it receives. +""" + +import base64 +import hashlib +import socket +import struct +import threading + +from loopstructural.main.geoviewer import protocol + +_GUID = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11' + + +def _read(conn, count, buffer): + while len(buffer) < count: + chunk = conn.recv(65536) + if not chunk: + raise ConnectionError + buffer += chunk + return buffer[:count], buffer[count:] + + +class FakeViewer: + def __init__(self, token='secret', origin=(10.0, 20.0, 0.0), inventory=(), reply_errors=None): + self.token = token + self.origin = list(origin) + self.inventory = list(inventory) + # object id -> error text; the viewer refuses these objects + self.reply_errors = dict(reply_errors or {}) + self.received = [] # (header, frame bytes) + self.pings = 0 + self._server = socket.socket() + self._server.bind(('127.0.0.1', 0)) + self._server.listen(1) + self.port = self._server.getsockname()[1] + self._thread = threading.Thread(target=self._serve, daemon=True) + self._thread.start() + + def close(self): + self._server.close() + + # -- server ------------------------------------------------------------- + + def _serve(self): + try: + conn, _ = self._server.accept() + except OSError: + return + with conn: + try: + self._connection(conn) + except (ConnectionError, OSError): + pass + + def _connection(self, conn): + data = b'' + while b'\r\n\r\n' not in data: + chunk = conn.recv(4096) + if not chunk: + return + data += chunk + head, _, buffer = data.partition(b'\r\n\r\n') + lines = head.decode().split('\r\n')[1:] + headers = {k.strip().lower(): v.strip() for k, _, v in (l.partition(':') for l in lines)} + if headers.get('authorization') != f'Bearer {self.token}': + conn.sendall(b'HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n') + return + accept = base64.b64encode( + hashlib.sha1(headers['sec-websocket-key'].encode() + _GUID).digest() + ).decode() + conn.sendall( + ( + 'HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n' + f'Connection: Upgrade\r\nSec-WebSocket-Accept: {accept}\r\n\r\n' + ).encode() + ) + while True: + opcode, payload, buffer = self._read_frame(conn, buffer) + if opcode == 0x8: + conn.sendall(bytes([0x88, 0])) + return + if opcode != 0x2: + continue + header = protocol.decode_header(payload) + self.received.append((header, payload)) + if header['type'] == 'hello': + self._send( + conn, + protocol.encode_frame( + { + 'type': 'hello', + 'application': {'name': 'fake', 'version': '0'}, + 'crs': header['crs'], + 'origin': self.origin, + } + ), + ) + self._send( + conn, + protocol.encode_frame( + { + 'type': 'scene.inventory', + 'entries': [ + {'id': i, 'object_type': 'Surface', 'revision': 1} + for i in self.inventory + ], + } + ), + ) + # a ping, to see that the client answers + conn.sendall(bytes([0x89, 0])) + elif header.get('id') in self.reply_errors: + self._send( + conn, + protocol.encode_frame( + { + 'type': 'error', + 'code': 'invalid_object', + 'object': header['id'], + 'message': self.reply_errors[header['id']], + } + ), + ) + + @staticmethod + def _read_frame(conn, buffer): + head, buffer = _read(conn, 2, buffer) + opcode, length = head[0] & 0x0F, head[1] & 0x7F + assert head[1] & 0x80, 'a client frame must be masked' + if length == 126: + raw, buffer = _read(conn, 2, buffer) + (length,) = struct.unpack('>H', raw) + elif length == 127: + raw, buffer = _read(conn, 8, buffer) + (length,) = struct.unpack('>Q', raw) + key, buffer = _read(conn, 4, buffer) + payload, buffer = _read(conn, length, buffer) + payload = bytes(b ^ key[i % 4] for i, b in enumerate(payload)) + return opcode, payload, buffer + + @staticmethod + def _send(conn, payload): + length = len(payload) + if length < 126: + head = bytes([0x82, length]) + elif length < 1 << 16: + head = bytes([0x82, 126]) + struct.pack('>H', length) + else: + head = bytes([0x82, 127]) + struct.pack('>Q', length) + conn.sendall(head + payload) + + # -- helpers for the tests --------------------------------------------- + + def types(self): + return [header['type'] for header, _ in self.received] + + def headers(self, kind): + return [header for header, _ in self.received if header['type'] == kind] diff --git a/tests/unit/test_geoviewer_link.py b/tests/unit/test_geoviewer_link.py new file mode 100644 index 00000000..f358799c --- /dev/null +++ b/tests/unit/test_geoviewer_link.py @@ -0,0 +1,94 @@ +"""Tests for the WebSocket client and the push session (main/geoviewer/link.py). + +A fake viewer on 127.0.0.1 plays the part of GeoViewer. +""" + +import pytest + +from loopstructural.main.geoviewer import protocol +from loopstructural.main.geoviewer.link import GeoViewerLink, LinkError +from loopstructural.main.geoviewer.websocket_client import WebSocketClient, WebSocketError + +from .geoviewer_fake_server import FakeViewer + + +@pytest.fixture +def viewer(): + fake = FakeViewer() + yield fake + fake.close() + + +def triangle(object_id): + return protocol.surface_message( + object_id, object_id, [[0, 0, 0], [1, 0, 0], [0, 1, 0]], [[0, 1, 2]] + ) + + +def test_wrong_token_is_refused(viewer): + with pytest.raises(WebSocketError, match='token'): + WebSocketClient.connect('127.0.0.1', viewer.port, token='wrong') + + +def test_nothing_listens(): + fake = FakeViewer() + port = fake.port + fake.close() + with pytest.raises(LinkError, match='Cannot connect'): + GeoViewerLink(port, 'secret').push('EPSG:1', (0, 0, 0), lambda origin: []) + + +def test_push_says_hello_and_sends_the_objects(viewer): + seen_origin = [] + + def build(origin): + seen_origin.append(list(origin)) + yield 'a', triangle('a') + yield 'b', triangle('b') + + result = GeoViewerLink(viewer.port, 'secret').push('EPSG:28350', (1, 2, 3), build) + assert viewer.types() == ['hello', 'scene.upsert', 'scene.upsert'] + assert viewer.headers('hello')[0]['origin'] == [1.0, 2.0, 3.0] + assert viewer.headers('hello')[0]['crs'] == 'EPSG:28350' + # the viewer origin applies, not the origin that the plugin proposed + assert seen_origin == [[10.0, 20.0, 0.0]] + assert result.sent == 2 and result.removed == 0 and not result.refused + + +def test_push_removes_old_objects_with_the_prefix_only(): + fake = FakeViewer(inventory=['loopstructural/old', 'loopstructural/keep', 'other/file']) + try: + result = GeoViewerLink(fake.port, 'secret').push( + 'EPSG:1', + (0, 0, 0), + lambda origin: [('loopstructural/keep', triangle('loopstructural/keep'))], + replace_prefix='loopstructural/', + ) + finally: + fake.close() + assert [h['id'] for h in fake.headers('scene.remove')] == ['loopstructural/old'] + assert result.removed == 1 + + +def test_refused_object_is_reported(): + fake = FakeViewer(reply_errors={'a': 'the surface is not valid'}) + try: + result = GeoViewerLink(fake.port, 'secret').push( + 'EPSG:1', (0, 0, 0), lambda origin: [('a', triangle('a')), ('b', triangle('b'))] + ) + finally: + fake.close() + assert result.sent == 2 + assert result.refused == [('a', 'the surface is not valid')] + assert '1 messages' in result.summary() + + +def test_large_message_is_sent_whole(viewer): + import numpy as np + + count = 40000 + vertices = np.random.default_rng(0).random((count, 3)) * 100 + triangles = np.arange(count - count % 3).reshape(-1, 3) + frame = protocol.surface_message('big', 'big', vertices, triangles) + GeoViewerLink(viewer.port, 'secret').push('EPSG:1', (0, 0, 0), lambda origin: [('big', frame)]) + assert viewer.received[-1][1] == frame diff --git a/tests/unit/test_geoviewer_protocol.py b/tests/unit/test_geoviewer_protocol.py new file mode 100644 index 00000000..c04ac3bc --- /dev/null +++ b/tests/unit/test_geoviewer_protocol.py @@ -0,0 +1,120 @@ +"""Tests for the frames of the GeoViewer link (main/geoviewer/protocol.py).""" + +import json +import struct + +import numpy as np +import pytest + +from loopstructural.main.geoviewer import protocol + + +def split(frame): + """Return the header and the bytes after it.""" + (length,) = struct.unpack_from('