This repository has been archived by the owner on Jan 12, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 7
/
wsgirefserver.py
147 lines (125 loc) · 4.74 KB
/
wsgirefserver.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
# -*- coding: utf-8 -*-
__doc__ = """
Add WebSocket support to the built-in WSGI server
provided by the :py:mod:`wsgiref`. This is clearly not
meant to be a production server so please consider this
only for testing purpose.
Mostly, this module overrides bits and pieces of
the built-in classes so that it supports the WebSocket
workflow.
.. code-block:: python
from wsgiref.simple_server import make_server
from ws4py.websocket import EchoWebSocket
from ws4py.server.wsgirefserver import WSGIServer, WebSocketWSGIRequestHandler
from ws4py.server.wsgiutils import WebSocketWSGIApplication
server = make_server('', 9000, server_class=WSGIServer,
handler_class=WebSocketWSGIRequestHandler,
app=WebSocketWSGIApplication(handler_cls=EchoWebSocket))
server.initialize_websockets_manager()
server.serve_forever()
.. note::
For some reason this server may fail against autobahntestsuite.
"""
import logging
import sys
from wsgiref.handlers import SimpleHandler
from wsgiref.simple_server import WSGIRequestHandler, WSGIServer as _WSGIServer
from wsgiref import util
util._hoppish = {}.__contains__
from ws4py.manager import WebSocketManager
from ws4py import format_addresses
from ws4py.server.wsgiutils import WebSocketWSGIApplication
from ws4py.compat import get_connection
__all__ = ['WebSocketWSGIHandler', 'WebSocketWSGIRequestHandler',
'WSGIServer']
logger = logging.getLogger('ws4py')
class WebSocketWSGIHandler(SimpleHandler):
def setup_environ(self):
"""
Setup the environ dictionary and add the
`'ws4py.socket'` key. Its associated value
is the real socket underlying socket.
"""
SimpleHandler.setup_environ(self)
self.environ['ws4py.socket'] = get_connection(self.environ['wsgi.input'])
def finish_response(self):
"""
Completes the response and performs the following tasks:
- Remove the `'ws4py.socket'` and `'ws4py.websocket'`
environ keys.
- Attach the returned websocket, if any, to the WSGI server
using its ``link_websocket_to_server`` method.
"""
ws = None
if self.environ:
self.environ.pop('ws4py.socket', None)
ws = self.environ.pop('ws4py.websocket', None)
try:
SimpleHandler.finish_response(self)
except:
if ws:
ws.close(1011, reason='Something broke')
raise
else:
if ws:
self.request_handler.server.link_websocket_to_server(ws)
class WebSocketWSGIRequestHandler(WSGIRequestHandler):
def handle(self):
"""
Unfortunately the base class forces us
to override the whole method to actually provide our wsgi handler.
"""
self.raw_requestline = self.rfile.readline()
if not self.parse_request(): # An error code has been sent, just exit
return
# next line is where we'd have expect a configuration key somehow
handler = WebSocketWSGIHandler(
self.rfile, self.wfile, self.get_stderr(), self.get_environ()
)
handler.request_handler = self # backpointer for logging
handler.run(self.server.get_app())
class WSGIServer(_WSGIServer):
def initialize_websockets_manager(self):
"""
Call thos to start the underlying websockets
manager. Make sure to call it once your server
is created.
"""
self.manager = WebSocketManager()
self.manager.start()
def shutdown_request(self, request):
"""
The base class would close our socket
if we didn't override it.
"""
pass
def link_websocket_to_server(self, ws):
"""
Call this from your WSGI handler when a websocket
has been created.
"""
self.manager.add(ws)
def server_close(self):
"""
Properly initiate closing handshakes on
all websockets when the WSGI server terminates.
"""
if hasattr(self, 'manager'):
self.manager.close_all()
self.manager.stop()
self.manager.join()
delattr(self, 'manager')
_WSGIServer.server_close(self)
if __name__ == '__main__':
from ws4py import configure_logger
configure_logger()
from wsgiref.simple_server import make_server
from ws4py.websocket import EchoWebSocket
server = make_server('', 9000, server_class=WSGIServer,
handler_class=WebSocketWSGIRequestHandler,
app=WebSocketWSGIApplication(handler_cls=EchoWebSocket))
server.initialize_websockets_manager()
try:
server.serve_forever()
except KeyboardInterrupt:
server.server_close()