Kaynağa Gözat

feat: add crate features: tls, server, client

Yujia Qiao 4 yıl önce
ebeveyn
işleme
dc04befa5a
6 değiştirilmiş dosya ile 64 ekleme ve 16 silme
  1. 11 1
      Cargo.toml
  2. 13 7
      src/client.rs
  3. 8 0
      src/helper.rs
  4. 19 4
      src/lib.rs
  5. 10 3
      src/server.rs
  6. 3 1
      src/transport/mod.rs

+ 11 - 1
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"

+ 13 - 7
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<bool>) -> Result<()> {
     let config = match &config.client {
@@ -32,8 +33,13 @@ pub async fn run_client(config: &Config, shutdown_rx: broadcast::Receiver<bool>)
             client.run(shutdown_rx).await
         }
         TransportType::Tls => {
-            let mut client = Client::<TlsTransport>::from(config).await?;
-            client.run(shutdown_rx).await
+            #[cfg(feature = "tls")]
+            {
+                let mut client = Client::<TlsTransport>::from(config).await?;
+                client.run(shutdown_rx).await
+            }
+            #[cfg(not(feature = "tls"))]
+            crate::helper::feature_not_compile("tls")
         }
     }
 }

+ 8 - 0
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
+    )
+}

+ 19 - 4
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<bool>) -> Result<()> {
@@ -27,8 +32,18 @@ pub async fn run(args: &Cli, shutdown_rx: broadcast::Receiver<bool>) -> 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
+        }
     }
 }
 

+ 10 - 3
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<bool>)
             server.run(shutdown_rx).await?;
         }
         TransportType::Tls => {
-            let mut server = Server::<TlsTransport>::from(config).await?;
-            server.run(shutdown_rx).await?;
+            #[cfg(feature = "tls")]
+            {
+                let mut server = Server::<TlsTransport>::from(config).await?;
+                server.run(shutdown_rx).await?;
+            }
+            #[cfg(not(feature = "tls"))]
+            crate::helper::feature_not_compile("tls")
         }
     }
 

+ 3 - 1
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;