feat: add crate features: tls, server, client

This commit is contained in:
Yujia Qiao 2021-12-20 12:40:02 +08:00 committed by Yujia Qiao
parent 531428f838
commit dc04befa5a
6 changed files with 64 additions and 16 deletions

View File

@ -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"

View File

@ -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")
}
}
}

View File

@ -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
)
}

View File

@ -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
}
}
}

View File

@ -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")
}
}

View File

@ -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;