fix infinite loop bug

This commit is contained in:
clowwindy 2012-12-30 14:43:18 +08:00
parent 004e9292f4
commit 480c9ec51e
2 changed files with 17 additions and 9 deletions

View file

@ -72,15 +72,19 @@ class Socks5Server(SocketServer.StreamRequestHandler):
r, w, e = select.select(fdset, [], []) r, w, e = select.select(fdset, [], [])
if sock in r: if sock in r:
data = sock.recv(4096) data = sock.recv(4096)
if data <= 0: if len(data) <= 0:
break 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: if remote in r:
data = remote.recv(4096) data = remote.recv(4096)
if data <= 0: if len(data) <= 0:
break 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: finally:
sock.close() sock.close()
remote.close() remote.close()

View file

@ -60,7 +60,6 @@ def send_all(sock, data):
if bytes_sent == len(data): if bytes_sent == len(data):
return bytes_sent return bytes_sent
class ThreadingTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer): class ThreadingTCPServer(SocketServer.ThreadingMixIn, SocketServer.TCPServer):
allow_reuse_address = True allow_reuse_address = True
@ -73,14 +72,19 @@ class Socks5Server(SocketServer.StreamRequestHandler):
r, w, e = select.select(fdset, [], []) r, w, e = select.select(fdset, [], [])
if sock in r: if sock in r:
data = sock.recv(4096) data = sock.recv(4096)
if data <= 0: if len(data) <= 0:
break 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: if remote in r:
data = remote.recv(4096) data = remote.recv(4096)
if data <= 0: if len(data) <= 0:
break 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: finally:
sock.close() sock.close()
remote.close() remote.close()