feat(network/onion): detect liveness via TcpStream close

instead of polling

Author: jetjinser

test/src/specs/tor/mod.rs

@@ -2,7 +2,7 @@ mod tor_basic;
 mod tor_connect;
 mod tor_hash_password;
 mod tor_reconnect;
-use ckb_async_runtime::Runtime;
+use ckb_async_runtime::{Runtime, new_global_runtime};
 use ckb_logger::{error, info};
 use std::{path::Path, process::Child};
 use tempfile::{TempDir, tempdir};
@@ -52,9 +52,15 @@ impl TorServer {
     pub fn tor_wait_bootstrap_done(&self) {
         let tor_controller_url = format!("127.0.0.1:{}", self.control_port);
         let controller_password = self.controller_password.clone();
-        Runtime::new().unwrap().block_on(async {
-            let tor_controller =
-                ckb_onion::TorController::new(tor_controller_url, controller_password, None).await;
+        let (handle, _handle_stop_rx, _runtime) = new_global_runtime(None);
+        handle.block_on(async {
+            let tor_controller = ckb_onion::TorController::new(
+                tor_controller_url,
+                controller_password,
+                None,
+                handle.clone(),
+            )
+            .await;
             let mut tor_controller = tor_controller.unwrap();
             if let Err(err) = tor_controller.wait_tor_server_bootstrap_done().await {
                 error!("wait tor server bootstrap done error: {:?}", err);

util/onion/src/onion_service.rs

@@ -104,7 +104,8 @@ impl OnionService {
         let tor_controller = self.config.tor_controller.to_string();
         let tor_password = self.config.tor_password.clone();
 
-        let mut tor_controller = TorController::new(tor_controller, tor_password, None).await?;
+        let mut tor_controller =
+            TorController::new(tor_controller, tor_password, None, self.handle.clone()).await?;
 
         tor_controller.wait_tor_server_bootstrap_done().await?;
 
@@ -128,22 +129,15 @@ impl OnionService {
         );
 
         self.handle.spawn(async move {
-            let mut ticker = tokio::time::interval(tokio::time::Duration::from_secs(3));
-            loop {
-                tokio::select! {
-                    _ = ticker.tick() => {
-                        let uptime = tor_controller.get_uptime().await;
-                        if let Err(err) = uptime {
-                            error!("Failed to get tor server uptime: {:?}", err);
-                            drop(tor_server_alive_tx);
-                            return;
-                        }
-                    }
-                    _ = stop_rx.cancelled() => {
-                        info!("OnionService received stop signal, exiting...");
-                        drop(tor_server_alive_tx);
-                        return;
-                    }
+            tokio::select! {
+                _ = tor_controller.wait_for_disconnect() => {
+                    error!("not alive");
+                    drop(tor_server_alive_tx);
+
+                }
+                _ = stop_rx.cancelled() => {
+                    info!("OnionService received stop signal, exiting...");
+                    drop(tor_server_alive_tx);
                 }
             }
         });

util/onion/src/tor_controller.rs

@@ -1,12 +1,14 @@
+use ckb_async_runtime::Handle;
 use ckb_error::{Error, InternalErrorKind};
 use ckb_logger::{debug, error, info, warn};
 use futures::future::BoxFuture;
 use std::borrow::Cow;
 use std::net::SocketAddr;
 use std::time::Duration;
 use tokio::fs::File;
-use tokio::io::AsyncReadExt;
+use tokio::io::{duplex, AsyncRead, AsyncReadExt, AsyncWrite, DuplexStream};
 use tokio::net::TcpStream;
+use tokio::sync::mpsc;
 use torut::control::{TorAuthData, TorAuthMethod, UnauthenticatedConn, COOKIE_LENGTH};
 use torut::{
     control::{AsyncEvent, AuthenticatedConn, ConnError},
@@ -15,12 +17,18 @@ use torut::{
 
 use crate::TorEventHandlerFn;
 
-type TorAuthenticatedConn =
-    AuthenticatedConn<TcpStream, fn(AsyncEvent<'_>) -> BoxFuture<'static, Result<(), ConnError>>>;
+type TorAuthenticatedConn = AuthenticatedConn<
+    DuplexStream,
+    fn(AsyncEvent<'_>) -> BoxFuture<'static, Result<(), ConnError>>,
+>;
 
 /// A controller for a Tor server.
 pub struct TorController {
     inner: TorAuthenticatedConn,
+    /// Notified when the underlying TCP connection to the Tor control port is
+    /// closed or encounters a fatal I/O error. Receiving a value indicates
+    /// the connection is dead.
+    _disconnect_rx: mpsc::UnboundedReceiver<()>,
 }
 
 impl TorController {
@@ -30,8 +38,13 @@ impl TorController {
         tor_controller_url: String,
         tor_password: Option<String>,
         event_handler: Option<TorEventHandlerFn>,
+        handle: Handle,
     ) -> Result<Self, Error> {
-        let s = TcpStream::connect(tor_controller_url.clone())
+        let (client, server) = duplex(1024);
+
+        let (disconnect_tx, disconnect_rx) = mpsc::unbounded_channel();
+
+        let raw_stream = TcpStream::connect(tor_controller_url.clone())
             .await
             .map_err(|err| {
                 InternalErrorKind::Other.other(format!(
@@ -40,15 +53,45 @@ impl TorController {
                 ))
             })?;
 
-        let mut utc: UnauthenticatedConn<TcpStream> = UnauthenticatedConn::new(s);
+        handle.spawn(async move {
+            let (mut tcp_read, mut tcp_write) = raw_stream.into_split();
+            let (mut server_read, mut server_write) = tokio::io::split(server);
+
+            tokio::select! {
+                // TCP -> Duplex
+                res = tokio::io::copy(&mut tcp_read, &mut server_write) => {
+                    match res {
+                        Ok(n) => debug!("Tor TCP connection closed after {} bytes", n),
+                        Err(e) => error!("Tor TCP read error: {}", e),
+                    }
+                    _ = disconnect_tx.send(());
+                }
+
+                // Duplex -> TCP
+                res = tokio::io::copy(&mut server_read, &mut tcp_write) => {
+                    match res {
+                        Ok(n) => debug!("Tor Duplex closed after {} bytes", n),
+                        Err(e) => error!("Tor TCP write error: {}", e),
+                    }
+                    _ = disconnect_tx.send(());
+                }
+            }
+
+            debug!("Tor DuplexStream exited.");
+        });
+
+        let mut utc: UnauthenticatedConn<_> = UnauthenticatedConn::new(client);
 
         authenticate(tor_password, &mut utc).await?;
 
         let mut ac = utc.into_authenticated().await;
 
         ac.set_async_event_handler(event_handler);
 
-        Ok(TorController { inner: ac })
+        Ok(TorController {
+            inner: ac,
+            _disconnect_rx: disconnect_rx,
+        })
     }
 
     /// get tor server's status
@@ -112,6 +155,16 @@ impl TorController {
         Ok(())
     }
 
+    /// Waits asynchronously until the underlying TCP connection to the Tor
+    /// controller is severed (either cleanly closed or due to an I/O error).
+    ///
+    /// Returns `Some(())` when the connection is lost, or `None` if the internal
+    /// notification channel was closed unexpectedly (which should not occur during
+    /// normal operation).
+    pub async fn wait_for_disconnect(&mut self) -> Option<()> {
+        self._disconnect_rx.recv().await
+    }
+
     /// Add a new v3 onion service to the Tor server.
     pub async fn add_onion_v3(
         &mut self,
@@ -125,10 +178,13 @@ impl TorController {
 }
 
 /// Authenticates with the Tor controller using the given password or cookie.
-pub async fn authenticate(
+pub async fn authenticate<S>(
     tor_password: Option<String>,
-    utc: &mut UnauthenticatedConn<TcpStream>,
-) -> Result<(), Error> {
+    utc: &mut UnauthenticatedConn<S>,
+) -> Result<(), Error>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+{
     let proto_info = utc.load_protocol_info().await.map_err(|err| {
         InternalErrorKind::Other.other(format!("Failed to load protocol info: {:?}", err))
     })?;