Flatten ServerMaster into Master

This commit is contained in:
Aldo Cortesi 2016-05-29 12:01:46 +12:00
parent ed415877d4
commit e4f043f69c
3 changed files with 11 additions and 38 deletions

View File

@ -12,17 +12,18 @@ class ControlError(Exception):
class Master(object):
"""
The master handles mitmproxy's main event loop.
"""
def __init__(self):
self.event_queue = queue.Queue()
self.should_exit = threading.Event()
self.servers = []
def start(self):
self.should_exit.clear()
for server in self.servers:
ServerThread(server).start()
def run(self):
self.start()
@ -57,38 +58,18 @@ class Master(object):
return changed
def shutdown(self):
for server in self.servers:
server.shutdown()
self.should_exit.set()
class ServerMaster(Master):
"""
The ServerMaster adds server thread support to the master.
"""
def __init__(self):
super(ServerMaster, self).__init__()
self.servers = []
def add_server(self, server):
# We give a Channel to the server which can be used to communicate with the master
channel = Channel(self.event_queue, self.should_exit)
server.set_channel(channel)
self.servers.append(server)
def start(self):
super(ServerMaster, self).start()
for server in self.servers:
ServerThread(server).start()
def shutdown(self):
for server in self.servers:
server.shutdown()
super(ServerMaster, self).shutdown()
class ServerThread(threading.Thread):
def __init__(self, server):
self.server = server
super(ServerThread, self).__init__()
@ -100,12 +81,10 @@ class ServerThread(threading.Thread):
class Channel(object):
"""
The only way for the proxy server to communicate with the master
is to use the channel it has been given.
"""
def __init__(self, q, should_exit):
self.q = q
self.should_exit = should_exit
@ -141,12 +120,10 @@ class Channel(object):
class DummyReply(object):
"""
A reply object that does nothing. Useful when we need an object to seem
like it has a channel, and during testing.
"""
def __init__(self):
self.acked = False
self.taken = False
@ -192,7 +169,6 @@ def handler(f):
class Reply(object):
"""
Messages sent through a channel are decorated with a "reply" attribute.
This object is used to respond to the message through the return

View File

@ -638,7 +638,7 @@ class State(object):
self.flows.kill_all(master)
class FlowMaster(controller.ServerMaster):
class FlowMaster(controller.Master):
@property
def server(self):

View File

@ -16,7 +16,6 @@ class TMsg:
class TestMaster(object):
def test_simple(self):
class DummyMaster(controller.Master):
@controller.handler
def handle_panic(self, _):
@ -34,10 +33,8 @@ class TestMaster(object):
m.run()
assert m.should_exit.is_set()
class TestServerMaster(object):
def test_simple(self):
m = controller.ServerMaster()
def test_server_simple(self):
m = controller.Master()
s = DummyServer(None)
m.add_server(s)
m.start()