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

refactor tcp handling, fix #280

Browse Source
This commit is contained in:
Maximilian Hils 2014-07-27 02:39:17 +02:00
parent 83536e7e53
commit 2b31a9c49e
1 changed files with 35 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
libmproxy/protocol/tcp.py
View File

@ -1,16 +1,21 @@
from __future__ import absolute_import
import select, socket
from cStringIO import StringIO
from .primitives import ProtocolHandler
class TCPHandler(ProtocolHandler):
"""
TCPHandler acts as a generic TCP forwarder.
Data will be .log()ed, but not stored any further.
"""
chunk_size = 4096
def handle_messages(self):
self.c.establish_server_connection()
server = "%s:%s" % self.c.server_conn.address()[:2]
buf = memoryview(bytearray(self.chunk_size))
conns = [self.c.client_conn.rfile, self.c.server_conn.rfile]
while not self.c.close:
r, _, _ = select.select(conns, [], [], 10)
@ -18,43 +23,43 @@ class TCPHandler(ProtocolHandler):
if self.c.client_conn.rfile == rfile:
src, dst = self.c.client_conn, self.c.server_conn
direction = "-> tcp ->"
dst_str = "%s:%s" % self.c.server_conn.address()[:2]
src_str, dst_str = "client", server
else:
dst, src = self.c.client_conn, self.c.server_conn
direction = "<- tcp <-"
dst_str = "client"
dst_str, src_str = "client", server
data = StringIO()
while range(4096):
# Do non-blocking select() to see if there is further data on in the buffer.
r, _, _ = select.select([rfile], [], [], 0)
if len(r):
d = rfile.read(1)
if d == "": # connection closed
break
data.write(d)
# OpenSSL Connections have an internal buffer that might
# contain data altough everything is read from the socket.
# Thankfully, connection.pending() returns the amount of
# bytes in this buffer, so we can read it completely at
# once.
if src.ssl_established:
data.write(rfile.read(src.connection.pending()))
else: # no data left, but not closed yet
break
data = data.getvalue()
closed = False
if src.ssl_established:
# Unfortunately, pyOpenSSL lacks a recv_into function.
contents = src.rfile.read(1) # We need to read a single byte before .pending() becomes usable
contents += src.rfile.read(src.connection.pending())
if not contents:
closed = True
else:
size = src.connection.recv_into(buf)
if not size:
closed = True
if data == "": # no data received, rfile is closed
self.c.log("Close writing connection to %s" % dst_str, "debug")
conns.remove(rfile)
if closed:
conns.remove(src.rfile)
# Shutdown connection to the other peer
if dst.ssl_established:
dst.connection.shutdown()
else:
dst.connection.shutdown(socket.SHUT_WR)
if len(conns) == 0:
self.c.close = True
break
continue
self.c.log("%s %s\r\n%s" % (direction, dst_str, data), "debug")
dst.wfile.write(data)
dst.wfile.flush()
if src.ssl_established or dst.ssl_established:
# if one of the peers is over SSL, we need to send bytes/strings
if not src.ssl_established: # only ssl to dst, i.e. we revc'd into buf but need bytes/string now.
contents = buf[:size].tobytes()
self.c.log("%s %s\r\n%s" % (direction, dst_str, contents[:100]), "debug")
dst.connection.send(contents)
else:
# socket.socket.send supports raw bytearrays/memoryviews
self.c.log("%s %s\r\n%s" % (direction, dst_str, buf[:100]), "debug")
dst.connection.send(buf[:size])