Skip to content

Instantly share code, notes, and snippets.

@tebeka
Created September 12, 2012 15:58
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save tebeka/3707665 to your computer and use it in GitHub Desktop.
Save tebeka/3707665 to your computer and use it in GitHub Desktop.
TCP reverse proxy with epoll (not working)
#!/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