[tor-commits] [stem/master] Remove bytes normalization from recv_message

[atagar at torproject.org]

commit cf3c55ef6aec640711e26be4a4f985850612c29c
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Sep 18 10:21:41 2017 -0700

    Remove bytes normalization from recv_message
    
    Sockets always provide bytes, but we still normalized the type because tests
    and a few methods like ControlMessage.from_str() provided it a StringIO
    instead. This is silly. We can provide BytesIOs instead. Our recv_message is in
    the critical path so it's important to keep it as minimal as possible.
    
    This might break our tests for python 2.6. Not sure - don't have that python
    verison and getting it on modern systems is a pita. I'll check later. 2.6 has
    the io module, but unclear if it has BytesIO. If so we might need to tweak our
    imports later.
---
 stem/control.py                       |  4 ++--
 stem/response/__init__.py             |  2 +-
 stem/socket.py                        |  6 +-----
 test/unit/response/control_message.py | 15 ++++++---------
 4 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/stem/control.py b/stem/control.py
index 2f1268ef..fd5f9f87 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -3298,7 +3298,7 @@ class Controller(BaseController):
     response = self.get_info('circuit-status')
 
     for circ in response.splitlines():
-      circ_message = stem.socket.recv_message(io.StringIO(stem.util.str_tools._to_unicode('650 CIRC ' + circ + '\r\n')))
+      circ_message = stem.socket.recv_message(io.BytesIO(stem.util.str_tools._to_unicode('650 CIRC %s\r\n' % circ)))
       stem.response.convert('EVENT', circ_message, arrived_at = 0)
       circuits.append(circ_message)
 
@@ -3481,7 +3481,7 @@ class Controller(BaseController):
     response = self.get_info('stream-status')
 
     for stream in response.splitlines():
-      message = stem.socket.recv_message(io.StringIO(stem.util.str_tools._to_unicode('650 STREAM ' + stream + '\r\n')))
+      message = stem.socket.recv_message(io.BytesIO(stem.util.str_tools._to_bytes('650 STREAM %s\r\n' % stream)))
       stem.response.convert('EVENT', message, arrived_at = 0)
       streams.append(message)
 
diff --git a/stem/response/__init__.py b/stem/response/__init__.py
index eadb25c7..e0de0f5c 100644
--- a/stem/response/__init__.py
+++ b/stem/response/__init__.py
@@ -156,7 +156,7 @@ class ControlMessage(object):
 
       content = re.sub('([\r]?)\n', '\r\n', content)
 
-    msg = stem.socket.recv_message(io.StringIO(stem.util.str_tools._to_unicode(content)))
+    msg = stem.socket.recv_message(io.BytesIO(stem.util.str_tools._to_bytes(content)))
 
     if msg_type is not None:
       convert(msg_type, msg, **kwargs)
diff --git a/stem/socket.py b/stem/socket.py
index ca56ea69..9817c129 100644
--- a/stem/socket.py
+++ b/stem/socket.py
@@ -530,11 +530,7 @@ def recv_message(control_file):
 
   while True:
     try:
-      # From a real socket readline() would always provide bytes, but during
-      # tests we might be given a StringIO in which case it's unicode under
-      # python 3.x.
-
-      line = stem.util.str_tools._to_bytes(control_file.readline())
+      line = control_file.readline()
     except AttributeError:
       # if the control_file has been closed then we will receive:
       # AttributeError: 'NoneType' object has no attribute 'recv'
diff --git a/test/unit/response/control_message.py b/test/unit/response/control_message.py
index c0fb2f32..8b05d6a8 100644
--- a/test/unit/response/control_message.py
+++ b/test/unit/response/control_message.py
@@ -2,17 +2,14 @@
 Unit tests for the stem.response.ControlMessage parsing and class.
 """
 
+import io
 import socket
 import unittest
 
-try:
-  from StringIO import StringIO
-except:
-  from io import StringIO
-
 import stem.socket
 import stem.response
 import stem.response.getinfo
+import stem.util.str_tools
 
 OK_REPLY = '250 OK\r\n'
 
@@ -128,7 +125,7 @@ class TestControlMessage(unittest.TestCase):
     for index, line in enumerate(infonames_lines):
       # replace the CRLF for the line
       infonames_lines[index] = line.rstrip('\r\n') + '\n'
-      test_socket_file = StringIO(''.join(infonames_lines))
+      test_socket_file = io.BytesIO(stem.util.str_tools._to_bytes(''.join(infonames_lines)))
       self.assertRaises(stem.ProtocolError, stem.socket.recv_message, test_socket_file)
 
       # puts the CRLF back
@@ -154,8 +151,8 @@ class TestControlMessage(unittest.TestCase):
         # - this is part of the message prefix
         # - this is disrupting the line ending
 
-        self.assertRaises(stem.ProtocolError, stem.socket.recv_message, StringIO(removal_test_input))
-        self.assertRaises(stem.ProtocolError, stem.socket.recv_message, StringIO(replacement_test_input))
+        self.assertRaises(stem.ProtocolError, stem.socket.recv_message, io.BytesIO(stem.util.str_tools._to_bytes(removal_test_input)))
+        self.assertRaises(stem.ProtocolError, stem.socket.recv_message, io.BytesIO(stem.util.str_tools._to_bytes(replacement_test_input)))
       else:
         # otherwise the data will be malformed, but this goes undetected
         self._assert_message_parses(removal_test_input)
@@ -179,7 +176,7 @@ class TestControlMessage(unittest.TestCase):
       stem.response.ControlMessage for the given input
     """
 
-    message = stem.socket.recv_message(StringIO(controller_reply))
+    message = stem.socket.recv_message(io.BytesIO(stem.util.str_tools._to_bytes(controller_reply)))
 
     # checks that the raw_content equals the input value
     self.assertEqual(controller_reply, message.raw_content())