Fix Lawouach#167 hopefully by reading the internal SSL socket object's buffer

Lawouach · Lawouach · commit 75b88bdecd98 · 2015-09-14T22:40:36.000+02:00
diff --git a/example/bug167_client.py b/example/bug167_client.py
@@ -0,0 +1,65 @@
+def run_threaded():
+    from ws4py.client.threadedclient import WebSocketClient
+    class EchoClient(WebSocketClient):
+        def opened(self):
+            self.send("hello")
+            
+        def closed(self, code, reason=None):
+            print(("Closed down", code, reason))
+
+        def received_message(self, m):
+            print(m)
+            self.close()
+            
+    try:
+        ws = EchoClient('wss://localhost:9000/ws')
+        ws.connect()
+        ws.run_forever()
+    except KeyboardInterrupt:
+        ws.close()
+
+def run_tornado():
+    from tornado import ioloop
+    from ws4py.client.tornadoclient import TornadoWebSocketClient
+    class MyClient(TornadoWebSocketClient):
+        def opened(self):
+            self.send("hello")
+            
+        def closed(self, code, reason=None):
+            print(("Closed down", code, reason))
+            ioloop.IOLoop.instance().stop()
+
+        def received_message(self, m):
+            print(m)
+            self.close()
+
+    ws = MyClient('wss://localhost:9000/ws')
+    ws.connect()
+
+    ioloop.IOLoop.instance().start()
+
+def run_gevent():
+    from gevent import monkey; monkey.patch_all()
+    import gevent
+    from ws4py.client.geventclient import WebSocketClient
+    
+    ws = WebSocketClient('wss://localhost:9000/ws')
+    ws.connect()
+
+    ws.send("hello")
+    
+    def incoming():
+        while True:
+            m = ws.receive()
+            if m is not None:
+                print(m)
+            else:
+                break
+            
+        ws.close()
+
+    gevent.joinall([gevent.spawn(incoming)])
+
+#run_gevent()
+run_threaded()
+run_tornado()
diff --git a/ws4py/client/__init__.py b/ws4py/client/__init__.py
@@ -205,7 +205,8 @@ def connect(self):
         if self.scheme == "wss":
             # default port is now 443; upgrade self.sender to send ssl
             self.sock = ssl.wrap_socket(self.sock, **self.ssl_options)
-
+            self._is_secure = True
+            
         self.sock.connect(self.bind_addr)
 
         self._write(self.handshake_request)
diff --git a/ws4py/client/tornadoclient.py b/ws4py/client/tornadoclient.py
@@ -35,6 +35,7 @@ def closed(self, code, reason=None):
                                      ssl_options=ssl_options, headers=headers)
         if self.scheme == "wss":
             self.sock = ssl.wrap_socket(self.sock, do_handshake_on_connect=False, **self.ssl_options)
+            self._is_secure = True
             self.io = iostream.SSLIOStream(self.sock, io_loop, ssl_options=self.ssl_options)
         else:
             self.io = iostream.IOStream(self.sock, io_loop)
diff --git a/ws4py/websocket.py b/ws4py/websocket.py
@@ -1,10 +1,17 @@
 # -*- coding: utf-8 -*-
 import logging
 import socket
+import ssl
 import time
 import threading
 import types
 
+try:
+    from OpenSSL.SSL import Error as pyOpenSSLError
+except ImportError:
+    class pyOpenSSLError(Exception):
+        pass
+    
 from ws4py import WS_KEY, WS_VERSION
 from ws4py.exc import HandshakeError, StreamClosed
 from ws4py.streaming import Stream
@@ -99,7 +106,12 @@ def __init__(self, sock, protocols=None, extensions=None, environ=None, heartbea
         """
         Underlying connection.
         """
-
+        
+        self._is_secure = hasattr(sock, '_ssl') or hasattr(sock, '_sslobj')
+        """
+        Tell us if the socket is secure or not.
+        """
+        
         self.client_terminated = False
         """
         Indicates if the client has been marked as terminated.
@@ -301,6 +313,50 @@ def send(self, payload, binary=False):
         else:
             raise ValueError("Unsupported type '%s' passed to send()" % type(payload))
 
+    def _get_from_pending(self):
+        """
+        The SSL socket object provides the same interface
+        as the socket interface but behaves differently.
+
+        When data is sent over a SSL connection
+        more data may be read than was requested from by 
+        the ws4py websocket object.
+
+        In that case, the data may have been indeed read 
+        from the underlying real socket, but not read by the 
+        application which will expect another trigger from the 
+        manager's polling mechanism as if more data was still on the
+        wire. This will happen only when new data is 
+        sent by the other peer which means there will be
+        some delay before the initial read data is handled
+        by the application.
+
+        Due to this, we have to rely on a non-public method
+        to query the internal SSL socket buffer if it has indeed
+        more data pending in its buffer.
+
+        Now, some people in the Python community 
+        `discourage <https://bugs.python.org/issue21430>`_
+        this usage of the ``pending()`` method because it's not
+        the right way of dealing with such use case. They advise
+        `this approach <https://docs.python.org/dev/library/ssl.html#notes-on-non-blocking-sockets>`_
+        instead. Unfortunately, this applies only if the
+        application can directly control the poller which is not
+        the case with the WebSocket abstraction here.
+
+        We therefore rely on this `technic <http://stackoverflow.com/questions/3187565/select-and-ssl-in-python>`_
+        which seems to be valid anyway.
+
+        This is a bit of a shame because we have to process 
+        more data than what wanted initially.
+        """
+        data = b""
+        pending = self.sock.pending()
+        while pending:
+            data += self.sock.recv(pending)
+            pending = self.sock.pending()
+        return data
+        
     def once(self):
         """
         Performs the operation of reading from the underlying
@@ -322,7 +378,10 @@ def once(self):
 
         try:
             b = self.sock.recv(self.reading_buffer_size)
-        except (socket.error, OSError) as e:
+            # This will only make sense with secure sockets.
+            if self._is_secure:
+                b += self._get_from_pending()
+        except (socket.error, OSError, pyOpenSSLError) as e:
             self.unhandled_error(e)
             return False
         else:
