Merge pull request #67 from TheTorProject/tor-conn-fixes

Refactor Tor connection and configuration code

Author: aagbsn

bwscanner/attacher.py

@@ -1,19 +1,9 @@
-import sys
-import itertools
-
-from twisted.internet import defer, reactor
+import txtorcon
+from twisted.internet import defer, reactor, endpoints
 from txtorcon.interface import CircuitListenerMixin, IStreamAttacher, StreamListenerMixin
-from txtorcon import TorState, launch_tor
-from txtorcon.util import available_tcp_port
 from zope.interface import implementer
 
-
-FETCH_ALL_DESCRIPTOR_OPTIONS = {
-    'UseMicroDescriptors': 0,
-    'FetchUselessDescriptors': 1,
-    'FetchDirInfoEarly': 1,
-    'FetchDirInfoExtraEarly': 1,
-}
+from bwscanner.logger import log
 
 
 @implementer(IStreamAttacher)
@@ -118,47 +108,80 @@ def stream_closed(self, *args, **kw):
         self.circ.close(ifUnused=True)
 
 
-def start_tor(config):
+def options_need_new_consensus(tor_config, new_options):
     """
-    Launches tor with random TCP ports chosen for SocksPort and ControlPort,
-    and other options specified by a txtorcon.torconfig.TorConfig instance.
-
-    Returns a deferred that calls back with a txtorcon.torstate.TorState
-    instance.
+    Check if we need to wait for a new consensus after updating
+    the Tor config with the new options.
     """
-    def get_random_tor_ports():
-        d2 = available_tcp_port(reactor)
-        d2.addCallback(lambda port: config.__setattr__('SocksPort', port))
-        d2.addCallback(lambda _: available_tcp_port(reactor))
-        d2.addCallback(lambda port: config.__setattr__('ControlPort', port))
-        return d2
+    if "UseMicroDescriptors" in new_options:
+        if tor_config.UseMicroDescriptors != new_options["UseMicroDescriptors"]:
+            log.debug("Changing UseMicroDescriptors from {current} to {new}.",
+                      current=tor_config.UseMicroDescriptors,
+                      new=new_options["UseMicroDescriptors"])
+            return True
+    return False
+
 
-    def launch_and_get_state(ignore):
-        d2 = launch_tor(config, reactor, stdout=sys.stdout)
-        d2.addCallback(lambda tpp: TorState(tpp.tor_protocol).post_bootstrap)
-        return d2
-    return get_random_tor_ports().addCallback(launch_and_get_state)
+def wait_for_newconsensus(tor_state):
+    got_consensus = defer.Deferred()
 
+    def got_newconsensus(event):
+        log.debug("Got NEWCONSENSUS event: {event}", event=event)
+        got_consensus.callback(event)
+        tor_state.protocol.remove_event_listener('NEWCONSENSUS', got_newconsensus)
 
-def update_tor_config(tor, config):
+    tor_state.protocol.add_event_listener('NEWCONSENSUS', got_newconsensus)
+    return got_consensus
+
+
+@defer.inlineCallbacks
+def connect_to_tor(launch_tor, circuit_build_timeout, control_port=None,
+                   tor_overrides=None):
     """
-    Update the Tor config from a dict of config key: value pairs.
+    Launch or connect to a Tor instance
+
+    Configure Tor with the passed options and return a Deferred
     """
-    config_pairs = [(key, value) for key, value in config.items()]
-    d = tor.protocol.set_conf(*itertools.chain.from_iterable(config_pairs))
-    return d.addCallback(lambda result: tor)
-
-
-def setconf_singleport_exit(tor):
-    port = available_tcp_port(reactor)
-
-    def add_single_port_exit(port):
-        tor.protocol.set_conf('PublishServerDescriptor', '0',
-                              'PortForwarding', '1',
-                              'AssumeReachable', '1',
-                              'ClientRejectInternalAddresses', '0',
-                              'OrPort', 'auto',
-                              'ExitPolicyRejectPrivate', '0',
-                              'ExitPolicy', 'accept 127.0.0.1:{}, reject *:*'.format(port))
-    return port.addCallback(add_single_port_exit).addCallback(
-        lambda ign: tor.routers[tor.protocol.get_info("fingerprint")])
+    # Options for spawned or running Tor to load the correct descriptors.
+    tor_options = {
+        'LearnCircuitBuildTimeout': 0,  # Disable adaptive circuit timeouts.
+        'CircuitBuildTimeout': circuit_build_timeout,
+        'UseEntryGuards': 0,  # Disable UseEntryGuards to avoid PathBias warnings.
+        'UseMicroDescriptors': 0,
+        'FetchUselessDescriptors': 1,
+        'FetchDirInfoEarly': 1,
+        'FetchDirInfoExtraEarly': 1,
+    }
+
+    if tor_overrides:
+        tor_options.update(tor_overrides)
+
+    if launch_tor:
+        log.info("Spawning a new Tor instance.")
+        # TODO: Pass in data_dir directory so consensus can be cached
+        tor = yield txtorcon.launch(reactor)
+    else:
+        log.info("Trying to connect to a running Tor instance.")
+        if control_port:
+            endpoint = endpoints.TCP4ClientEndpoint(reactor, "localhost", control_port)
+        else:
+            endpoint = None
+        tor = yield txtorcon.connect(reactor, endpoint)
+
+    # Get Tor state first to avoid a race conditions where CONF_CHANGED
+    # messages are received while Txtorcon is reading the consensus.
+    tor_state = yield tor.create_state()
+
+    # Get current TorConfig object
+    tor_config = yield tor.get_config()
+    wait_for_consensus = options_need_new_consensus(tor_config, tor_options)
+
+    # Update Tor config options from dictionary
+    for key, value in tor_options.items():
+        setattr(tor_config, key, value)
+    yield tor_config.save()  # Send updated options to Tor
+
+    if wait_for_consensus:
+        yield wait_for_newconsensus(tor_state)
+
+    defer.returnValue(tor_state)

bwscanner/circuit.py

@@ -31,7 +31,7 @@ def random_path_to_exit(exit_relay, relays):
 class CircuitGenerator(object):
     def __init__(self, state):
         self.state = state
-        self.relays = list(set(state.routers.values()))
+        self.relays = list(set(r for r in state.routers.values() if r))
         self.exits = [relay for relay in self.relays if is_valid_exit(relay)]
 
     def __iter__(self):

bwscanner/fetcher.py

@@ -12,7 +12,7 @@
 from bwscanner.logger import log
 
 
-def get_orport_endpoint(tor_state):
+def get_tor_socks_endpoint(tor_state):
     proxy_endpoint = tor_state.protocol.get_conf("SocksPort")
 
     def extract_port_value(result):
@@ -54,23 +54,23 @@ def __init__(self, host, port, state, path):
         self.path = path
         self.state = state
 
-        self.or_endpoint = get_orport_endpoint(state)
+        self.tor_socks_endpoint = get_tor_socks_endpoint(state)
 
     def connect(self, protocol_factory):
         """
         Implements L{IStreamClientEndpoint.connect} to connect via TCP, after
         SOCKS5 negotiation and Tor circuit construction is done.
         """
         proxy_factory = SOCKS5ClientFactory(self.host, self.port, protocol_factory)
-        self.or_endpoint.addCallback(lambda end: end.connect(proxy_factory))
+        self.tor_socks_endpoint.addCallback(lambda end: end.connect(proxy_factory))
 
         def _create_circ(proto):
             hp = proto.transport.getHost()
-            d = self.state.attacher.create_circuit(hp.host, hp.port, self.path)
+            d = self.state._attacher.create_circuit(hp.host, hp.port, self.path)
             d.addErrback(proxy_factory.deferred.errback)
             return proxy_factory.deferred
 
-        return self.or_endpoint.addCallback(_create_circ)
+        return self.tor_socks_endpoint.addCallback(_create_circ)
 
 
 class OnionRoutedAgent(Agent):

bwscanner/scanner.py

@@ -4,54 +4,16 @@
 
 import click
 from twisted.internet import reactor
-from txtorcon import build_local_tor_connection, TorConfig
 
+from bwscanner.attacher import connect_to_tor
 from bwscanner.logger import setup_logging, log
-from bwscanner.attacher import start_tor, update_tor_config, FETCH_ALL_DESCRIPTOR_OPTIONS
 from bwscanner.measurement import BwScan
 from bwscanner.aggregate import write_aggregate_data
 
 
 BWSCAN_VERSION = '0.0.1'
 
 
-def connect_to_tor(launch_tor, circuit_build_timeout, circuit_idle_timeout):
-    """
-    Launch or connect to a Tor instance
-
-    Configure Tor with the passed options and return a Deferred
-    """
-    # Options for spawned or running Tor to load the correct descriptors.
-    tor_options = {
-        'LearnCircuitBuildTimeout': 0,  # Disable adaptive circuit timeouts.
-        'CircuitBuildTimeout': circuit_build_timeout,
-        'CircuitIdleTimeout': circuit_idle_timeout,
-        'UseEntryGuards': 0,  # Disable UseEntryGuards to avoid PathBias warnings.
-    }
-
-    def tor_status(tor):
-        log.info("Connected successfully to Tor.")
-        return tor
-
-    if launch_tor:
-        log.info("Spawning a new Tor instance.")
-        c = TorConfig()
-        # Update Tor config before launching a new Tor.
-        c.config.update(tor_options)
-        c.config.update(FETCH_ALL_DESCRIPTOR_OPTIONS)
-        tor = start_tor(c)
-
-    else:
-        log.info("Trying to connect to a running Tor instance.")
-        tor = build_local_tor_connection(reactor)
-        # Update the Tor config on a running Tor.
-        tor.addCallback(update_tor_config, tor_options)
-        tor.addCallback(update_tor_config, FETCH_ALL_DESCRIPTOR_OPTIONS)
-
-    tor.addCallback(tor_status)
-    return tor
-
-
 class ScanInstance(object):
     """
     Store the configuration and state for the CLI tool.
@@ -80,11 +42,9 @@ def __repr__(self):
               help='Launch Tor or try to connect to an existing Tor instance.')
 @click.option('--circuit-build-timeout', default=20,
               help='Option passed when launching Tor.')
-@click.option('--circuit-idle-timeout', default=20,
-              help='Option passed when launching Tor.')
 @click.version_option(BWSCAN_VERSION)
 @click.pass_context
-def cli(ctx, data_dir, loglevel, logfile, launch_tor, circuit_build_timeout, circuit_idle_timeout):
+def cli(ctx, data_dir, loglevel, logfile, launch_tor, circuit_build_timeout):
     """
     The bwscan tool measures Tor relays and calculates their bandwidth. These
     bandwidth measurements can then be aggregate to create the bandwidth
@@ -98,7 +58,7 @@ def cli(ctx, data_dir, loglevel, logfile, launch_tor, circuit_build_timeout, cir
         os.makedirs(ctx.obj.measurement_dir)
 
     # Create a connection to a Tor instance
-    ctx.obj.tor = connect_to_tor(launch_tor, circuit_build_timeout, circuit_idle_timeout)
+    ctx.obj.tor_state = connect_to_tor(launch_tor, circuit_build_timeout)
 
     # Set up the logger to only output log lines of level `loglevel` and above.
     setup_logging(log_level=loglevel, log_name=logfile)
@@ -131,12 +91,14 @@ def rename_finished_scan(deferred):
         click.echo(deferred)
         os.rename(scan_data_dir, os.path.join(scan.measurement_dir, scan_time))
 
-    scan.tor.addCallback(BwScan, reactor, scan_data_dir, request_timeout=timeout,
-                         request_limit=request_limit, partitions=partitions,
-                         this_partition=current_partition)
-    scan.tor.addCallback(lambda scanner: scanner.run_scan())
-    scan.tor.addCallback(lambda _: reactor.stop())
-    scan.tor.addCallback(rename_finished_scan)
+    scan.tor_state.addCallback(BwScan, reactor, scan_data_dir,
+                               request_timeout=timeout,
+                               request_limit=request_limit,
+                               partitions=partitions,
+                               this_partition=current_partition)
+    scan.tor_state.addCallback(lambda scanner: scanner.run_scan())
+    scan.tor_state.addCallback(lambda _: reactor.stop())
+    scan.tor_state.addCallback(rename_finished_scan)
 
     reactor.run()
 
@@ -192,7 +154,7 @@ def aggregate(scan, scan_name, previous):
         scan_data_dirs = [os.path.join(scan.measurement_dir, name) for name in recent_scan_names]
         log.info("Aggregating data from past {count} scans.", count=len(scan_data_dirs))
 
-    scan.tor.addCallback(lambda tor: write_aggregate_data(tor, scan_data_dirs))
-    scan.tor.addErrback(lambda failure: log.failure("Unexpected error"))
-    scan.tor.addCallback(lambda _: reactor.stop())
+    scan.tor_state.addCallback(lambda tor_state: write_aggregate_data(tor_state, scan_data_dirs))
+    scan.tor_state.addErrback(lambda failure: log.failure("Unexpected error"))
+    scan.tor_state.addCallback(lambda _: reactor.stop())
     reactor.run()

requirements.txt

@@ -4,4 +4,4 @@ service-identity==16.0.0
 stem>=1.4.0
 Twisted==16.2.0
 txsocksx==1.15.0.2
-txtorcon==0.15.0
+txtorcon==0.19.3

test/template.py

@@ -2,28 +2,28 @@
 import random
 
 from twisted.internet import defer, reactor
-from twisted.internet.endpoints import TCP4ClientEndpoint
 from twisted.trial import unittest
-from txtorcon.torstate import build_tor_connection
 
 from bwscanner import circuit
-from bwscanner.attacher import SOCKSClientStreamAttacher
+from bwscanner.attacher import SOCKSClientStreamAttacher, connect_to_tor
 
 
 class TorTestCase(unittest.TestCase):
 
     @defer.inlineCallbacks
     def setUp(self):
-        self.tor = yield build_tor_connection(
-            TCP4ClientEndpoint(reactor, 'localhost', int(
-                os.environ.get('CHUTNEY_CONTROL_PORT'))))
+        self.tor_state = yield connect_to_tor(
+                launch_tor=False,
+                control_port=int(os.environ.get('CHUTNEY_CONTROL_PORT')),
+                circuit_build_timeout=30,
+        )
 
-        self.attacher = SOCKSClientStreamAttacher(self.tor)
-        yield self.tor.set_attacher(self.attacher, reactor)
+        self.attacher = SOCKSClientStreamAttacher(self.tor_state)
+        yield self.tor_state.set_attacher(self.attacher, reactor)
 
     @property
     def routers(self):
-        return list(set(self.tor.routers.values()))
+        return list(set(self.tor_state.routers.values()))
 
     @property
     def exits(self):
@@ -35,7 +35,7 @@ def random_path(self):
 
     @defer.inlineCallbacks
     def tearDown(self):
-        yield self.tor.set_attacher(None, reactor)
-        yield self.tor.protocol.quit()
+        yield self.tor_state.set_attacher(None, reactor)
+        yield self.tor_state.protocol.quit()
         # seems to leave dirty reactor otherwise?
-        yield self.tor.protocol.transport.loseConnection()
+        yield self.tor_state.protocol.transport.loseConnection()

test/test_attacher.py

@@ -54,7 +54,7 @@ class TestSOCKSClientStreamAttacher(TorTestCase):
     def setUp(self):
         yield super(TestSOCKSClientStreamAttacher, self).setUp()
         # do not attach circuits automatically
-        yield self.tor.set_attacher(None, reactor)
+        yield self.tor_state.set_attacher(None, reactor)
 
     @defer.inlineCallbacks
     def test_create_circuit(self):

test/test_circuit.py

@@ -8,7 +8,7 @@ def test_exit_scan(self):
         all_exits = set(self.exits)
         num_circuits = 0
         seen = set()
-        for circuit in ExitScan(self.tor):
+        for circuit in ExitScan(self.tor_state):
             assert len(circuit) == 3
             assert 'exit' in circuit[-1].flags
             seen.add(circuit[-1])
@@ -20,7 +20,7 @@ def test_two_hop(self):
         all_r = set(self.routers)
         num_circuits = 0
         seen = set()
-        for circuit in TwoHop(self.tor):
+        for circuit in TwoHop(self.tor_state):
             assert len(circuit) == 2
             assert 'exit' in circuit[-1].flags
             num_circuits = num_circuits + 1