diff --git a/examples/websocket.html b/examples/websocket.html new file mode 100644 index 00000000..a3a98f35 --- /dev/null +++ b/examples/websocket.html @@ -0,0 +1,29 @@ + + + + WebSocket demo + + + + + diff --git a/examples/websocket.py b/examples/websocket.py new file mode 100644 index 00000000..ecc14f8d --- /dev/null +++ b/examples/websocket.py @@ -0,0 +1,23 @@ +from sanic import Sanic +from sanic.response import file + +app = Sanic(__name__) + + +@app.route('/') +async def index(request): + return await file('websocket.html') + + +@app.ws('/feed') +async def feed(request, ws): + while True: + data = 'hello!' + print('Sending: ' + data) + await ws.send(data) + data = await ws.recv() + print('Received: ' + data) + + +if __name__ == '__main__': + app.run() diff --git a/sanic/app.py b/sanic/app.py index 63700e9d..b4f7fcb7 100644 --- a/sanic/app.py +++ b/sanic/app.py @@ -19,6 +19,7 @@ from sanic.server import serve, serve_multiple, HttpProtocol from sanic.static import register as static_register from sanic.testing import TestClient from sanic.views import CompositionView +from sanic.ws import WebSocketProtocol, ConnectionClosed class Sanic: @@ -51,6 +52,7 @@ class Sanic: self.sock = None self.listeners = defaultdict(list) self.is_running = False + self.needs_websocket = False # Register alternative method names self.go_fast = self.run @@ -168,6 +170,40 @@ class Sanic: self.route(uri=uri, methods=methods, host=host)(handler) return handler + # Decorator + def ws(self, uri, host=None): + """Decorate a function to be registered as a websocket route + :param uri: path of the URL + :param host: + :return: decorated function + """ + self.needs_websocket = True + + # Fix case where the user did not prefix the URL with a / + # and will probably get confused as to why it's not working + if not uri.startswith('/'): + uri = '/' + uri + + def response(handler): + async def websocket_handler(request, *args, **kwargs): + protocol = request.transport.get_protocol() + ws = await protocol.websocket_handshake(request) + try: + # invoke the application handler + await handler(request, ws, *args, **kwargs) + except ConnectionClosed: + pass + await ws.close() + + self.router.add(uri=uri, handler=websocket_handler, + methods=frozenset({'GET'}), host=host) + return handler + + return response + + def add_ws_route(self, handler, uri, host=None): + return self.ws(uri, host=host)(handler) + def remove_route(self, uri, clean_cache=True, host=None): self.router.remove(uri, clean_cache, host) @@ -437,7 +473,7 @@ class Sanic: def run(self, host="127.0.0.1", port=8000, debug=False, before_start=None, after_start=None, before_stop=None, after_stop=None, ssl=None, - sock=None, workers=1, loop=None, protocol=HttpProtocol, + sock=None, workers=1, loop=None, protocol=None, backlog=100, stop_event=None, register_sys_signals=True): """Run the HTTP Server and listen until keyboard interrupt or term signal. On termination, drain connections before closing. @@ -464,6 +500,9 @@ class Sanic: :param protocol: Subclass of asyncio protocol class :return: Nothing """ + if protocol is None: + protocol = WebSocketProtocol if self.needs_websocket \ + else HttpProtocol server_settings = self._helper( host=host, port=port, debug=debug, before_start=before_start, after_start=after_start, before_stop=before_stop, @@ -491,13 +530,16 @@ class Sanic: async def create_server(self, host="127.0.0.1", port=8000, debug=False, before_start=None, after_start=None, before_stop=None, after_stop=None, ssl=None, - sock=None, loop=None, protocol=HttpProtocol, + sock=None, loop=None, protocol=None, backlog=100, stop_event=None): """Asynchronous version of `run`. NOTE: This does not support multiprocessing and is not the preferred way to run a Sanic application. """ + if protocol is None: + protocol = WebSocketProtocol if self.needs_websocket \ + else HttpProtocol server_settings = self._helper( host=host, port=port, debug=debug, before_start=before_start, after_start=after_start, before_stop=before_stop, diff --git a/sanic/ws.py b/sanic/ws.py new file mode 100644 index 00000000..8d0ddc70 --- /dev/null +++ b/sanic/ws.py @@ -0,0 +1,53 @@ +from sanic.server import HttpProtocol +from httptools import HttpParserUpgrade +from websockets import handshake, WebSocketCommonProtocol +from websockets import ConnectionClosed # noqa + + +class WebSocketProtocol(HttpProtocol): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.ws = None + + def data_received(self, data): + if self.ws is not None: + # pass the data to the websocket protocol + self.ws.data_received(data) + else: + try: + super().data_received(data) + except HttpParserUpgrade: + # this is okay, it just indicates we've got an upgrade request + pass + + def write_response(self, response): + if self.ws is not None: + # websocket requests do not write a response + self.transport.close() + else: + super().write_response(response) + + async def websocket_handshake(self, request): + # let the websockets package do the handshake with the client + headers = [] + + def get_header(k): + return request.headers.get(k, '') + + def set_header(k, v): + headers.append((k, v)) + + key = handshake.check_request(get_header) + handshake.build_response(set_header, key) + + # write the 101 response back to the client + rv = b'HTTP/1.1 101 Switching Protocols\r\n' + for k, v in headers: + rv += k.encode('utf-8') + b': ' + v.encode('utf-8') + b'\r\n' + rv += b'\r\n' + request.transport.write(rv) + + # hook up the websocket protocol + self.ws = WebSocketCommonProtocol() + self.ws.connection_made(request.transport) + return self.ws diff --git a/setup.py b/setup.py index c73f3848..52e4a56e 100644 --- a/setup.py +++ b/setup.py @@ -19,6 +19,7 @@ install_requires = [ 'httptools>=0.0.9', 'ujson>=1.35', 'aiofiles>=0.3.0', + 'websockets>=3.2', ] if os.name != 'nt':