From 16f6ee4b54d8a952079e529a3b5f9681f0165112 Mon Sep 17 00:00:00 2001
From: Nick Mathewson <nickm@torproject.org>
Date: Mon, 25 Oct 2021 12:02:44 -0400
Subject: [PATCH] Add an isolate_client() function to create an isolated
 TorClient.

When two TorClients are isolated, their streams shouldn't share
circuits, even though they share internal circuit and guard state.
---
 crates/arti-client/src/client.rs | 49 ++++++++++++++++++++++----------
 1 file changed, 34 insertions(+), 15 deletions(-)

diff --git a/crates/arti-client/src/client.rs b/crates/arti-client/src/client.rs
index c2e5af8e0..225347489 100644
--- a/crates/arti-client/src/client.rs
+++ b/crates/arti-client/src/client.rs
@@ -36,6 +36,8 @@ use tracing::{debug, error, info, warn};
 pub struct TorClient<R: Runtime> {
     /// Asynchronous runtime object.
     runtime: R,
+    /// Default isolation token for streams through this client.
+    client_isolation: IsolationToken,
     /// Circuit manager for keeping our circuits up to date and building
     /// them on-demand.
     circmgr: Arc<tor_circmgr::CircMgr<R>>,
@@ -46,12 +48,12 @@ pub struct TorClient<R: Runtime> {
 }
 
 /// Preferences for how to route a stream over the Tor network.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default)]
 pub struct ConnectPrefs {
     /// What kind of IPv6/IPv4 we'd prefer, and how strongly.
     ip_ver_pref: IpVersionPreference,
     /// Id of the isolation group the connection should be part of
-    isolation_group: IsolationToken,
+    isolation_group: Option<IsolationToken>,
 }
 
 impl ConnectPrefs {
@@ -114,28 +116,19 @@ impl ConnectPrefs {
     /// Indicate which other connections might use the same circuit
     /// as this one.
     pub fn set_isolation_group(&mut self, isolation_group: IsolationToken) -> &mut Self {
-        self.isolation_group = isolation_group;
+        self.isolation_group = Some(isolation_group);
         self
     }
 
-    /// Return a u64 to describe which connections might use
+    /// Return a token to describe which connections might use
     /// the same circuit as this one.
-    fn isolation_group(&self) -> IsolationToken {
+    fn isolation_group(&self) -> Option<IsolationToken> {
         self.isolation_group
     }
 
     // TODO: Add some way to be IPFlexible, and require exit to support both.
 }
 
-impl Default for ConnectPrefs {
-    fn default() -> Self {
-        ConnectPrefs {
-            ip_ver_pref: Default::default(),
-            isolation_group: IsolationToken::no_isolation(),
-        }
-    }
-}
-
 impl<R: Runtime> TorClient<R> {
     /// Bootstrap a network connection configured by `dir_cfg` and `circ_cfg`.
     ///
@@ -190,14 +183,35 @@ impl<R: Runtime> TorClient<R> {
             Arc::downgrade(&dirmgr),
         ))?;
 
+        let client_isolation = IsolationToken::new();
+
         Ok(TorClient {
             runtime,
+            client_isolation,
             circmgr,
             dirmgr,
             addrcfg: addr_cfg,
         })
     }
 
+    /// Return a new isolated `TorClient` instance.
+    ///
+    /// The two `TorClient`s will share some internal state, but their
+    /// streams will haver share circuits with one another.
+    ///
+    /// Use this function when you want separate parts of your program to
+    /// each have a TorClient handle, but where you don't want their
+    /// activities to be linkable to one another over the Tor network.
+    ///
+    /// Calling this function is usually preferable to creating a
+    /// completely separate TorClient instance, since it can share its
+    /// internals with the existing `TorClient`.
+    pub fn isolated_client(&self) -> TorClient<R> {
+        let mut result = self.clone();
+        result.client_isolation = IsolationToken::new();
+        result
+    }
+
     /// Launch an anonymized connection to the provided address and
     /// port over the Tor network.
     ///
@@ -303,9 +317,14 @@ impl<R: Runtime> TorClient<R> {
         flags: &ConnectPrefs,
     ) -> Result<Arc<ClientCirc>> {
         let dir = self.dirmgr.netdir();
+
+        // XXXX: this isn't what we really want.  We'd like to have _both_
+        // of these tokens considered.
+        let isolation = flags.isolation_group().unwrap_or(self.client_isolation);
+
         let circ = self
             .circmgr
-            .get_or_launch_exit(dir.as_ref().into(), exit_ports, flags.isolation_group())
+            .get_or_launch_exit(dir.as_ref().into(), exit_ports, isolation)
             .await
             .map_err(|_| Error::Internal("Unable to launch circuit"))?;
         drop(dir); // This decreases the refcount on the netdir.