Merge branch 'isolation-example' into 'main'

add example for Isolation

Closes #414

See merge request tpo/core/arti!524

crates/arti-client/examples/advanced_isolation.rs

@@ -0,0 +1,198 @@
+// This example showcase how to use the trait IsolationHelper to build complexe isolation rules.
+// For most usages, using a combination of TorClient::isolated_client and IsolationToken should
+// be enough.
+use std::collections::HashSet;
+
+use anyhow::Result;
+use arti_client::isolation::IsolationHelper;
+use arti_client::{IsolationToken, StreamPrefs, TorClient, TorClientConfig};
+use tokio_crate as tokio;
+
+use futures::io::{AsyncReadExt, AsyncWriteExt};
+
+/// Example Isolation which isolate streams deemed sensitive from each other, but won't isolate
+/// Innocent streams from other Innoncent streams or from Sensitive streams.
+///
+/// More formally, for two streams to share the same circuit, it's required that either:
+/// - at least one stream is Innocent
+/// - both are Sensitive, with the same inner IsolationToken
+#[derive(Debug, Clone, Copy)]
+enum IsolateSensitive {
+    Sensitive(IsolationToken),
+    Innocent,
+}
+
+impl IsolateSensitive {
+    fn new_sensitive() -> Self {
+        IsolateSensitive::Sensitive(IsolationToken::new())
+    }
+}
+
+impl IsolationHelper for IsolateSensitive {
+    fn compatible_same_type(&self, other: &Self) -> bool {
+        match (self, other) {
+            (IsolateSensitive::Sensitive(i), IsolateSensitive::Sensitive(j)) => {
+                i.compatible_same_type(j)
+            }
+            _ => true,
+        }
+    }
+    fn join_same_type(&self, other: &Self) -> Option<Self> {
+        match (self, other) {
+            (IsolateSensitive::Sensitive(i), IsolateSensitive::Sensitive(j)) => {
+                i.join_same_type(j).map(IsolateSensitive::Sensitive)
+            }
+            (res @ IsolateSensitive::Sensitive(_), _)
+            | (_, res @ IsolateSensitive::Sensitive(_)) => Some(*res),
+            _ => Some(IsolateSensitive::Innocent),
+        }
+    }
+}
+
+/// Example Isolation which limit how many different purposes a single circuit can have.
+///
+/// Note that two IsolateOverused with different MAX_USAGE or different T are different types, so they
+/// would get isolated from each other, even if neither has reached its MAX_USAGE.
+#[derive(Debug, Clone)]
+struct IsolateOverused<T, const MAX_USAGE: usize> {
+    usages: HashSet<T>,
+}
+
+impl<T: Eq + std::hash::Hash, const MAX_USAGE: usize> IsolateOverused<T, MAX_USAGE> {
+    fn new_with_usage(usage: T) -> Self {
+        let mut usages = HashSet::new();
+        usages.insert(usage);
+
+        IsolateOverused { usages }
+    }
+}
+
+impl<T: Eq + std::hash::Hash + Clone, const MAX_USAGE: usize> IsolationHelper
+    for IsolateOverused<T, MAX_USAGE>
+{
+    fn compatible_same_type(&self, other: &Self) -> bool {
+        self.usages.union(&other.usages).count() <= MAX_USAGE
+    }
+
+    fn join_same_type(&self, other: &Self) -> Option<Self> {
+        let usages: HashSet<_> = self.usages.union(&other.usages).cloned().collect();
+        if usages.len() <= MAX_USAGE {
+            Some(IsolateOverused { usages })
+        } else {
+            None
+        }
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    tracing_subscriber::fmt::init();
+    let config = TorClientConfig::default();
+
+    eprintln!("connecting to Tor...");
+    let tor_client = TorClient::create_bootstrapped(config).await?;
+
+    // requests using this won't be isolated from each others, or from "sensitive" requests
+    let innocent = {
+        let mut pref = StreamPrefs::new();
+        pref.set_isolation(IsolateSensitive::Innocent);
+        pref
+    };
+
+    // requests using this will be isolated from "sensitive" requests using a different
+    // IsolateSensitive::Sensitive. They won't be isolated from each other or from Innocent requests
+    let sensitive_1 = {
+        let mut pref = StreamPrefs::new();
+        pref.set_isolation(IsolateSensitive::new_sensitive());
+        pref
+    };
+
+    let sensitive_2 = {
+        let mut pref = StreamPrefs::new();
+        pref.set_isolation(IsolateSensitive::new_sensitive());
+        pref
+    };
+
+    eprintln!("sending a bunch of sensitive and innocent requests");
+
+    let mut stream = tor_client
+        .connect_with_prefs(("example.net", 80), &innocent)
+        .await?;
+    send_request(&mut stream, "example.net").await?;
+
+    let mut stream = tor_client
+        .connect_with_prefs(("example.com", 80), &sensitive_1)
+        .await?;
+    send_request(&mut stream, "example.com").await?;
+
+    let mut stream = tor_client
+        .connect_with_prefs(("example.com", 80), &sensitive_2)
+        .await?;
+    send_request(&mut stream, "example.com").await?;
+
+    let mut stream = tor_client
+        .connect_with_prefs(("example.com", 80), &sensitive_1)
+        .await?;
+    send_request(&mut stream, "example.com").await?;
+
+    // Each of these requests can share a circuit with at most 2 other usages (3 including itself).
+    // As there are 4 different usages, there will be at least 2 circuits used. Also as
+    // IsolateOverused and IsolateSensitive are not the same type, they won't share any circuits.
+    let first_usage = {
+        let mut pref = StreamPrefs::new();
+        pref.set_isolation(IsolateOverused::<_, 3>::new_with_usage("first usage"));
+        pref
+    };
+    let second_usage = {
+        let mut pref = StreamPrefs::new();
+        pref.set_isolation(IsolateOverused::<_, 3>::new_with_usage("second usage"));
+        pref
+    };
+    let third_usage = {
+        let mut pref = StreamPrefs::new();
+        pref.set_isolation(IsolateOverused::<_, 3>::new_with_usage("third usage"));
+        pref
+    };
+    let fourth_usage = {
+        let mut pref = StreamPrefs::new();
+        pref.set_isolation(IsolateOverused::<_, 3>::new_with_usage("fourth usage"));
+        pref
+    };
+
+    let mut stream = tor_client
+        .connect_with_prefs(("example.net", 80), &first_usage)
+        .await?;
+    send_request(&mut stream, "example.net").await?;
+    let mut stream = tor_client
+        .connect_with_prefs(("example.org", 80), &second_usage)
+        .await?;
+    send_request(&mut stream, "example.org").await?;
+    let mut stream = tor_client
+        .connect_with_prefs(("example.com", 80), &third_usage)
+        .await?;
+    send_request(&mut stream, "example.com").await?;
+    let mut stream = tor_client
+        .connect_with_prefs(("example.net", 80), &fourth_usage)
+        .await?;
+    send_request(&mut stream, "example.net").await?;
+
+    Ok(())
+}
+
+async fn send_request(stream: &mut arti_client::DataStream, host: &str) -> Result<()> {
+    stream
+        .write_all(
+            format!(
+                "GET / HTTP/1.1\r\nHost: {}\r\nConnection: close\r\n\r\n",
+                host
+            )
+            .as_bytes(),
+        )
+        .await?;
+    stream.flush().await?;
+
+    let mut buf = Vec::new();
+    stream.read_to_end(&mut buf).await?;
+    println!("{}", String::from_utf8_lossy(&buf));
+    Ok(())
+}