Kill deadlock that sometimes occurred on shutdown.

4 files changed, 18 insertions, 15 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index d0213249..fbb8a925 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -714,6 +714,9 @@ class ConsoleMaster(controller.Master):
 
         self.ui.run_wrapper(self.loop)
         # If True, quit just pops out to connection list view.
+        print >> sys.stderr, "Shutting down..."
+        sys.stderr.flush()
+        self.shutdown()
         self.nested = False
 
     def make_view(self):
@@ -1015,7 +1018,6 @@ class ConsoleMaster(controller.Master):
                         self.view.keypress(size, k)
         except (Stop, KeyboardInterrupt):
             pass
-        self.shutdown()
 
     def shutdown(self):
         for i in self.state.flow_list:
diff --git a/libmproxy/controller.py b/libmproxy/controller.py
index 29e2f7a8..ddfbe7f1 100644
--- a/libmproxy/controller.py
+++ b/libmproxy/controller.py
@@ -29,8 +29,11 @@ class Msg:
 
     def send(self, masterq):
         self.acked = False
-        masterq.put(self)
-        return self.q.get()
+        try:
+            masterq.put(self, timeout=3)
+            return self.q.get(timeout=3)
+        except (Queue.Empty, Queue.Full):
+            return None
 
 
 class Slave(threading.Thread):
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 8eb7342a..296b6738 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -78,14 +78,13 @@ class Flow:
         return isinstance(self.connection, ReplayConnection)
 
     def kill(self):
-        if self.intercepting:
-            if not self.request.acked:
-                self.request.kill = True
-                self.request.ack()
-            elif self.response and not self.response.acked:
-                self.response.kill = True
-                self.response.ack()
-            self.intercepting = False
+        if self.request and not self.request.acked:
+            self.request.kill = True
+            self.request.ack()
+        elif self.response and not self.response.acked:
+            self.response.kill = True
+            self.response.ack()
+        self.intercepting = False
 
     def intercept(self):
         self.intercepting = True
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index f3b9fe53..5e8da3e4 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -85,7 +85,6 @@ class Request(controller.Msg):
         self.connection = connection
         self.host, self.port, self.scheme = host, port, scheme
         self.method, self.path, self.headers, self.content = method, path, headers, content
-        self.kill = False
         controller.Msg.__init__(self)
 
     def get_state(self):
@@ -161,7 +160,6 @@ class Response(controller.Msg):
         self.request = request
         self.code, self.proto, self.msg = code, proto, msg
         self.headers, self.content = headers, content
-        self.kill = False
         controller.Msg.__init__(self)
 
     def get_state(self):
@@ -341,13 +339,13 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
         try:
             request = self.read_request(bc)
             request = request.send(self.mqueue)
-            if request.kill:
+            if request is None:
                 self.finish()
                 return
             server = ServerConnection(request)
             response = server.read_response()
             response = response.send(self.mqueue)
-            if response.kill:
+            if response is None:
                 server.terminate()
                 self.finish()
                 return
@@ -405,6 +403,7 @@ class ProxyHandler(SocketServer.StreamRequestHandler):
         self.wfile.flush()
 
     def terminate(self, connection, wfile, rfile):
+        self.request.close()
         try:
             if not getattr(wfile, "closed", False):
                 wfile.flush()