├── .gitignore ├── Makefile ├── README.md ├── conf.py ├── flask_app.py ├── message_proxy.py ├── nginx.conf ├── requirements.txt ├── static └── js │ ├── cookies.js │ ├── user.js │ └── websocket.js ├── supervisord.conf ├── templates └── index.html └── websocket_server.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | log 3 | *~ 4 | \#* 5 | run 6 | __pycache__ 7 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: all paths supervisord 2 | 3 | SHELL=/bin/bash 4 | 5 | # Supervisord only runs under Python 2, but that is no problem, even 6 | # if everything else runs under Python 3. 7 | SUPERVISORD=supervisord 8 | 9 | all: supervisord 10 | 11 | paths: 12 | mkdir -p log run 13 | mkdir -p log/sv_child 14 | 15 | supervisord: paths 16 | $(SUPERVISORD) -c supervisord.conf 17 | 18 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Message Flow Demo 2 | 3 | Demonstrates how to set up a Python web service with WebSockets for 4 | pushing messages from server to front-end. 5 | 6 | ## Installation 7 | 8 | - Install nginx 9 | - Install supervisord 10 | - Edit the Makefile to point to supervisord 11 | - ``pip install -r requirements.txt`` 12 | 13 | ## Running 14 | 15 | - ``make`` 16 | - Connect to ``http://localhost:5000`` 17 | 18 | ## System outline 19 | 20 | 21 | 22 | 23 | -------------------------------------------------------------------------------- /conf.py: -------------------------------------------------------------------------------- 1 | # Secret used to authenticate websockets. This can be anything you want. 2 | secret = 'kljasdlkahsdlkajhsdlkajhsd71826381623123' 3 | -------------------------------------------------------------------------------- /flask_app.py: -------------------------------------------------------------------------------- 1 | from flask import Flask, render_template, request, jsonify, escape 2 | 3 | import json 4 | import threading 5 | import hashlib 6 | import uuid 7 | import collections 8 | import conf 9 | import jwt 10 | 11 | 12 | app = Flask(__name__) 13 | app.secret_key = conf.secret 14 | 15 | 16 | import zmq 17 | ctx = zmq.Context() 18 | pub = ctx.socket(zmq.PUB) 19 | pub.connect('ipc:///tmp/message_flow_in') 20 | 21 | 22 | # Task IDs 23 | class TID: 24 | OK = 'OK' 25 | ERROR = 'ERROR' 26 | DONE = 'TASK DONE' 27 | 28 | 29 | def push(user, task_id, data): 30 | """Push message to `user` over websocket. 31 | 32 | """ 33 | pub.send(b"0 " + json.dumps({'username': user, 34 | 'id': task_id, 35 | 'data': data}).encode('utf-8')) 36 | 37 | 38 | def long_task(user, message): 39 | """This is an example of a task that takes long to execute. 40 | 41 | """ 42 | import time 43 | try: 44 | from time import monotonic 45 | except ImportError: 46 | from time import time as monotonic 47 | import random 48 | 49 | m0 = monotonic() 50 | time.sleep(2 + random.random()) 51 | m1 = monotonic() 52 | 53 | processed_message = 'Message processed, "{}" in {:.2f} seconds'.format(escape(message), m1 - m0) 54 | push(user, TID.DONE, processed_message) 55 | 56 | 57 | @app.route('/send', methods=['POST']) 58 | def send(): 59 | message = request.form.get('message', '') 60 | 61 | if not message: 62 | return jsonify(status=TID.ERROR, data='No message submitted') 63 | else: 64 | threading.Thread(target=long_task, args=(get_username(), message)).start() 65 | return jsonify(status=TID.OK, 66 | data='Message submitted for processing') 67 | 68 | 69 | # Modify this for your specific application 70 | # 71 | # !! Certainly don't get usernames from the client like I do here--that would be highly 72 | # insecure 73 | def get_username(): 74 | return request.cookies.get('username') 75 | 76 | 77 | # !!! 78 | # This API call should **only be callable by logged in users** 79 | # !!! 80 | @app.route('/socket_auth_token', methods=['GET']) 81 | def socket_auth_token(): 82 | return jwt.encode({'username': get_username()}, app.secret_key) 83 | 84 | 85 | @app.route('/') 86 | def index(): 87 | return render_template('index.html', username=get_username()) 88 | -------------------------------------------------------------------------------- /message_proxy.py: -------------------------------------------------------------------------------- 1 | # http://zguide.zeromq.org/page:all#The-Dynamic-Discovery-Problem 2 | 3 | import zmq 4 | 5 | IN = 'ipc:///tmp/message_flow_in' 6 | OUT = 'ipc:///tmp/message_flow_out' 7 | 8 | context = zmq.Context() 9 | 10 | feed_in = context.socket(zmq.XSUB) 11 | feed_in.bind(IN) 12 | 13 | feed_out = context.socket(zmq.XPUB) 14 | feed_out.bind(OUT) 15 | 16 | print('[message_proxy] Forwarding messages between {} and {}'.format(IN, OUT)) 17 | zmq.proxy(feed_in, feed_out) 18 | -------------------------------------------------------------------------------- /nginx.conf: -------------------------------------------------------------------------------- 1 | http { 2 | upstream websocket_server { 3 | server localhost:4567; 4 | } 5 | 6 | server { 7 | listen 5000; 8 | 9 | location / { 10 | proxy_pass http://unix:run/flask_app.sock:/; 11 | 12 | proxy_set_header Host $http_host; 13 | proxy_set_header X-Real-IP $remote_addr; 14 | proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; 15 | proxy_set_header X-Forwarded-Proto $scheme; 16 | } 17 | 18 | location /websocket { 19 | proxy_pass http://websocket_server/websocket; 20 | proxy_http_version 1.1; 21 | proxy_set_header Upgrade $http_upgrade; 22 | proxy_set_header Connection "upgrade"; 23 | proxy_read_timeout 60s; 24 | } 25 | 26 | error_log log/error.log; 27 | access_log log/nginx-access.log; 28 | 29 | } 30 | 31 | } 32 | 33 | events { 34 | } 35 | 36 | pid run/nginx.pid; 37 | error_log log/error.log; 38 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | tornado>=3.0 2 | supervisor 3 | pyzmq 4 | pyjwt 5 | -------------------------------------------------------------------------------- /static/js/cookies.js: -------------------------------------------------------------------------------- 1 | // From http://www.quirksmode.org/js/cookies.html 2 | 3 | function createCookie(name, value, minutes) { 4 | var expires = ""; 5 | if (minutes) { 6 | var date = new Date(); 7 | date.setTime(date.getTime() + (minutes * 60 * 1000)); 8 | expires = "; expires=" + date.toGMTString(); 9 | } 10 | document.cookie = name + "=" + value + expires + "; path=/"; 11 | } 12 | 13 | function readCookie(name) { 14 | var nameEQ = name + "="; 15 | var ca = document.cookie.split(';'); 16 | for(var i=0; i < ca.length; i++) { 17 | var c = ca[i]; 18 | while (c.charAt(0) == ' ') c = c.substring(1, c.length); 19 | if (c.indexOf(nameEQ) == 0) return c.substring(nameEQ.length, c.length); 20 | } 21 | return null; 22 | } 23 | 24 | function eraseCookie(name) { 25 | createCookie(name, "", -1); 26 | } 27 | -------------------------------------------------------------------------------- /static/js/user.js: -------------------------------------------------------------------------------- 1 | function generateUserCookie() { 2 | var username = "test" + 3 | Math.floor(Math.random() * 9000) + 1000 + "@myservice.org"; 4 | createCookie("username", username); 5 | 6 | return username; 7 | } 8 | -------------------------------------------------------------------------------- /static/js/websocket.js: -------------------------------------------------------------------------------- 1 | function getTime() { 2 | var date = new Date(); 3 | var n = date.toDateString(); 4 | return date.toLocaleTimeString(); 5 | }; 6 | 7 | function webSocketAuth() { 8 | var auth_token = new Promise( 9 | function(resolve, reject) { 10 | // First, try and read the authentication token from a cookie 11 | var cookie_token = readCookie('auth_token'); 12 | var no_token = "no_auth_token_user bad_token"; 13 | 14 | if (cookie_token) { 15 | resolve(cookie_token); 16 | } else { 17 | // If not found, ask the server for a new one 18 | $.ajax({ 19 | url: $SCRIPT_ROOT + "/socket_auth_token", 20 | statusCode: { 21 | // If we get a gateway error, it probably means nginx is being restarted. 22 | // Not much we can do, other than wait a bit and continue with a 23 | // fake token. 24 | 405: function() { 25 | setTimeout(function() { resolve(no_token); }, 1000); 26 | } 27 | }}) 28 | .done(function(data) { 29 | createCookie('auth_token', data); 30 | resolve(data); 31 | }) 32 | .fail(function() { 33 | // Same situation as with 405. 34 | setTimeout(function() { resolve(no_token); }, 1000); 35 | }); 36 | } 37 | } 38 | ); 39 | 40 | auth_token.then( 41 | function(token) { 42 | ws.send(token); 43 | }); 44 | }; 45 | 46 | 47 | function createSocket(url, messageHandler) { 48 | var ws = new ReconnectingWebSocket(url); 49 | 50 | ws.onopen = function (event) { 51 | $("#connected").empty().append("Online"); 52 | $("#authenticated").empty().append("No"); 53 | }; 54 | 55 | ws.onmessage = function (event) { 56 | var message = event.data; 57 | 58 | // Ignore heartbeat signals 59 | if (message === '<3') { 60 | return; 61 | } 62 | 63 | var data = JSON.parse(message); 64 | var id = data["id"]; 65 | 66 | switch (id) { 67 | case "AUTH REQUEST": 68 | webSocketAuth(); 69 | break; 70 | case "AUTH FAILED": 71 | eraseCookie('auth_token'); 72 | break; 73 | case "AUTH OK": 74 | $("#connected").empty().append("Online"); 75 | $("#authenticated").empty().append("Yes"); 76 | break; 77 | default: 78 | messageHandler(data); 79 | } 80 | }; 81 | 82 | ws.onclose = function (event) { 83 | $("#connected").empty().append("Offline"); 84 | $("#authenticated").empty().append("No"); 85 | }; 86 | 87 | return ws; 88 | }; 89 | -------------------------------------------------------------------------------- /supervisord.conf: -------------------------------------------------------------------------------- 1 | [supervisord] 2 | logfile=log/supervisord.log 3 | pidfile=run/supervisord.pid 4 | nodaemon=true 5 | childlogdir=log/sv_child 6 | 7 | [program:waitress] 8 | command=waitress-serve --unix-socket=run/flask_app.sock flask_app:app 9 | environment=PYTHONUNBUFFERED=1 10 | stdout_logfile=log/waitress.log 11 | redirect_stderr=true 12 | 13 | [program:nginx] 14 | command=nginx -c nginx.conf -p . -g "daemon off;" 15 | 16 | [program:message_proxy] 17 | command=/usr/bin/env python message_proxy.py 18 | environment=PYTHONUNBUFFERED=1 19 | stdout_logfile=log/message_proxy.log 20 | redirect_stderr=true 21 | 22 | [program:websocket] 23 | command=/usr/bin/env python websocket_server.py 24 | environment=PYTHONUNBUFFERED=1 25 | stdout_logfile=log/websocket_server.log 26 | redirect_stderr=true 27 | 28 | -------------------------------------------------------------------------------- /templates/index.html: -------------------------------------------------------------------------------- 1 | 2 | Index 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 15 | 16 | 17 | 18 |

WebSocket Demo

19 | 20 | 25 | 26 |
27 | Message: 28 | 29 |
30 | 31 |

API Result:

32 |
33 |
    34 |
35 |
36 | 37 |

Incoming:

38 |
39 |
    40 |
41 |
42 | 43 | 97 | 98 | 99 | 100 | -------------------------------------------------------------------------------- /websocket_server.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | 3 | from tornado import websocket, web, ioloop 4 | import json 5 | import zmq 6 | import jwt 7 | 8 | from conf import secret 9 | 10 | ctx = zmq.Context() 11 | 12 | 13 | # Could also use: http://aaugustin.github.io/websockets/ 14 | 15 | 16 | class WebSocket(websocket.WebSocketHandler): 17 | participants = set() 18 | 19 | def __init__(self, *args, **kwargs): 20 | websocket.WebSocketHandler.__init__(self, *args, **kwargs) 21 | 22 | self.authenticated = False 23 | self.auth_failures = 0 24 | self.username = None 25 | 26 | def check_origin(self, origin): 27 | return True 28 | 29 | def open(self): 30 | if self not in self.participants: 31 | self.participants.add(self) 32 | self.request_auth() 33 | 34 | def on_close(self): 35 | if self in self.participants: 36 | self.participants.remove(self) 37 | 38 | def on_message(self, auth_token): 39 | self.authenticate(auth_token) 40 | if not self.authenticated and self.auth_failures < 3: 41 | self.request_auth() 42 | 43 | def request_auth(self): 44 | self.auth_failures += 1 45 | self.send_json(id="AUTH REQUEST") 46 | 47 | def send_json(self, **kwargs): 48 | self.write_message(json.dumps(kwargs)) 49 | 50 | def authenticate(self, auth_token): 51 | try: 52 | token_payload = jwt.decode(auth_token, secret) 53 | self.username = token_payload['username'] 54 | self.authenticated = True 55 | self.send_json(id='AUTH OK') 56 | except DecodeError: 57 | self.send_json(id='AUTH FAILED') 58 | 59 | @classmethod 60 | def heartbeat(cls): 61 | for p in cls.participants: 62 | p.write_message(b'<3') 63 | 64 | # http://mrjoes.github.io/2013/06/21/python-realtime.html 65 | @classmethod 66 | def broadcast(cls, data): 67 | channel, data = data[0].decode('utf-8').split(" ", 1) 68 | user = json.loads(data)["username"] 69 | 70 | for p in cls.participants: 71 | if p.authenticated and p.username == user: 72 | p.write_message(data) 73 | 74 | 75 | if __name__ == "__main__": 76 | PORT = 4567 77 | LOCAL_OUTPUT = 'ipc:///tmp/message_flow_out' 78 | 79 | import zmq 80 | 81 | # https://zeromq.github.io/pyzmq/eventloop.html 82 | from zmq.eventloop import ioloop, zmqstream 83 | 84 | ioloop.install() 85 | 86 | sub = ctx.socket(zmq.SUB) 87 | sub.connect(LOCAL_OUTPUT) 88 | sub.setsockopt(zmq.SUBSCRIBE, b'') 89 | 90 | print('[websocket_server] Broadcasting {} to all websockets'.format(LOCAL_OUTPUT)) 91 | stream = zmqstream.ZMQStream(sub) 92 | stream.on_recv(WebSocket.broadcast) 93 | 94 | server = web.Application([ 95 | (r'/websocket', WebSocket), 96 | ]) 97 | server.listen(PORT) 98 | 99 | # We send a heartbeat every 45 seconds to make sure that nginx 100 | # proxy does not time out and close the connection 101 | ioloop.PeriodicCallback(WebSocket.heartbeat, 45000).start() 102 | 103 | print('[websocket_server] Listening for incoming websocket connections on port {}'.format(PORT)) 104 | ioloop.IOLoop.instance().start() 105 | --------------------------------------------------------------------------------