1# -*- coding: utf-8 -*-
2__doc__ = """
3WSGI entities to support WebSocket from within gevent.
4
5Its usage is rather simple:
6
7.. code-block: python
8
9    from gevent import monkey; monkey.patch_all()
10    from ws4py.websocket import EchoWebSocket
11    from ws4py.server.geventserver import WSGIServer
12    from ws4py.server.wsgiutils import WebSocketWSGIApplication
13
14    server = WSGIServer(('localhost', 9000), WebSocketWSGIApplication(handler_cls=EchoWebSocket))
15    server.serve_forever()
16
17"""
18import logging
19import sys
20
21import gevent
22from gevent.pywsgi import WSGIHandler, WSGIServer as _WSGIServer
23from gevent.pool import Pool
24
25from ws4py import format_addresses
26from ws4py.server.wsgiutils import WebSocketWSGIApplication
27
28logger = logging.getLogger('ws4py')
29
30__all__ = ['WebSocketWSGIHandler', 'WSGIServer',
31           'GEventWebSocketPool']
32
33class WebSocketWSGIHandler(WSGIHandler):
34    """
35    A WSGI handler that will perform the :rfc:`6455`
36    upgrade and handshake before calling the WSGI application.
37
38    If the incoming request doesn't have a `'Upgrade'` header,
39    the handler will simply fallback to the gevent builtin's handler
40    and process it as per usual.
41    """
42
43    def run_application(self):
44        upgrade_header = self.environ.get('HTTP_UPGRADE', '').lower()
45        if upgrade_header:
46            try:
47                # Build and start the HTTP response
48                self.environ['ws4py.socket'] = self.socket or self.environ['wsgi.input'].rfile._sock
49                self.result = self.application(self.environ, self.start_response) or []
50                self.process_result()
51            except:
52                raise
53            else:
54                del self.environ['ws4py.socket']
55                self.socket = None
56                self.rfile.close()
57
58                ws = self.environ.pop('ws4py.websocket')
59                if ws:
60                    self.server.pool.track(ws)
61        else:
62            gevent.pywsgi.WSGIHandler.run_application(self)
63
64class GEventWebSocketPool(Pool):
65    """
66    Simple pool of bound websockets.
67    Internally it uses a gevent group to track
68    the websockets. The server should call the ``clear``
69    method to initiate the closing handshake when the
70    server is shutdown.
71    """
72
73    def track(self, websocket):
74        logger.info("Managing websocket %s" % format_addresses(websocket))
75        return self.spawn(websocket.run)
76
77    def clear(self):
78        logger.info("Terminating server and all connected websockets")
79        for greenlet in self:
80            try:
81                websocket = greenlet._run.im_self
82                if websocket:
83                    websocket.close(1001, 'Server is shutting down')
84            except:
85                pass
86            finally:
87                self.discard(greenlet)
88
89class WSGIServer(_WSGIServer):
90    handler_class = WebSocketWSGIHandler
91
92    def __init__(self, *args, **kwargs):
93        """
94        WSGI server that simply tracks websockets
95        and send them a proper closing handshake
96        when the server terminates.
97
98        Other than that, the server is the same
99        as its :class:`gevent.pywsgi.WSGIServer`
100        base.
101        """
102        _WSGIServer.__init__(self, *args, **kwargs)
103        self.pool = GEventWebSocketPool()
104
105    def stop(self, *args, **kwargs):
106        self.pool.clear()
107        _WSGIServer.stop(self, *args, **kwargs)
108
109if __name__ == '__main__':
110    import os
111
112    from ws4py import configure_logger
113    configure_logger()
114
115    from ws4py.websocket import EchoWebSocket
116    server = WSGIServer(('127.0.0.1', 9000),
117                        WebSocketWSGIApplication(handler_cls=EchoWebSocket))
118    server.serve_forever()
119