Merge pull request #124 from TheSpeedM/main

Add websocket header support to the ROS-client

CHANGELOG.rst

@@ -12,6 +12,8 @@ Unreleased
 
 **Added**
 
+* Added websocket header support to the ROS-client.
+
 **Changed**
 
 **Fixed**

requirements-dev.txt

@@ -10,4 +10,5 @@ pydocstyle
 pytest>=6.0
 sphinx >=3.4
 twine
+websockets >= 12.0
 -e .

src/roslibpy/ros.py

@@ -32,12 +32,13 @@ class Ros(object):
         host (:obj:`str`): Name or IP address of the ROS bridge host, e.g. ``127.0.0.1``.
         port (:obj:`int`): ROS bridge port, e.g. ``9090``.
         is_secure (:obj:`bool`): ``True`` to use a secure web sockets connection, otherwise ``False``.
+        headers (:obj:`dict`): Additional headers to include in the WebSocket connection.
     """
 
-    def __init__(self, host, port=None, is_secure=False):
+    def __init__(self, host, port=None, is_secure=False, headers=None):
         self._id_counter = 0
         url = RosBridgeClientFactory.create_url(host, port, is_secure)
-        self.factory = RosBridgeClientFactory(url)
+        self.factory = RosBridgeClientFactory(url, headers=headers)
         self.is_connecting = False
         self.connect()
 

tests/test_ws_headers.py

@@ -0,0 +1,66 @@
+from __future__ import print_function
+
+import asyncio
+import threading
+import time
+
+import websockets
+
+from roslibpy import Ros
+
+headers = {
+    'cookie': 'token=rosbridge',
+    'authorization': 'Some auth'
+}
+
+
+async def websocket_handler(websocket, path):
+    request_headers = websocket.request_headers
+    for key, value in headers.items():
+        assert request_headers.get(key) == value, f"Header {key} did not match expected value {value}"
+    await websocket.close()
+
+
+async def start_server(stop_event):
+    server = await websockets.serve(websocket_handler, '127.0.0.1', 9000)
+    await stop_event.wait()
+    server.close()
+    await server.wait_closed()
+
+
+def run_server(stop_event):
+    asyncio.run(start_server(stop_event))
+
+
+def run_client():
+    client = Ros('127.0.0.1', 9000, headers=headers)
+    client.run()
+    client.close()
+
+
+def test_websocket_headers():
+    server_stop_event = asyncio.Event()
+    stop_event = threading.Event()
+
+    server_thread = threading.Thread(target=run_server, args=(server_stop_event,))
+    server_thread.start()
+
+    time.sleep(1)  # Give the server time to start
+
+    client_thread = threading.Thread(target=run_client)
+    client_thread.start()
+
+    # Wait for the client thread to finish or timeout after 10 seconds
+    client_thread.join(timeout=10)
+
+    if client_thread.is_alive():
+        raise Exception("Client did not terminate as expected")
+
+    # Signal the server to stop
+    server_stop_event.set()
+    server_thread.join(timeout=10)
+
+    if server_thread.is_alive():
+        raise Exception("Server did not stop as expected")
+
+    stop_event.set()