[tor-commits] [stem/master] Accepting ORPort and DirPort endpoints

[atagar at torproject.org]

commit 60f99f2b2056b8784d80599350c2be9158d2abac
Author: Damian Johnson <atagar at torproject.org>
Date:   Fri Apr 20 12:15:20 2018 -0700

    Accepting ORPort and DirPort endpoints
    
    We do *not* yet support actually using ORPort endpoints, but our constructor
    now takes them. ORPorts are treated like DirPorts (clearly wrong), but this
    will be followed up by actual support.
---
 stem/descriptor/remote.py      | 24 ++++++++++++++++++++----
 test/unit/descriptor/remote.py | 13 +++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/stem/descriptor/remote.py b/stem/descriptor/remote.py
index 89019c2d..5b942787 100644
--- a/stem/descriptor/remote.py
+++ b/stem/descriptor/remote.py
@@ -333,13 +333,18 @@ class Query(object):
      this was called httplib.HTTPMessage, whereas in python3 the class was
      renamed to http.client.HTTPMessage.
 
+  .. versionchanged:: 1.7.0
+     Endpoints are now expected to be :class:`~stem.DirPort` or
+     :class:`~stem.ORPort` instances. Usage of tuples for this
+     argument is deprecated and will be removed in the future.
+
   :var str resource: resource being fetched, such as '/tor/server/all'
   :var str descriptor_type: type of descriptors being fetched (for options see
     :func:`~stem.descriptor.__init__.parse_file`), this is guessed from the
     resource if **None**
 
-  :var list endpoints: (address, dirport) tuples of the authority or mirror
-    we're querying, this uses authorities if undefined
+  :var list endpoints: :class:`~stem.DirPort` or :class:`~stem.ORPort` of the
+    authority or mirror we're querying, this uses authorities if undefined
   :var list compression: list of :data:`stem.descriptor.remote.Compression`
     we're willing to accept, when none are mutually supported downloads fall
     back to Compression.PLAINTEXT
@@ -401,9 +406,19 @@ class Query(object):
     else:
       self.descriptor_type = _guess_descriptor_type(resource)
 
+    self.endpoints = []
+
+    if endpoints:
+      for endpoint in endpoints:
+        if isinstance(endpoint, tuple) and len(endpoint) == 2:
+          self.endpoints.append(stem.DirPort(endpoint[0], endpoint[1]))
+        elif isinstance(endpoint, (stem.ORPort, stem.DirPort)):
+          self.endpoints.append(endpoint)
+        else:
+          raise ValueError("Endpoints must be an stem.ORPort, stem.DirPort, or two value tuple. '%s' is a %s." % (endpoint, type(endpoint).__name__))
+
     self.resource = resource
     self.compression = compression
-    self.endpoints = endpoints if endpoints else []
     self.retries = retries
     self.fall_back_to_authority = fall_back_to_authority
 
@@ -526,7 +541,8 @@ class Query(object):
       picked = random.choice(list(directories))
       address, dirport = picked.address, picked.dir_port
     else:
-      address, dirport = random.choice(self.endpoints)
+      picked = random.choice(self.endpoints)
+      address, dirport = picked.address, picked.port
 
     return 'http://%s:%i/%s' % (address, dirport, self.resource.lstrip('/'))
 
diff --git a/test/unit/descriptor/remote.py b/test/unit/descriptor/remote.py
index a033d1c1..deb10060 100644
--- a/test/unit/descriptor/remote.py
+++ b/test/unit/descriptor/remote.py
@@ -3,6 +3,7 @@ Unit tests for stem.descriptor.remote.
 """
 
 import io
+import re
 import socket
 import tempfile
 import time
@@ -343,6 +344,18 @@ class TestDescriptorDownloader(unittest.TestCase):
     self.assertRaises(socket.timeout, query.run)
     self.assertEqual(2, urlopen_mock.call_count)
 
+  def test_query_with_invalid_endpoints(self):
+    invalid_endpoints = {
+      'hello': "'h' is a str.",
+      ('hello',): "'hello' is a str.",
+      (15,): "'15' is a int.",
+      (('12.34.56.78', 15, 'third arg'),): "'('12.34.56.78', 15, 'third arg')' is a tuple.",
+    }
+
+    for endpoints, error_suffix in invalid_endpoints.items():
+      expected_error = re.escape('Endpoints must be an stem.ORPort, stem.DirPort, or two value tuple. ' + error_suffix)
+      self.assertRaisesRegexp(ValueError, expected_error, stem.descriptor.remote.Query, TEST_RESOURCE, 'server-descriptor 1.0', endpoints = endpoints)
+
   @patch(URL_OPEN, _urlopen_mock(TEST_DESCRIPTOR))
   def test_can_iterate_multiple_times(self):
     query = stem.descriptor.remote.Query(