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

asyncio: fix channel interface and tests

Browse Source 
We now need to synthesize a tick event when changing addons in tests, because
tick is asynchronously called on the event loop.
This commit is contained in:
Aldo Cortesi 2018-04-01 11:37:35 +12:00
parent b6d943cfa3
commit 3cc5d81a4a
4 changed files with 33 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mitmproxy/controller.py
View File

@ -8,8 +8,9 @@ class Channel:
The only way for the proxy server to communicate with the master
is to use the channel it has been given.
"""
def __init__(self, loop, q):
def __init__(self, loop, q, should_exit):
self.loop = loop
self.should_exit = should_exit
self._q = q
def ask(self, mtype, m):

12
mitmproxy/master.py
View File

@ -1,7 +1,6 @@
import threading
import contextlib
import asyncio
import signal
from mitmproxy import addonmanager
from mitmproxy import options
@ -37,11 +36,16 @@ class Master:
"""
def __init__(self, opts):
self.event_queue = asyncio.Queue()
self.should_exit = threading.Event()
self.channel = controller.Channel(
asyncio.get_event_loop(),
self.event_queue,
self.should_exit,
)
self.options = opts or options.Options() # type: options.Options
self.commands = command.CommandManager(self)
self.addons = addonmanager.AddonManager(self)
self.should_exit = threading.Event()
self._server = None
self.first_tick = True
self.waiting_flows = []
@ -52,7 +56,7 @@ class Master:
@server.setter
def server(self, server):
server.set_channel(controller.Channel(asyncio.get_event_loop(), self.event_queue))
server.set_channel(self.channel)
self._server = server
@contextlib.contextmanager
@ -202,7 +206,7 @@ class Master:
host = f.request.headers.pop(":authority")
f.request.headers.insert(0, "host", host)
rt = http_replay.RequestReplayThread(self.options, f, self.server.channel)
rt = http_replay.RequestReplayThread(self.options, f, self.channel)
rt.start() # pragma: no cover
if block:
rt.join()

34
test/mitmproxy/proxy/test_server.py
View File

@ -276,10 +276,9 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin):
s = script.Script(
tutils.test_data.path("mitmproxy/data/addonscripts/stream_modify.py")
)
self.master.addons.add(s)
self.set_addons(s)
d = self.pathod('200:b"foo"')
assert d.content == b"bar"
self.master.addons.remove(s)
def test_first_line_rewrite(self):
"""
@ -583,12 +582,11 @@ class TestTransparent(tservers.TransparentProxyTest, CommonMixin, TcpMixin):
s = script.Script(
tutils.test_data.path("mitmproxy/data/addonscripts/tcp_stream_modify.py")
)
self.master.addons.add(s)
self.set_addons(s)
self._tcpproxy_on()
d = self.pathod('200:b"foo"')
self._tcpproxy_off()
assert d.content == b"bar"
self.master.addons.remove(s)
class TestTransparentSSL(tservers.TransparentProxyTest, CommonMixin, TcpMixin):
@ -731,7 +729,7 @@ class TestRedirectRequest(tservers.HTTPProxyTest):
This test verifies that the original destination is restored for the third request.
"""
self.proxy.tmaster.addons.add(ARedirectRequest(self.server2.port))
self.set_addons(ARedirectRequest(self.server2.port))
p = self.pathoc()
with p.connect():
@ -770,7 +768,7 @@ class AStreamRequest:
class TestStreamRequest(tservers.HTTPProxyTest):
def test_stream_simple(self):
self.proxy.tmaster.addons.add(AStreamRequest())
self.set_addons(AStreamRequest())
p = self.pathoc()
with p.connect():
# a request with 100k of data but without content-length
@ -779,7 +777,7 @@ class TestStreamRequest(tservers.HTTPProxyTest):
assert len(r1.content) > 100000
def test_stream_multiple(self):
self.proxy.tmaster.addons.add(AStreamRequest())
self.set_addons(AStreamRequest())
p = self.pathoc()
with p.connect():
# simple request with streaming turned on
@ -791,7 +789,7 @@ class TestStreamRequest(tservers.HTTPProxyTest):
assert r1.status_code == 201
def test_stream_chunked(self):
self.proxy.tmaster.addons.add(AStreamRequest())
self.set_addons(AStreamRequest())
connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
connection.connect(("127.0.0.1", self.proxy.port))
fconn = connection.makefile("rb")
@ -820,7 +818,7 @@ class AFakeResponse:
class TestFakeResponse(tservers.HTTPProxyTest):
def test_fake(self):
self.proxy.tmaster.addons.add(AFakeResponse())
self.set_addons(AFakeResponse())
f = self.pathod("200")
assert "header-response" in f.headers
@ -836,7 +834,7 @@ class TestServerConnect(tservers.HTTPProxyTest):
def test_unnecessary_serverconnect(self):
"""A replayed/fake response with no upstream_cert should not connect to an upstream server"""
self.proxy.tmaster.addons.add(AFakeResponse())
self.set_addons(AFakeResponse())
assert self.pathod("200").status_code == 200
assert not self.proxy.tmaster.has_log("serverconnect")
@ -849,7 +847,7 @@ class AKillRequest:
class TestKillRequest(tservers.HTTPProxyTest):
def test_kill(self):
self.proxy.tmaster.addons.add(AKillRequest())
self.set_addons(AKillRequest())
with pytest.raises(exceptions.HttpReadDisconnect):
self.pathod("200")
# Nothing should have hit the server
@ -863,7 +861,7 @@ class AKillResponse:
class TestKillResponse(tservers.HTTPProxyTest):
def test_kill(self):
self.proxy.tmaster.addons.add(AKillResponse())
self.set_addons(AKillResponse())
with pytest.raises(exceptions.HttpReadDisconnect):
self.pathod("200")
# The server should have seen a request
@ -886,7 +884,7 @@ class AIncomplete:
class TestIncompleteResponse(tservers.HTTPProxyTest):
def test_incomplete(self):
self.proxy.tmaster.addons.add(AIncomplete())
self.set_addons(AIncomplete())
assert self.pathod("200").status_code == 502
@ -969,7 +967,7 @@ class TestUpstreamProxySSL(
def test_change_upstream_proxy_connect(self):
# skip chain[0].
self.proxy.tmaster.addons.add(
self.set_addons(
UpstreamProxyChanger(
("127.0.0.1", self.chain[1].port)
)
@ -988,8 +986,8 @@ class TestUpstreamProxySSL(
Client <- HTTPS -> Proxy <- HTTP -> Proxy <- HTTPS -> Server
"""
self.proxy.tmaster.addons.add(RewriteToHttp())
self.chain[1].tmaster.addons.add(RewriteToHttps())
self.set_addons(RewriteToHttp())
self.set_addons(RewriteToHttps())
p = self.pathoc()
with p.connect():
resp = p.request("get:'/p/418'")
@ -1063,8 +1061,8 @@ class TestProxyChainingSSLReconnect(tservers.HTTPUpstreamProxyTest):
http1obj.server_conn.wfile.write(headers)
http1obj.server_conn.wfile.flush()
self.chain[0].tmaster.addons.add(RequestKiller([1, 2]))
self.chain[1].tmaster.addons.add(RequestKiller([1]))
self.chain[0].set_addons(RequestKiller([1, 2]))
self.chain[1].set_addons(RequestKiller([1]))
p = self.pathoc()
with p.connect():

8
test/mitmproxy/tservers.py
View File

@ -92,6 +92,7 @@ class ProxyThread(threading.Thread):
self.masterclass = masterclass
self.options = options
self.tmaster = None
self.event_loop = None
controller.should_exit = False
@property
@ -106,7 +107,8 @@ class ProxyThread(threading.Thread):
self.tmaster.shutdown()
def run(self):
asyncio.set_event_loop(asyncio.new_event_loop())
self.event_loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.event_loop)
self.tmaster = self.masterclass(self.options)
self.tmaster.addons.add(core.Core())
self.name = "ProxyThread (%s:%s)" % (
@ -177,6 +179,10 @@ class ProxyTestBase:
ssl_insecure=True,
)
def set_addons(self, *addons):
self.proxy.tmaster.reset(addons)
self.proxy.tmaster.addons.trigger("tick")
def addons(self):
"""
Can be over-ridden to add a standard set of addons to tests.