mirror/rustdesk-server
1
0
Fork 0
mirror of https://github.com/rustdesk/rustdesk-server.git synced 2024-11-27 23:29:00 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

refactor stop

Browse Source
This commit is contained in:
open-trade 2021-03-25 17:55:44 +08:00
parent 3d1260c8ad
commit 9132d85f74
2 changed files with 47 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
src/relay_server.rs
View File

@ -6,6 +6,7 @@ use hbb_common::{
tcp::{new_listener, FramedStream},
tokio::{
self,
net::TcpListener,
time::{interval, Duration},
},
ResultType,
@ -26,8 +27,19 @@ pub const DEFAULT_PORT: &'static str = "21117";
pub async fn start(port: &str, license: &str, stop: Arc<Mutex<bool>>) -> ResultType<()> {
let addr = format!("0.0.0.0:{}", port);
log::info!("Listening on tcp {}", addr);
let mut timer = interval(Duration::from_millis(300));
let mut listener = new_listener(addr, false).await?;
loop {
if *stop.lock().unwrap() {
sleep(0.1).await;
continue;
}
log::info!("Start");
io_loop(&mut listener, license, stop.clone()).await;
}
}
async fn io_loop(listener: &mut TcpListener, license: &str, stop: Arc<Mutex<bool>>) {
let mut timer = interval(Duration::from_millis(100));
loop {
tokio::select! {
Ok((stream, addr)) = listener.accept() => {
@ -44,7 +56,6 @@ pub async fn start(port: &str, license: &str, stop: Arc<Mutex<bool>>) -> ResultT
}
}
}
Ok(())
}
async fn make_pair(stream: FramedStream, addr: SocketAddr, license: &str) -> ResultType<()> {

40
src/rendezvous_server.rs
View File

@ -9,11 +9,12 @@ use hbb_common::{
log,
protobuf::{Message as _, MessageField},
rendezvous_proto::*,
sleep,
tcp::{new_listener, FramedStream},
timeout,
tokio::{
self,
net::TcpStream,
net::{TcpListener, TcpStream},
sync::mpsc,
time::{interval, Duration},
},
@ -137,6 +138,7 @@ impl PeerMap {
const REG_TIMEOUT: i32 = 30_000;
type Sink = SplitSink<Framed<TcpStream, BytesCodec>, Bytes>;
type Sender = mpsc::UnboundedSender<(RendezvousMessage, SocketAddr)>;
type Receiver = mpsc::UnboundedReceiver<(RendezvousMessage, SocketAddr)>;
static mut ROTATION_RELAY_SERVER: usize = 0;
#[derive(Clone)]
@ -184,7 +186,34 @@ impl RendezvousServer {
};
let mut listener = new_listener(addr, false).await?;
let mut listener2 = new_listener(addr2, false).await?;
let mut timer = interval(Duration::from_millis(300));
loop {
if *stop.lock().unwrap() {
sleep(0.1).await;
continue;
}
log::info!("Start");
rs.io_loop(
&mut rx,
&mut listener,
&mut listener2,
&mut socket,
license,
stop.clone(),
)
.await;
}
}
async fn io_loop(
&mut self,
rx: &mut Receiver,
listener: &mut TcpListener,
listener2: &mut TcpListener,
socket: &mut FramedSocket,
license: &str,
stop: Arc<Mutex<bool>>,
) {
let mut timer = interval(Duration::from_millis(100));
loop {
tokio::select! {
_ = timer.tick() => {
@ -197,7 +226,7 @@ impl RendezvousServer {
allow_err!(socket.send(&msg, addr).await);
}
Some(Ok((bytes, addr))) = socket.next() => {
allow_err!(rs.handle_msg(&bytes, addr, &mut socket, license).await);
allow_err!(self.handle_msg(&bytes, addr, socket, license).await);
}
Ok((stream, addr)) = listener2.accept() => {
let stream = FramedStream::from(stream);
@ -220,8 +249,8 @@ impl RendezvousServer {
Ok((stream, addr)) = listener.accept() => {
log::debug!("Tcp connection from {:?}", addr);
let (a, mut b) = Framed::new(stream, BytesCodec::new()).split();
let tcp_punch = rs.tcp_punch.clone();
let mut rs = rs.clone();
let tcp_punch = self.tcp_punch.clone();
let mut rs = self.clone();
let license = license.to_owned();
tokio::spawn(async move {
let mut sender = Some(a);
@ -305,7 +334,6 @@ impl RendezvousServer {
}
}
}
Ok(())
}
#[inline]