diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..feb1f2f --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,17 @@ +# Repository Instructions + +## Rust Rules + +- In Rust code, do not introduce `unwrap()` or `expect()`. +- Allowed exceptions: +- Tests may use `unwrap()` or `expect()` when it keeps the test focused and readable. +- Lock acquisition may use `unwrap()` only when the locking API makes that the practical option and the failure mode is poison handling rather than normal control flow. +- Outside those exceptions, propagate errors, handle them explicitly, or use safer fallbacks instead of `unwrap()` and `expect()`. + +## Editing Hygiene + +- Do not introduce formatting-only changes. +- Do not run repository-wide formatters or reflow unrelated code unless the + user explicitly asks for formatting. +- Keep diffs limited to semantic changes required for the task. + diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..c317064 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +AGENTS.md diff --git a/README.md b/README.md index 4d7590f..887cee6 100644 --- a/README.md +++ b/README.md @@ -43,9 +43,10 @@ The most common options: | Option | Flag | Env var | Applies to | Purpose | | --- | --- | --- | --- | --- | -| Key | `-k` | `KEY` | hbbs, hbbr | Server key; `-k _` requires encryption | +| Key | `-k` | `KEY` | hbbs, hbbr | `hbbs` loads/generates one by default | +| Bind address | `-b` | `BIND` | hbbs, hbbr | Local IP address to listen on (default: all interfaces; requires 1.1.17+) | | Port | `-p` | `PORT` | hbbs, hbbr | Listening port (hbbs `21116`, hbbr `21117`) | -| Relay servers | `-r` | `RELAY-SERVERS` | hbbs | Relay address given to clients, e.g. your domain | +| Relay servers | `-r` | `RELAY-SERVERS` | hbbs | Override when the relay uses a different address or a non-standard port | | Force relay | — | `ALWAYS_USE_RELAY` | hbbs | `Y` disables direct connections | | Log level | — | `RUST_LOG` | hbbs, hbbr | e.g. `debug` (default `info`) | diff --git a/db_v2.sqlite3 b/db_v2.sqlite3 index 93d9801..5e1c781 100644 Binary files a/db_v2.sqlite3 and b/db_v2.sqlite3 differ diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 8531f52..488921f 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -25,70 +25,31 @@ the hood every source is turned into a process environment variable, and the code then reads that variable — so "flag", "config file" and "env var" are just three ways to set the same thing. -For **`hbbr`** the precedence is: **flag** (`-p`, `-k`) → **`.env`** → +For **`hbbr`** the precedence is: **flag** (`-b`, `-p`, `-k`) → **`.env`** → **inherited environment**. `RUST_LOG` is an exception to these rules. Both binaries initialize logging before loading `.env` (or `hbbs`'s `--config` file), so `RUST_LOG` must be set in the inherited process environment. -### ⚠️ The `.env` naming gotcha (read this) - -`hbbs` and `hbbr` do **not** parse `.env` the same way: - -| | Key written in `.env` | Variable the code sees | -|---|---|---| -| **`hbbs`** | `relay_servers` **or** `relay-servers` | `RELAY-SERVERS` (upper‑cased, `_`→`-`) | -| **`hbbr`** | `downgrade_threshold` | `DOWNGRADE_THRESHOLD` (used verbatim) | - -* `hbbs` rewrites every `.env`/`--config` key to **UPPERCASE** and replaces - underscores with dashes. So a key with an underscore in `.env` (for example - `DB_URL` or `TEST_HBBS`) becomes `DB-URL` / `TEST-HBBS` and will **not** match - the `DB_URL` / `TEST_HBBS` the code looks for. Those "direct" variables - (marked 🅴 in the tables below) can therefore **only** be set as a real - environment variable, not through the `hbbs` `.env`/`--config` file. -* `hbbr` uses `.env` keys verbatim, so names must match the documented - uppercase spelling exactly. This works for variables read after `.env` is - loaded; `RUST_LOG` is the exception described above. - -**Recommendation:** -* Set multi-word `hbbs` flag options via the **command line** or the **`.env` - file** (using the dashed lowercase name, e.g. `relay-servers`). -* In a `.env` file shared by both binaries, spell shared names such as `KEY` and - `PORT` in uppercase. `hbbs` accepts that spelling after normalization, and - `hbbr` requires it. -* Set the 🅴 "direct" tuning variables as **real environment variables** - (docker‑compose `environment:`, systemd `Environment=`, or `export`). - -Multi‑word options such as `relay-servers` have an internal env‑var name that -contains a dash (`RELAY-SERVERS`). Most shells cannot `export RELAY-SERVERS=…`, -so for those prefer the CLI flag or the `.env` file. Single‑word options -(`PORT`, `KEY`, `MASK`, `SERIAL`, `RMEM`) are ordinary identifiers and work fine -as exported environment variables. - --- ## `hbbs` — ID / rendezvous server | Variable | CLI flag | Default | Description | |---|---|---|---| -| `KEY` | `-k`, `--key` | `-` | Public key clients must use, a base64 secret key, or `-` / `_` to load or generate a key pair (`id_ed25519`, `id_ed25519.pub`). Use `_` to require encryption. An explicitly empty value disables key validation; see [Keys](#keys-and-encryption). | +| `KEY` | `-k`, `--key` | `-` | Public key clients must use, a base64 secret key, or `-` / `_` to load or generate a key pair (`id_ed25519`, `id_ed25519.pub`). `-` and `_` have the same behavior, so explicitly passing `-k _` to `hbbs` is unnecessary. An explicitly empty value disables key validation; see [Keys](#keys-and-encryption). | +| `BIND` | `-b`, `--bind` | all interfaces | **Available since 1.1.17.** Local IPv4 or IPv6 address on which all `hbbs` TCP, UDP, and WebSocket listeners bind. This does not change the addresses advertised to clients. Supported by `--config`, `.env`, and the inherited environment. | | `PORT` | `-p`, `--port` | `21116` | Main TCP/UDP listening port. `hbbs` also binds `PORT-1` (NAT type test) and `PORT+2` (WebSocket). | -| `RELAY-SERVERS` | `-r`, `--relay-servers` | *(empty)* | Default relay server(s) handed to clients, comma‑separated `host` or `host:port`. Usually your public IP / domain. | -| `RENDEZVOUS-SERVERS` | `-R`, `--rendezvous-servers` | *(empty)* | Peer rendezvous servers to forward to, comma‑separated. For multi‑server setups; leave empty for a single server. | -| `MASK` | `--mask` | *(none)* | CIDR that marks a client as "LAN", e.g. `192.168.0.0/16`. When set, LAN peers get `LOCAL-IP` instead of their public address. | -| `LOCAL-IP` | *(none — env/`.env` only)* | auto‑detected | LAN address advertised to peers matched by `MASK`. Defaults to the machine's primary local IP. | -| `SERIAL` | `-s`, `--serial` | `0` | Config update serial. Bump it to push updated relay/rendezvous lists to clients. | +| `RELAY-SERVERS` | `-r`, `--relay-servers` | *(empty)* | Optional relay server override handed to clients, as comma-separated `host` or `host:port` values. Leave empty when `hbbr` uses the same address as `hbbs` and the standard port `21117`; clients derive it automatically. Set this only when the relay uses a different IP/hostname or a non-standard port. | | `RMEM` | `-M`, `--rmem` | `0` (system default) | UDP receive‑buffer size in bytes. Raise the OS limit first: `sudo sysctl -w net.core.rmem_max=52428800`. | -| `SOFTWARE-URL` | `-u`, `--software-url` | *(empty)* | Download URL of the newest RustDesk client; the version is parsed from it and offered to clients. | | *(config file)* | `-c`, `--config` | *(none)* | Path to an extra INI config file (see precedence above). | | `TEST_HBBS` 🅴 | *(none)* | *(auto)* | UDP self‑test target checked at start‑up. Set to `no` to skip the check (useful behind some NATs/proxies), or to an explicit `host:port`. | | `ALWAYS_USE_RELAY` 🅴 | *(none)* | `N` | `Y` forces every session through a relay (disables direct/hole‑punched connections). At runtime, send `always-use-relay Y` or `always-use-relay N` to the `hbbs` [loopback console](#runtime-console). | | `DB_URL` 🅴 | *(none)* | `./db_v2.sqlite3` | Path/URL of the SQLite database file. See [Database](#database). | | `MAX_DATABASE_CONNECTIONS` 🅴 | *(none)* | `1` | Size of the SQLite connection pool. | -🅴 = read directly from the environment; **cannot** be set through the `hbbs` -`.env`/`--config` file (see the gotcha above). +🅴 = set through the inherited process environment. > `PORT_FOR_API` / `KEY_FOR_API` are only used by RustDesk Server **Pro** and its > API; they have no effect in the open‑source server. @@ -99,10 +60,11 @@ as exported environment variables. | Variable | CLI flag | Default | Description | |---|---|---|---| -| `KEY` | `-k`, `--key` | *(empty)* | Must match the non-empty key `hbbs` uses. `-` / `_` load or generate a key pair. The empty default disables key validation and leaves the relay unauthenticated; do not leave it empty on an exposed server. | +| `KEY` | `-k`, `--key` | *(empty)* | The empty default intentionally disables relay key validation, avoiding key-pair setup and mismatch failures. To enable relay key validation, use the same non-empty key as `hbbs`; `-` / `_` have the same behavior and load or generate a key pair. An empty key allows clients without a matching key to use the relay, so choose this tradeoff deliberately on an exposed server. | +| `BIND` | `-b`, `--bind` | all interfaces | **Available since 1.1.17.** Local IPv4 or IPv6 address on which the relay TCP and WebSocket listeners bind. Supported by `.env` and the inherited environment; `hbbr` does not support `--config`. | | `PORT` | `-p`, `--port` | `21117` | Relay listening port. `hbbr` also binds `PORT+2` for WebSocket relay. **Note:** when set via the `PORT` env var (not `-p`), `hbbr` listens on `PORT + 1`, so a shared `PORT=21116` makes `hbbs`=21116 and `hbbr`=21117. | -### Relay bandwidth / QoS (all 🅴, set as environment variables) +### Relay bandwidth / QoS These have no CLI flag and can also be changed through the `hbbr` [loopback console](#runtime-console) (`tb`, `sb`, `ls`, `dt`, `t`, …; send `h` @@ -123,7 +85,7 @@ capped to `LIMIT_SPEED` once its average throughput since it started exceeds downgraded even when the relay is otherwise idle. `TOTAL_BANDWIDTH` is a separate aggregate cap. -`hbbr` reads its `.env` verbatim, so these may also be placed in `.env` +These may also be placed in `.env` using the uppercase spellings shown above (e.g. `SINGLE_BANDWIDTH=256`). ### Blocklists / blacklists (files, not env vars) @@ -176,7 +138,7 @@ before launching the binary. A value in `.env` or `hbbs`'s `--config` file is loaded too late and has no effect on logging. ```bash -RUST_LOG=debug hbbs -r example.com +RUST_LOG=debug hbbs ``` --- @@ -194,11 +156,16 @@ The `KEY` / `-k` value can be: * **empty** — key validation is disabled. `hbbs` still loads or generates key files for signing but deliberately leaves its active validation key empty; `hbbr` neither loads nor generates a key. Both services then accept clients - without validating a key. Do not use an empty value on an exposed server. + without validating a key. This is the intentional `hbbr` default; use a + non-empty value when relay key validation is required. -By convention `-k _` is used to run an **encryption‑only** server (the official -supervisor Docker image exposes this as `ENCRYPTED_ONLY=1`). `hbbs` and `hbbr` -must be started with the **same non-empty** key. +`hbbs` defaults to `-`, so it already loads or generates a key pair without an +explicit `-k _`. `hbbr` intentionally defaults to an empty key to avoid +key-pair setup and mismatch failures. Leave it empty for the default mode +without key validation. To enable relay key validation, give it the same +non-empty key as `hbbs`; both services can reuse key material from a shared +working directory. The `_` value is not a stricter mode than `-` in the current +implementation. To supply your own key pair, place `id_ed25519` and `id_ed25519.pub` in the process's **current working directory** before first start. That directory may @@ -216,7 +183,7 @@ by `hbbs`/`hbbr` directly: | Variable | Default | Description | |---|---|---| | `RELAY` | `relay.example.com` | Passed to `hbbs` as `-r $RELAY` (your public address). | -| `ENCRYPTED_ONLY` | `0` | `1` adds `-k _` to both servers, forcing encryption. | +| `ENCRYPTED_ONLY` | `0` | `1` adds `-k _` to both servers. This is redundant for `hbbs`, whose default is `-`, and opts `hbbr` into key validation instead of its intentional empty default. | | `KEY_PUB` | *(unset)* | If set, written to `/data/id_ed25519.pub` on first start. | | `KEY_PRIV` | *(unset)* | If set, written to `/data/id_ed25519` on first start. Provide **both** `KEY_PUB` and `KEY_PRIV`, or neither. | @@ -231,28 +198,22 @@ binaries and does **not** implement `RELAY`, `ENCRYPTED_ONLY`, `KEY_PUB`, or ## Examples -### Command line +### Command line — non-standard ports ```bash -# ID server: relay clients to this host, LAN detection, force encryption -hbbs -r rustdesk.example.com:21117 --mask 192.168.0.0/16 -k _ - -# Relay server, same key -hbbr -k _ +# Tell clients where the relay listens because it is not using port 21117. +hbbs -p 22116 -r rustdesk.example.com:22117 +hbbr -p 22117 ``` ### `.env` file (working directory) ```ini -# Shared by both binaries. hbbr requires exact uppercase KEY and PORT names. -relay-servers=rustdesk.example.com:21117 -KEY=_ -PORT=21116 +# Non-standard ports shared by both binaries; hbbr listens on PORT+1. +relay-servers=rustdesk.example.com:22117 +PORT=22116 ``` -> Reminder: put the 🅴 variables (`DB_URL`, `TEST_HBBS`, `ALWAYS_USE_RELAY`, -> `MAX_DATABASE_CONNECTIONS`) in the real environment, not in the `hbbs` `.env`. - ### docker-compose ```yaml @@ -261,7 +222,6 @@ services: image: rustdesk/rustdesk-server-s6:latest environment: - RELAY=rustdesk.example.com:21117 - - ENCRYPTED_ONLY=1 - ALWAYS_USE_RELAY=Y - RUST_LOG=info - SINGLE_BANDWIDTH=256 @@ -282,7 +242,7 @@ services: [Service] Environment=ALWAYS_USE_RELAY=Y Environment=RUST_LOG=info -ExecStart=/usr/bin/hbbs -r rustdesk.example.com:21117 -k _ +ExecStart=/usr/bin/hbbs ``` --- diff --git a/src/common.rs b/src/common.rs index 9c76f41..80fc73a 100644 --- a/src/common.rs +++ b/src/common.rs @@ -7,10 +7,33 @@ use sodiumoxide::crypto::sign; use std::{ io::prelude::*, io::Read, - net::SocketAddr, + net::{IpAddr, SocketAddr}, time::{Instant, SystemTime}, }; +pub fn parse_bind_address(value: &str) -> Result> { + let value = value.trim(); + if value.is_empty() { + Ok(None) + } else { + value + .parse() + .with_context(|| format!("Invalid bind address: {value}")) + .map(Some) + } +} + +pub async fn listen_tcp( + bind_addr: Option, + port: u16, +) -> ResultType { + if let Some(bind_addr) = bind_addr { + hbb_common::tcp::new_listener(SocketAddr::new(bind_addr, port), true).await + } else { + hbb_common::tcp::listen_any(port).await + } +} + #[allow(dead_code)] pub(crate) fn get_expired_time() -> Instant { let now = Instant::now(); @@ -52,6 +75,12 @@ fn arg_name(name: &str) -> String { name.to_uppercase().replace('_', "-") } +#[allow(dead_code)] +#[inline] +pub fn set_arg(name: &str, value: &str) { + std::env::set_var(arg_name(name), value); +} + #[allow(dead_code)] pub fn init_args(args: &str, name: &str, about: &str) { let matches = App::new(name) @@ -64,7 +93,7 @@ pub fn init_args(args: &str, name: &str, about: &str) { if let Some(section) = v.section(None::) { section .iter() - .for_each(|(k, v)| std::env::set_var(arg_name(k), v)); + .for_each(|(k, v)| set_arg(k, v)); } } if let Some(config) = matches.value_of("config") { @@ -72,17 +101,42 @@ pub fn init_args(args: &str, name: &str, about: &str) { if let Some(section) = v.section(None::) { section .iter() - .for_each(|(k, v)| std::env::set_var(arg_name(k), v)); + .for_each(|(k, v)| set_arg(k, v)); } } } for (k, v) in matches.args { if let Some(v) = v.vals.first() { - std::env::set_var(arg_name(k), v.to_string_lossy().to_string()); + set_arg(k, &v.to_string_lossy()); } } } +#[allow(dead_code)] +pub fn get_arg_opt(name: &str) -> Option { + let dashed = arg_name(name); + let underscored = dashed.replace('-', "_"); + let lower_dashed = dashed.to_lowercase(); + let lower_underscored = underscored.to_lowercase(); + for alias in [&dashed, &underscored, &lower_dashed, &lower_underscored] { + if let Ok(value) = std::env::var(alias) { + return Some(value); + } + } + let mut aliases = std::env::vars_os() + .filter_map(|(key, value)| { + let key = key.into_string().ok()?; + if arg_name(&key) == dashed { + Some((key, value.into_string().ok()?)) + } else { + None + } + }) + .collect::>(); + aliases.sort_by(|a, b| a.0.cmp(&b.0)); + aliases.into_iter().next().map(|(_, value)| value) +} + #[allow(dead_code)] #[inline] pub fn get_arg(name: &str) -> String { @@ -92,7 +146,7 @@ pub fn get_arg(name: &str) -> String { #[allow(dead_code)] #[inline] pub fn get_arg_or(name: &str, default: String) -> String { - std::env::var(arg_name(name)).unwrap_or(default) + get_arg_opt(name).unwrap_or(default) } #[allow(dead_code)] @@ -215,4 +269,60 @@ async fn check_software_update_() -> hbb_common::ResultType<()> { log::info!("new version is available: {}", latest_release_version); } Ok(()) -} \ No newline at end of file +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::{Ipv4Addr, Ipv6Addr}; + + #[test] + fn argument_names_ignore_case_and_separator() { + let aliases = [ + "RUSTDESK-CONFIG-ALIAS-TEST", + "RUSTDESK_CONFIG_ALIAS_TEST", + "rustdesk-config-alias-test", + "rustdesk_config_alias_test", + "RustDesk_Config-Alias_Test", + ]; + for alias in aliases { + std::env::remove_var(alias); + } + for alias in aliases { + std::env::set_var(alias, alias); + assert_eq!(get_arg("RUSTDESK_CONFIG_ALIAS_TEST"), alias); + std::env::remove_var(alias); + } + set_arg("rustdesk_config_alias_test", "normalized"); + assert_eq!( + std::env::var("RUSTDESK-CONFIG-ALIAS-TEST").unwrap(), + "normalized" + ); + std::env::set_var("RUSTDESK_CONFIG_ALIAS_TEST", "inherited"); + set_arg("rustdesk-config-alias-test", "higher-priority"); + assert_eq!(get_arg("rustdesk_config_alias_test"), "higher-priority"); + std::env::remove_var("RUSTDESK-CONFIG-ALIAS-TEST"); + std::env::remove_var("RUSTDESK_CONFIG_ALIAS_TEST"); + } + + #[test] + fn parses_bind_address() { + assert_eq!(parse_bind_address("").unwrap(), None); + assert_eq!( + parse_bind_address("127.0.0.1").unwrap(), + Some(IpAddr::V4(Ipv4Addr::LOCALHOST)) + ); + assert_eq!( + parse_bind_address("::1").unwrap(), + Some(IpAddr::V6(Ipv6Addr::LOCALHOST)) + ); + assert!(parse_bind_address("not-an-ip").is_err()); + } + + #[hbb_common::tokio::test] + async fn tcp_listener_uses_bind_address() { + let bind_addr = IpAddr::V4(Ipv4Addr::LOCALHOST); + let listener = listen_tcp(Some(bind_addr), 0).await.unwrap(); + assert_eq!(listener.local_addr().unwrap().ip(), bind_addr); + } +} diff --git a/src/database.rs b/src/database.rs index fa1b6ed..1ea016e 100644 --- a/src/database.rs +++ b/src/database.rs @@ -51,8 +51,7 @@ impl Database { if !std::path::Path::new(url).exists() { std::fs::File::create(url).ok(); } - let n: usize = std::env::var("MAX_DATABASE_CONNECTIONS") - .unwrap_or_else(|_| "1".to_owned()) + let n: usize = crate::common::get_arg_or("MAX_DATABASE_CONNECTIONS", "1".to_owned()) .parse() .unwrap_or(1); log::debug!("MAX_DATABASE_CONNECTIONS={}", n); diff --git a/src/hbbr.rs b/src/hbbr.rs index 4c8a784..e6763fa 100644 --- a/src/hbbr.rs +++ b/src/hbbr.rs @@ -13,7 +13,8 @@ fn main() -> ResultType<()> { .write_mode(WriteMode::Async) .start()?; let args = format!( - "-p, --port=[NUMBER(default={RELAY_PORT})] 'Sets the listening port' + "-b, --bind=[IP] 'Sets the IP address to bind to (default: all interfaces)' + -p, --port=[NUMBER(default={RELAY_PORT})] 'Sets the listening port' -k, --key=[KEY] 'Only allow the client with the same key' ", ); @@ -25,21 +26,29 @@ fn main() -> ResultType<()> { .get_matches(); if let Ok(v) = ini::Ini::load_from_file(".env") { if let Some(section) = v.section(None::) { - section.iter().for_each(|(k, v)| std::env::set_var(k, v)); + section.iter().for_each(|(k, v)| common::set_arg(k, v)); } } let mut port = RELAY_PORT; - if let Ok(v) = std::env::var("PORT") { + if let Some(v) = common::get_arg_opt("PORT") { let v: i32 = v.parse().unwrap_or_default(); if v > 0 { port = v + 1; } } - start( + let bind = matches + .value_of("bind") + .map(str::to_owned) + .unwrap_or_else(|| common::get_arg("BIND")); + let bind_addr = common::parse_bind_address(&bind)?; + let key = matches + .value_of("key") + .map(str::to_owned) + .unwrap_or_else(|| common::get_arg("KEY")); + start_with_bind( + bind_addr, matches.value_of("port").unwrap_or(&port.to_string()), - matches - .value_of("key") - .unwrap_or(&std::env::var("KEY").unwrap_or_default()), + &key, )?; Ok(()) } diff --git a/src/main.rs b/src/main.rs index 8a0ff24..d422c3c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,13 +15,14 @@ fn main() -> ResultType<()> { .start()?; let args = format!( "-c --config=[FILE] +takes_value 'Sets a custom config file' + -b, --bind=[IP] 'Sets the IP address to bind to (default: all interfaces)' -p, --port=[NUMBER(default={RENDEZVOUS_PORT})] 'Sets the listening port' - -s, --serial=[NUMBER(default=0)] 'Sets configure update serial number' - -R, --rendezvous-servers=[HOSTS] 'Sets rendezvous servers, separated by comma' - -u, --software-url=[URL] 'Sets download url of RustDesk software of newest version' + -s, --serial=[NUMBER(default=0)] '[DEPRECATED] Sets configure update serial number' + -R, --rendezvous-servers=[HOSTS] '[DEPRECATED] Sets rendezvous servers, separated by comma' + -u, --software-url=[URL] '[DEPRECATED] Sets download url of RustDesk software of newest version' -r, --relay-servers=[HOST] 'Sets the default relay servers, separated by comma' -M, --rmem=[NUMBER(default={RMEM})] 'Sets UDP recv buffer size, set system rmem_max first, e.g., sudo sysctl -w net.core.rmem_max=52428800. vi /etc/sysctl.conf, net.core.rmem_max=52428800, sudo sysctl –p' - , --mask=[MASK] 'Determine if the connection comes from LAN, e.g. 192.168.0.0/16' + , --mask=[MASK] '[DEPRECATED] Determine if the connection comes from LAN, e.g. 192.168.0.0/16' -k, --key=[KEY] 'Only allow the client with the same key'", ); init_args(&args, "hbbs", "RustDesk ID/Rendezvous Server"); @@ -29,9 +30,16 @@ fn main() -> ResultType<()> { if port < 3 { bail!("Invalid port"); } + let bind_addr = parse_bind_address(&get_arg("bind"))?; let rmem = get_arg("rmem").parse::().unwrap_or(RMEM); let serial: i32 = get_arg("serial").parse().unwrap_or(0); crate::common::check_software_update(); - RendezvousServer::start(port, serial, &get_arg_or("key", "-".to_owned()), rmem)?; + RendezvousServer::start_with_bind( + bind_addr, + port, + serial, + &get_arg_or("key", "-".to_owned()), + rmem, + )?; Ok(()) } diff --git a/src/peer.rs b/src/peer.rs index 4ca87cf..0f218df 100644 --- a/src/peer.rs +++ b/src/peer.rs @@ -67,7 +67,7 @@ pub(crate) struct PeerMap { impl PeerMap { pub(crate) async fn new() -> ResultType { - let db = std::env::var("DB_URL").unwrap_or({ + let db = get_arg_opt("DB_URL").unwrap_or_else(|| { let mut db = "db_v2.sqlite3".to_owned(); #[cfg(all(windows, not(debug_assertions)))] { diff --git a/src/relay_server.rs b/src/relay_server.rs index 72c21e4..70c3d3c 100644 --- a/src/relay_server.rs +++ b/src/relay_server.rs @@ -8,7 +8,7 @@ use hbb_common::{ protobuf::Message as _, rendezvous_proto::*, sleep, - tcp::{listen_any, FramedStream}, + tcp::FramedStream, timeout, tokio::{ self, @@ -24,7 +24,7 @@ use std::{ collections::{HashMap, HashSet}, io::prelude::*, io::Error, - net::SocketAddr, + net::{IpAddr, SocketAddr}, sync::atomic::{AtomicUsize, Ordering}, }; @@ -46,7 +46,11 @@ const BLACKLIST_FILE: &str = "blacklist.txt"; const BLOCKLIST_FILE: &str = "blocklist.txt"; #[tokio::main(flavor = "multi_thread")] -pub async fn start(port: &str, key: &str) -> ResultType<()> { +pub async fn start_with_bind( + bind_addr: Option, + port: &str, + key: &str, +) -> ResultType<()> { let key = get_server_sk(key); if let Ok(mut file) = std::fs::File::open(BLACKLIST_FILE) { let mut contents = String::new(); @@ -85,7 +89,12 @@ pub async fn start(port: &str, key: &str) -> ResultType<()> { let main_task = async move { loop { log::info!("Start"); - io_loop(listen_any(port).await?, listen_any(port2).await?, &key).await; + io_loop( + crate::common::listen_tcp(bind_addr, port).await?, + crate::common::listen_tcp(bind_addr, port2).await?, + &key, + ) + .await; } }; let listen_signal = crate::common::listen_signal(); @@ -96,8 +105,8 @@ pub async fn start(port: &str, key: &str) -> ResultType<()> { } fn check_params() { - let tmp = std::env::var("DOWNGRADE_THRESHOLD") - .map(|x| x.parse::().unwrap_or(0.)) + let tmp = crate::common::get_arg("DOWNGRADE_THRESHOLD") + .parse::() .unwrap_or(0.); if tmp > 0. { DOWNGRADE_THRESHOLD_100.store((tmp * 100.) as _, Ordering::SeqCst); @@ -106,8 +115,8 @@ fn check_params() { "DOWNGRADE_THRESHOLD: {}", DOWNGRADE_THRESHOLD_100.load(Ordering::SeqCst) as f64 / 100. ); - let tmp = std::env::var("DOWNGRADE_START_CHECK") - .map(|x| x.parse::().unwrap_or(0)) + let tmp = crate::common::get_arg("DOWNGRADE_START_CHECK") + .parse::() .unwrap_or(0); if tmp > 0 { DOWNGRADE_START_CHECK.store(tmp * 1000, Ordering::SeqCst); @@ -116,8 +125,8 @@ fn check_params() { "DOWNGRADE_START_CHECK: {}s", DOWNGRADE_START_CHECK.load(Ordering::SeqCst) / 1000 ); - let tmp = std::env::var("LIMIT_SPEED") - .map(|x| x.parse::().unwrap_or(0.)) + let tmp = crate::common::get_arg("LIMIT_SPEED") + .parse::() .unwrap_or(0.); if tmp > 0. { LIMIT_SPEED.store((tmp * 1024. * 1024.) as usize, Ordering::SeqCst); @@ -126,8 +135,8 @@ fn check_params() { "LIMIT_SPEED: {}Mb/s", LIMIT_SPEED.load(Ordering::SeqCst) as f64 / 1024. / 1024. ); - let tmp = std::env::var("TOTAL_BANDWIDTH") - .map(|x| x.parse::().unwrap_or(0.)) + let tmp = crate::common::get_arg("TOTAL_BANDWIDTH") + .parse::() .unwrap_or(0.); if tmp > 0. { TOTAL_BANDWIDTH.store((tmp * 1024. * 1024.) as usize, Ordering::SeqCst); @@ -137,8 +146,8 @@ fn check_params() { "TOTAL_BANDWIDTH: {}Mb/s", TOTAL_BANDWIDTH.load(Ordering::SeqCst) as f64 / 1024. / 1024. ); - let tmp = std::env::var("SINGLE_BANDWIDTH") - .map(|x| x.parse::().unwrap_or(0.)) + let tmp = crate::common::get_arg("SINGLE_BANDWIDTH") + .parse::() .unwrap_or(0.); if tmp > 0. { SINGLE_BANDWIDTH.store((tmp * 1024. * 1024.) as usize, Ordering::SeqCst); diff --git a/src/rendezvous_server.rs b/src/rendezvous_server.rs index 0de0878..b194ae0 100644 --- a/src/rendezvous_server.rs +++ b/src/rendezvous_server.rs @@ -16,7 +16,7 @@ use hbb_common::{ register_pk_response::Result::{TOO_FREQUENT, UUID_MISMATCH}, *, }, - tcp::{listen_any, FramedStream}, + tcp::FramedStream, timeout, tokio::{ self, @@ -98,18 +98,25 @@ enum LoopFailure { } impl RendezvousServer { + pub fn start(port: i32, serial: i32, key: &str, rmem: usize) -> ResultType<()> { + Self::start_with_bind(None, port, serial, key, rmem) + } + #[tokio::main(flavor = "multi_thread")] - pub async fn start(port: i32, serial: i32, key: &str, rmem: usize) -> ResultType<()> { + pub async fn start_with_bind( + bind_addr: Option, + port: i32, + serial: i32, + key: &str, + rmem: usize, + ) -> ResultType<()> { let (key, sk) = Self::get_server_sk(key); let nat_port = port - 1; let ws_port = port + 2; let pm = PeerMap::new().await?; log::info!("serial={}", serial); let rendezvous_servers = get_servers(&get_arg("rendezvous-servers"), "rendezvous-servers"); - log::info!("Listening on tcp/udp :{}", port); - log::info!("Listening on tcp :{}, extra port for NAT test", nat_port); - log::info!("Listening on websocket :{}", ws_port); - let mut socket = create_udp_listener(port, rmem).await?; + let mut socket = create_udp_listener(bind_addr, port, rmem).await?; let (tx, mut rx) = mpsc::unbounded_channel::(); let software_url = get_arg("software-url"); let version = hbb_common::get_version_from_url(&software_url); @@ -147,15 +154,17 @@ impl RendezvousServer { log::info!("local-ip: {:?}", rs.inner.local_ip); std::env::set_var("PORT_FOR_API", port.to_string()); rs.parse_relay_servers(&get_arg("relay-servers")); - let mut listener = create_tcp_listener(port).await?; - let mut listener2 = create_tcp_listener(nat_port).await?; - let mut listener3 = create_tcp_listener(ws_port).await?; - let test_addr = std::env::var("TEST_HBBS").unwrap_or_default(); - if std::env::var("ALWAYS_USE_RELAY") - .unwrap_or_default() - .to_uppercase() - == "Y" - { + let mut listener = create_tcp_listener(bind_addr, port).await?; + let mut listener2 = create_tcp_listener(bind_addr, nat_port).await?; + let mut listener3 = create_tcp_listener(bind_addr, ws_port).await?; + log::info!("Listening on tcp/udp {}", listener.local_addr()?); + log::info!( + "Listening on tcp {}, extra port for NAT test", + listener2.local_addr()? + ); + log::info!("Listening on websocket {}", listener3.local_addr()?); + let test_addr = get_arg("TEST_HBBS"); + if get_arg("ALWAYS_USE_RELAY").to_uppercase() == "Y" { ALWAYS_USE_RELAY.store(true, Ordering::SeqCst); } log::info!( @@ -204,19 +213,19 @@ impl RendezvousServer { { LoopFailure::UdpSocket => { drop(socket); - socket = create_udp_listener(port, rmem).await?; + socket = create_udp_listener(bind_addr, port, rmem).await?; } LoopFailure::Listener => { drop(listener); - listener = create_tcp_listener(port).await?; + listener = create_tcp_listener(bind_addr, port).await?; } LoopFailure::Listener2 => { drop(listener2); - listener2 = create_tcp_listener(nat_port).await?; + listener2 = create_tcp_listener(bind_addr, nat_port).await?; } LoopFailure::Listener3 => { drop(listener3); - listener3 = create_tcp_listener(ws_port).await?; + listener3 = create_tcp_listener(bind_addr, ws_port).await?; } } } @@ -1352,7 +1361,15 @@ async fn send_rk_res( socket.send(&msg_out, addr).await } -async fn create_udp_listener(port: i32, rmem: usize) -> ResultType { +async fn create_udp_listener( + bind_addr: Option, + port: i32, + rmem: usize, +) -> ResultType { + if let Some(bind_addr) = bind_addr { + let addr = SocketAddr::new(bind_addr, port as _); + return FramedSocket::new_reuse(&addr, true, rmem).await; + } let addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), port as _); if let Ok(s) = FramedSocket::new_reuse(&addr, true, rmem).await { log::debug!("listen on udp {:?}", s.local_addr()); @@ -1365,8 +1382,20 @@ async fn create_udp_listener(port: i32, rmem: usize) -> ResultType } #[inline] -async fn create_tcp_listener(port: i32) -> ResultType { - let s = listen_any(port as _).await?; +async fn create_tcp_listener(bind_addr: Option, port: i32) -> ResultType { + let s = listen_tcp(bind_addr, port as _).await?; log::debug!("listen on tcp {:?}", s.local_addr()); Ok(s) } + +#[cfg(test)] +mod tests { + use super::*; + + #[hbb_common::tokio::test] + async fn udp_listener_uses_bind_address() { + let bind_addr = IpAddr::V4(Ipv4Addr::LOCALHOST); + let socket = create_udp_listener(Some(bind_addr), 0, 0).await.unwrap(); + assert_eq!(socket.local_addr().unwrap().ip(), bind_addr); + } +}