summaryrefslogtreecommitdiff
path: root/src/server/rz_mesh.py
blob: b20aa64b4cea0795b905b550cc176e642a96016f (plain)
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
from flask import Response
from flask import request
from geventwebsocket.handler import WebSocketHandler
import logging
from socketio import socketio_manage
from socketio.mixins import BroadcastMixin
from socketio.namespace import BaseNamespace
from socketio.server import SocketIOHandler
from socketio.server import SocketIOServer
from functools import wraps

from rz_kernel import RZ_Kernel
from model.graph import Attr_Diff, Topo_Diff

log = logging.getLogger('rhizi')

class WebSocket_Graph_NS(BaseNamespace, BroadcastMixin):
    """
    Rhizi '/graph' websocket namespace
    """
    def multicast_msg(self, msg_name, *args):
        self.socket.server.log_multicast(msg_name)
        super(WebSocket_Graph_NS, self).broadcast_event_not_me(msg_name, *args)

    def _log_conn(self, prefix_msg):
        rmt_addr = self.environ['REMOTE_ADDR']
        rmt_port = self.environ['REMOTE_PORT']
        sid = self.environ['socketio'].sessid
        log.info('ws: %s: sid: %s, remote-socket: %s:%s' % (prefix_msg, sid, rmt_addr, rmt_port))

    def recv_connect(self):
        self._log_conn('conn open')

    def recv_disconnect(self):
        self._log_conn('conn close')

    def on_diff_commit__topo(self, topo_diff):
        log.info('ws: rx: topo diff: ' + str(topo_diff))
        self.multicast_msg('diff_commit__topo', topo_diff)

    def on_diff_commit__attr(self, attr_diff):
        log.info('ws: rx: attr diff: ' + str(attr_diff))
        self.multicast_msg('diff_commit__attr', attr_diff)

class RZ_WebSocket_Server(SocketIOServer):
    """
    Rhizi customized SocketIOServer:
        - allow response header injection on websocket connections
    """

    class WebSocketHandlerExt(SocketIOHandler):

        def start_response(self, status, headers, exc_info=None):
            headers['Access-Control-Allow-Origin'] = '*'
            return WebSocketHandler.start_response(self, status, headers, exc_info)

        def handle_one_response(self):
            return WebSocketHandler.handle_one_response(self)

        def upgrade_websocket(self):
            return WebSocketHandler.upgrade_websocket(self)

    def __init__ (self, cfg, webapp):
        # Thread.__init__(self)
        SocketIOServer.__init__(self,
                                (cfg.listen_address, cfg.listen_port),
                                webapp,
                                resource='socket.io',
                                policy_server=False)

    def log_multicast(self, msg_name):
        multicast_size = len(self.sockets) - 1  # subtract self socket
        log.info('ws: multicast: msg: \'%s\', cast-size ~= %d' % (msg_name, multicast_size))  # ~=: as race conditions apply

def init_ws_interface(cfg, kernel, flask_webapp):
    """
    Initialize websocket interface:
       - apply websocket route handlers
    
    @return: an initialized RZ_WebSocket_Server object
    """

    rz_request_env = {}  # available in namespace functions via 'self.request' - see socketio_manage
    def socketio_route_handler(url_path):
        try:
            socketio_manage(request.environ, {'/graph': WebSocket_Graph_NS}, rz_request_env)
        except:
            flask_webapp.logger.error("Exception while handling socketio connection",
                             exc_info=True)
        return Response()

    # connect socketio route
    route_dec = flask_webapp.route('/socket.io/<path:url_path>')
    f = route_dec(socketio_route_handler)
    flask_webapp.f = f

    # init ws server
    ws_srv = RZ_WebSocket_Server(cfg, flask_webapp)

    # link ws hooks: multicast on topo_diff, attr_diff
    def decorator__ws_multicast(ws_srv, f, f_multicast):
        """
        @param f: [!] wrapped function, name used to derive socket message name
        """

        @wraps(f)
        def wrapped_function(*args, **kw):
            f_ret = f(*args, **kw)

            pkt_data = f_ret
            if isinstance(f_ret, Topo_Diff):
                pkt_data = f_ret.to_json_dict()

            msg_name = f.__name__
            pkt = dict(type="event",
                       name=msg_name,
                       args=[pkt_data],
                       endpoint='/graph')

            ws_srv.log_multicast(msg_name)
            for sessid, socket in ws_srv.sockets.iteritems():
                socket.send_packet(pkt)

            return f_ret

        return wrapped_function

    kernel.diff_commit__topo = decorator__ws_multicast(ws_srv,
                                                       kernel.diff_commit__topo,
                                                       f_multicast=WebSocket_Graph_NS.on_diff_commit__topo)
    kernel.diff_commit__attr = decorator__ws_multicast(ws_srv,
                                                       kernel.diff_commit__attr,
                                                       f_multicast=WebSocket_Graph_NS.on_diff_commit__attr)

    return ws_srv;