Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2025-02-07 02:28:50 +00:00
Code Issues Projects Releases Wiki Activity

fix #113, make Reader.peek() work on Python 3

Browse Source
This commit is contained in:
Maximilian Hils 2016-02-01 19:38:14 +01:00
parent 7c83a709ea
commit bda49dd178
2 changed files with 26 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
netlib/tcp.py
View File

@ -25,6 +25,10 @@ from netlib.exceptions import InvalidCertificateException, TcpReadIncomplete, Tl
version_check.check_pyopenssl_version()
if six.PY2:
socket_fileobject = socket._fileobject
else:
socket_fileobject = socket.SocketIO
EINTR = 4
@ -270,7 +274,7 @@ class Reader(_FileLike):
TlsException if there was an error with pyOpenSSL.
NotImplementedError if the underlying file object is not a (pyOpenSSL) socket
"""
if isinstance(self.o, socket._fileobject):
if isinstance(self.o, socket_fileobject):
try:
return self.o._sock.recv(length, socket.MSG_PEEK)
except socket.error as e:
@ -423,8 +427,17 @@ class _Connection(object):
def __init__(self, connection):
if connection:
self.connection = connection
self.rfile = Reader(self.connection.makefile('rb', self.rbufsize))
self.wfile = Writer(self.connection.makefile('wb', self.wbufsize))
# Ideally, we would use the Buffered IO in Python 3 by default.
# Unfortunately, the implementation of .peek() is broken for n>1 bytes,
# as it may just return what's left in the buffer and not all the bytes we want.
# As a workaround, we just use unbuffered sockets directly.
# https://mail.python.org/pipermail/python-dev/2009-June/089986.html
if six.PY2:
self.rfile = Reader(self.connection.makefile('rb', self.rbufsize))
self.wfile = Writer(self.connection.makefile('wb', self.wbufsize))
else:
self.rfile = Reader(socket.SocketIO(self.connection, "rb"))
self.wfile = Writer(socket.SocketIO(self.connection, "wb"))
else:
self.connection = None
self.rfile = None
@ -663,8 +676,15 @@ class TCPClient(_Connection):
connection.connect(self.address())
if not self.source_address:
self.source_address = Address(connection.getsockname())
self.rfile = Reader(connection.makefile('rb', self.rbufsize))
self.wfile = Writer(connection.makefile('wb', self.wbufsize))
# See _Connection.__init__ why we do this dance.
if six.PY2:
self.rfile = Reader(connection.makefile('rb', self.rbufsize))
self.wfile = Writer(connection.makefile('wb', self.wbufsize))
else:
self.rfile = Reader(socket.SocketIO(connection, "rb"))
self.wfile = Writer(socket.SocketIO(connection, "wb"))
except (socket.error, IOError) as err:
raise TcpException(
'Error connecting to "%s": %s' %

2
test/test_tcp.py
View File

@ -723,7 +723,7 @@ class TestPeek(tservers.ServerTestBase):
c.wfile.write(testval)
c.wfile.flush()
assert c.rfile.peek(4) == "peek"[:4]
assert c.rfile.peek(4) == b"peek"[:4]
assert c.rfile.peek(6) == testval