From 93890e9ccc8dc01b303524c3ae4e55fdefb3582e Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Fri, 20 Apr 2012 23:09:55 +0800
Subject: [PATCH] close remote fd

---
 local.py  | 43 ++++++++++++++++++++++++-------------------
 server.py | 19 ++++++++++++-------
 2 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/local.py b/local.py
index d4ca21a..9753915 100755
--- a/local.py
+++ b/local.py
@@ -52,7 +52,7 @@ my_lock = threading.Lock()
 def lock_print(msg):
     my_lock.acquire()
     try:
-        print "[%s]%s" % (time.ctime(), msg)
+        print "[%s] %s" % (time.ctime(), msg)
     finally:
         my_lock.release()
 
@@ -69,28 +69,33 @@ class Socks5Server(SocketServer.StreamRequestHandler):
         return data.translate(decrypt_table)
 
     def handle_tcp(self, sock, remote):
-        fdset = [sock, remote]
-        counter = 0
-        while True:
-            r, w, e = select.select(fdset, [], [])
-            if sock in r:
-                r_data = sock.recv(4096)
-                if counter == 1:
-                    try:
-                        lock_print("Connecting " + r_data[5:5 + ord(r_data[4])])
-                    except Exception:
-                        pass
-                if counter < 2:
-                    counter += 1
-                if remote.send(self.encrypt(r_data)) <= 0: break
-            if remote in r:
-                if sock.send(self.decrypt(remote.recv(4096))) <= 0:
-                    break
+        try:
+            fdset = [sock, remote]
+            counter = 0
+            while True:
+                r, w, e = select.select(fdset, [], [])
+                if sock in r:
+                    r_data = sock.recv(4096)
+                    if counter == 1:
+                        try:
+                            lock_print(
+                                "Connecting " + r_data[5:5 + ord(r_data[4])])
+                        except Exception:
+                            pass
+                    if counter < 2:
+                        counter += 1
+                    if remote.send(self.encrypt(r_data)) <= 0:
+                        break
+                if remote in r:
+                    if sock.send(self.decrypt(remote.recv(4096))) <= 0:
+                        break
+        finally:
+            remote.close()
 
     def handle(self):
         try:
             sock = self.connection
-            remote = socket.socket()
+            remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
             remote.connect((SERVER, REMOTE_PORT))
             self.handle_tcp(sock, remote)
         except socket.error:
diff --git a/server.py b/server.py
index 95f36b6..6a6043f 100755
--- a/server.py
+++ b/server.py
@@ -47,13 +47,18 @@ class ThreadingTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
 
 class Socks5Server(SocketServer.StreamRequestHandler):
     def handle_tcp(self, sock, remote):
-        fdset = [sock, remote]
-        while True:
-            r, w, e = select.select(fdset, [], [])
-            if sock in r:
-                if remote.send(self.decrypt(sock.recv(4096))) <= 0: break
-            if remote in r:
-                if sock.send(self.encrypt(remote.recv(4096))) <= 0: break
+        try:
+            fdset = [sock, remote]
+            while True:
+                r, w, e = select.select(fdset, [], [])
+                if sock in r:
+                    if remote.send(self.decrypt(sock.recv(4096))) <= 0:
+                        break
+                if remote in r:
+                    if sock.send(self.encrypt(remote.recv(4096))) <= 0:
+                        break
+        finally:
+            remote.close()
 
     def encrypt(self, data):
         return data.translate(encrypt_table)