Source code for ws4py.server.wsgi.middleware

# -*- coding: utf-8 -*-
import copy
import base64
from hashlib import sha1
import types
import socket

from ws4py import WS_KEY, WS_VERSION
from ws4py.exc import HandshakeError, StreamClosed
from ws4py.streaming import Stream
from ws4py.websocket import WebSocket

[docs]class WebSocketUpgradeMiddleware(object): def __init__(self, app, fallback_app=None, protocols=None, extensions=None, websocket_class=WebSocket, versions=WS_VERSION): """ WSGI middleware that performs the WebSocket upgrade handshake. .. code-block:: python :linenos: def ws_handler(websocket): ... app = WebSocketUpgradeMiddleware(ws_handler) If the handshake succeeds, it calls ``app`` with an instance of ``websocket_class`` with a copy of the environ dictionary. If an error occurs and ``fallback_app`` is provided, it must be a WSGI application which will be called. Otherwise it returns a simple error through the inner ``start_response``. One interesting aspect is that wsgiref fails with this middleware due to the ``Upgrade`` hop-by-hop header which is not allowed. Make sure that your server does not close the underlying socket for you since it would close the whole WebSocket connection as well. You may provide your own representation of the socket by setting the environ key: ``'upgrade.socket'``. Otherwise, ``'wsgi.input'._sock`` will be used. """ self.app = app self.fallback_app = fallback_app self.protocols = protocols self.extensions = extensions self.websocket_class = websocket_class self.versions = versions self.supported_versions = ', '.join([str(v) for v in versions]) def __call__(self, environ, start_response): # Initial handshake validation try: if 'websocket' not in environ.get('upgrade.protocol', environ.get('HTTP_UPGRADE', '')).lower(): raise HandshakeError("Upgrade protocol is not websocket") if environ.get('REQUEST_METHOD') != 'GET': raise HandshakeError('Method is not GET') key = environ.get('HTTP_SEC_WEBSOCKET_KEY') if key: ws_key = base64.b64decode(key) if len(ws_key) != 16: raise HandshakeError("WebSocket key's length is invalid") else: raise HandshakeError("Not a valid HyBi WebSocket request") version = environ.get('HTTP_SEC_WEBSOCKET_VERSION') version_is_valid = False if version: try: version = int(version) except: pass else: version_is_valid = version in self.versions if not version_is_valid: raise HandshakeError('Unsupported WebSocket version: %s' % version) environ['websocket.version'] = str(version) except HandshakeError as e: if self.fallback_app: return self.fallback_app(environ, start_response) else: start_response("400 Bad Handshake", [('Sec-WebSocket-Version', self.supported_versions)]) return [str(e)] # Collect supported subprotocols protocols = self.protocols or [] subprotocols = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL') ws_protocols = [] if subprotocols: for s in subprotocols.split(','): s = s.strip() if s in protocols: ws_protocols.append(s) # Collect supported extensions exts = self.extensions or [] ws_extensions = [] extensions = environ.get('HTTP_SEC_WEBSOCKET_EXTENSIONS') if extensions: for ext in extensions.split(','): ext = ext.strip() if ext in exts: ws_extensions.append(ext) # Build and start the HTTP response headers = [ ('Upgrade', 'websocket'), ('Connection', 'Upgrade'), ('Sec-WebSocket-Version', environ['websocket.version']), ('Sec-WebSocket-Accept', base64.b64encode(sha1(key + WS_KEY).digest())), ] if ws_protocols: headers.append(('Sec-WebSocket-Protocol', ', '.join(ws_protocols))) if ws_extensions: headers.append(('Sec-WebSocket-Extensions', ','.join(ws_extensions))) start_response("101 Web Socket Hybi Handshake", headers) if 'upgrade.socket' in environ: upgrade_socket = environ['upgrade.socket'] else: upgrade_socket = environ['wsgi.input']._sock return self.app(self.websocket_class(upgrade_socket, ws_protocols, ws_extensions, environ.copy()))