Add headers to ws connection

Author: TheSpeedM

src/roslibpy/comm/comm_autobahn.py

@@ -23,10 +23,17 @@
 class AutobahnRosBridgeProtocol(RosBridgeProtocol, WebSocketClientProtocol):
     def __init__(self, *args, **kwargs):
         super(AutobahnRosBridgeProtocol, self).__init__(*args, **kwargs)
+        self.headers = {}
 
     def onConnect(self, response):
         LOGGER.debug("Server connected: %s", response.peer)
 
+    def getHandshakeRequestHeaders(self):
+        headers = super(AutobahnRosBridgeProtocol, self).getHandshakeRequestHeaders()
+        for key, value in self.headers.items():
+            headers.append((key, value))
+        return headers
+
     def onOpen(self):
         LOGGER.info("Connection to ROS ready.")
         self._manual_disconnect = False
@@ -62,13 +69,20 @@ class AutobahnRosBridgeClientFactory(EventEmitterMixin, ReconnectingClientFactor
 
     protocol = AutobahnRosBridgeProtocol
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, headers=None, **kwargs):
         super(AutobahnRosBridgeClientFactory, self).__init__(*args, **kwargs)
+        self.headers = headers or {}
         self._proto = None
         self._manager = None
         self.connector = None
         self.setProtocolOptions(closeHandshakeTimeout=5)
 
+    def buildProtocol(self, addr):
+        proto = self.protocol()
+        proto.factory = self
+        proto.headers = self.headers
+        return proto
+
     def connect(self):
         """Establish WebSocket connection to the ROS server defined for this factory."""
         self.connector = connectWS(self)

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()