Correct handling of cookies in the client (Fixes #162)

engineio/asyncio_client.py

@@ -243,18 +243,28 @@ async def _connect_websocket(self, url, headers, engineio_path):
         if self.http is None or self.http.closed:  # pragma: no cover
             self.http = aiohttp.ClientSession()
 
+        # extract any new cookies passed in a header so that they can also be
+        # sent the the WebSocket route
+        cookies = {}
+        for header, value in headers.items():
+            if header.lower() == 'cookie':
+                cookies = dict(
+                    [cookie.split('=') for cookie in value.split('; ')])
+                del headers[header]
+                break
+
         try:
             if not self.ssl_verify:
                 ssl_context = ssl.create_default_context()
                 ssl_context.check_hostname = False
                 ssl_context.verify_mode = ssl.CERT_NONE
                 ws = await self.http.ws_connect(
                     websocket_url + self._get_url_timestamp(),
-                    headers=headers, ssl=ssl_context)
+                    headers=headers, cookies=cookies, ssl=ssl_context)
             else:
                 ws = await self.http.ws_connect(
                     websocket_url + self._get_url_timestamp(),
-                    headers=headers)
+                    headers=headers, cookies=cookies)
         except (aiohttp.client_exceptions.WSServerHandshakeError,
                 aiohttp.client_exceptions.ServerConnectionError):
             if upgrade:

engineio/client.py

@@ -360,6 +360,13 @@ def _connect_websocket(self, url, headers, engineio_path):
         if self.http:
             cookies = '; '.join(["{}={}".format(cookie.name, cookie.value)
                                  for cookie in self.http.cookies])
+            for header, value in headers.items():
+                if header.lower() == 'cookie':
+                    if cookies:
+                        cookies += '; '
+                    cookies += value
+                    del headers[header]
+                    break
 
         try:
             if not self.ssl_verify:

tests/asyncio/test_asyncio_client.py

@@ -464,7 +464,7 @@ def test_websocket_connection_failed(self, _time):
                       headers={'Foo': 'Bar'}))
         c.http.ws_connect.mock.assert_called_once_with(
             'ws://foo/engine.io/?transport=websocket&EIO=3&t=123.456',
-            headers={'Foo': 'Bar'})
+            headers={'Foo': 'Bar'}, cookies={})
 
     @mock.patch('engineio.client.time.time', return_value=123.456)
     def test_websocket_upgrade_failed(self, _time):
@@ -477,7 +477,7 @@ def test_websocket_upgrade_failed(self, _time):
             'http://foo', transports=['websocket'])))
         c.http.ws_connect.mock.assert_called_once_with(
             'ws://foo/engine.io/?transport=websocket&EIO=3&sid=123&t=123.456',
-            headers={})
+            headers={}, cookies={})
 
     def test_websocket_connection_no_open_packet(self):
         c = asyncio_client.AsyncClient()
@@ -528,7 +528,7 @@ def test_websocket_connection_successful(self, _time):
         self.assertEqual(c.ws, ws)
         c.http.ws_connect.mock.assert_called_once_with(
             'ws://foo/engine.io/?transport=websocket&EIO=3&t=123.456',
-            headers={})
+            headers={}, cookies={})
 
     @mock.patch('engineio.client.time.time', return_value=123.456)
     def test_websocket_https_noverify_connection_successful(self, _time):
@@ -596,7 +596,63 @@ def test_websocket_connection_with_cookies(self, _time):
         time.sleep(0.1)
         c.http.ws_connect.mock.assert_called_once_with(
             'ws://foo/engine.io/?transport=websocket&EIO=3&t=123.456',
-            headers={})
+            headers={}, cookies={})
+
+    @mock.patch('engineio.client.time.time', return_value=123.456)
+    def test_websocket_connection_with_cookie_header(self, _time):
+        c = asyncio_client.AsyncClient()
+        c.http = mock.MagicMock(closed=False)
+        c.http.ws_connect = AsyncMock()
+        ws = c.http.ws_connect.mock.return_value
+        ws.receive = AsyncMock()
+        ws.receive.mock.return_value.data = packet.Packet(packet.OPEN, {
+            'sid': '123', 'upgrades': [], 'pingInterval': 1000,
+            'pingTimeout': 2000
+        }).encode()
+        c.http._cookie_jar = []
+        c._ping_loop = AsyncMock()
+        c._read_loop_polling = AsyncMock()
+        c._read_loop_websocket = AsyncMock()
+        c._write_loop = AsyncMock()
+        on_connect = mock.MagicMock()
+        c.on('connect', on_connect)
+        _run(c.connect('ws://foo',
+                       headers={'Cookie': 'key=value; key2=value2'},
+                       transports=['websocket']))
+        time.sleep(0.1)
+        c.http.ws_connect.mock.assert_called_once_with(
+            'ws://foo/engine.io/?transport=websocket&EIO=3&t=123.456',
+            headers={}, cookies={'key': 'value', 'key2': 'value2'})
+
+    @mock.patch('engineio.client.time.time', return_value=123.456)
+    def test_websocket_connection_with_cookies_and_headers(self, _time):
+        c = asyncio_client.AsyncClient()
+        c.http = mock.MagicMock(closed=False)
+        c.http.ws_connect = AsyncMock()
+        ws = c.http.ws_connect.mock.return_value
+        ws.receive = AsyncMock()
+        ws.receive.mock.return_value.data = packet.Packet(packet.OPEN, {
+            'sid': '123', 'upgrades': [], 'pingInterval': 1000,
+            'pingTimeout': 2000
+        }).encode()
+        c.http._cookie_jar = [mock.MagicMock(), mock.MagicMock()]
+        c.http._cookie_jar[0].key = 'key'
+        c.http._cookie_jar[0].value = 'value'
+        c.http._cookie_jar[1].key = 'key2'
+        c.http._cookie_jar[1].value = 'value2'
+        c._ping_loop = AsyncMock()
+        c._read_loop_polling = AsyncMock()
+        c._read_loop_websocket = AsyncMock()
+        c._write_loop = AsyncMock()
+        on_connect = mock.MagicMock()
+        c.on('connect', on_connect)
+        _run(c.connect('ws://foo',
+                       headers={'Foo': 'Bar', 'Cookie': 'key3=value3'},
+                       transports=['websocket']))
+        time.sleep(0.1)
+        c.http.ws_connect.mock.assert_called_once_with(
+            'ws://foo/engine.io/?transport=websocket&EIO=3&t=123.456',
+            headers={'Foo': 'Bar'}, cookies={'key3': 'value3'})
 
     def test_websocket_upgrade_no_pong(self):
         c = asyncio_client.AsyncClient()

tests/common/test_client.py

@@ -608,6 +608,60 @@ def test_websocket_connection_with_cookies(self, create_connection):
         self.assertEqual(create_connection.call_args[1],
                          {'header': {}, 'cookie': 'key=value; key2=value2'})
 
+    @mock.patch('engineio.client.websocket.create_connection')
+    def test_websocket_connection_with_cookie_header(self, create_connection):
+        create_connection.return_value.recv.return_value = packet.Packet(
+            packet.OPEN, {
+                'sid': '123', 'upgrades': [], 'pingInterval': 1000,
+                'pingTimeout': 2000
+            }).encode()
+        c = client.Client()
+        c.http = mock.MagicMock()
+        c.http.cookies = []
+        c._ping_loop = mock.MagicMock()
+        c._read_loop_polling = mock.MagicMock()
+        c._read_loop_websocket = mock.MagicMock()
+        c._write_loop = mock.MagicMock()
+        on_connect = mock.MagicMock()
+        c.on('connect', on_connect)
+        c.connect('ws://foo', headers={'Foo': 'bar', 'Cookie': 'key=value'},
+                  transports=['websocket'])
+        time.sleep(0.1)
+
+        self.assertEqual(len(create_connection.call_args_list), 1)
+        self.assertEqual(create_connection.call_args[1],
+                         {'header': {'Foo': 'bar'}, 'cookie': 'key=value'})
+
+    @mock.patch('engineio.client.websocket.create_connection')
+    def test_websocket_connection_with_cookies_and_headers(
+            self, create_connection):
+        create_connection.return_value.recv.return_value = packet.Packet(
+            packet.OPEN, {
+                'sid': '123', 'upgrades': [], 'pingInterval': 1000,
+                'pingTimeout': 2000
+            }).encode()
+        c = client.Client()
+        c.http = mock.MagicMock()
+        c.http.cookies = [mock.MagicMock(), mock.MagicMock()]
+        c.http.cookies[0].name = 'key'
+        c.http.cookies[0].value = 'value'
+        c.http.cookies[1].name = 'key2'
+        c.http.cookies[1].value = 'value2'
+        c._ping_loop = mock.MagicMock()
+        c._read_loop_polling = mock.MagicMock()
+        c._read_loop_websocket = mock.MagicMock()
+        c._write_loop = mock.MagicMock()
+        on_connect = mock.MagicMock()
+        c.on('connect', on_connect)
+        c.connect('ws://foo', headers={'Cookie': 'key3=value3'},
+                  transports=['websocket'])
+        time.sleep(0.1)
+
+        self.assertEqual(len(create_connection.call_args_list), 1)
+        self.assertEqual(
+            create_connection.call_args[1],
+            {'header': {}, 'cookie': 'key=value; key2=value2; key3=value3'})
+
     @mock.patch('engineio.client.websocket.create_connection')
     def test_websocket_upgrade_no_pong(self, create_connection):
         create_connection.return_value.recv.return_value = packet.Packet(