tentative tcp reserve proxy

epoll_tcprp.py

import socket
import select

BACKEND_HOST = 'example'
BACKEND_PORT = 1234
BUF_SIZE = 1024
LISTEN_PORT = 1234

proxy_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
proxy_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
proxy_socket.bind(('0.0.0.0', LISTEN_PORT))
proxy_socket.listen(5)
proxy_socket.setblocking(0)


epoll = select.epoll()
epoll.register(proxy_socket.fileno(), select.EPOLLIN)


def switch_fileno(fileno, filenos):
    a, b = zip(*filenos)
    return b[a.index(fileno)] if fileno in a else a[b.index(fileno)]


def delete_fileno_pair(fileno, filenos):
    a, b = zip(*filenos)
    if fileno in a:
        del filenos[a.index(fileno)]
    elif fileno in b:
        del filenos[b.index(fileno)]


client_conns = {}
backend_conns = {}
filenos = []
client_reqs = {}
client_resps = {}
while True:
    try:
        events = epoll.poll(1)
        for fileno, event in events:
            if fileno == proxy_socket.fileno():
                conn, addr = proxy_socket.accept()
                conn.setblocking(0)
                epoll.register(conn.fileno(), select.EPOLLIN)
                client_conns[conn.fileno()] = conn
                backend_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                backend_socket.connect((socket.gethostbyname(BACKEND_HOST), BACKEND_PORT))
                backend_socket.setblocking(0)
                backend_conns[conn.fileno()] = backend_socket
                backend_conns[backend_socket.fileno()] = backend_socket
                epoll.register(backend_socket.fileno(), select.EPOLLIN)
                filenos.append((conn.fileno(), backend_socket.fileno()))
                client_reqs[backend_socket.fileno()] = b''
                client_resps[conn.fileno()] = b''
            elif event & select.EPOLLIN:
                if fileno in client_conns:
                    conn = client_conns[fileno]
                    data = conn.recv(BUF_SIZE)
                    print('Receive from client %d' % fileno)
                    if data and len(data) <= BUF_SIZE:
                        client_reqs[switch_fileno(fileno, filenos)] += data
                        print('Receive all from client %d' % fileno)
                        epoll.modify(switch_fileno(fileno, filenos), select.EPOLLOUT)
                    if not data:
                        epoll.modify(fileno, 0)
                        conn.shutdown(socket.SHUT_RDWR)
                elif fileno in backend_conns:
                    conn = backend_conns[fileno]
                    data = conn.recv(BUF_SIZE)
                    print('Receive from backend %d' % fileno)
                    if data and len(data) <= BUF_SIZE:
                        client_resps[switch_fileno(fileno, filenos)] += data
                        print('Receive all from backend %d' % fileno)
                        epoll.modify(switch_fileno(fileno, filenos), select.EPOLLOUT)
                    if not data:
                        epoll.modify(fileno, 0)
                        conn.shutdown(socket.SHUT_RDWR)
            elif event & select.EPOLLOUT:
                if fileno in client_conns:
                    conn = client_conns[fileno]
                    bytes = conn.send(client_resps[fileno])
                    print('Send to client %d' % fileno)
                    client_resps[fileno] = client_resps[fileno][bytes:]
                    if len(client_resps[fileno]) == 0:
                        print('Send all to client %d' % fileno)
                        epoll.modify(fileno, select.EPOLLIN)
                elif fileno in backend_conns:
                    conn = backend_conns[fileno]
                    bytes = conn.send(client_reqs[fileno])
                    print('Send to backend %d' % fileno)
                    client_reqs[fileno] = client_reqs[fileno][bytes:]
                    if len(client_reqs[fileno]) == 0:
                        print('Send all to backend %d' % fileno)
                        epoll.modify(fileno, select.EPOLLIN)
            elif event & select.EPOLLHUP:
                print("Hang up")
                epoll.unregister(fileno)
                epoll.unregister(switch_fileno(fileno, filenos))
                if fileno in client_conns:
                    client_conns[fileno].close()
                    del client_conns[fileno]
                elif fileno in backend_conns:
                    backend_conns[fileno].close()
                    del backend_conns[fileno]
                delete_fileno_pair(fileno, filenos)
    except KeyboardInterrupt:
        print("\nTCP reserve proxy exited..")
        break
    except:
        continue

epoll.unregister(proxy_socket.fileno())
epoll.close()
proxy_socket.close()