From b7472de43f42b7559813c044efe407f671d9516a Mon Sep 17 00:00:00 2001
From: Henry de Valence <hdevalence@hdevalence.ca>
Date: Tue, 1 Sep 2020 20:41:43 -0700
Subject: [PATCH] network: add a zebra_network::connect_isolated() method.

The peer set provides an automatically managed connection pool, abstracting
away all the details of handling individual peer connections.  However, it's
also useful to be able to create completely isolated and
minimally-distinguishable connections to individual peers, in order to be able
to send specific messages over Tor, or to implement some custom network crawler
logic.
---
 zebra-network/src/isolated.rs | 78 +++++++++++++++++++++++++++++++++++
 zebra-network/src/lib.rs      |  2 +
 2 files changed, 80 insertions(+)
 create mode 100644 zebra-network/src/isolated.rs

diff --git a/zebra-network/src/isolated.rs b/zebra-network/src/isolated.rs
new file mode 100644
index 00000000..149311f0
--- /dev/null
+++ b/zebra-network/src/isolated.rs
@@ -0,0 +1,78 @@
+//! Code for creating isolated connections to specific peers.
+
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use futures::future::{FutureExt, TryFutureExt};
+
+use tokio::net::TcpStream;
+use tower::{
+    util::{BoxService, Oneshot},
+    Service,
+};
+
+use crate::{peer, BoxedStdError, Config, Request, Response};
+
+/// Use the provided TCP connection to create a Zcash connection completely
+/// isolated from all other node state.
+///
+/// The connection pool returned by `init` should be used for all requests that
+/// don't require isolated state or use of an existing TCP connection. However,
+/// this low-level API is useful for custom network crawlers or Tor connections.
+///
+/// In addition to being completely isolated from all other node state, this
+/// method also aims to be minimally distinguishable from other clients.
+///
+/// Note that this method does not implement any timeout behavior, so callers may
+/// want to layer it with a timeout as appropriate for their application.
+///
+/// # Inputs
+///
+/// - `conn`: an existing TCP connection to use. Passing an existing TCP
+/// connection allows this method to be used with clearnet or Tor transports.
+///
+/// - `user_agent`: a valid BIP14 user-agent, e.g., the empty string.
+pub async fn connect_isolated(
+    conn: TcpStream,
+    user_agent: String,
+) -> Result<BoxService<Request, Response, BoxedStdError>, BoxedStdError> {
+    let handshake = peer::Handshake::builder()
+        .with_config(Config::default())
+        .with_inbound_service(tower::service_fn(|_req| async move {
+            Ok::<Response, BoxedStdError>(Response::Nil)
+        }))
+        .with_user_agent(user_agent)
+        .finish()
+        .expect("provided mandatory builder parameters");
+
+    // We can't get the remote addr from conn, because it might be a tcp
+    // connection through a socks proxy, not directly to the remote. But it
+    // doesn't seem like zcashd cares if we give a bogus one, and Zebra doesn't
+    // touch it at all.
+    let remote_addr = "0.0.0.0:8233".parse().unwrap();
+
+    let client = Oneshot::new(handshake, (conn, remote_addr)).await?;
+
+    Ok(BoxService::new(Wrapper(client)))
+}
+
+// This can be deleted when a new version of Tower with map_err is released.
+struct Wrapper(peer::Client);
+
+impl Service<Request> for Wrapper {
+    type Response = Response;
+    type Error = BoxedStdError;
+    type Future =
+        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.0.poll_ready(cx).map_err(Into::into)
+    }
+
+    fn call(&mut self, req: Request) -> Self::Future {
+        self.0.call(req).map_err(Into::into).boxed()
+    }
+}
diff --git a/zebra-network/src/lib.rs b/zebra-network/src/lib.rs
index e99506c8..3924e7ff 100644
--- a/zebra-network/src/lib.rs
+++ b/zebra-network/src/lib.rs
@@ -58,6 +58,7 @@ pub type BoxedStdError = Box<dyn std::error::Error + Send + Sync + 'static>;
 mod address_book;
 mod config;
 mod constants;
+mod isolated;
 mod meta_addr;
 mod peer;
 mod peer_set;
@@ -68,6 +69,7 @@ mod timestamp_collector;
 pub use crate::{
     address_book::AddressBook,
     config::Config,
+    isolated::connect_isolated,
     peer_set::init,
     policies::{RetryErrors, RetryLimit},
     protocol::internal::{Request, Response},