Feat: listen for unix signal

5 files changed, 98 insertions, 41 deletions

diff --git a/src/common.rs b/src/common.rs
index a57485b..e5171b7 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -1,5 +1,8 @@
 use clap::App;
-use hbb_common::{anyhow::Context, log, ResultType};
+use hbb_common::{
+    anyhow::{Context, Result},
+    log, ResultType,
+};
 use ini::Ini;
 use sodiumoxide::crypto::sign;
 use std::{
@@ -138,3 +141,43 @@ pub fn gen_sk(wait: u64) -> (String, Option<sign::SecretKey>) {
     }
     ("".to_owned(), None)
 }
+
+#[cfg(unix)]
+pub async fn listen_signal() -> Result<()> {
+    use hbb_common::tokio;
+    use hbb_common::tokio::signal::unix::{signal, SignalKind};
+
+    tokio::spawn(async {
+        let mut s = signal(SignalKind::hangup())?;
+        let hangup = s.recv();
+        let mut s = signal(SignalKind::terminate())?;
+        let terminate = s.recv();
+        let mut s = signal(SignalKind::interrupt())?;
+        let interrupt = s.recv();
+        let mut s = signal(SignalKind::quit())?;
+        let quit = s.recv();
+
+        tokio::select! {
+            _ = hangup => {
+                log::info!("signal hangup");
+            }
+            _ = terminate => {
+                log::info!("signal terminate");
+            }
+            _ = interrupt => {
+                log::info!("signal interrupt");
+            }
+            _ = quit => {
+                log::info!("signal quit");
+            }
+        }
+        Ok(())
+    })
+    .await?
+}
+
+#[cfg(not(unix))]
+pub async fn listen_signal() -> Result<()> {
+    let () = std::future::pending().await;
+    unreachable!();
+}
diff --git a/src/peer.rs b/src/peer.rs
index 49d440d..1ce5a2e 100644
--- a/src/peer.rs
+++ b/src/peer.rs
@@ -1,10 +1,10 @@
 use crate::common::*;
 use crate::database;
 use hbb_common::{
+    bytes::Bytes,
     log,
     rendezvous_proto::*,
     tokio::sync::{Mutex, RwLock},
-    bytes::Bytes,
     ResultType,
 };
 use serde_derive::{Deserialize, Serialize};
diff --git a/src/relay_server.rs b/src/relay_server.rs
index 7b198d6..be8a69d 100644
--- a/src/relay_server.rs
+++ b/src/relay_server.rs
@@ -81,15 +81,22 @@ pub async fn start(port: &str, key: &str) -> ResultType<()> {
     log::info!("Listening on tcp {}", addr);
     let addr2 = format!("0.0.0.0:{}", port.parse::<u16>().unwrap() + 2);
     log::info!("Listening on websocket {}", addr2);
-    loop {
-        log::info!("Start");
-        io_loop(
-            new_listener(&addr, false).await?,
-            new_listener(&addr2, false).await?,
-            &key,
-        )
-        .await;
-    }
+    let main_task = async move {
+        loop {
+            log::info!("Start");
+            io_loop(
+                new_listener(&addr, false).await?,
+                new_listener(&addr2, false).await?,
+                &key,
+            )
+            .await;
+        }
+    };
+    let listen_signal = crate::common::listen_signal();
+    tokio::select!(
+        res = main_task => res,
+        res = listen_signal => res,
+    )
 }
 
 fn check_params() {
diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs
index b93e04f..da69540 100644
--- a/src/rendezvous_server.rs
+++ b/src/rendezvous_server.rs
@@ -170,37 +170,44 @@ impl RendezvousServer {
                 allow_err!(test_hbbs(test_addr).await);
             });
         };
-        loop {
-            log::info!("Start");
-            match rs
-                .io_loop(
-                    &mut rx,
-                    &mut listener,
-                    &mut listener2,
-                    &mut listener3,
-                    &mut socket,
-                    &key,
-                )
-                .await
-            {
-                LoopFailure::UdpSocket => {
-                    drop(socket);
-                    socket = FramedSocket::new_with_buf_size(&addr, rmem).await?;
-                }
-                LoopFailure::Listener => {
-                    drop(listener);
-                    listener = new_listener(&addr, false).await?;
-                }
-                LoopFailure::Listener2 => {
-                    drop(listener2);
-                    listener2 = new_listener(&addr2, false).await?;
-                }
-                LoopFailure::Listener3 => {
-                    drop(listener3);
-                    listener3 = new_listener(&addr3, false).await?;
+        let main_task = async move {
+            loop {
+                log::info!("Start");
+                match rs
+                    .io_loop(
+                        &mut rx,
+                        &mut listener,
+                        &mut listener2,
+                        &mut listener3,
+                        &mut socket,
+                        &key,
+                    )
+                    .await
+                {
+                    LoopFailure::UdpSocket => {
+                        drop(socket);
+                        socket = FramedSocket::new_with_buf_size(&addr, rmem).await?;
+                    }
+                    LoopFailure::Listener => {
+                        drop(listener);
+                        listener = new_listener(&addr, false).await?;
+                    }
+                    LoopFailure::Listener2 => {
+                        drop(listener2);
+                        listener2 = new_listener(&addr2, false).await?;
+                    }
+                    LoopFailure::Listener3 => {
+                        drop(listener3);
+                        listener3 = new_listener(&addr3, false).await?;
+                    }
                 }
             }
-        }
+        };
+        let listen_signal = listen_signal();
+        tokio::select!(
+            res = main_task => res,
+            res = listen_signal => res,
+        )
     }
 
     async fn io_loop(
diff --git a/src/version.rs b/src/version.rs
index a68926b..b54facc 100644
--- a/src/version.rs
+++ b/src/version.rs
@@ -1 +1 @@
-pub const VERSION: &str = "1.1.6";
\ No newline at end of file
+pub const VERSION: &str = "1.1.6";