rendezvous_servers

src/main.rs

@@ -4,14 +4,18 @@
 use clap::App;
 use hbb_common::{env_logger::*, log, tokio, ResultType};
 use hbbs::*;
+use ini::Ini;
 const DEFAULT_PORT: &'static str = "21116";
 
 #[tokio::main]
 async fn main() -> ResultType<()> {
     init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info"));
     let args = format!(
-        "-p, --port=[default={}] 'Sets the listening port'
-    -r, --relay-server=[] 'Sets the default relay server'",
+        "-c --config=[FILE] +takes_value 'Sets a custom config file'
+        -p, --port=[NUMBER(default={})] 'Sets the listening port'
+        -s, --serial=[NUMBER(default={0})] 'Sets configure update serial number'
+        -R, --rendezvous-servers=[HOSTS] 'Sets rendezvous servers, seperated by colon'
+    -r, --relay-server=[HOST] 'Sets the default relay server'",
         DEFAULT_PORT
     );
     let matches = App::new("hbbs")
@@ -20,15 +24,49 @@ async fn main() -> ResultType<()> {
         .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)
-    );
+    let mut section = None;
+    let conf; // for holding section
+    if let Some(config) = matches.value_of("config") {
+        if let Ok(v) = Ini::load_from_file(config) {
+            conf = v;
+            section = conf.section(None::<String>);
+        }
+    }
+    let get_arg = |name: &str, default: &str| -> String {
+        if let Some(v) = matches.value_of(name) {
+            return v.to_owned();
+        } else if let Some(section) = section {
+            if let Some(v) = section.get(name) {
+                return v.to_owned();
+            }
+        }
+        return default.to_owned();
+    };
+    let port = get_arg("port", DEFAULT_PORT);
+    let mut relay_server = get_arg("relay-server", "");
+    if !relay_server.contains(":") {
+        relay_server = format!("{}:21117", relay_server);
+    }
+    if !relay_server.parse::<std::net::SocketAddr>().is_ok() {
+        relay_server = "".to_owned();
+    }
+    let serial: i32 = get_arg("serial", "").parse().unwrap_or(0);
+    let rendezvous_servers: Vec<String> = get_arg("rendezvous-servers", "")
+        .split(",")
+        .map(|x| {
+            if !x.contains(":") {
+                format!("{}:21116", x)
+            } else {
+                x.to_owned()
+            }
+        })
+        .filter(|x| x.parse::<std::net::SocketAddr>().is_ok())
+        .collect();
+    let addr = format!("0.0.0.0:{}", port);
     log::info!("Listening on {}", addr);
-    RendezvousServer::start(
-        &addr,
-        matches.value_of("relay-server").unwrap_or("").to_owned(),
-    )
-    .await?;
+    log::info!("relay-server={}", relay_server);
+    log::info!("serial={}", serial);
+    log::info!("rendzvous-servers={:?}", rendezvous_servers);
+    RendezvousServer::start(&addr, relay_server, serial, rendezvous_servers).await?;
     Ok(())
 }

src/rendezvous_server.rs

@@ -126,10 +126,17 @@ pub struct RendezvousServer {
     pm: PeerMap,
     tx: Sender,
     relay_server: String,
+    serial: i32,
+    rendezvous_servers: Vec<String>,
 }
 
 impl RendezvousServer {
-    pub async fn start(addr: &str, relay_server: String) -> ResultType<()> {
+    pub async fn start(
+        addr: &str,
+        relay_server: String,
+        serial: i32,
+        rendezvous_servers: Vec<String>,
+    ) -> ResultType<()> {
         let mut socket = FramedSocket::new(addr).await?;
         let (tx, mut rx) = mpsc::unbounded_channel::<(RendezvousMessage, SocketAddr)>();
         let mut rs = Self {
@@ -137,6 +144,8 @@ impl RendezvousServer {
             pm: PeerMap::new()?,
             tx: tx.clone(),
             relay_server,
+            serial,
+            rendezvous_servers,
         };
         let mut listener = new_listener(addr, true).await?;
         loop {
@@ -219,6 +228,17 @@ impl RendezvousServer {
                     if rp.id.len() > 0 {
                         log::debug!("New peer registered: {:?} {:?}", &rp.id, &addr);
                         self.update_addr(rp.id, addr, socket).await?;
+                        if self.serial != rp.serial {
+                            let mut msg_out = RendezvousMessage::new();
+                            let mut mi = MiscInfo::new();
+                            mi.set_configure_update(ConfigUpdate {
+                                serial: self.serial,
+                                rendezvous_servers: self.rendezvous_servers.clone(),
+                                ..Default::default()
+                            });
+                            msg_out.set_misc_info(mi);
+                            socket.send(&msg_out, addr).await?;
+                        }
                     }
                 }
                 Some(rendezvous_message::Union::register_pk(rk)) => {
@@ -267,9 +287,6 @@ impl RendezvousServer {
                 Some(rendezvous_message::Union::local_addr(la)) => {
                     self.handle_local_addr(&la, addr, Some(socket)).await?;
                 }
-                Some(rendezvous_message::Union::system_info(info)) => {
-                    log::info!("{}", info.value);
-                }
                 _ => {}
             }
         }