X-Git-Url: https://code.delx.au/monosys/blobdiff_plain/2451a634205ef6a673f8ea76b9ad67f2afafc73f..9969880cb9004f8a93848fae9ad645030a834ec7:/scripts/proxy.py

diff --git a/scripts/proxy.py b/scripts/proxy.py
index 0f314ae..66220ef 100755
--- a/scripts/proxy.py
+++ b/scripts/proxy.py
@@ -84,6 +84,7 @@ import os
 import socket
 import struct
 import sys
+import traceback
 
 
 if sys.platform == "linux2":
@@ -122,8 +123,17 @@ elif sys.platform == "darwin":
 
 
 class Proxy(asyncore.dispatcher):
-	def __init__(self, sock):
-		asyncore.dispatcher.__init__(self, sock)
+	def __init__(self, arg):
+		if isinstance(arg, tuple):
+			asyncore.dispatcher.__init__(self)
+			self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
+			self.connect(arg)
+		else:
+			asyncore.dispatcher.__init__(self, arg)
+		self.init()
+
+	def init(self):
+		self.end = False
 		self.other = None
 		self.buffer = ""
 
@@ -132,8 +142,8 @@ class Proxy(asyncore.dispatcher):
 		other.other = self
 
 	def handle_error(self):
-		print >>sys.stderr, "Proxy error:", sys.exc_info()
-		self.handle_close()
+		print >>sys.stderr, "Proxy error:", traceback.format_exc()
+		self.close()
 
 	def handle_read(self):
 		data = self.recv(8192)
@@ -143,6 +153,8 @@ class Proxy(asyncore.dispatcher):
 	def handle_write(self):
 		sent = self.send(self.buffer)
 		self.buffer = self.buffer[sent:]
+		if len(self.buffer) == 0 and self.end:
+			self.close()
 
 	def writable(self):
 		return len(self.buffer) > 0
@@ -151,8 +163,10 @@ class Proxy(asyncore.dispatcher):
 		if not self.other:
 			return
 		print >>sys.stderr, "Proxy closed"
-		self.other.close()
 		self.close()
+		if len(self.other.buffer) == 0:
+			self.other.close()
+		self.other.end = True
 		self.other = None
 
 class ConnectProxy(asyncore.dispatcher):
@@ -161,8 +175,8 @@ class ConnectProxy(asyncore.dispatcher):
 		self.buffer = ""
 
 	def handle_error(self):
-		print >>sys.stderr, "ConnectProxy error:", sys.exc_info()
-		self.handle_close()
+		print >>sys.stderr, "ConnectProxy error:", traceback.format_exc()
+		self.close()
 
 	def handle_read(self):
 		self.buffer += self.recv(8192)
@@ -191,15 +205,13 @@ class ConnectProxy(asyncore.dispatcher):
 		print >>sys.stderr, "Forwarding connection", host, port
 
 		# Create server proxy
-		server_connection = socket.socket()
-		server_connection.connect((host, port))
-		server = Proxy(server_connection)
+		server = Proxy((host, port))
+		server.buffer = self.buffer
 
 		# Morph and connect
 		self.__class__ = Proxy
+		self.init()
 		server.meet(self)
-		server.buffer = self.buffer
-		self.buffer = ""
 
 
 class BasicForwarder(asyncore.dispatcher):
@@ -209,11 +221,13 @@ class BasicForwarder(asyncore.dispatcher):
 		self.port = port
 		self.allowed = allowed
 		self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
+		self.set_reuse_addr()
 		self.bind(("", listen_port))
-		self.listen(5)
+		self.listen(50)
+		print >>sys.stderr, "BasicForwarder bound on", listen_port
 
 	def handle_error(self):
-		print >>sys.stderr, "BasicForwarder error:", sys.exc_info()
+		print >>sys.stderr, "BasicForwarder error:", traceback.format_exc()
 
 	def handle_accept(self):
 		client_connection, source_addr = self.accept()
@@ -223,12 +237,10 @@ class BasicForwarder(asyncore.dispatcher):
 			return
 
 		print >>sys.stderr, "Accepted connection from", source_addr
-		server_connection = socket.socket()
-		server_connection.connect((self.host, self.port))
 
 		# Hook the sockets up to the event loop
 		client = Proxy(client_connection)
-		server = Proxy(server_connection)
+		server = Proxy((self.host, self.port))
 		server.meet(client)
 
 class Forwarder(asyncore.dispatcher):
@@ -236,11 +248,13 @@ class Forwarder(asyncore.dispatcher):
 		asyncore.dispatcher.__init__(self)
 		self.allowed = allowed
 		self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
+		self.set_reuse_addr()
 		self.bind(("", listen_port))
-		self.listen(5)
+		self.listen(50)
+		print >>sys.stderr, "Forwarder bound on", listen_port
 
 	def handle_error(self):
-		print >>sys.stderr, "Forwarder error:", sys.exc_info()
+		print >>sys.stderr, "Forwarder error:", traceback.format_exc()
 
 	def handle_accept(self):
 		client_connection, source_addr = self.accept()
@@ -258,23 +272,23 @@ class Interceptor(asyncore.dispatcher):
 		self.host = host
 		self.port = port
 		self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
+		self.set_reuse_addr()
 		self.bind(("0.0.0.0", listen_port))
-		self.listen(5)
+		self.listen(50)
+		print >>sys.stderr, "Interceptor bound on", listen_port
 
 	def handle_error(self):
-		print >>sys.stderr, "Interceptor error!", sys.exc_info()
+		print >>sys.stderr, "Interceptor error!", traceback.format_exc()
 
 	def handle_accept(self):
 		# Get sockets
 		client_connection, source_addr = self.accept()
 		dest = get_original_dest(client_connection)
 		print >>sys.stderr, "Accepted connection from", source_addr
-		server_connection = socket.socket()
-		server_connection.connect((self.host, self.port))
 
 		# Hook them up to the event loop
 		client = Proxy(client_connection)
-		server = Proxy(server_connection)
+		server = Proxy((self.host, self.port))
 		server.buffer += "%s\n%d\n" % dest
 		server.meet(client)
 
@@ -357,7 +371,7 @@ if __name__ == "__main__":
 		if os.fork() == 0:
 			# We are the child
 			try:
-				sys.exit(main(listen_port, host, port, allowed))
+				sys.exit(main(listen_port, host, port, mode, allowed))
 			except KeyboardInterrupt:
 				print
 			sys.exit(0)