[tor-commits] [stem/master] Replacing protocolinfo lookups with ControlSocket

[atagar at torproject.org]

commit e313bb35c827fdb75a382d7c5e837692baeead73
Author: Damian Johnson <atagar at torproject.org>
Date:   Fri Nov 25 22:36:57 2011 -0800

    Replacing protocolinfo lookups with ControlSocket
    
    Replacing raw socket use in the protocolinfo lookup functions with the
    ControlSocket class, and attaching it to the responses instead.
---
 stem/connection.py |   41 ++++++++++++++---------------------------
 1 files changed, 14 insertions(+), 27 deletions(-)

diff --git a/stem/connection.py b/stem/connection.py
index ba851a2..453b1d7 100644
--- a/stem/connection.py
+++ b/stem/connection.py
@@ -101,37 +101,24 @@ def _get_protocolinfo_impl(control_socket, connection_args, keep_alive):
   connects the given socket and issues a PROTOCOLINFO query with it.
   """
   
-  control_socket_file = control_socket.makefile()
-  protocolinfo_response, raised_exc = None, None
-  
   try:
-    # initiates connection
     control_socket.connect(connection_args)
-    
-    # issues the PROTOCOLINFO query
-    stem.socket.send_message(control_socket_file, "PROTOCOLINFO 1")
-    
-    protocolinfo_response = stem.socket.recv_message(control_socket_file)
-    ProtocolInfoResponse.convert(protocolinfo_response)
+    control_socket = stem.socket.ControlSocket(control_socket)
   except socket.error, exc:
-    raised_exc = stem.socket.SocketError(exc)
-  except (stem.socket.ProtocolError, stem.socket.SocketError), exc:
-    raised_exc = exc
-  
-  control_socket_file.close() # done with the linked file
+    raise stem.socket.SocketError(exc)
   
-  if not keep_alive or raised_exc:
-    # shut down the socket we were using
-    try: control_socket.shutdown(socket.SHUT_RDWR)
-    except socket.error: pass
+  try:
+    control_socket.send("PROTOCOLINFO 1")
+    protocolinfo_response = control_socket.recv()
+    ProtocolInfoResponse.convert(protocolinfo_response)
     
+    if keep_alive: protocolinfo_response.socket = control_socket
+    else: control_socket.close()
+    
+    return protocolinfo_response
+  except stem.socket.ControllerError, exc:
     control_socket.close()
-  else:
-    # if we're keeping the socket open then attach it to the response
-    protocolinfo_response.socket = control_socket
-  
-  if raised_exc: raise raised_exc
-  else: return protocolinfo_response
+    raise exc
 
 def _expand_cookie_path(cookie_path, pid_resolver, pid_resolution_arg):
   """
@@ -173,7 +160,7 @@ class ProtocolInfoResponse(stem.socket.ControlMessage):
   correct this.
   
   The protocol_version is the only mandatory data for a valid PROTOCOLINFO
-  response, so all other values are None if undefined or empty if a collecion.
+  response, so all other values are None if undefined or empty if a collection.
   
   Attributes:
     protocol_version (int)             - protocol version of the response
@@ -181,7 +168,7 @@ class ProtocolInfoResponse(stem.socket.ControlMessage):
     auth_methods (tuple)               - AuthMethod types that tor will accept
     unknown_auth_methods (tuple)       - strings of unrecognized auth methods
     cookie_path (str)                  - path of tor's authentication cookie
-    socket (socket.socket)             - socket used to make the query
+    socket (stem.socket.ControlSocket) - socket used to make the query
   """
   
   def convert(control_message):