[tor-commits] [stem/master] Transparently allowing multiple PROTOCOLINFO calls

[atagar at torproject.org]

commit b4ca09871c7650ad6469c3b998ebc174605f4a7d
Author: Damian Johnson <atagar at torproject.org>
Date:   Tue Dec 20 09:41:29 2011 -0800

    Transparently allowing multiple PROTOCOLINFO calls
    
    Tor hangs up on the socket for multiple reasons prior to authentication. One
    of them is if multiple PROTOCOLINFO queries are made. This is annoying and does
    not make sense to controller callers, so transparently re-establishing those
    socket connections.
---
 stem/connection.py                    |   14 ++++++++++----
 test/integ/connection/protocolinfo.py |   21 +++++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/stem/connection.py b/stem/connection.py
index ddfed89..39be31d 100644
--- a/stem/connection.py
+++ b/stem/connection.py
@@ -521,9 +521,6 @@ def get_protocolinfo(control_socket):
   Issues a PROTOCOLINFO query to a control socket, getting information about
   the tor process running on it.
   
-  Tor hangs up on sockets after receiving a PROTOCOLINFO query if it isn't next
-  followed by authentication.
-  
   Arguments:
     control_socket (stem.socket.ControlSocket) - connected tor control socket
   
@@ -538,6 +535,15 @@ def get_protocolinfo(control_socket):
   
   control_socket.send("PROTOCOLINFO 1")
   protocolinfo_response = control_socket.recv()
+  
+  # Tor hangs up on sockets after receiving a PROTOCOLINFO query if it isn't
+  # next followed by authentication. Transparently reconnect if that happens.
+  
+  if str(protocolinfo_response) == "Authentication required.":
+    control_socket.connect()
+    control_socket.send("PROTOCOLINFO 1")
+    protocolinfo_response = control_socket.recv()
+  
   ProtocolInfoResponse.convert(protocolinfo_response)
   
   # attempt ot expand relative cookie paths via the control port or socket file
@@ -640,7 +646,7 @@ class ProtocolInfoResponse(stem.socket.ControlMessage):
     
     # sanity check that we're a PROTOCOLINFO response
     if not list(self)[0].startswith("PROTOCOLINFO"):
-      msg = "Message is not a PROTOCOLINFO response"
+      msg = "Message is not a PROTOCOLINFO response (%s)" % self
       raise stem.socket.ProtocolError(msg)
     
     for line in self:
diff --git a/test/integ/connection/protocolinfo.py b/test/integ/connection/protocolinfo.py
index dfc1c01..e37f63b 100644
--- a/test/integ/connection/protocolinfo.py
+++ b/test/integ/connection/protocolinfo.py
@@ -118,6 +118,27 @@ class TestProtocolInfo(unittest.TestCase):
       # we don't have a control socket
       self.assertRaises(stem.socket.SocketError, stem.socket.ControlSocketFile, test.runner.CONTROL_SOCKET_PATH)
   
+  def test_multiple_protocolinfo_calls(self):
+    """
+    Tests making repeated PROTOCOLINFO queries. This use case is interesting
+    because tor will shut down the socket and stem should transparently
+    re-establish it.
+    """
+    
+    runner = test.runner.get_runner()
+    connection_type = runner.get_connection_type()
+    
+    if connection_type == test.runner.TorConnection.NONE:
+      self.skipTest("(no connection)")
+    
+    control_socket = runner.get_tor_socket(False)
+    
+    for i in range(5):
+      protocolinfo_response = stem.connection.get_protocolinfo(control_socket)
+      self.assert_protocolinfo_attr(protocolinfo_response, connection_type)
+    
+    control_socket.close()
+  
   def assert_protocolinfo_attr(self, protocolinfo_response, connection_type):
     """
     Makes assertions that the protocolinfo response's attributes match those of