Created
September 12, 2012 15:58
-
-
Save tebeka/3707665 to your computer and use it in GitHub Desktop.
TCP reverse proxy with epoll (not working)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python | |
# TCP Proxy | |
# Each socket has a read buffers | |
# if we proxy s1 <-> s2 then | |
# s1 writes to s2 read buffer | |
# s2 writes to s1 read buffer | |
# Using epool (kqueue?) | |
# When new event on server socket: | |
# - Accept | |
# - Create new socket to backend | |
# - Register read/write on both sockets | |
# Based on http://scotdoyle.com/python-epoll-howto.html | |
import select | |
import socket | |
def new_socket(): | |
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | |
sock.setblocking(0) | |
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) | |
return sock | |
def match_maker(poll, conn, connections, buffers, mapping): | |
conn.setblocking(0) | |
backend = new_socket() | |
backend.connect(('localhost', 8080)) | |
buffers[backend.fileno()] = '' | |
connections[backend.fileno()] = backend | |
poll.register(backend.fileno(), | |
select.EPOLLIN | select.EPOLLOUT) | |
buffers[conn.fileno()] = '' | |
connections[conn.fileno()] = conn | |
poll.register(conn.fileno(), select.EPOLLIN | select.EPOLLOUT) | |
mapping[conn.fileno()] = backend.fileno() | |
mapping[backend.fileno()] = conn.fileno() | |
def main(argv=None): | |
import sys | |
from argparse import ArgumentParser | |
argv = argv or sys.argv | |
parser = ArgumentParser(description='') | |
parser.add_argument('port', help='port to listen on', type=int) | |
args = parser.parse_args(argv[1:]) | |
server = new_socket() | |
server.bind(('localhost', args.port)) | |
server.listen(1) | |
connections, buffers, mapping = {}, {}, {} | |
poll = select.epoll() | |
poll.register(server.fileno(), select.EPOLLIN) | |
try: | |
while True: | |
for fileno, event in poll.poll(1): | |
if fileno == server.fileno(): | |
# New connection | |
conn, address = server.accept() | |
match_maker(poll, conn, connections, buffers, mapping) | |
elif event & select.EPOLLIN: | |
other = mapping[fileno] | |
buffers[other] += connections[fileno].recv(1024) | |
elif event & select.EPOLLOUT: | |
size = connections[fileno].send(buffers[fileno]) | |
if size == 0: | |
connections[fileno].shutdown(socket.SHUT_RDWR) | |
poll.modify(fileno, 0) | |
other = connections[mapping[fileno]] | |
other.shutdown(socket.SHUT_RDWR) | |
poll.modify(other.fileno(), 0) | |
elif event & select.EPOLLHUP: | |
poll.unregister(fileno) | |
connections.pop(fileno) | |
buffers.pop(fileno) | |
other = connections[mapping[fileno]] | |
other.shutdown(socket.SHUT_RDWR) | |
poll.modify(other.fileno(), 0) | |
finally: | |
poll.unregister(server.fileno()) | |
poll.close() | |
server.close() | |
if __name__ == '__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment