fix

qwj · qwj · commit ef163aaf5bfc · 2021-02-19T15:02:07.000+08:00
diff --git a/pproxy/server.py b/pproxy/server.py
@@ -6,22 +6,12 @@
 UDP_LIMIT = 30
 DUMMY = lambda s: s
 
-class ProxyReader(asyncio.StreamReader):
- def __init__(self, o=None):
- if o:
- for n in dir(o):
- if not n.startswith('__'):
- setattr(self, n, getattr(o, n))
- else:
- super(self).__init__()
- def read_w(self, n):
- return asyncio.wait_for(self.read(n), timeout=SOCKET_TIMEOUT)
- def read_n(self, n):
- return asyncio.wait_for(self.readexactly(n), timeout=SOCKET_TIMEOUT)
- def read_until(self, s):
- return asyncio.wait_for(self.readuntil(s), timeout=SOCKET_TIMEOUT)
- def rollback(self, s):
- self._buffer.__setitem__(slice(0, 0), s)
+def patch_StreamReader(c=asyncio.StreamReader):
+ c.read_w = lambda self, n: asyncio.wait_for(self.read(n), timeout=SOCKET_TIMEOUT)
+ c.read_n = lambda self, n: asyncio.wait_for(self.readexactly(n), timeout=SOCKET_TIMEOUT)
+ c.read_until = lambda self, s: asyncio.wait_for(self.readuntil(s), timeout=SOCKET_TIMEOUT)
+ c.rollback = lambda self, s: self._buffer.__setitem__(slice(0, 0), s)
+patch_StreamReader()
 
 class AuthTable(object):
 _auth ={}
@@ -68,7 +58,7 @@ def schedule(rserver, salgorithm, host_name, port):
 
 async def stream_handler(reader, writer, unix, lbind, protos, rserver, cipher, sslserver, debug=0, authtime=86400*30, block=None, salgorithm='fa', verbose=DUMMY, modstat=lambda u,r,h:lambda i:DUMMY, **kwargs):
 try:
- reader, writer = proto.sslwrap(ProxyReader(reader), writer, sslserver, True, None, verbose)
+ reader, writer = proto.sslwrap(reader, writer, sslserver, True, None, verbose)
 if unix:
 remote_ip, server_ip, remote_text = 'local', None, 'unix_local'
 else:
@@ -231,7 +221,7 @@ async def open_connection(self, host, port, local_addr, lbind, timeout=SOCKET_TI
 reader, writer = await asyncio.wait_for(wait, timeout=timeout)
 except Exception as ex:
 raise
- return ProxyReader(reader), writer
+ return reader, writer
 async def prepare_connection(self, reader_remote, writer_remote, host, port):
 return reader_remote, writer_remote
 async def tcp_connect(self, host, port, local_addr=None, lbind=None):
@@ -532,7 +522,7 @@ async def handler(reader, writer, **kw):
 auth = b'\x01'+auth
 if auth:
 try:
- assert auth == (await asyncio.wait_for(reader.readexactly(len(auth)), timeout=SOCKET_TIMEOUT))
+ assert auth == (await reader.read_n(len(auth)))
 except Exception:
 return
 await self.conn.put((reader, writer))

-Original file line number
+Diff line change
 UDP_LIMIT=30
 DUMMY=lambdas: s
 -classProxyReader(asyncio.StreamReader):
 -def__init__(self, o=None):
 -ifo:
 -fornindir(o):
 -ifnotn.startswith('__'):
 -setattr(self, n, getattr(o, n))
 -else:
 -super(self).__init__()
 -defread_w(self, n):
 -returnasyncio.wait_for(self.read(n), timeout=SOCKET_TIMEOUT)
 -defread_n(self, n):
 -returnasyncio.wait_for(self.readexactly(n), timeout=SOCKET_TIMEOUT)
 -defread_until(self, s):
 -returnasyncio.wait_for(self.readuntil(s), timeout=SOCKET_TIMEOUT)
 -defrollback(self, s):
 -self._buffer.__setitem__(slice(0, 0), s)
 +defpatch_StreamReader(c=asyncio.StreamReader):
 +c.read_w=lambdaself, n: asyncio.wait_for(self.read(n), timeout=SOCKET_TIMEOUT)
 +c.read_n=lambdaself, n: asyncio.wait_for(self.readexactly(n), timeout=SOCKET_TIMEOUT)
 +c.read_until=lambdaself, s: asyncio.wait_for(self.readuntil(s), timeout=SOCKET_TIMEOUT)
 +c.rollback=lambdaself, s: self._buffer.__setitem__(slice(0, 0), s)
 +patch_StreamReader()
 classAuthTable(object):
 _auth={}
 asyncdefstream_handler(reader, writer, unix, lbind, protos, rserver, cipher, sslserver, debug=0, authtime=86400*30, block=None, salgorithm='fa', verbose=DUMMY, modstat=lambdau,r,h:lambdai:DUMMY, **kwargs):
 try:
 -reader, writer=proto.sslwrap(ProxyReader(reader), writer, sslserver, True, None, verbose)
 +reader, writer=proto.sslwrap(reader, writer, sslserver, True, None, verbose)
 ifunix:
 remote_ip, server_ip, remote_text='local', None, 'unix_local'
 else:
 reader, writer=awaitasyncio.wait_for(wait, timeout=timeout)
 exceptExceptionasex:
 raise
 -returnProxyReader(reader), writer
 +returnreader, writer
 asyncdefprepare_connection(self, reader_remote, writer_remote, host, port):
 returnreader_remote, writer_remote
 asyncdeftcp_connect(self, host, port, local_addr=None, lbind=None):
 auth=b'\x01'+auth
 ifauth:
 try:
 -assertauth== (awaitasyncio.wait_for(reader.readexactly(len(auth)), timeout=SOCKET_TIMEOUT))
 +assertauth== (awaitreader.read_n(len(auth)))
 exceptException:
 return
 awaitself.conn.put((reader, writer))