diff --git a/crates/arti-client/examples/advanced_isolation.rs b/crates/arti-client/examples/advanced_isolation.rs new file mode 100644 index 000000000..a386a38c8 --- /dev/null +++ b/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 { + 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 { + usages: HashSet, +} + +impl IsolateOverused { + fn new_with_usage(usage: T) -> Self { + let mut usages = HashSet::new(); + usages.insert(usage); + + IsolateOverused { usages } + } +} + +impl IsolationHelper + for IsolateOverused +{ + fn compatible_same_type(&self, other: &Self) -> bool { + self.usages.union(&other.usages).count() <= MAX_USAGE + } + + fn join_same_type(&self, other: &Self) -> Option { + 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(()) +}