From cf19630d15d7af536f96d8018ca320c738433985 Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Sat, 25 May 2013 13:32:39 +0800
Subject: [PATCH] add pause and resume

---
 local.py  | 18 +++++++++++++++---
 server.py | 17 +++++++++++++++--
 ssloop    |  2 +-
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/local.py b/local.py
index c855aa9..29f08e7 100755
--- a/local.py
+++ b/local.py
@@ -40,6 +40,7 @@ class RemoteHandler(object):
         self.local_handler = local_handler
         conn.on('connect', self.on_connect)
         conn.on('data', self.on_data)
+        conn.on('drain', self.on_drain)
         conn.on('close', self.on_close)
         conn.on('end', self.on_end)
         conn.connect((SERVER, REMOTE_PORT))
@@ -53,7 +54,12 @@ class RemoteHandler(object):
 
     def on_data(self, s, data):
         data = self.local_handler.encryptor.decrypt(data)
-        self.local_handler.conn.write(data)
+        if not self.local_handler.conn.write(data):
+            self.conn.pause()
+
+    def on_drain(self, s):
+        if self.local_handler:
+            self.local_handler.conn.resume()
 
     def on_close(self, s):
         # self.local_handler.conn.end()
@@ -67,7 +73,8 @@ class LocalHandler(object):
     def on_data(self, s, data):
         if self.stage == 5:
             data = self.encryptor.encrypt(data)
-            self.remote_handler.conn.write(data)
+            if not self.remote_handler.conn.write(data):
+                self.conn.pause()
             return
         if self.stage == 0:
             self.conn.write('\x05\00')
@@ -117,6 +124,10 @@ class LocalHandler(object):
         if self.stage == 4:
             self.cached_pieces.append(data)
 
+    def on_drain(self, s):
+        if self.remote_handler:
+            self.remote_handler.conn.resume()
+
     def on_end(self, s):
         if self.remote_handler:
             self.remote_handler.conn.end()
@@ -127,14 +138,15 @@ class LocalHandler(object):
 
     def __init__(self, conn):
         self.stage = 0
-        self.remote = None
         self.addr_len = 0
         self.addr_to_send = ''
         self.conn = conn
         self.cached_pieces = []
         self.encryptor = encrypt.Encryptor(KEY, METHOD)
+        self.remote_handler = None
 
         conn.on('data', self.on_data)
+        conn.on('drain', self.on_drain)
         conn.on('end', self.on_end)
         conn.on('close', self.on_close)
 
diff --git a/server.py b/server.py
index e79f8b0..ee0e9ab 100755
--- a/server.py
+++ b/server.py
@@ -40,6 +40,7 @@ class RemoteHandler(object):
         self.local_handler = local_handler
         conn.on('connect', self.on_connect)
         conn.on('data', self.on_data)
+        conn.on('drain', self.on_drain)
         conn.on('close', self.on_close)
         conn.on('end', self.on_end)
         conn.connect(local_handler.remote_addr_pair)
@@ -53,7 +54,12 @@ class RemoteHandler(object):
 
     def on_data(self, s, data):
         data = self.local_handler.encryptor.encrypt(data)
-        self.local_handler.conn.write(data)
+        if not self.local_handler.conn.write(data):
+            self.conn.pause()
+
+    def on_drain(self, s):
+        if self.local_handler:
+            self.local_handler.conn.resume()
 
     def on_close(self, s):
         # self.local_handler.conn.end()
@@ -67,7 +73,8 @@ class LocalHandler(object):
     def on_data(self, s, data):
         data = self.encryptor.decrypt(data)
         if self.stage == 5:
-            self.remote_handler.conn.write(data)
+            if not self.remote_handler.conn.write(data):
+                self.conn.pause()
             return
         if self.stage == 0:
             try:
@@ -108,6 +115,10 @@ class LocalHandler(object):
         if self.stage == 4:
             self.cached_pieces.append(data)
 
+    def on_drain(self, s):
+        if self.remote_handler:
+            self.remote_handler.conn.resume()
+
     def on_end(self, s):
         if self.remote_handler:
             self.remote_handler.conn.end()
@@ -124,8 +135,10 @@ class LocalHandler(object):
         self.conn = conn
         self.cached_pieces = []
         self.encryptor = encrypt.Encryptor(KEY, METHOD)
+        self.remote_handler = None
 
         conn.on('data', self.on_data)
+        conn.on('drain', self.on_drain)
         conn.on('end', self.on_end)
         conn.on('close', self.on_close)
 
diff --git a/ssloop b/ssloop
index 39beb21..20caed1 160000
--- a/ssloop
+++ b/ssloop
@@ -1 +1 @@
-Subproject commit 39beb21517a7adb28408c01bd258d0843ea731b5
+Subproject commit 20caed14e2a9c842c28b76efa6d23c26dd2c5db9