From dc04befa5a1b357ba0e17adec6512cc22eedc24b Mon Sep 17 00:00:00 2001 From: Yujia Qiao Date: Mon, 20 Dec 2021 12:40:02 +0800 Subject: [PATCH] feat: add crate features: tls, server, client --- Cargo.toml | 12 +++++++++++- src/client.rs | 20 +++++++++++++------- src/helper.rs | 8 ++++++++ src/lib.rs | 23 +++++++++++++++++++---- src/server.rs | 13 ++++++++++--- src/transport/mod.rs | 4 +++- 6 files changed, 64 insertions(+), 16 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2f01c16..4b4b397 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,10 +8,20 @@ license = "Apache-2.0" repository = "https://github.com/rapiz1/rathole" readme = "README.md" +[features] +default = ["tls", "server", "client"] +tls = ["tokio-native-tls"] +server = [] +client = [] + # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [profile.bench] debug = 1 +[profile.minimal] +inherits = "release" +opt-level = "s" + [dependencies] tokio = { version = "1", features = ["full"] } bytes = { version = "1"} @@ -29,5 +39,5 @@ tracing = "0.1" tracing-subscriber = "0.2" socket2 = "0.4" fdlimit = "0.2.1" -tokio-native-tls = "0.3.0" +tokio-native-tls = { version = "0.3.0", optional = true } async-trait = "0.1.52" diff --git a/src/client.rs b/src/client.rs index 32f6c50..a48abfa 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,22 +1,23 @@ -use std::collections::HashMap; -use std::sync::Arc; - use crate::config::{ClientConfig, ClientServiceConfig, Config, TransportType}; use crate::protocol::Hello::{self, *}; use crate::protocol::{ self, read_ack, read_control_cmd, read_data_cmd, read_hello, Ack, Auth, ControlChannelCmd, DataChannelCmd, CURRENT_PROTO_VRESION, HASH_WIDTH_IN_BYTES, }; -use crate::transport::{TcpTransport, TlsTransport, Transport}; +use crate::transport::{TcpTransport, Transport}; use anyhow::{anyhow, bail, Context, Result}; use backoff::ExponentialBackoff; - +use std::collections::HashMap; +use std::sync::Arc; use tokio::io::{copy_bidirectional, AsyncWriteExt}; use tokio::net::TcpStream; use tokio::sync::{broadcast, oneshot}; use tokio::time::{self, Duration}; use tracing::{debug, error, info, instrument, Instrument, Span}; +#[cfg(feature = "tls")] +use crate::transport::TlsTransport; + // The entrypoint of running a client pub async fn run_client(config: &Config, shutdown_rx: broadcast::Receiver) -> Result<()> { let config = match &config.client { @@ -32,8 +33,13 @@ pub async fn run_client(config: &Config, shutdown_rx: broadcast::Receiver) client.run(shutdown_rx).await } TransportType::Tls => { - let mut client = Client::::from(config).await?; - client.run(shutdown_rx).await + #[cfg(feature = "tls")] + { + let mut client = Client::::from(config).await?; + client.run(shutdown_rx).await + } + #[cfg(not(feature = "tls"))] + crate::helper::feature_not_compile("tls") } } } diff --git a/src/helper.rs b/src/helper.rs index ddca7ad..a136213 100644 --- a/src/helper.rs +++ b/src/helper.rs @@ -13,3 +13,11 @@ pub fn set_tcp_keepalive(conn: &TcpStream) -> Result<()> { s.set_tcp_keepalive(&keepalive) .with_context(|| "Failed to set keepalive") } + +#[allow(dead_code)] +pub fn feature_not_compile(feature: &str) -> ! { + panic!( + "The feature '{}' is not compiled in this binary. Please re-compile rathole", + feature + ) +} diff --git a/src/lib.rs b/src/lib.rs index de79348..e6c7d68 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,10 +1,8 @@ mod cli; -mod client; mod config; mod helper; mod multi_map; mod protocol; -mod server; mod transport; pub use cli::Cli; @@ -14,7 +12,14 @@ use anyhow::{anyhow, Result}; use tokio::sync::broadcast; use tracing::debug; +#[cfg(feature = "client")] +mod client; +#[cfg(feature = "client")] use client::run_client; + +#[cfg(feature = "server")] +mod server; +#[cfg(feature = "server")] use server::run_server; pub async fn run(args: &Cli, shutdown_rx: broadcast::Receiver) -> Result<()> { @@ -27,8 +32,18 @@ pub async fn run(args: &Cli, shutdown_rx: broadcast::Receiver) -> Result<( match determine_run_mode(&config, args) { RunMode::Undetermine => Err(anyhow!("Cannot determine running as a server or a client")), - RunMode::Client => run_client(&config, shutdown_rx).await, - RunMode::Server => run_server(&config, shutdown_rx).await, + RunMode::Client => { + #[cfg(not(feature = "client"))] + crate::helper::feature_not_compile("client"); + #[cfg(feature = "client")] + run_client(&config, shutdown_rx).await + } + RunMode::Server => { + #[cfg(not(feature = "server"))] + crate::helper::feature_not_compile("server"); + #[cfg(feature = "server")] + run_server(&config, shutdown_rx).await + } } } diff --git a/src/server.rs b/src/server.rs index a9c9630..3d843d9 100644 --- a/src/server.rs +++ b/src/server.rs @@ -4,7 +4,9 @@ use crate::protocol::Hello::{ControlChannelHello, DataChannelHello}; use crate::protocol::{ self, read_auth, read_hello, Ack, ControlChannelCmd, DataChannelCmd, Hello, HASH_WIDTH_IN_BYTES, }; -use crate::transport::{TcpTransport, TlsTransport, Transport}; +#[cfg(feature = "tls")] +use crate::transport::TlsTransport; +use crate::transport::{TcpTransport, Transport}; use anyhow::{anyhow, bail, Context, Result}; use backoff::backoff::Backoff; use backoff::ExponentialBackoff; @@ -41,8 +43,13 @@ pub async fn run_server(config: &Config, shutdown_rx: broadcast::Receiver) server.run(shutdown_rx).await?; } TransportType::Tls => { - let mut server = Server::::from(config).await?; - server.run(shutdown_rx).await?; + #[cfg(feature = "tls")] + { + let mut server = Server::::from(config).await?; + server.run(shutdown_rx).await?; + } + #[cfg(not(feature = "tls"))] + crate::helper::feature_not_compile("tls") } } diff --git a/src/transport/mod.rs b/src/transport/mod.rs index 777ef9f..9d447c7 100644 --- a/src/transport/mod.rs +++ b/src/transport/mod.rs @@ -19,6 +19,8 @@ pub trait Transport: Debug + Send + Sync { } mod tcp; -mod tls; pub use tcp::TcpTransport; +#[cfg(feature = "tls")] +mod tls; +#[cfg(feature = "tls")] pub use tls::TlsTransport;