From c02fdb24639c8c0fe89534205d48933de832287a Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Mon, 2 Apr 2012 13:24:51 +1200
Subject: [PATCH] Refactor proxy.Server to fix a crash when replaying with -n

---
 libmproxy/console/__init__.py    |  8 +++-----
 libmproxy/controller.py          |  4 +---
 libmproxy/proxy.py               | 17 +++++++++++++++++
 mitmdump                         |  2 +-
 mitmproxy                        |  2 +-
 test/test_console_contentview.py |  4 ++--
 6 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py
index 6940b4826..c0700b56a 100644
--- a/libmproxy/console/__init__.py
+++ b/libmproxy/console/__init__.py
@@ -150,7 +150,7 @@ class StatusBar(common.WWrap):
             r.append("[")
             r.append(("heading_key", "u"))
             r.append(":%s]"%self.master.stickyauth_txt)
-        if self.master.server and self.master.server.config.reverse_proxy:
+        if self.master.server.config.reverse_proxy:
             r.append("[")
             r.append(("heading_key", "P"))
             r.append(":%s]"%utils.unparse_url(*self.master.server.config.reverse_proxy))
@@ -187,7 +187,7 @@ class StatusBar(common.WWrap):
         t = [
                 ('heading', ("[%s]"%self.master.state.flow_count()).ljust(7)),
             ]
-        if self.master.server:
+        if self.master.server.bound:
             boundaddr = "[%s:%s]"%(self.master.server.address or "*", self.master.server.port)
         else:
             boundaddr = ""
@@ -510,9 +510,7 @@ class ConsoleMaster(flow.FlowMaster):
 
         self.view_flowlist()
 
-        if self.server:
-            slave = controller.Slave(self.masterq, self.server)
-            slave.start()
+        self.server.start_slave(controller.Slave, self.masterq)
 
         if self.options.rfile:
             ret = self.load_flows(self.options.rfile)
diff --git a/libmproxy/controller.py b/libmproxy/controller.py
index dac2c82aa..174a3c6e5 100644
--- a/libmproxy/controller.py
+++ b/libmproxy/controller.py
@@ -83,9 +83,7 @@ class Master:
     def run(self):
         global should_exit
         should_exit = False
-        if self.server:
-            slave = Slave(self.masterq, self.server)
-            slave.start()
+        self.server.start_slave(Slave, self.masterq)
         while not should_exit:
             self.tick(self.masterq)
         self.shutdown()
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 33e50890b..c6a68ba33 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -478,6 +478,7 @@ ServerBase.daemon_threads = True        # Terminate workers when main thread ter
 class ProxyServer(ServerBase):
     request_queue_size = 20
     allow_reuse_address = True
+    bound = True
     def __init__(self, config, port, address=''):
         """
             Raises ProxyServerError if there's a startup problem.
@@ -491,6 +492,10 @@ class ProxyServer(ServerBase):
         self.certdir = tempfile.mkdtemp(prefix="mitmproxy")
         config.certdir = self.certdir
 
+    def start_slave(self, klass, masterq):
+        slave = klass(masterq, self)
+        slave.start()
+
     def set_mqueue(self, q):
         self.masterq = q
 
@@ -505,6 +510,18 @@ class ProxyServer(ServerBase):
             pass
 
 
+class DummyServer:
+    bound = False
+    def __init__(self, config):
+        self.config = config
+
+    def start_slave(self, klass, masterq):
+        pass
+
+    def shutdown(self):
+        pass
+
+
 # Command-line utils
 def certificate_option_group(parser):
     group = optparse.OptionGroup(parser, "SSL")
diff --git a/mitmdump b/mitmdump
index 92113f269..83748c711 100755
--- a/mitmdump
+++ b/mitmdump
@@ -40,7 +40,7 @@ if __name__ == '__main__':
 
     proxyconfig = proxy.process_proxy_options(parser, options)
     if options.no_server:
-        server = None
+        server = proxy.DummyServer(proxyconfig)
     else:
         try:
             server = proxy.ProxyServer(proxyconfig, options.port, options.addr)
diff --git a/mitmproxy b/mitmproxy
index 93a2facc2..e7e0db225 100755
--- a/mitmproxy
+++ b/mitmproxy
@@ -46,7 +46,7 @@ if __name__ == '__main__':
     config = proxy.process_proxy_options(parser, options)
 
     if options.no_server:
-        server = None
+        server = proxy.DummyServer(config)
     else:
         try:
             server = proxy.ProxyServer(config, options.port, options.addr)
diff --git a/test/test_console_contentview.py b/test/test_console_contentview.py
index 28265dec2..babe59ea9 100644
--- a/test/test_console_contentview.py
+++ b/test/test_console_contentview.py
@@ -131,7 +131,7 @@ Larry
                 [["content-type", "application/json"]],
                 "[1, 2"
               )
-        assert r[0] == "Raw"
+        assert "Raw" in r[0]
 
         r = cv.get_content_view(
                 cv.VIEW_AUTO,
@@ -153,7 +153,7 @@ Larry
                 encoding.encode('gzip', "[1, 2, 3]")
               )
         assert "decoded gzip" in r[0]
-        assert "forced" in r[0]
+        assert "XML" in r[0]
 
 
 tests = [