From f200dff600c8bc4070090fc5b8f40391b8670c87 Mon Sep 17 00:00:00 2001 From: zonyitoo Date: Thu, 15 Aug 2024 23:55:47 +0800 Subject: [PATCH] feat: unified ServerHandle --- crates/shadowsocks-service/src/lib.rs | 1 + crates/shadowsocks-service/src/local/mod.rs | 28 ++------------- crates/shadowsocks-service/src/server/mod.rs | 34 ++---------------- .../shadowsocks-service/src/server/server.rs | 30 ++-------------- crates/shadowsocks-service/src/utils.rs | 36 +++++++++++++++++++ 5 files changed, 45 insertions(+), 84 deletions(-) create mode 100644 crates/shadowsocks-service/src/utils.rs diff --git a/crates/shadowsocks-service/src/lib.rs b/crates/shadowsocks-service/src/lib.rs index 13116b107c69..8514b82f6834 100644 --- a/crates/shadowsocks-service/src/lib.rs +++ b/crates/shadowsocks-service/src/lib.rs @@ -70,6 +70,7 @@ pub mod net; #[cfg(feature = "server")] pub mod server; mod sys; +mod utils; /// Default UDP association's expire duration #[allow(dead_code)] diff --git a/crates/shadowsocks-service/src/local/mod.rs b/crates/shadowsocks-service/src/local/mod.rs index 7b5791ed43c4..7ece6e16dcee 100644 --- a/crates/shadowsocks-service/src/local/mod.rs +++ b/crates/shadowsocks-service/src/local/mod.rs @@ -1,27 +1,24 @@ //! Shadowsocks Local Server use std::{ - future::Future, io::{self, ErrorKind}, - pin::Pin, sync::Arc, - task::{Context, Poll}, time::Duration, }; -use futures::{future, ready}; +use futures::future; use log::trace; use shadowsocks::{ config::Mode, net::{AcceptOpts, ConnectOpts}, }; -use tokio::task::JoinHandle; #[cfg(feature = "local-flow-stat")] use crate::{config::LocalFlowStatAddress, net::FlowStat}; use crate::{ config::{Config, ConfigType, ProtocolType}, dns::build_dns_resolver, + utils::ServerHandle, }; use self::{ @@ -70,27 +67,6 @@ pub mod utils; /// This is borrowed from Go's `net` library's default setting pub(crate) const LOCAL_DEFAULT_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(15); -struct ServerHandle(JoinHandle>); - -impl Drop for ServerHandle { - #[inline] - fn drop(&mut self) { - self.0.abort(); - } -} - -impl Future for ServerHandle { - type Output = io::Result<()>; - - #[inline] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match ready!(Pin::new(&mut self.0).poll(cx)) { - Ok(res) => res.into(), - Err(err) => Err(io::Error::new(ErrorKind::Other, err)).into(), - } - } -} - /// Local Server instance pub struct Server { balancer: PingBalancer, diff --git a/crates/shadowsocks-service/src/server/mod.rs b/crates/shadowsocks-service/src/server/mod.rs index d58ab450134f..e019658fc04d 100644 --- a/crates/shadowsocks-service/src/server/mod.rs +++ b/crates/shadowsocks-service/src/server/mod.rs @@ -1,22 +1,15 @@ //! Shadowsocks server -use std::{ - future::Future, - io::{self, ErrorKind}, - pin::Pin, - sync::Arc, - task::{Context, Poll}, - time::Duration, -}; +use std::{io, sync::Arc, time::Duration}; -use futures::{future, ready}; +use futures::future; use log::trace; use shadowsocks::net::{AcceptOpts, ConnectOpts}; -use tokio::task::JoinHandle; use crate::{ config::{Config, ConfigType}, dns::build_dns_resolver, + utils::ServerHandle, }; pub use self::{ @@ -170,24 +163,3 @@ pub async fn run(config: Config) -> io::Result<()> { let (res, ..) = future::select_all(vfut).await; res } - -struct ServerHandle(JoinHandle>); - -impl Drop for ServerHandle { - #[inline] - fn drop(&mut self) { - self.0.abort(); - } -} - -impl Future for ServerHandle { - type Output = io::Result<()>; - - #[inline] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match ready!(Pin::new(&mut self.0).poll(cx)) { - Ok(res) => res.into(), - Err(err) => Err(io::Error::new(ErrorKind::Other, err)).into(), - } - } -} diff --git a/crates/shadowsocks-service/src/server/server.rs b/crates/shadowsocks-service/src/server/server.rs index ae9102bdf848..fffbefd003af 100644 --- a/crates/shadowsocks-service/src/server/server.rs +++ b/crates/shadowsocks-service/src/server/server.rs @@ -2,15 +2,12 @@ use std::{ collections::HashMap, - future::Future, io::{self, ErrorKind}, - pin::Pin, sync::Arc, - task::{Context, Poll}, time::Duration, }; -use futures::{future, ready}; +use futures::future; use log::{error, trace}; use shadowsocks::{ config::{ManagerAddr, ServerConfig}, @@ -19,9 +16,9 @@ use shadowsocks::{ plugin::{Plugin, PluginMode}, ManagerClient, }; -use tokio::{task::JoinHandle, time}; +use tokio::time; -use crate::{acl::AccessControl, config::SecurityConfig, net::FlowStat}; +use crate::{acl::AccessControl, config::SecurityConfig, net::FlowStat, utils::ServerHandle}; use super::{context::ServiceContext, tcprelay::TcpServer, udprelay::UdpServer}; @@ -162,27 +159,6 @@ impl ServerBuilder { } } -struct ServerHandle(JoinHandle>); - -impl Drop for ServerHandle { - #[inline] - fn drop(&mut self) { - self.0.abort(); - } -} - -impl Future for ServerHandle { - type Output = io::Result<()>; - - #[inline] - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match ready!(Pin::new(&mut self.0).poll(cx)) { - Ok(res) => res.into(), - Err(err) => Err(io::Error::new(ErrorKind::Other, err)).into(), - } - } -} - /// Shadowsocks Server instance pub struct Server { context: Arc, diff --git a/crates/shadowsocks-service/src/utils.rs b/crates/shadowsocks-service/src/utils.rs new file mode 100644 index 000000000000..4da3e04d12a9 --- /dev/null +++ b/crates/shadowsocks-service/src/utils.rs @@ -0,0 +1,36 @@ +//! Service Utilities + +use std::{ + future::Future, + io::{self, ErrorKind}, + pin::Pin, + task::{Context, Poll}, +}; + +use futures::ready; +use tokio::task::JoinHandle; + +/// Wrapper of `tokio::task::JoinHandle`, which links to a server instance. +/// +/// `ServerHandle` implements `Future` which will join the `JoinHandle` and get the result. +/// When `ServerHandle` drops, it will abort the task. +pub struct ServerHandle(pub JoinHandle>); + +impl Drop for ServerHandle { + #[inline] + fn drop(&mut self) { + self.0.abort(); + } +} + +impl Future for ServerHandle { + type Output = io::Result<()>; + + #[inline] + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + match ready!(Pin::new(&mut self.0).poll(cx)) { + Ok(res) => res.into(), + Err(err) => Err(io::Error::new(ErrorKind::Other, err)).into(), + } + } +}