refine VPS

fungchou · Jan 5, 2015 · c9d6722 · c9d6722
1 parent aa686f9
commit c9d6722
Showing 1 changed file with 38 additions and 1 deletion.
diff --git a/local/proxy.py b/local/proxy.py
@@ -87,6 +87,7 @@
 import urllib2
 import urlparse
 import zlib
+import select
 
 import gevent
 import gevent.server
@@ -627,6 +628,42 @@ def __init__(self, *args, **kwargs):
         self.fetchservers = kwargs.pop('fetchservers')
         gevent.server.StreamServer.__init__(self, *args, **kwargs)
 
+    def forward_socket(self, local, remote, timeout, bufsize):
+        """forward socket"""
+        try:
+            tick = 1
+            timecount = timeout
+            while 1:
+                timecount -= tick
+                if timecount <= 0:
+                    break
+                (ins, _, errors) = select.select([local, remote], [], [local, remote], tick)
+                if errors:
+                    break
+                for sock in ins:
+                    data = sock.recv(bufsize)
+                    if not data:
+                        break
+                    if sock is remote:
+                        local.sendall(data)
+                        timecount = timeout
+                    else:
+                        remote.sendall(data)
+                        timecount = timeout
+        except socket.timeout:
+            pass
+        except (socket.error, ssl.SSLError, OpenSSL.SSL.Error) as e:
+            if e.args[0] not in (errno.ECONNABORTED, errno.ECONNRESET, errno.ENOTCONN, errno.EPIPE):
+                raise
+            if e.args[0] in (errno.EBADF,):
+                return
+        finally:
+            for sock in (remote, local):
+                try:
+                    sock.close()
+                except StandardError:
+                    pass
+
     def handle(self, sock, addr):
         request_data = data = ''
         while True:
@@ -652,7 +689,7 @@ def handle(self, sock, addr):
         remote = self.net2.create_ssl_connection(host, port, 8, cache_key=netloc)
         request_data = '%s\r\nProxy-Authorization: Baisic %s\r\n%s' % (request_line, base64.b64encode('%s:%s' % (username, password)).strip(), header_data)
         remote.sendall(request_data)
-        forward_socket(sock, remote, 60, bufsize=256*1024)
+        self.forward_socket(sock, remote, 60, bufsize=256*1024)
 
 
 class GAEFetchFilter(BaseProxyHandlerFilter):