Jeffrey Gelens avatar Jeffrey Gelens committed b66e3d8

Refactored websocket passthrough. WebSocket can now be found using
environ['wsgi.websocket'].
Handler will fallback to default PyWSGI functionality in case of a normal HTTP
REquest

Comments (0)

Files changed (5)

examples/gunicorn_websocket.py

 # demo app
 import os
 import random
-def app(environ, start_response, ws=None):
-    if ws.path == '/echo':
-        while True:
-            m = ws.wait()
-            if m is None:
-                break
-            ws.send(m)
-
-    elif ws.path == '/data':
+def app(environ, start_response):
+    if environ['PATH_INFO'] == '/test':
+        start_response("200 OK", [('Content-Type', 'text/plain')])
+        return ["blaat"]
+    elif environ['PATH_INFO'] == "/data":
+        ws = environ['wsgi.websocket']
         for i in xrange(10000):
             ws.send("0 %s %s\n" % (i, random.random()))
             gevent.sleep(1)
+    else:
+        start_response("404 Not Found", [])
+        return []
+
 # demo app
 import os
 import random
-def handle(environ, start_response, ws):
+def handle(ws):
     """  This is the websocket handler function.  Note that we
     can dispatch based on path in here, too."""
     if ws.path == '/echo':
             #print "0 %s %s\n" % (i, random.random())
             gevent.sleep(0.1)
 
-server = pywsgi.WSGIServer(('0.0.0.0', 9999), handle,
+
+def app(environ, start_response):
+    if environ['PATH_INFO'] == '/test':
+        start_response("200 OK", [('Content-Type', 'text/plain')])
+        return ["blaat"]
+    elif environ['PATH_INFO'] == "/data":
+        handle(environ['wsgi.websocket'])
+    else:
+        start_response("404 Not Found", [])
+        return []
+
+
+
+server = pywsgi.WSGIServer(('0.0.0.0', 8000), app,
         handler_class=WebSocketHandler)
 server.serve_forever()

geventwebsocket/handler.py

 from geventwebsocket import WebSocket
 
 class WebSocketHandler(WSGIHandler):
+    def __init__(self, *args, **kwargs):
+        self.websocket_connection = False
+        super(WebSocketHandler, self).__init__(*args, **kwargs)
+
     def handle_one_response(self):
-        self.time_start = time.time()
-        self.status = None
-        self.response_length = 0
-
+        # In case the client doesn't want to initialize a WebSocket connection
+        # we will proceed with the default PyWSGI functionality.
         if self.environ.get("HTTP_CONNECTION") != "Upgrade" or \
            self.environ.get("HTTP_UPGRADE") != "WebSocket" or \
            not self.environ.get("HTTP_ORIGIN"):
-            message = "Websocket connection expected"
-            headers = [("Content-Length", str(len(message))),]
-            self.start_response("HTTP/1.1 400 Bad Request", headers, message)
-            self.close_connection = True
-            return
+            return super(WebSocketHandler, self).handle_one_response()
+        else:
+            self.websocket_connection = True
 
-        ws = WebSocket(self.rfile, self.wfile, self.socket, self.environ)
+        self.websocket = WebSocket(self.rfile, self.wfile, self.socket, self.environ)
+        self.environ['wsgi.websocket'] = self.websocket
         challenge = self._get_challenge()
 
         headers = [
             ("Upgrade", "WebSocket"),
             ("Connection", "Upgrade"),
-            ("Sec-WebSocket-Origin", ws.origin),
-            ("Sec-WebSocket-Protocol", ws.protocol),
-            ("Sec-WebSocket-Location", "ws://" + self.environ.get('HTTP_HOST') + ws.path),
+            ("Sec-WebSocket-Origin", self.websocket.origin),
+            ("Sec-WebSocket-Protocol", self.websocket.protocol),
+            ("Sec-WebSocket-Location", "ws://" + self.environ.get('HTTP_HOST') + self.websocket.path),
         ]
 
-        self.start_response(
-            "HTTP/1.1 101 Web Socket Protocol Handshake", headers, challenge
-        )
+        self.start_response("101 Web Socket Protocol Handshake", headers)
+        self.write([challenge])
 
-        try:
-            self.application(self.environ, self.start_response, ws)
-        except Exception:
-            traceback.print_exc()
-            sys.exc_clear()
-            try:
-                args = (getattr(self, 'server', ''),
-                        getattr(self, 'requestline', ''),
-                        getattr(self, 'client_address', ''),
-                        getattr(self, 'application', ''))
-                msg = '%s: Failed to handle request:\n  request = %s from %s\n  application = %s\n\n' % args
-                sys.stderr.write(msg)
-            except Exception:
-                sys.exc_clear()
-        finally:
-            self.wsgi_input._discard()
-            self.time_finish = time.time()
-            self.log_request()
+        return self.application(self.environ, self.start_response)
 
-    def start_response(self, status, headers, body=None):
-        towrite = [status]
-        for header in headers:
-            towrite.append(": ".join(header))
+    def write(self, data):
+        if self.websocket_connection:
+            self.wfile.writelines(data)
+        else:
+            super(WebSocketHandler, self).write(data)
 
-        if body is not None:
-            towrite.append("")
-            towrite.append(body)
+    def start_response(self, status, headers, exc_info=None):
+        if self.websocket_connection:
+            self.status = status
 
-        self.wfile.write("\r\n".join(towrite))
+            towrite = []
+            towrite.append('%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.headers_sent = True
+        else:
+            super(WebSocketHandler, self).start_response(status, headers, exc_info)
 
     def _get_key_value(self, key_value):
         key_number = int(re.sub("\\D", "", key_value))
         if not (key1 and key2):
             message = "Client using old protocol implementation"
             headers = [("Content-Length", str(len(message))),]
-            self.start_response("HTTP/1.1 400 Bad Request", headers, message)
+            self.start_response("400 Bad Request", headers)
+            self.write([message])
             self.close_connection = True
             return
 

geventwebsocket/websocket.py

+# This class implements the Websocket protocol draft version as of May 23, 2010
+# The version as of August 6, 2010 will be implementend once Firefox or
+# Webkit-trunk support this version.
+
 class WebSocket(object):
     def __init__(self, rfile, wfile, sock, environ):
         self.rfile = rfile

tests/test__websocket.py

 class TestWebSocket(TestCase):
     message = "\x00Hello world\xff"
 
-    def application(self, environ, start_response, ws):
+    def application(self, environ, start_response):
         if environ['PATH_INFO'] == "/echo":
+            try:
+                ws = environ['wsgi.websocket']
+            except KeyError:
+                start_response("400 Bad Request", [])
+                return []
+
             while True:
                 message = ws.wait()
                 if message is None:
 
     def test_badrequest(self):
         fd = self.connect().makefile(bufsize=1)
-        fd.write('GET / HTTP/1.1\r\nHost: localhost\r\n\r\n')
-        read_http(fd, code=400, reason='Bad Request', body='Websocket connection expected')
+        fd.write('GET /echo HTTP/1.1\r\nHost: localhost\r\n\r\n')
+        read_http(fd, code=400, reason='Bad Request')
         fd.close()
 
     def test_oldprotocol_version(self):
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.