Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-26 18:18:25 +00:00
Code Issues Projects Releases Wiki Activity

Make the tcp connection closer cancellable

Browse Source 
And use this to make pathoc error handling more sophisticated
This commit is contained in:
Aldo Cortesi 2016-06-12 11:17:05 +12:00
parent 9bea616441
commit dc545ca0f6
2 changed files with 54 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
netlib/tcp.py
View File

@ -6,7 +6,6 @@ import sys
import threading import threading
import time import time
import traceback import traceback
import contextlib
import binascii import binascii
from six.moves import range from six.moves import range
@ -582,12 +581,24 @@ class _Connection(object):
return context return context
@contextlib.contextmanager class ConnectionCloser(object):
def _closer(client): def __init__(self, conn):
try: self.conn = conn
yield self._canceled = False
finally:
client.close() def pop(self):
"""
Cancel the current closer, and return a fresh one.
"""
self._canceled = True
return ConnectionCloser(self.conn)
def __enter__(self):
return self
def __exit__(self, *args):
if not self._canceled:
self.conn.close()
class TCPClient(_Connection): class TCPClient(_Connection):
@ -717,11 +728,12 @@ class TCPClient(_Connection):
except (socket.error, IOError) as err: except (socket.error, IOError) as err:
raise exceptions.TcpException( raise exceptions.TcpException(
'Error connecting to "%s": %s' % 'Error connecting to "%s": %s' %
(self.address.host, err)) (self.address.host, err)
)
self.connection = connection self.connection = connection
self.ip_address = Address(connection.getpeername()) self.ip_address = Address(connection.getpeername())
self._makefile() self._makefile()
return _closer(self) return ConnectionCloser(self)
def settimeout(self, n): def settimeout(self, n):
self.connection.settimeout(n) self.connection.settimeout(n)

5
pathod/pathoc.py
View File

@ -291,7 +291,7 @@ class Pathoc(tcp.TCPClient):
if self.use_http2 and not self.ssl: if self.use_http2 and not self.ssl:
raise NotImplementedError("HTTP2 without SSL is not supported.") raise NotImplementedError("HTTP2 without SSL is not supported.")
ret = tcp.TCPClient.connect(self) with tcp.TCPClient.connect(self) as closer:
if connect_to: if connect_to:
self.http_connect(connect_to) self.http_connect(connect_to)
@ -328,7 +328,8 @@ class Pathoc(tcp.TCPClient):
if self.timeout: if self.timeout:
self.settimeout(self.timeout) self.settimeout(self.timeout)
return ret
return closer.pop()
def stop(self): def stop(self):
if self.ws_framereader: if self.ws_framereader: