diff options
| -rw-r--r-- | db_v2.sqlite3 | bin | 24576 -> 24576 bytes | |||
| -rw-r--r-- | libs/hbb_common/build.rs | 5 | ||||
| -rw-r--r-- | libs/hbb_common/src/fs.rs | 1 | ||||
| -rw-r--r-- | src/common.rs | 45 | ||||
| -rw-r--r-- | src/peer.rs | 2 | ||||
| -rw-r--r-- | src/relay_server.rs | 25 | ||||
| -rw-r--r-- | src/rendezvous_server.rs | 65 | ||||
| -rw-r--r-- | src/version.rs | 2 |
8 files changed, 99 insertions, 46 deletions
diff --git a/db_v2.sqlite3 b/db_v2.sqlite3 Binary files differindex 3d9350d..c95a2f3 100644 --- a/db_v2.sqlite3 +++ b/db_v2.sqlite3 diff --git a/libs/hbb_common/build.rs b/libs/hbb_common/build.rs index 225ec34..5feda67 100644 --- a/libs/hbb_common/build.rs +++ b/libs/hbb_common/build.rs @@ -5,10 +5,7 @@ fn main() { .out_dir("src/protos") .inputs(&["protos/rendezvous.proto", "protos/message.proto"]) .include("protos") - .customize( - protobuf_codegen::Customize::default() - .tokio_bytes(true) - ) + .customize(protobuf_codegen::Customize::default().tokio_bytes(true)) .run() .expect("Codegen failed."); } diff --git a/libs/hbb_common/src/fs.rs b/libs/hbb_common/src/fs.rs index 69cd348..a5bc6ad 100644 --- a/libs/hbb_common/src/fs.rs +++ b/libs/hbb_common/src/fs.rs @@ -565,4 +565,3 @@ pub fn transform_windows_path(entries: &mut Vec<FileEntry>) { entry.name = entry.name.replace("\\", "/"); } } - 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"; |