tendermint/example/python/tmsp/server.py

import socket
import select
import sys

from wire import decode_varint, encode
from reader import BytesBuffer
from msg import RequestDecoder, message_types

# hold the asyncronous state of a connection
# ie. we may not get enough bytes on one read to decode the message

class Connection():

    def __init__(self, fd, app):
        self.fd = fd
        self.app = app
        self.recBuf = BytesBuffer(bytearray())
        self.resBuf = BytesBuffer(bytearray())
        self.msgLength = 0
        self.decoder = RequestDecoder(self.recBuf)
        self.inProgress = False  # are we in the middle of a message

    def recv(this):
        data = this.fd.recv(1024)
        if not data:  # what about len(data) == 0
            raise IOError("dead connection")
        this.recBuf.write(data)

# TMSP server responds to messges by calling methods on the app

class TMSPServer():

    def __init__(self, app, port=5410):
        self.app = app
        # map conn file descriptors to (app, reqBuf, resBuf, msgDecoder)
        self.appMap = {}

        self.port = port
        self.listen_backlog = 10

        self.listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.listener.setblocking(0)
        self.listener.bind(('', port))

        self.listener.listen(self.listen_backlog)

        self.shutdown = False

        self.read_list = [self.listener]
        self.write_list = []

    def handle_new_connection(self, r):
        new_fd, new_addr = r.accept()
        new_fd.setblocking(0)  # non-blocking
        self.read_list.append(new_fd)
        self.write_list.append(new_fd)
        print 'new connection to', new_addr

        self.appMap[new_fd] = Connection(new_fd, self.app)

    def handle_conn_closed(self, r):
        self.read_list.remove(r)
        self.write_list.remove(r)
        r.close()
        print "connection closed"

    def handle_recv(self, r):
        #  app, recBuf, resBuf, conn
        conn = self.appMap[r]
        while True:
            try:
                print "recv loop"
                # check if we need more data first
                if conn.inProgress:
                    if (conn.msgLength == 0 or conn.recBuf.size() < conn.msgLength):
                        conn.recv()
                else:
                    if conn.recBuf.size() == 0:
                        conn.recv()

                conn.inProgress = True

                # see if we have enough to get the message length
                if conn.msgLength == 0:
                    ll = conn.recBuf.peek()
                    if conn.recBuf.size() < 1 + ll:
                        # we don't have enough bytes to read the length yet
                        return
                    print "decoding msg length"
                    conn.msgLength = decode_varint(conn.recBuf)

                # see if we have enough to decode the message
                if conn.recBuf.size() < conn.msgLength:
                    return

                # now we can decode the message

                # first read the request type and get the particular msg
                # decoder
                typeByte = conn.recBuf.read(1)
                typeByte = int(typeByte[0])
                resTypeByte = typeByte + 0x10
                req_type = message_types[typeByte]

                if req_type == "flush":
                    # messages are length prefixed
                    conn.resBuf.write(encode(1))
                    conn.resBuf.write([resTypeByte])
                    conn.fd.send(str(conn.resBuf.buf))
                    conn.msgLength = 0
                    conn.inProgress = False
                    conn.resBuf = BytesBuffer(bytearray())
                    return

                decoder = getattr(conn.decoder, req_type)

                print "decoding args"
                req_args = decoder()
                print "got args", req_args

                # done decoding message
                conn.msgLength = 0
                conn.inProgress = False

                req_f = getattr(conn.app, req_type)
                if req_args is None:
                    res = req_f()
                elif isinstance(req_args, tuple):
                    res = req_f(*req_args)
                else:
                    res = req_f(req_args)

                if isinstance(res, tuple):
                    res, ret_code = res
                else:
                    ret_code = res
                    res = None

                print "called", req_type, "ret code:", ret_code
                if ret_code != 0:
                    print "non-zero retcode:", ret_code

                if req_type in ("echo", "info"):  # these dont return a ret code
                    enc = encode(res)
                    # messages are length prefixed
                    conn.resBuf.write(encode(len(enc) + 1))
                    conn.resBuf.write([resTypeByte])
                    conn.resBuf.write(enc)
                else:
                    enc, encRet = encode(res), encode(ret_code)
                    # messages are length prefixed
                    conn.resBuf.write(encode(len(enc) + len(encRet) + 1))
                    conn.resBuf.write([resTypeByte])
                    conn.resBuf.write(encRet)
                    conn.resBuf.write(enc)
            except TypeError as e:
                print "TypeError on reading from connection:", e
                self.handle_conn_closed(r)
                return
            except ValueError as e:
                print "ValueError on reading from connection:", e
                self.handle_conn_closed(r)
                return
            except IOError as e:
                print "IOError on reading from connection:", e
                self.handle_conn_closed(r)
                return
            except Exception as e:
                # sys.exc_info()[0] # TODO better
                print "error reading from connection", str(e)
                self.handle_conn_closed(r)
                return

    def main_loop(self):
        while not self.shutdown:
            r_list, w_list, _ = select.select(
                self.read_list, self.write_list, [], 2.5)

            for r in r_list:
                if (r == self.listener):
                    try:
                        self.handle_new_connection(r)
                        # undo adding to read list ...
                    except NameError as e:
                        print "Could not connect due to NameError:", e
                    except TypeError as e:
                        print "Could not connect due to TypeError:", e
                    except:
                        print "Could not connect due to unexpected error:", sys.exc_info()[0]
                else:
                    self.handle_recv(r)

    def handle_shutdown(self):
        for r in self.read_list:
            r.close()
        for w in self.write_list:
            try:
                w.close()
            except Exception as e:
                print(e)  # TODO: add logging
        self.shutdown = True