From a185e1671b52c12148ab7989e4faf756ab650876 Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Fri, 6 Jun 2014 22:52:02 +0800
Subject: [PATCH] add some status protection

---
 setup.py                 |  2 +-
 shadowsocks/eventloop.py |  6 ++++--
 shadowsocks/tcprelay.py  | 14 +++++++++++++-
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/setup.py b/setup.py
index 0263846..efdb366 100644
--- a/setup.py
+++ b/setup.py
@@ -6,7 +6,7 @@ with open('README.rst') as f:
 
 setup(
     name="shadowsocks",
-    version="2.0",
+    version="2.0.1",
     license='MIT',
     description="A fast tunnel proxy that help you get through firewalls",
     author='clowwindy',
diff --git a/shadowsocks/eventloop.py b/shadowsocks/eventloop.py
index d75905d..c9de6ef 100644
--- a/shadowsocks/eventloop.py
+++ b/shadowsocks/eventloop.py
@@ -204,9 +204,11 @@ class EventLoop(object):
                     logging.error(e)
                     continue
             for handler in self._handlers:
-                # no exceptions should be raised by users
                 # TODO when there are a lot of handlers
-                handler(events)
+                try:
+                    handler(events)
+                except (OSError, IOError) as e:
+                    logging.error(e)
 
 
 # from tornado
diff --git a/shadowsocks/tcprelay.py b/shadowsocks/tcprelay.py
index 8f12218..5f66962 100644
--- a/shadowsocks/tcprelay.py
+++ b/shadowsocks/tcprelay.py
@@ -57,6 +57,7 @@ STAGE_HELLO = 1
 STAGE_UDP_ASSOC = 2
 STAGE_REPLY = 4
 STAGE_STREAM = 5
+STAGE_DESTROYED = -1
 
 # stream direction
 STREAM_UP = 0
@@ -137,7 +138,7 @@ class TCPRelayHandler(object):
 
     def _write_to_sock(self, data, sock):
         if not data or not sock:
-            return
+            return False
         uncomplete = False
         try:
             l = len(data)
@@ -152,6 +153,7 @@ class TCPRelayHandler(object):
             else:
                 logging.error(e)
                 self.destroy()
+                return False
         if uncomplete:
             if sock == self._local_sock:
                 self._data_to_write_to_local.append(data)
@@ -168,6 +170,7 @@ class TCPRelayHandler(object):
                 self._update_stream(STREAM_UP, WAIT_STATUS_READING)
             else:
                 logging.error('write_all_to_sock:unknown socket')
+        return True
 
     def _handle_stage_reply(self, data):
         if self._is_local:
@@ -367,10 +370,14 @@ class TCPRelayHandler(object):
         self.destroy()
 
     def handle_event(self, sock, event):
+        if self._stage == STAGE_DESTROYED:
+            return
         # order is important
         if sock == self._remote_sock:
             if event & eventloop.POLL_IN:
                 self._on_remote_read()
+            if self._stage == STAGE_DESTROYED:
+                return
             if event & eventloop.POLL_OUT:
                 self._on_remote_write()
             if event & eventloop.POLL_ERR:
@@ -378,6 +385,8 @@ class TCPRelayHandler(object):
         elif sock == self._local_sock:
             if event & eventloop.POLL_IN:
                 self._on_local_read()
+            if self._stage == STAGE_DESTROYED:
+                return
             if event & eventloop.POLL_OUT:
                 self._on_local_write()
             if event & eventloop.POLL_ERR:
@@ -386,6 +395,9 @@ class TCPRelayHandler(object):
             logging.warn('unknown socket')
 
     def destroy(self):
+        if self._stage == STAGE_DESTROYED:
+            return
+        self._stage = STAGE_DESTROYED
         if self._remote_address:
             logging.debug('destroy: %s:%d' %
                           self._remote_address)