Thursday, April 22, 2010

Here is a really nice piece of code which implements a TCP relay (tunnel) using Python's asyncore module. It's really interesting because the way I implemented it makes a lot of sense, in contrary to the same code written using select.

There's also something about using asyncore.dispatcher_with_send on Unix systems which might send an EWOULDBLOCK sometimes, but I didn't dig deep enough.

The 3 sockets playing a role here are:
* A Relay server socket - accepts connections from clients
* Relay clients, which trigger a new connection to the tunnel destination
* Relay connection for each relay client connected.

The cool thing about asyncore is that you can also decide if handle_read will be called or not, even if data is available, by overriding the "readable" function. In this implementation, a relay client does not start reading from a socket until the relay connection has been established successfully.




class RelayConnection(asyncore.dispatcher):
def __init__(self, client, address):
asyncore.dispatcher.__init__(self)
self.client = client
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
print "connecting to %s..." % str(address)
self.connect(address)

def handle_connect(self):
print "connected."
# Allow reading once the connection
# on the other side is open.
self.client.is_readable = True

def handle_read(self):
self.client.send(self.recv(1024))

class RelayClient(asyncore.dispatcher):
def __init__(self, server, client, address):
asyncore.dispatcher.__init__(self, client)
self.is_readable = False
self.server = server
self.relay = RelayConnection(self, address)

def handle_read(self):
self.relay.send(self.recv(1024))

def handle_close(self):
print "Closing relay..."
# If the client disconnects, close the
# relay connection as well.
self.relay.close()
self.close()

def readable(self):
return self.is_readable

class RelayServer(asyncore.dispatcher):
def __init__(self, bind_address, dest_address):
asyncore.dispatcher.__init__(self)
self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
self.bind(bind_address)
self.dest_address = dest_address
self.listen(10)

def handle_accept(self):
conn, addr = self.accept()
RelayClient(self, conn, self.dest_address)


RelayServer(("0.0.0.0", 8080), ("127.0.0.1", 1234))
asyncore.loop()

No comments: