progrium · denik · Oct 2, 2011 · Oct 8, 2011 · Oct 8, 2011 · Oct 8, 2011
diff --git a/example/echo_gevent_server.py b/example/echo_gevent_server.py
@@ -0,0 +1,84 @@
+# -*- coding: utf-8 -*-
+from ws4py.server.geventserver import WebSocketServer
+from gevent import spawn
+from socket import error
+from collections import deque
+
+
+HTML = """<html>
+    <head>
+      <script type='application/javascript' src='https://ajax.googleapis.com/ajax/libs/jquery/1.4.2/jquery.min.js'> </script>
+      <script type='application/javascript'>
+
+        try { WebSocket } catch(err) { WebSocket = MozWebSocket; }
+
+        $(document).ready(function() {
+          var ws = new WebSocket('ws://localhost:9000/ws');
+          $(window).unload(function() {
+             ws.close();
+          });
+          ws.onmessage = function (evt) {
+             $('#chat').val($('#chat').val() + evt.data + '\\n');
+          };
+          ws.onopen = function() {
+             ws.send("%(username)s entered the room");
+          };
+          $('#chatform').submit(function() {
+             ws.send('%(username)s: ' + $('#message').val());
+             $('#message').val("");
+             return false;
+          });
+        });
+      </script>
+    </head>
+    <body>
+    <form action='/echo' id='chatform' method='get'>
+      <textarea id='chat' cols='35' rows='10'></textarea>
+      <br />
+      <label for='message'>%(username)s: </label><input type='text' id='message' />
+      <input type='submit' value='Send' />
+      </form>
+    </body>
+    </html>
+    """
+
+
+sockets = set()
+messagelog = deque(maxlen=20)
+
+
+def _send(socket, message):
+    try:
+        socket.send(message)
+    except error:
+        pass
+
+
+def broadcast(message):
+    messagelog.append(message)
+    for socket in sockets:
+        spawn(_send, socket, message)
+
+
+def application(environ, start_response):
+    websocket = environ.get('wsgi.websocket')
+    if websocket is None:
+        start_response('200 OK', [])
+        return [HTML % {'username': environ['REMOTE_ADDR']}]
+    else:
+        for msg in messagelog:
+            websocket.send(msg)
+        sockets.add(websocket)
+        try:
+            while True:
+                message = websocket.receive()
+                if not message:
+                    broadcast('%s left the room' % environ['REMOTE_ADDR'])
+                    break
+                broadcast(message)
+        finally:
+            sockets.discard(websocket)
+
+
+if __name__ == '__main__':
+    WebSocketServer(('0.0.0.0', 9000), application).serve_forever()
diff --git a/ws4py/server/geventserver.py b/ws4py/server/geventserver.py
@@ -2,81 +2,57 @@
 
 from ws4py.server.wsgi.middleware import WebSocketUpgradeMiddleware
 
+
+# This should probably be built-in in gevent:
+# environ['wsgi.socket'] is the connection that server puts here for applications like WebSocket
+# environ['wsgi.socket.detach'] is a flag that application can set if it wants to use socket separately
+# if it's True gevent server won't do anything about the connection anymore
+
+
 class UpgradableWSGIHandler(gevent.pywsgi.WSGIHandler):
-    """Upgradable version of gevent.pywsgi.WSGIHandler class
-
-    This is a drop-in replacement for gevent.pywsgi.WSGIHandler that supports
-    protocol upgrades via WSGI environment. This means you can create upgraders
-    as WSGI apps or WSGI middleware.
-
-    If an HTTP request comes in that includes the Upgrade header, it will add
-    to the environment two items:
-
-    `upgrade.protocol` 
-      The protocol to upgrade to. Checking for this lets you know the request
-      wants to be upgraded and the WSGI server supports this interface. 
-
-    `upgrade.socket`
-      The raw Python socket object for the connection. From this you can do any
-      upgrade negotiation and hand it off to the proper protocol handler.
-
-    The upgrade must be signalled by starting a response using the 101 status
-    code. This will inform the server to flush the headers and response status
-    immediately, not to expect the normal WSGI app return value, and not to 
-    look for more HTTP requests on this connection. 
-
-    To use this handler with gevent.pywsgi.WSGIServer, you can pass it to the
-    constructor:
-
-    server = WSGIServer(('127.0.0.1', 80), app, 
-                            handler_class=UpgradableWSGIHandler)
-
-    Alternatively, you can specify it as a class variable for a WSGIServer 
-    subclass:
-
-    class UpgradableWSGIServer(gevent.pywsgi.WSGIServer):
-        handler_class = UpgradableWSGIHandler
-
-    """
+
+    def start_response_for_upgrade(self, status, headers, exc_info=None):
+        write = self.start_response(status, headers, exc_info)
+        if self.code == 101:
+            # flushes headers now
+            towrite = ['%s %s\r\n' % (self.request_version, self.status)]
+            for header in headers:
+                towrite.append('%s: %s\r\n' % header)
+            towrite.append('\r\n')
+            self.wfile.writelines(towrite)
+            self.response_length += sum(len(x) for x in towrite)
+        return write
+
     def run_application(self):
-        upgrade_header = self.environ.get('HTTP_UPGRADE', '').lower()
-        if upgrade_header:
-            self.environ['upgrade.protocol'] = upgrade_header
-            self.environ['upgrade.socket'] = self.socket
-            def start_response_for_upgrade(status, headers, exc_info=None):
-                write = self.start_response(status, headers, exc_info)
-                if self.code == 101:
-                    # flushes headers now
-                    towrite = ['%s %s\r\n' % (self.request_version, self.status)]
-                    for header in headers:
-                        towrite.append('%s: %s\r\n' % header)
-                    towrite.append('\r\n')
-                    self.wfile.writelines(towrite)
-                    self.response_length += sum(len(x) for x in towrite)
-                return write
-            try:
-                self.result = self.application(self.environ, start_response_for_upgrade)
-                if self.code != 101:
-                    self.process_result()
-            finally:
-                if hasattr(self, 'code') and self.code == 101:
-                    self.rfile.close() # makes sure we stop processing requests
-        else:
-            gevent.pywsgi.WSGIHandler.run_application(self)
+        self.environ['wsgi.socket'] = self.socket
+        try:
+            self.result = self.application(self.environ, self.start_response_for_upgrade)
+            if not self.environ.get('wsgi.socket.detach'):
+                self.process_result()
+        finally:
+            if self.environ.get('wsgi.socket.detach'):
+                self.rfile.close() # makes sure we stop processing requests
+                self.socket = None  # otherwise gevent server would close the connection
+
 
 class WebSocketServer(gevent.pywsgi.WSGIServer):
     handler_class = UpgradableWSGIHandler
-    
+
     def __init__(self, *args, **kwargs):
         gevent.pywsgi.WSGIServer.__init__(self, *args, **kwargs)
         protocols = kwargs.pop('websocket_protocols', [])
         extensions = kwargs.pop('websocket_extensions', [])
-        self.application = WebSocketUpgradeMiddleware(self.application, 
+        self.application = WebSocketUpgradeMiddleware(self.application,
                             protocols=protocols,
-                            extensions=extensions)    
+                            extensions=extensions)
 
 if __name__ == '__main__':
-    def echo_handler(websocket, environ):
+
+    def echo_handler(environ, start_response):
+        if start_response is not None:
+            start_response('400 Bad Request', [])
+            return []
+        websocket = environ['wsgi.websocket']
         try:
             while True:
                 msg = websocket.receive(msg_obj=True)
@@ -86,6 +62,6 @@ def echo_handler(websocket, environ):
                     break
         finally:
             websocket.close()
-    
+
     server = WebSocketServer(('127.0.0.1', 9000), echo_handler)
-    server.serve_forever()
+    server.serve_forever()
diff --git a/ws4py/server/wsgi/middleware.py b/ws4py/server/wsgi/middleware.py
@@ -61,12 +61,12 @@ def __init__(self, handle, fallback_app=None, protocols=None, extensions=None,
         self.extensions = extensions
         self.websocket_class = websocket_class
 
-    def __call__(self, environ, start_response):        
+    def __call__(self, environ, start_response):
+        if 'websocket' not in environ.get('HTTP_UPGRADE', '').lower():
+            return self.handle(environ, start_response)
+
         # Initial handshake validation
         try:
-            if 'websocket' not in environ.get('upgrade.protocol', ''):
-                raise HandshakeError("Upgrade protocol is not websocket")
-
             if environ.get('REQUEST_METHOD') != 'GET':
                 raise HandshakeError('Method is not GET')
 
@@ -125,12 +125,10 @@ def __call__(self, environ, start_response):
             headers.append(('Sec-WebSocket-Extensions', ','.join(ws_extensions)))
 
         start_response("101 Web Socket Hybi Handshake", headers)
-
-        # Build a websocket object and pass it to the handler
-        self.handle(
-            self.websocket_class(
-                environ.get('upgrade.socket'), 
-                ws_protocols, 
-                ws_extensions, 
-                environ), 
-            environ)
+        environ['wsgi.socket.detach'] = True
+        websocket = self.websocket_class(environ['wsgi.socket'],
+                                         ws_protocols,
+                                         ws_extensions,
+                                         environ)
+        environ['wsgi.websocket'] = websocket
+        self.handle(environ, None)