clap

src/main.rs

@@ -1,14 +1,34 @@
 // https://tools.ietf.org/rfc/rfc5128.txt
 // https://blog.csdn.net/bytxl/article/details/44344855
 
+use clap::App;
 use hbb_common::{env_logger::*, log, tokio, ResultType};
 use hbbs::*;
+const DEFAULT_PORT: &'static str = "21116";
 
 #[tokio::main]
 async fn main() -> ResultType<()> {
     init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info"));
-    let addr = "0.0.0.0:21116";
+    let args = format!(
+        "-p, --port=[default={}] 'Sets the listening port'
+    -r, --relay-server=[] 'Sets the default relay server'",
+        DEFAULT_PORT
+    );
+    let matches = App::new("hbbs")
+        .version("1.0")
+        .author("Zhou Huabing <info@rustdesk.com>")
+        .about("RustDesk Rendezvous Server")
+        .args_from_usage(&args)
+        .get_matches();
+    let addr = format!(
+        "0.0.0.0:{}",
+        matches.value_of("port").unwrap_or(DEFAULT_PORT)
+    );
     log::info!("Listening on {}", addr);
-    RendezvousServer::start(&addr).await?;
+    RendezvousServer::start(
+        &addr,
+        matches.value_of("relay-server").unwrap_or("").to_owned(),
+    )
+    .await?;
     Ok(())
 }

src/rendezvous_server.rs

@@ -125,16 +125,18 @@ pub struct RendezvousServer {
     tcp_punch: Arc<Mutex<HashMap<SocketAddr, Sink>>>,
     pm: PeerMap,
     tx: Sender,
+    relay_server: String,
 }
 
 impl RendezvousServer {
-    pub async fn start(addr: &str) -> ResultType<()> {
+    pub async fn start(addr: &str, relay_server: String) -> ResultType<()> {
         let mut socket = FramedSocket::new(addr).await?;
         let (tx, mut rx) = mpsc::unbounded_channel::<(RendezvousMessage, SocketAddr)>();
         let mut rs = Self {
             tcp_punch: Arc::new(Mutex::new(HashMap::new())),
             pm: PeerMap::new()?,
             tx: tx.clone(),
+            relay_server,
         };
         let mut listener = new_listener(addr, true).await?;
         loop {
@@ -347,7 +349,10 @@ impl RendezvousServer {
             Some(peer) => peer.pk,
             _ => Vec::new(),
         };
-        let relay_server = phs.relay_server;
+        let mut relay_server = phs.relay_server;
+        if relay_server.is_empty() {
+            relay_server = self.relay_server.clone();
+        }
         msg_out.set_punch_hole_response(PunchHoleResponse {
             socket_addr: AddrMangle::encode(addr),
             pk,