From 6caeabd03de4da5e10298c64fa5586651958eb1f Mon Sep 17 00:00:00 2001 From: clowwindy Date: Sat, 10 May 2014 01:22:26 +0800 Subject: [PATCH] fix a potential RST problem --- shadowsocks/local.py | 24 ++++++++++++++++-------- shadowsocks/server.py | 24 ++++++++++++++++-------- 2 files changed, 32 insertions(+), 16 deletions(-) diff --git a/shadowsocks/local.py b/shadowsocks/local.py index 5ff9615..1bb9cec 100755 --- a/shadowsocks/local.py +++ b/shadowsocks/local.py @@ -100,6 +100,7 @@ class Socks5Server(SocketServer.StreamRequestHandler): else: fdset = [sock, remote] while True: + should_break = False r, w, e = select.select(fdset, [], []) if sock in r: if not connected and config_fast_open: @@ -123,18 +124,25 @@ class Socks5Server(SocketServer.StreamRequestHandler): pending_data = None data = encryptor.encrypt(data) if len(data) <= 0: - break - result = send_all(remote, data) - if result < len(data): - raise Exception('failed to send all data') + should_break = True + else: + result = send_all(remote, data) + if result < len(data): + raise Exception('failed to send all data') if remote in r: data = encryptor.decrypt(remote.recv(4096)) if len(data) <= 0: - break - result = send_all(sock, data) - if result < len(data): - raise Exception('failed to send all data') + should_break = True + else: + result = send_all(sock, data) + if result < len(data): + raise Exception('failed to send all data') + if should_break: + # make sure all data are read before we close the sockets + # TODO: we haven't read ALL the data, actually + # http://cs.ecs.baylor.edu/~donahoo/practical/CSockets/TCPRST.pdf + break finally: sock.close() remote.close() diff --git a/shadowsocks/server.py b/shadowsocks/server.py index aa0affc..8493b0a 100755 --- a/shadowsocks/server.py +++ b/shadowsocks/server.py @@ -79,21 +79,29 @@ class Socks5Server(SocketServer.StreamRequestHandler): try: fdset = [sock, remote] while True: + should_break = False r, w, e = select.select(fdset, [], []) if sock in r: data = self.decrypt(sock.recv(4096)) if len(data) <= 0: - break - result = send_all(remote, data) - if result < len(data): - raise Exception('failed to send all data') + should_break = True + else: + result = send_all(remote, data) + if result < len(data): + raise Exception('failed to send all data') if remote in r: data = self.encrypt(remote.recv(4096)) if len(data) <= 0: - break - result = send_all(sock, data) - if result < len(data): - raise Exception('failed to send all data') + should_break = True + else: + result = send_all(sock, data) + if result < len(data): + raise Exception('failed to send all data') + if should_break: + # make sure all data are read before we close the sockets + # TODO: we haven't read ALL the data, actually + # http://cs.ecs.baylor.edu/~donahoo/practical/CSockets/TCPRST.pdf + break finally: sock.close()