This commit is contained in:
rustdesk 2021-04-03 22:32:14 +08:00
parent 6a3903df2c
commit a10bb34f6e
5 changed files with 36 additions and 31 deletions

2
Cargo.lock generated
View File

@ -528,7 +528,7 @@ dependencies = [
[[package]]
name = "hbbs"
version = "1.0.0"
version = "1.1.3"
dependencies = [
"cc",
"clap",

View File

@ -1,6 +1,6 @@
[package]
name = "hbbs"
version = "1.0.0"
version = "1.1.3"
authors = ["open-trade <info@opentradesolutions.com>"]
edition = "2018"
build= "build.rs"

10
mod.rs
View File

@ -8,12 +8,12 @@ lazy_static::lazy_static! {
static ref STOP: Arc<Mutex<bool>> = Arc::new(Mutex::new(true));
}
pub fn bootstrap(license: &str, host: &str) {
pub fn bootstrap(key: &str, host: &str) {
let port = rendezvous_server::DEFAULT_PORT;
let addr = format!("0.0.0.0:{}", port);
let addr2 = format!("0.0.0.0:{}", port.parse::<i32>().unwrap_or(0) - 1);
let relay_servers: Vec<String> = vec![format!("{}:{}", host, relay_server::DEFAULT_PORT)];
let tmp_license = license.to_owned();
let tmp_key = key.to_owned();
std::thread::spawn(move || {
allow_err!(rendezvous_server::RendezvousServer::start(
&addr,
@ -22,15 +22,15 @@ pub fn bootstrap(license: &str, host: &str) {
0,
Default::default(),
Default::default(),
&tmp_license,
&tmp_key,
STOP.clone(),
));
});
let tmp_license = license.to_owned();
let tmp_key = key.to_owned();
std::thread::spawn(move || {
allow_err!(relay_server::start(
relay_server::DEFAULT_PORT,
&tmp_license,
&tmp_key,
STOP.clone()
));
});

View File

@ -24,7 +24,10 @@ lazy_static::lazy_static! {
pub const DEFAULT_PORT: &'static str = "21117";
#[tokio::main(basic_scheduler)]
pub async fn start(port: &str, license: &str, stop: Arc<Mutex<bool>>) -> ResultType<()> {
pub async fn start(port: &str, key: &str, stop: Arc<Mutex<bool>>) -> ResultType<()> {
if !key.is_empty() {
log::info!("Key: {}", key);
}
let addr = format!("0.0.0.0:{}", port);
log::info!("Listening on tcp {}", addr);
let mut listener = new_listener(addr, false).await?;
@ -34,18 +37,18 @@ pub async fn start(port: &str, license: &str, stop: Arc<Mutex<bool>>) -> ResultT
continue;
}
log::info!("Start");
io_loop(&mut listener, license, stop.clone()).await;
io_loop(&mut listener, key, stop.clone()).await;
}
}
async fn io_loop(listener: &mut TcpListener, license: &str, stop: Arc<Mutex<bool>>) {
async fn io_loop(listener: &mut TcpListener, key: &str, stop: Arc<Mutex<bool>>) {
let mut timer = interval(Duration::from_millis(100));
loop {
tokio::select! {
Ok((stream, addr)) = listener.accept() => {
let license = license.to_owned();
let key = key.to_owned();
tokio::spawn(async move {
make_pair(FramedStream::from(stream), addr, &license).await.ok();
make_pair(FramedStream::from(stream), addr, &key).await.ok();
});
}
_ = timer.tick() => {
@ -58,12 +61,12 @@ async fn io_loop(listener: &mut TcpListener, license: &str, stop: Arc<Mutex<bool
}
}
async fn make_pair(stream: FramedStream, addr: SocketAddr, license: &str) -> ResultType<()> {
async fn make_pair(stream: FramedStream, addr: SocketAddr, key: &str) -> ResultType<()> {
let mut stream = stream;
if let Some(Ok(bytes)) = stream.next_timeout(30_000).await {
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
if let Some(rendezvous_message::Union::request_relay(rf)) = msg_in.union {
if !license.is_empty() && rf.licence_key != license {
if !key.is_empty() && rf.licence_key != key {
return Ok(());
}
if !rf.uuid.is_empty() {

View File

@ -162,9 +162,12 @@ impl RendezvousServer {
serial: i32,
rendezvous_servers: Vec<String>,
software_url: String,
license: &str,
key: &str,
stop: Arc<Mutex<bool>>,
) -> ResultType<()> {
if !key.is_empty() {
log::info!("Key: {}", key);
}
log::info!("Listening on tcp/udp {}", addr);
log::info!("Listening on tcp {}, extra port for NAT test", addr2);
log::info!("relay-servers={:?}", relay_servers);
@ -197,7 +200,7 @@ impl RendezvousServer {
&mut listener,
&mut listener2,
&mut socket,
license,
key,
stop.clone(),
)
.await;
@ -210,7 +213,7 @@ impl RendezvousServer {
listener: &mut TcpListener,
listener2: &mut TcpListener,
socket: &mut FramedSocket,
license: &str,
key: &str,
stop: Arc<Mutex<bool>>,
) {
let mut timer = interval(Duration::from_millis(100));
@ -226,7 +229,7 @@ impl RendezvousServer {
allow_err!(socket.send(&msg, addr).await);
}
Some(Ok((bytes, addr))) = socket.next() => {
allow_err!(self.handle_msg(&bytes, addr, socket, license).await);
allow_err!(self.handle_msg(&bytes, addr, socket, key).await);
}
Ok((stream, addr)) = listener2.accept() => {
let stream = FramedStream::from(stream);
@ -251,7 +254,7 @@ impl RendezvousServer {
let (a, mut b) = Framed::new(stream, BytesCodec::new()).split();
let tcp_punch = self.tcp_punch.clone();
let mut rs = self.clone();
let license = license.to_owned();
let key = key.to_owned();
tokio::spawn(async move {
let mut sender = Some(a);
while let Ok(Some(Ok(bytes))) = timeout(30_000, b.next()).await {
@ -262,7 +265,7 @@ impl RendezvousServer {
if let Some(sender) = sender.take() {
tcp_punch.lock().unwrap().insert(addr, sender);
}
allow_err!(rs.handle_tcp_punch_hole_request(addr, ph, &license).await);
allow_err!(rs.handle_tcp_punch_hole_request(addr, ph, &key).await);
}
Some(rendezvous_message::Union::request_relay(mut rf)) => {
// there maybe several attempt, so sender can be none
@ -342,7 +345,7 @@ impl RendezvousServer {
bytes: &BytesMut,
addr: SocketAddr,
socket: &mut FramedSocket,
license: &str,
key: &str,
) -> ResultType<()> {
if let Ok(msg_in) = RendezvousMessage::parse_from_bytes(&bytes) {
match msg_in.union {
@ -392,14 +395,13 @@ impl RendezvousServer {
}
Some(rendezvous_message::Union::punch_hole_request(ph)) => {
if self.pm.is_in_memory(&ph.id) {
self.handle_udp_punch_hole_request(addr, ph, license)
.await?;
self.handle_udp_punch_hole_request(addr, ph, key).await?;
} else {
// not in memory, fetch from db with spawn in case blocking me
let mut me = self.clone();
let license = license.to_owned();
let key = key.to_owned();
tokio::spawn(async move {
allow_err!(me.handle_udp_punch_hole_request(addr, ph, &license).await);
allow_err!(me.handle_udp_punch_hole_request(addr, ph, &key).await);
});
}
}
@ -572,9 +574,9 @@ impl RendezvousServer {
&mut self,
addr: SocketAddr,
ph: PunchHoleRequest,
license: &str,
key: &str,
) -> ResultType<(RendezvousMessage, Option<SocketAddr>)> {
if !license.is_empty() && ph.licence_key != license {
if !key.is_empty() && ph.licence_key != key {
let mut msg_out = RendezvousMessage::new();
msg_out.set_punch_hole_response(PunchHoleResponse {
failure: punch_hole_response::Failure::LICENCE_MISMATCH.into(),
@ -686,9 +688,9 @@ impl RendezvousServer {
&mut self,
addr: SocketAddr,
ph: PunchHoleRequest,
license: &str,
key: &str,
) -> ResultType<()> {
let (msg, to_addr) = self.handle_punch_hole_request(addr, ph, license).await?;
let (msg, to_addr) = self.handle_punch_hole_request(addr, ph, key).await?;
if let Some(addr) = to_addr {
self.tx.send((msg, addr))?;
} else {
@ -702,9 +704,9 @@ impl RendezvousServer {
&mut self,
addr: SocketAddr,
ph: PunchHoleRequest,
license: &str,
key: &str,
) -> ResultType<()> {
let (msg, to_addr) = self.handle_punch_hole_request(addr, ph, license).await?;
let (msg, to_addr) = self.handle_punch_hole_request(addr, ph, key).await?;
self.tx.send((
msg,
match to_addr {