Skip to content

Commit 423c076

Browse files
mhilsmhils
committed
cleanup mitmproxy.controller, raise Kill in Channel (mitmproxy#1085)
1 parent bc60c26 commit 423c076

16 files changed

Lines changed: 265 additions & 178 deletions

File tree

‎mitmproxy/console/__init__.py‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -451,7 +451,7 @@ def set_palette(self, name):
451451
self.ui.clear()
452452

453453
def ticker(self, *userdata):
454-
changed = self.tick(self.masterq, timeout=0)
454+
changed = self.tick(timeout=0)
455455
if changed:
456456
self.loop.draw_screen()
457457
signals.update_settings.send()
@@ -467,11 +467,6 @@ def run(self):
467467
handle_mouse = not self.options.no_mouse,
468468
)
469469

470-
self.server.start_slave(
471-
controller.Slave,
472-
controller.Channel(self.masterq, self.should_exit)
473-
)
474-
475470
if self.options.rfile:
476471
ret = self.load_flows_path(self.options.rfile)
477472
if ret and self.state.flow_count():
@@ -507,6 +502,7 @@ def exit(s, f):
507502
lambda *args: self.view_flowlist()
508503
)
509504

505+
self.start()
510506
try:
511507
self.loop.run()
512508
except Exception:

‎mitmproxy/controller.py‎

Lines changed: 107 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -2,53 +2,105 @@
22
from six.moves import queue
33
import threading
44

5+
from .exceptions import Kill
56

6-
class DummyReply:
77

8+
class Master(object):
89
"""
9-
A reply object that does nothing. Useful when we need an object to seem
10-
like it has a channel, and during testing.
10+
The master handles mitmproxy's main event loop.
1111
"""
1212

1313
def __init__(self):
14-
self.acked = False
14+
self.event_queue = queue.Queue()
15+
self.should_exit = threading.Event()
1516

16-
def __call__(self, msg=False):
17-
self.acked = True
17+
def start(self):
18+
self.should_exit.clear()
1819

20+
def run(self):
21+
self.start()
22+
try:
23+
while not self.should_exit.is_set():
24+
# Don't choose a very small timeout in Python 2:
25+
# https://github.com/mitmproxy/mitmproxy/issues/443
26+
# TODO: Lower the timeout value if we move to Python 3.
27+
self.tick(0.1)
28+
finally:
29+
self.shutdown()
30+
31+
def tick(self, timeout):
32+
changed = False
33+
try:
34+
# This endless loop runs until the 'Queue.Empty'
35+
# exception is thrown.
36+
while True:
37+
mtype, obj = self.event_queue.get(timeout=timeout)
38+
handle_func = getattr(self, "handle_" + mtype)
39+
handle_func(obj)
40+
self.event_queue.task_done()
41+
changed = True
42+
except queue.Empty:
43+
pass
44+
return changed
45+
46+
def shutdown(self):
47+
self.should_exit.set()
1948

20-
class Reply:
2149

50+
class ServerMaster(Master):
2251
"""
23-
Messages sent through a channel are decorated with a "reply" attribute.
24-
This object is used to respond to the message through the return
25-
channel.
52+
The ServerMaster adds server thread support to the master.
2653
"""
2754

28-
def __init__(self, obj):
29-
self.obj = obj
30-
self.q = queue.Queue()
31-
self.acked = False
55+
def __init__(self):
56+
super(ServerMaster, self).__init__()
57+
self.servers = []
3258

33-
def __call__(self, msg=None):
34-
if not self.acked:
35-
self.acked = True
36-
if msg is None:
37-
self.q.put(self.obj)
38-
else:
39-
self.q.put(msg)
59+
def add_server(self, server):
60+
# We give a Channel to the server which can be used to communicate with the master
61+
channel = Channel(self.event_queue, self.should_exit)
62+
server.set_channel(channel)
63+
self.servers.append(server)
64+
65+
def start(self):
66+
super(ServerMaster, self).start()
67+
for server in self.servers:
68+
ServerThread(server).start()
69+
70+
def shutdown(self):
71+
for server in self.servers:
72+
server.shutdown()
73+
super(ServerMaster, self).shutdown()
4074

4175

42-
class Channel:
76+
class ServerThread(threading.Thread):
77+
def __init__(self, server):
78+
self.server = server
79+
super(ServerThread, self).__init__()
80+
address = getattr(self.server, "address", None)
81+
self.name = "ServerThread ({})".format(repr(address))
82+
83+
def run(self):
84+
self.server.serve_forever()
85+
86+
87+
class Channel(object):
88+
"""
89+
The only way for the proxy server to communicate with the master
90+
is to use the channel it has been given.
91+
"""
4392

4493
def __init__(self, q, should_exit):
4594
self.q = q
4695
self.should_exit = should_exit
4796

4897
def ask(self, mtype, m):
4998
"""
50-
Decorate a message with a reply attribute, and send it to the
51-
master. then wait for a response.
99+
Decorate a message with a reply attribute, and send it to the
100+
master. Then wait for a response.
101+
102+
Raises:
103+
Kill: All connections should be closed immediately.
52104
"""
53105
m.reply = Reply(m)
54106
self.q.put((mtype, m))
@@ -58,85 +110,54 @@ def ask(self, mtype, m):
58110
g = m.reply.q.get(timeout=0.5)
59111
except queue.Empty: # pragma: no cover
60112
continue
113+
if g == Kill:
114+
raise Kill()
61115
return g
62116

117+
raise Kill()
118+
63119
def tell(self, mtype, m):
64120
"""
65-
Decorate a message with a dummy reply attribute, send it to the
66-
master, then return immediately.
121+
Decorate a message with a dummy reply attribute, send it to the
122+
master, then return immediately.
67123
"""
68124
m.reply = DummyReply()
69125
self.q.put((mtype, m))
70126

71127

72-
class Slave(threading.Thread):
73-
128+
class DummyReply(object):
74129
"""
75-
Slaves get a channel end-point through which they can send messages to
76-
the master.
130+
A reply object that does nothing. Useful when we need an object to seem
131+
like it has a channel, and during testing.
77132
"""
78133

79-
def __init__(self, channel, server):
80-
self.channel, self.server = channel, server
81-
self.server.set_channel(channel)
82-
threading.Thread.__init__(self)
83-
self.name = "SlaveThread ({})".format(repr(self.server.address))
134+
def __init__(self):
135+
self.acked = False
84136

85-
def run(self):
86-
self.server.serve_forever()
137+
def __call__(self, msg=False):
138+
self.acked = True
87139

88140

89-
class Master(object):
141+
# Special value to distinguish the case where no reply was sent
142+
NO_REPLY = object()
143+
90144

145+
class Reply(object):
91146
"""
92-
Masters get and respond to messages from slaves.
147+
Messages sent through a channel are decorated with a "reply" attribute.
148+
This object is used to respond to the message through the return
149+
channel.
93150
"""
94151

95-
def __init__(self, server):
96-
"""
97-
server may be None if no server is needed.
98-
"""
99-
self.server = server
100-
self.masterq = queue.Queue()
101-
self.should_exit = threading.Event()
102-
103-
def tick(self, q, timeout):
104-
changed = False
105-
try:
106-
# This endless loop runs until the 'Queue.Empty'
107-
# exception is thrown. If more than one request is in
108-
# the queue, this speeds up every request by 0.1 seconds,
109-
# because get_input(..) function is not blocking.
110-
while True:
111-
msg = q.get(timeout=timeout)
112-
self.handle(*msg)
113-
q.task_done()
114-
changed = True
115-
except queue.Empty:
116-
pass
117-
return changed
118-
119-
def run(self):
120-
self.should_exit.clear()
121-
self.server.start_slave(Slave, Channel(self.masterq, self.should_exit))
122-
while not self.should_exit.is_set():
123-
124-
# Don't choose a very small timeout in Python 2:
125-
# https://github.com/mitmproxy/mitmproxy/issues/443
126-
# TODO: Lower the timeout value if we move to Python 3.
127-
self.tick(self.masterq, 0.1)
128-
self.shutdown()
129-
130-
def handle(self, mtype, obj):
131-
c = "handle_" + mtype
132-
m = getattr(self, c, None)
133-
if m:
134-
m(obj)
135-
else:
136-
obj.reply()
152+
def __init__(self, obj):
153+
self.obj = obj
154+
self.q = queue.Queue()
155+
self.acked = False
137156

138-
def shutdown(self):
139-
if not self.should_exit.is_set():
140-
self.should_exit.set()
141-
if self.server:
142-
self.server.shutdown()
157+
def __call__(self, msg=NO_REPLY):
158+
if not self.acked:
159+
self.acked = True
160+
if msg is NO_REPLY:
161+
self.q.put(self.obj)
162+
else:
163+
self.q.put(msg)

‎mitmproxy/dump.py‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -343,15 +343,8 @@ def handle_error(self, f):
343343
self._process_flow(f)
344344
return f
345345

346-
def shutdown(self): # pragma: no cover
347-
return flow.FlowMaster.shutdown(self)
348-
349346
def run(self): # pragma: no cover
350347
if self.o.rfile and not self.o.keepserving:
351348
self.shutdown()
352349
return
353-
try:
354-
return super(DumpMaster, self).run()
355-
except BaseException:
356-
self.shutdown()
357-
raise
350+
super(DumpMaster, self).run()

‎mitmproxy/exceptions.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,13 @@ def __init__(self, message=None):
1717
super(ProxyException, self).__init__(message)
1818

1919

20+
class Kill(ProxyException):
21+
"""
22+
Signal that both client and server connection(s) should be killed immediately.
23+
"""
24+
pass
25+
26+
2027
class ProtocolException(ProxyException):
2128
pass
2229

0 commit comments

Comments
 (0)