From aca48ddfffc529fe85e5b839b50d9e700bbd8ead Mon Sep 17 00:00:00 2001 From: Shirom Makkad Date: Thu, 20 Jun 2024 21:25:49 +0000 Subject: [PATCH] Add parsing tokens from files or environment variables --- Cargo.lock | 60 ++++++- Cargo.toml | 3 + README.md | 16 ++ src/config.rs | 165 ++++++++++++++---- src/config_watcher.rs | 14 +- src/helper.rs | 13 +- src/main.rs | 3 +- src/transport/tcp.rs | 6 +- tests/common/mod.rs | 8 +- .../config_test/valid_config/token_file.toml | 47 +++++ tests/integration_test.rs | 10 +- tests/token | 1 + 12 files changed, 277 insertions(+), 69 deletions(-) create mode 100644 tests/config_test/valid_config/token_file.toml create mode 100644 tests/token diff --git a/Cargo.lock b/Cargo.lock index 95fb8ac4..2b435c17 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -681,6 +681,7 @@ checksum = "23342abe12aba583913b2e62f22225ff9c950774065e4bfb61a19cd9770fec40" dependencies = [ "futures-channel", "futures-core", + "futures-executor", "futures-io", "futures-sink", "futures-task", @@ -703,6 +704,17 @@ version = "0.3.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4bca583b7e26f571124fe5b7561d49cb2868d79116cfa0eefce955557c6fee8c" +[[package]] +name = "futures-executor" +version = "0.3.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccecee823288125bd88b4d7f565c9e58e41858e47ab72e8ea2d64e93624386e0" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.28" @@ -738,10 +750,13 @@ version = "0.3.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26b01e40b772d54cf6c6d721c1d1abd0647a0106a12ecaa1c186273392a69533" dependencies = [ + "futures-channel", "futures-core", + "futures-io", "futures-macro", "futures-sink", "futures-task", + "memchr", "pin-project-lite", "pin-utils", "slab", @@ -1282,9 +1297,9 @@ dependencies = [ [[package]] name = "once_cell" -version = "1.18.0" +version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d" +checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" [[package]] name = "opaque-debug" @@ -1614,6 +1629,7 @@ dependencies = [ "rustls-native-certs", "rustls-pemfile", "serde", + "serial_test", "sha2", "snowstorm", "socket2 0.4.9", @@ -1808,6 +1824,15 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "scc" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76ad2bbb0ae5100a07b7a6f2ed7ab5fd0045551a4c507989b7a620046ea3efdc" +dependencies = [ + "sdd", +] + [[package]] name = "schannel" version = "0.1.22" @@ -1823,6 +1848,12 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sdd" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b84345e4c9bd703274a082fb80caaa99b7612be48dfaa1dd9266577ec412309d" + [[package]] name = "security-framework" version = "2.9.2" @@ -1883,6 +1914,31 @@ dependencies = [ "serde", ] +[[package]] +name = "serial_test" +version = "3.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b4b487fe2acf240a021cf57c6b2b4903b1e78ca0ecd862a71b71d2a51fed77d" +dependencies = [ + "futures", + "log", + "once_cell", + "parking_lot", + "scc", + "serial_test_derive", +] + +[[package]] +name = "serial_test_derive" +version = "3.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82fe9db325bcef1fbcde82e078a5cc4efdf787e96b3b9cf45b50b529f2083d67" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.37", +] + [[package]] name = "sha1" version = "0.10.6" diff --git a/Cargo.toml b/Cargo.toml index e7a944ec..1738f2e1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -136,3 +136,6 @@ vergen = { version = "7.4.2", default-features = false, features = [ "cargo", ] } anyhow = "1.0" + +[dev-dependencies] +serial_test = "3.1.1" diff --git a/README.md b/README.md index 25cf97f1..7bdca9c6 100644 --- a/README.md +++ b/README.md @@ -101,12 +101,25 @@ Before heading to the full configuration specification, it's recommend to skim [ See [Transport](./docs/transport.md) for more details about encryption and the `transport` block. +Tokens can also be set through environment variables. The variable `RATHOLE_{service name in uppercase}_TOKEN` can be set or `RATHOLE_DEFAULT_TOKEN` for all services. +Tokens are parsed in the following order for "servicex": +1. (client/server).services.servicex.token +2. (client/server).services.servicex.token_file +3. RATHOLE_SERVICEX_TOKEN +4. (client/server).default_token +5. (client/server).default_token_file +6. RATHOLE_DEFAULT_TOKEN + +Tokens should be generated by yourself (not on someone's website or on random.com) using a cryptographic pseudorandom generator. On Linux, use `openssl rand -hex 64 > /path/to/key`. Make sure to do this on a system with high entropy. +Most systems will have plenty of entropy. The random network delay between packets, using the computer and typing, access latency from your hdd all can be used to create entropy. Just use your system for anything other than extremely repetitive tasks and don't generate the key right after boot. + Here is the full configuration specification: ```toml [client] remote_addr = "example.com:2333" # Necessary. The address of the server default_token = "default_token_if_not_specify" # Optional. The default token of services, if they don't define their own ones +default_token_file = "/path/to/token" # Optional. This will pull the default token from the path specified heartbeat_timeout = 40 # Optional. Set to 0 to disable the application-layer heartbeat test. The value must be greater than `server.heartbeat_interval`. Default: 40 seconds retry_interval = 1 # Optional. The interval between retry to connect to the server. Default: 1 second @@ -134,6 +147,7 @@ tls = true # If `true` then it will use settings in `client.transport.tls` [client.services.service1] # A service that needs forwarding. The name `service1` can change arbitrarily, as long as identical to the name in the server's configuration type = "tcp" # Optional. The protocol that needs forwarding. Possible values: ["tcp", "udp"]. Default: "tcp" token = "whatever" # Necessary if `client.default_token` not set +token_file = "/path/to/token" # Necessary if token, default_token, the env var, and default_token_file are unset. local_addr = "127.0.0.1:1081" # Necessary. The address of the service that needs to be forwarded nodelay = true # Optional. Override the `client.transport.nodelay` per service retry_interval = 1 # Optional. The interval between retry to connect to the server. Default: inherits the global config @@ -144,6 +158,7 @@ local_addr = "127.0.0.1:1082" [server] bind_addr = "0.0.0.0:2333" # Necessary. The address that the server listens for clients. Generally only the port needs to be change. default_token = "default_token_if_not_specify" # Optional +default_token_file = "/path/to/token" # Optional. This will pull the default token from the path specified heartbeat_interval = 30 # Optional. The interval between two application-layer heartbeat. Set to 0 to disable sending heartbeat. Default: 30 seconds [server.transport] # Same as `[client.transport]` @@ -169,6 +184,7 @@ tls = true # If `true` then it will use settings in `server.transport.tls` [server.services.service1] # The service name must be identical to the client side type = "tcp" # Optional. Same as the client `[client.services.X.type] token = "whatever" # Necessary if `server.default_token` not set +token_file = "/path/to/token" # Necessary if token, default_token, and default_token_file are unset. bind_addr = "0.0.0.0:8081" # Necessary. The address of the service is exposed at. Generally only the port needs to be change. nodelay = true # Optional. Same as the client diff --git a/src/config.rs b/src/config.rs index ca85fc20..0fdeba20 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,7 @@ use anyhow::{anyhow, bail, Context, Result}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +use std::env; use std::fmt::{Debug, Formatter}; use std::ops::Deref; use std::path::Path; @@ -64,6 +65,7 @@ pub struct ClientServiceConfig { pub name: String, pub local_addr: String, pub token: Option, + pub token_file: Option, pub nodelay: Option, pub retry_interval: Option, } @@ -101,6 +103,7 @@ pub struct ServerServiceConfig { pub name: String, pub bind_addr: String, pub token: Option, + pub token_file: Option, pub nodelay: Option, } @@ -201,6 +204,7 @@ fn default_client_retry_interval() -> u64 { pub struct ClientConfig { pub remote_addr: String, pub default_token: Option, + pub default_token_file: Option, pub services: HashMap, #[serde(default)] pub transport: TransportConfig, @@ -219,6 +223,7 @@ fn default_heartbeat_interval() -> u64 { pub struct ServerConfig { pub bind_addr: String, pub default_token: Option, + pub default_token_file: Option, pub services: HashMap, #[serde(default)] pub transport: TransportConfig, @@ -234,15 +239,15 @@ pub struct Config { } impl Config { - fn from_str(s: &str) -> Result { + async fn from_str(s: &str) -> Result { let mut config: Config = toml::from_str(s).with_context(|| "Failed to parse the config")?; if let Some(server) = config.server.as_mut() { - Config::validate_server_config(server)?; + Config::validate_server_config(server).await?; } if let Some(client) = config.client.as_mut() { - Config::validate_client_config(client)?; + Config::validate_client_config(client).await?; } if config.server.is_none() && config.client.is_none() { @@ -252,15 +257,48 @@ impl Config { } } - fn validate_server_config(server: &mut ServerConfig) -> Result<()> { + async fn parse_token( + name: &str, + token: &Option, + token_file: &Option, + default_token: &Option, + ) -> Option { + if token.is_some() { + return token.clone(); + } + + if let Some(v) = token_file { + return fs::read_to_string(v).await.ok().map(MaskedString); + } + + if default_token.is_some() { + return default_token.clone(); + } + + if let Ok(v) = env::var(format!("RATHOLE_{}_TOKEN", name.to_uppercase())) { + return Some(MaskedString(v)); + } + + None + } + + async fn validate_server_config(server: &mut ServerConfig) -> Result<()> { + let default_token = Self::parse_token( + "default", + &server.default_token, + &server.default_token_file, + &None, + ) + .await; + // Validate services for (name, s) in &mut server.services { - s.name = name.clone(); + s.name.clone_from(name); + s.token = + Self::parse_token(name.as_str(), &s.token, &s.token_file, &default_token).await; + if s.token.is_none() { - s.token = server.default_token.clone(); - if s.token.is_none() { - bail!("The token of service {} is not set", name); - } + bail!("The token of service {} is not set", name); } } @@ -269,15 +307,23 @@ impl Config { Ok(()) } - fn validate_client_config(client: &mut ClientConfig) -> Result<()> { + async fn validate_client_config(client: &mut ClientConfig) -> Result<()> { + let default_token = Self::parse_token( + "default", + &client.default_token, + &client.default_token_file, + &None, + ) + .await; + // Validate services for (name, s) in &mut client.services { - s.name = name.clone(); + s.name.clone_from(name); + s.token = + Self::parse_token(name.as_str(), &s.token, &s.token_file, &default_token).await; + if s.token.is_none() { - s.token = client.default_token.clone(); - if s.token.is_none() { - bail!("The token of service {} is not set", name); - } + bail!("The token of service {} is not set", name); } if s.retry_interval.is_none() { s.retry_interval = Some(client.retry_interval); @@ -327,7 +373,7 @@ impl Config { let s: String = fs::read_to_string(path) .await .with_context(|| format!("Failed to read the config {:?}", path))?; - Config::from_str(&s).with_context(|| { + Config::from_str(&s).await.with_context(|| { "Configuration is invalid. Please refer to the configuration specification." }) } @@ -336,9 +382,12 @@ impl Config { #[cfg(test)] mod tests { use super::*; - use std::{fs, path::PathBuf}; + use std::fs; + use std::path::PathBuf; use anyhow::Result; + use serial_test::{parallel, serial}; + use tokio::runtime::Runtime; fn list_config_files>(root: T) -> Result> { let mut files = Vec::new(); @@ -361,38 +410,42 @@ mod tests { .collect()) } - #[test] - fn test_example_config() -> Result<()> { + #[tokio::test] + #[parallel] + async fn test_example_config() -> Result<()> { let paths = get_all_example_config()?; for p in paths { let s = fs::read_to_string(p)?; - Config::from_str(&s)?; + Config::from_str(&s).await?; } Ok(()) } - #[test] - fn test_valid_config() -> Result<()> { + #[tokio::test] + #[parallel] + async fn test_valid_config() -> Result<()> { let paths = list_config_files("tests/config_test/valid_config")?; for p in paths { let s = fs::read_to_string(p)?; - Config::from_str(&s)?; + Config::from_str(&s).await?; } Ok(()) } - #[test] - fn test_invalid_config() -> Result<()> { + #[tokio::test] + #[parallel] + async fn test_invalid_config() -> Result<()> { let paths = list_config_files("tests/config_test/invalid_config")?; for p in paths { let s = fs::read_to_string(p)?; - assert!(Config::from_str(&s).is_err()); + assert!(Config::from_str(&s).await.is_err()); } Ok(()) } - #[test] - fn test_validate_server_config() -> Result<()> { + #[tokio::test] + #[parallel] + async fn test_validate_server_config() -> Result<()> { let mut cfg = ServerConfig::default(); cfg.services.insert( @@ -407,11 +460,11 @@ mod tests { ); // Missing the token - assert!(Config::validate_server_config(&mut cfg).is_err()); + assert!(Config::validate_server_config(&mut cfg).await.is_err()); // Use the default token cfg.default_token = Some("123".into()); - assert!(Config::validate_server_config(&mut cfg).is_ok()); + assert!(Config::validate_server_config(&mut cfg).await.is_ok()); assert_eq!( cfg.services .get("foo1") @@ -426,7 +479,7 @@ mod tests { // The default token won't override the service token cfg.services.get_mut("foo1").unwrap().token = Some("4".into()); - assert!(Config::validate_server_config(&mut cfg).is_ok()); + assert!(Config::validate_server_config(&mut cfg).await.is_ok()); assert_eq!( cfg.services .get("foo1") @@ -441,8 +494,9 @@ mod tests { Ok(()) } - #[test] - fn test_validate_client_config() -> Result<()> { + #[tokio::test] + #[parallel] + async fn test_validate_client_config() -> Result<()> { let mut cfg = ClientConfig::default(); cfg.services.insert( @@ -457,11 +511,12 @@ mod tests { ); // Missing the token - assert!(Config::validate_client_config(&mut cfg).is_err()); + println!("{:?}", env::var("DEFAULT_TOKEN").ok()); + assert!(Config::validate_client_config(&mut cfg).await.is_err()); // Use the default token cfg.default_token = Some("123".into()); - assert!(Config::validate_client_config(&mut cfg).is_ok()); + assert!(Config::validate_client_config(&mut cfg).await.is_ok()); assert_eq!( cfg.services .get("foo1") @@ -476,7 +531,7 @@ mod tests { // The default token won't override the service token cfg.services.get_mut("foo1").unwrap().token = Some("4".into()); - assert!(Config::validate_client_config(&mut cfg).is_ok()); + assert!(Config::validate_client_config(&mut cfg).await.is_ok()); assert_eq!( cfg.services .get("foo1") @@ -490,4 +545,42 @@ mod tests { ); Ok(()) } + + #[serial(env_default_token)] + fn read_from_env_var() { + let mut cfg = ClientConfig::default(); + + cfg.services.insert( + "foo1".into(), + ClientServiceConfig { + service_type: ServiceType::Tcp, + name: "foo1".into(), + local_addr: "127.0.0.1:80".into(), + token: None, + ..Default::default() + }, + ); + + env::set_var("RATHOLE_DEFAULT_TOKEN", "test-token"); + + // Can't .await with tokio::test while env vars are set. There must be a block surrounding the futures. + let rt = Runtime::new().unwrap(); + rt.block_on(async { + Config::validate_client_config(&mut cfg).await.unwrap(); + }); + assert_eq!( + cfg.services + .get("foo1") + .as_ref() + .unwrap() + .token + .as_ref() + .unwrap() + .0 + .as_str(), + "test-token" + ); + + env::remove_var("RATHOLE_DEFAULT_TOKEN"); + } } diff --git a/src/config_watcher.rs b/src/config_watcher.rs index 993fdcce..61a6a2d5 100644 --- a/src/config_watcher.rs +++ b/src/config_watcher.rs @@ -1,13 +1,9 @@ -use crate::{ - config::{ClientConfig, ClientServiceConfig, ServerConfig, ServerServiceConfig}, - Config, -}; +use crate::config::{ClientConfig, ClientServiceConfig, ServerConfig, ServerServiceConfig}; +use crate::Config; use anyhow::{Context, Result}; -use std::{ - collections::HashMap, - env, - path::{Path, PathBuf}, -}; +use std::collections::HashMap; +use std::env; +use std::path::{Path, PathBuf}; use tokio::sync::{broadcast, mpsc}; use tracing::{error, info, instrument}; diff --git a/src/helper.rs b/src/helper.rs index a292969f..37669122 100644 --- a/src/helper.rs +++ b/src/helper.rs @@ -1,13 +1,14 @@ use anyhow::{anyhow, Context, Result}; use async_http_proxy::{http_connect_tokio, http_connect_tokio_with_basic_auth}; -use backoff::{backoff::Backoff, Notify}; +use backoff::backoff::Backoff; +use backoff::Notify; use socket2::{SockRef, TcpKeepalive}; -use std::{future::Future, net::SocketAddr, time::Duration}; +use std::future::Future; +use std::net::SocketAddr; +use std::time::Duration; use tokio::io::{AsyncWrite, AsyncWriteExt}; -use tokio::{ - net::{lookup_host, TcpStream, ToSocketAddrs, UdpSocket}, - sync::broadcast, -}; +use tokio::net::{lookup_host, TcpStream, ToSocketAddrs, UdpSocket}; +use tokio::sync::broadcast; use tracing::trace; use url::Url; diff --git a/src/main.rs b/src/main.rs index 92ab75c1..f4a45266 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,8 @@ use anyhow::Result; use clap::Parser; use rathole::{run, Cli}; -use tokio::{signal, sync::broadcast}; +use tokio::signal; +use tokio::sync::broadcast; use tracing_subscriber::EnvFilter; #[tokio::main] diff --git a/src/transport/tcp.rs b/src/transport/tcp.rs index 3c5e242e..4ee905c8 100644 --- a/src/transport/tcp.rs +++ b/src/transport/tcp.rs @@ -1,7 +1,5 @@ -use crate::{ - config::{TcpConfig, TransportConfig}, - helper::tcp_connect_with_proxy, -}; +use crate::config::{TcpConfig, TransportConfig}; +use crate::helper::tcp_connect_with_proxy; use super::{AddrMaybeCached, SocketOpts, Transport}; use anyhow::Result; diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 9e59f92d..af353d24 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -1,11 +1,9 @@ use std::path::PathBuf; use anyhow::Result; -use tokio::{ - io::{self, AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream, ToSocketAddrs}, - sync::broadcast, -}; +use tokio::io::{self, AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream, ToSocketAddrs}; +use tokio::sync::broadcast; pub const PING: &str = "ping"; pub const PONG: &str = "pong"; diff --git a/tests/config_test/valid_config/token_file.toml b/tests/config_test/valid_config/token_file.toml new file mode 100644 index 00000000..1937c4e1 --- /dev/null +++ b/tests/config_test/valid_config/token_file.toml @@ -0,0 +1,47 @@ +[client] +remote_addr = "example.com:2333" # Necessary. The address of the server +default_token_file = "tests/token" # Optional. The file that stores the token. Can be stored as plain text in the config and on each service + +[client.transport] +type = "tcp" # Optional. Possible values: ["tcp", "tls"]. Default: "tcp" + +[client.transport.tls] # Necessary if `type` is "tls" +trusted_root = "ca.pem" # Necessary. The certificate of CA that signed the server's certificate +hostname = "example.com" # Optional. The hostname that the client uses to validate the certificate. If not set, fallback to `client.remote_addr` + +[client.transport.noise] # Noise protocol. See `docs/transport.md` for further explanation +pattern = "Noise_NK_25519_ChaChaPoly_BLAKE2s" # Optional. Default value as shown +local_private_key = "key_encoded_in_base64" # Optional +remote_public_key = "key_encoded_in_base64" # Optional + +[client.services.service1] # A service that needs forwarding. The name `service1` can change arbitrarily, as long as identical to the name in the server's configuration +type = "tcp" # Optional. The protocol that needs forwarding. Possible values: ["tcp", "udp"]. Default: "tcp" +token = "whatever" # Necessary if `client.default_token` not set +local_addr = "127.0.0.1:1081" # Necessary. The address of the service that needs to be forwarded + +[client.services.service2] # Multiple services can be defined +local_addr = "127.0.0.1:1082" + +[server] +bind_addr = "0.0.0.0:2333" # Necessary. The address that the server listens for clients. Generally only the port needs to be change. +default_token = "default_token_if_not_specify" # Optional + +[server.transport] +type = "tcp" # Same as `[client.transport]` + +[server.transport.tls] # Necessary if `type` is "tls" +pkcs12 = "identify.pfx" # Necessary. pkcs12 file of server's certificate and private key +pkcs12_password = "password" # Necessary. Password of the pkcs12 file + +[server.transport.noise] # Same as `[client.transport.noise]` +pattern = "Noise_NK_25519_ChaChaPoly_BLAKE2s" +local_private_key = "key_encoded_in_base64" +remote_public_key = "key_encoded_in_base64" + +[server.services.service1] # The service name must be identical to the client side +type = "tcp" # Optional. Same as the client `[client.services.X.type] +token = "whatever" # Necesary if `server.default_token` not set +bind_addr = "0.0.0.0:8081" # Necessary. The address of the service is exposed at. Generally only the port needs to be change. + +[server.services.service2] +bind_addr = "0.0.0.1:8082" diff --git a/tests/integration_test.rs b/tests/integration_test.rs index 7b5d408d..e6fcfed9 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -2,12 +2,10 @@ use anyhow::{Ok, Result}; use common::{run_rathole_client, PING, PONG}; use rand::Rng; use std::time::Duration; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpStream, UdpSocket}, - sync::broadcast, - time, -}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpStream, UdpSocket}; +use tokio::sync::broadcast; +use tokio::time; use tracing::{debug, info, instrument}; use tracing_subscriber::EnvFilter; diff --git a/tests/token b/tests/token new file mode 100644 index 00000000..8c79029d --- /dev/null +++ b/tests/token @@ -0,0 +1 @@ +test-file-token