[tor-commits] [stem/master] Tesing and fix for disconnected socket reads

[atagar at torproject.org]

commit c74ac0f7f0c75d429a7d32f7d516b3518fc90ff0
Author: Damian Johnson <atagar at torproject.org>
Date:   Wed Oct 12 09:37:49 2011 -0700

    Tesing and fix for disconnected socket reads
    
    Adding a unit test and fix for when the input file is derived from a socket
    that's never been connected. I'm suspicious that this won't catch disconnects
    that accur a little while into the socket's use, but I'll need to implement
    integration tests for that. Guess that's next...
---
 stem/types.py        |    9 +++++++--
 test/unit/message.py |   11 +++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/stem/types.py b/stem/types.py
index 730accb..2339007 100644
--- a/stem/types.py
+++ b/stem/types.py
@@ -19,6 +19,7 @@ Version - Tor versioning information.
 """
 
 import re
+import socket
 
 from stem.util import log
 
@@ -51,7 +52,9 @@ def read_message(control_file):
   parsed_content, raw_content = [], ""
   
   while True:
-    line = control_file.readline()
+    try: line = control_file.readline()
+    except socket.error, exc: raise ControlSocketClosed(exc)
+    
     raw_content += line
     
     # Parses the tor control lines. These are of the form...
@@ -85,7 +88,9 @@ def read_message(control_file):
       # get a line with just a period
       
       while True:
-        line = control_file.readline()
+        try: line = control_file.readline()
+        except socket.error, exc: raise ControlSocketClosed(exc)
+        
         raw_content += line
         
         if not line.endswith("\r\n"):
diff --git a/test/unit/message.py b/test/unit/message.py
index e1fb622..851b394 100644
--- a/test/unit/message.py
+++ b/test/unit/message.py
@@ -2,6 +2,7 @@
 Unit tests for the types.ControlMessage parsing and class.
 """
 
+import socket
 import StringIO
 import unittest
 import stem.types
@@ -147,6 +148,16 @@ class TestMessageFunctions(unittest.TestCase):
         self.assert_message_parses(removal_test_input)
         self.assert_message_parses(replacement_test_input)
   
+  def test_disconnected_socket(self):
+    """
+    Tests when the read function is given a file derived from a disconnected
+    socket.
+    """
+    
+    control_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    control_socket_file = control_socket.makefile()
+    self.assertRaises(stem.types.ControlSocketClosed, stem.types.read_message, control_socket_file)
+  
   def assert_message_parses(self, controller_reply):
     """
     Performs some basic sanity checks that a reply mirrors its parsed result.