From 480c9ec51e7c6a049fc5de9c97f3fdb0de2b1982 Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Sun, 30 Dec 2012 14:43:18 +0800
Subject: [PATCH] fix infinite loop bug

---
 local.py  | 12 ++++++++----
 server.py | 14 +++++++++-----
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/local.py b/local.py
index aaa4a4d..25121b4 100755
--- a/local.py
+++ b/local.py
@@ -72,15 +72,19 @@ class Socks5Server(SocketServer.StreamRequestHandler):
                 r, w, e = select.select(fdset, [], [])
                 if sock in r:
                     data = sock.recv(4096)
-                    if data <= 0:
+                    if len(data) <= 0:
                         break
-                    send_all(remote, self.encrypt(data))
+                    result = send_all(remote, self.encrypt(data))
+                    if result < len(data):
+                        raise Exception('failed to send all data')
 
                 if remote in r:
                     data = remote.recv(4096)
-                    if data <= 0:
+                    if len(data) <= 0:
                         break
-                    send_all(sock, self.decrypt(data))
+                    result = send_all(sock, self.decrypt(data))
+                    if result < len(data):
+                        raise Exception('failed to send all data')
         finally:
             sock.close()
             remote.close()
diff --git a/server.py b/server.py
index b608c24..35f2a51 100755
--- a/server.py
+++ b/server.py
@@ -60,7 +60,6 @@ def send_all(sock, data):
         if bytes_sent == len(data):
             return bytes_sent
 
-
 class ThreadingTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
     allow_reuse_address = True
 
@@ -73,14 +72,19 @@ class Socks5Server(SocketServer.StreamRequestHandler):
                 r, w, e = select.select(fdset, [], [])
                 if sock in r:
                     data = sock.recv(4096)
-                    if data <= 0:
+                    if len(data) <= 0:
                         break
-                    send_all(remote, self.decrypt(data))
+                    result = send_all(remote, self.decrypt(data))
+                    if result < len(data):
+                        raise Exception('failed to send all data')
                 if remote in r:
                     data = remote.recv(4096)
-                    if data <= 0:
+                    if len(data) <= 0:
                         break
-                    send_all(sock, self.encrypt(data))
+                    result = send_all(sock, self.encrypt(data))
+                    if result < len(data):
+                        raise Exception('failed to send all data')
+
         finally:
             sock.close()
             remote.close()