diff --git a/Cargo.lock b/Cargo.lock index 199eb46..2663159 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3259,6 +3259,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "lz4_flex" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "373f5eceeeab7925e0c1098212f2fbc4d416adec9d35051a6ab251e824c1854a" + [[package]] name = "matchers" version = "0.2.0" @@ -4035,6 +4041,8 @@ dependencies = [ "hex", "hmac", "kcp-tokio", + "lz4_flex", + "pin-project-lite", "quinn", "rcgen", "rustls", diff --git a/Cargo.toml b/Cargo.toml index 515eab9..3dbbb1f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -63,3 +63,5 @@ rust-embed = "8" hostname = "0.4" chrono = { version = "0.4", default-features = false, features = ["clock", "std"] } tls-parser = "0.12" +lz4_flex = { version = "0.11", default-features = false, features = ["std"] } +pin-project-lite = "0.2" diff --git a/client/src/control/session.rs b/client/src/control/session.rs index c2c05a9..a39755e 100644 --- a/client/src/control/session.rs +++ b/client/src/control/session.rs @@ -738,6 +738,9 @@ fn new_tunnel_base( route_by_http_user: String::new(), bandwidth: transport.bandwidth, bandwidth_limit_side: omit_client_side(&transport.bandwidth_limit_side), + compression: orbien_core::compression::CompressionAlgo::parse(&transport.compression) + .unwrap_or_default() + .wire_str(), }; extra(&mut np); np diff --git a/client/src/tunnel/manager.rs b/client/src/tunnel/manager.rs index 16513da..fa91489 100644 --- a/client/src/tunnel/manager.rs +++ b/client/src/tunnel/manager.rs @@ -1,9 +1,10 @@ use super::udp::run_udp_session; use crate::plugin::{self, ConnectionInfo, Plugin, PluginContext}; use anyhow::{anyhow, Result}; +use orbien_core::compression::{wrap_data_conn, CompressionAlgo}; use orbien_core::config::{ClientConfig, TunnelConfig}; use orbien_core::io; -use orbien_core::limit::{self, maybe_limit, BandwidthLimitSide, BandwidthLimiter}; +use orbien_core::limit::{self, BandwidthLimitSide, BandwidthLimiter}; use orbien_core::msg::StartDataConn; use orbien_core::net::{ addrs_from_start_data_conn, build_proxy_protocol_header, parse_proxy_protocol_version, @@ -18,6 +19,7 @@ use tokio::sync::{oneshot, Mutex as AsyncMutex}; struct TunnelEntry { cfg: TunnelConfig, limiter: Option>, + compression: CompressionAlgo, plugin: Option>, proxy_protocol: Option<&'static str>, udp_cancel: AsyncMutex>>, @@ -89,7 +91,7 @@ impl TunnelManager { start: &StartDataConn, data: DynStream, ) -> Result<()> { - let data = maybe_limit(data, entry.limiter.clone()); + let data = wrap_data_conn(data, entry.limiter.clone(), entry.compression); if let Some(ref plugin) = entry.plugin { tracing::debug!( @@ -192,7 +194,7 @@ impl TunnelManager { "udp data conn; starting forwarder" ); - let data = maybe_limit(data, entry.limiter.clone()); + let data = wrap_data_conn(data, entry.limiter.clone(), entry.compression); run_udp_session( data, @@ -246,6 +248,15 @@ fn build_entry(tunnel: &TunnelConfig) -> Result> { ); } + let compression = CompressionAlgo::parse(&tunnel.transport.compression)?; + if !compression.is_none() { + tracing::info!( + tunnel = %tunnel.name, + algo = compression.as_str(), + "data connection compression enabled" + ); + } + let plugin = if let Some(ref pc) = tunnel.plugin { if pc.plugin_type.is_empty() { None @@ -273,6 +284,7 @@ fn build_entry(tunnel: &TunnelConfig) -> Result> { Ok(Arc::new(TunnelEntry { cfg: tunnel.clone(), limiter, + compression, plugin, proxy_protocol, udp_cancel: AsyncMutex::new(None), diff --git a/conf/orbien-full.toml b/conf/orbien-full.toml index fdedcc1..4b13e41 100644 --- a/conf/orbien-full.toml +++ b/conf/orbien-full.toml @@ -36,6 +36,7 @@ remotePort = 9000 bandwidth = 10 bandwidthLimitSide = "client" proxyProtocolVersion = "v2" +compression = "lz4" [[tunnels]] name = "dns" @@ -47,6 +48,7 @@ remotePort = 5353 bandwidth = 10 bandwidthLimitSide = "server" proxyProtocolVersion = "v1" +compression = "none" [[tunnels]] name = "web" diff --git a/conf/orbien.toml b/conf/orbien.toml index eee1f69..cb8674c 100644 --- a/conf/orbien.toml +++ b/conf/orbien.toml @@ -7,4 +7,4 @@ server = "127.0.0.1:9527" name = "ssh" protocol = "tcp" service = "127.0.0.1:22" -remotePort = 9000 +remotePort = 9000 \ No newline at end of file diff --git a/core/Cargo.toml b/core/Cargo.toml index 014c7a0..3addc12 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -32,3 +32,5 @@ futures-util = { workspace = true } kcp-tokio = { workspace = true } yamux = { workspace = true } tls-parser = { workspace = true } +lz4_flex = { workspace = true } +pin-project-lite = { workspace = true } diff --git a/core/src/compression/algo.rs b/core/src/compression/algo.rs new file mode 100644 index 0000000..42e1bd5 --- /dev/null +++ b/core/src/compression/algo.rs @@ -0,0 +1,39 @@ +use anyhow::{bail, Result}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub enum CompressionAlgo { + #[default] + None, + Lz4, +} + +impl CompressionAlgo { + pub fn parse(raw: &str) -> Result { + let s = raw.trim(); + if s.is_empty() || s.eq_ignore_ascii_case("none") { + return Ok(Self::None); + } + if s.eq_ignore_ascii_case("lz4") { + return Ok(Self::Lz4); + } + bail!("unsupported compression {raw:?}; use \"none\" | \"lz4\"") + } + + pub fn as_str(self) -> &'static str { + match self { + Self::None => "none", + Self::Lz4 => "lz4", + } + } + + pub fn is_none(self) -> bool { + matches!(self, Self::None) + } + + pub fn wire_str(self) -> String { + match self { + Self::None => String::new(), + Self::Lz4 => "lz4".into(), + } + } +} diff --git a/core/src/compression/mod.rs b/core/src/compression/mod.rs new file mode 100644 index 0000000..e7305fc --- /dev/null +++ b/core/src/compression/mod.rs @@ -0,0 +1,17 @@ +mod algo; +mod stream; + +use crate::limit::{maybe_limit, BandwidthLimiter}; +use crate::transport::DynStream; +use std::sync::Arc; + +pub use algo::CompressionAlgo; +pub use stream::{maybe_compress, Lz4Stream}; + +pub fn wrap_data_conn( + stream: DynStream, + limiter: Option>, + compression: CompressionAlgo, +) -> DynStream { + maybe_compress(maybe_limit(stream, limiter), compression) +} diff --git a/core/src/compression/stream.rs b/core/src/compression/stream.rs new file mode 100644 index 0000000..cb0f327 --- /dev/null +++ b/core/src/compression/stream.rs @@ -0,0 +1,364 @@ +use super::CompressionAlgo; +use crate::transport::{boxed_stream, DynStream}; +use pin_project_lite::pin_project; +use std::io; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +const MAX_PLAIN_CHUNK: usize = 64 * 1024; +const MAX_FRAME_PAYLOAD: usize = 1 + MAX_PLAIN_CHUNK + 64; +const MAX_PLAIN_DECODE: usize = MAX_PLAIN_CHUNK; + +const KIND_RAW: u8 = 0; +const KIND_LZ4: u8 = 1; + +pin_project! { + pub struct Lz4Stream { + #[pin] + inner: S, + write: WriteState, + read: ReadState, + } +} + +struct WriteState { + pending: Vec, + pending_off: usize, + acked: Option, + finished: bool, + compress_buf: Vec, +} + +impl Default for WriteState { + fn default() -> Self { + Self { + pending: Vec::new(), + pending_off: 0, + acked: None, + finished: false, + compress_buf: Vec::new(), + } + } +} + +enum ReadPhase { + Len { got: usize, buf: [u8; 4] }, + Body { need: usize, buf: Vec }, + Deliver { data: Vec, off: usize }, +} + +struct ReadState { + phase: ReadPhase, +} + +impl Default for ReadState { + fn default() -> Self { + Self { + phase: ReadPhase::Len { + got: 0, + buf: [0; 4], + }, + } + } +} + +impl Lz4Stream { + pub fn new(inner: S) -> Self { + Self { + inner, + write: WriteState::default(), + read: ReadState::default(), + } + } + + fn decode_frame(body: &[u8]) -> io::Result> { + if body.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "empty compression frame body", + )); + } + let kind = body[0]; + let payload = &body[1..]; + match kind { + KIND_RAW => { + if payload.len() > MAX_PLAIN_DECODE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "raw frame plaintext too large", + )); + } + Ok(payload.to_vec()) + } + KIND_LZ4 => { + if payload.len() < 4 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "truncated lz4 payload", + )); + } + let declared = u32::from_le_bytes(payload[..4].try_into().unwrap()) as usize; + if declared == 0 || declared > MAX_PLAIN_DECODE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("lz4 declared size {declared} out of range"), + )); + } + lz4_flex::block::decompress_size_prepended(payload).map_err(|e| { + io::Error::new(io::ErrorKind::InvalidData, format!("lz4 decompress: {e}")) + }) + } + other => Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("unknown compression frame kind {other}"), + )), + } + } +} + +impl WriteState { + fn pending_remaining(&self) -> bool { + self.pending_off < self.pending.len() + } + + fn clear_pending(&mut self) { + self.pending.clear(); + self.pending_off = 0; + } + + fn encode_into_pending(&mut self, plain: &[u8]) -> io::Result<()> { + debug_assert!(!plain.is_empty()); + debug_assert!(plain.len() <= MAX_PLAIN_CHUNK); + + self.compress_buf.clear(); + let max_comp = lz4_flex::block::get_maximum_output_size(plain.len()); + self.compress_buf.resize(4 + max_comp, 0); + let written = + lz4_flex::block::compress_into(plain, &mut self.compress_buf[4..]).map_err(|e| { + io::Error::new(io::ErrorKind::InvalidData, format!("lz4 compress: {e}")) + })?; + self.compress_buf[0..4].copy_from_slice(&(plain.len() as u32).to_le_bytes()); + self.compress_buf.truncate(4 + written); + + let use_lz4 = self.compress_buf.len() < plain.len(); + let (kind, payload_len) = if use_lz4 { + (KIND_LZ4, self.compress_buf.len()) + } else { + (KIND_RAW, plain.len()) + }; + + let body_len = 1 + payload_len; + if body_len > MAX_FRAME_PAYLOAD { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "compression frame payload too large", + )); + } + + self.pending.clear(); + self.pending.reserve(4 + body_len); + self.pending + .extend_from_slice(&(body_len as u32).to_le_bytes()); + self.pending.push(kind); + if use_lz4 { + self.pending.extend_from_slice(&self.compress_buf); + } else { + self.pending.extend_from_slice(plain); + } + self.pending_off = 0; + Ok(()) + } +} + +fn poll_drain_pending( + mut inner: Pin<&mut W>, + write: &mut WriteState, + cx: &mut Context<'_>, +) -> Poll> { + while write.pending_remaining() { + match inner + .as_mut() + .poll_write(cx, &write.pending[write.pending_off..]) + { + Poll::Ready(Ok(0)) => { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::WriteZero, + "write zero while sending compression frame", + ))); + } + Poll::Ready(Ok(n)) => write.pending_off += n, + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + Poll::Ready(Ok(())) +} + +impl AsyncRead for Lz4Stream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if buf.remaining() == 0 { + return Poll::Ready(Ok(())); + } + + let mut this = self.project(); + loop { + match &mut this.read.phase { + ReadPhase::Deliver { data, off } => { + let n = buf.remaining().min(data.len() - *off); + if n == 0 { + this.read.phase = ReadPhase::Len { + got: 0, + buf: [0; 4], + }; + continue; + } + buf.put_slice(&data[*off..*off + n]); + *off += n; + if *off >= data.len() { + this.read.phase = ReadPhase::Len { + got: 0, + buf: [0; 4], + }; + } + return Poll::Ready(Ok(())); + } + ReadPhase::Len { got, buf: len_buf } => { + while *got < 4 { + let mut tmp = [0u8; 4]; + let mut rb = ReadBuf::new(&mut tmp[..4 - *got]); + match this.inner.as_mut().poll_read(cx, &mut rb) { + Poll::Ready(Ok(())) => { + let filled = rb.filled(); + if filled.is_empty() { + if *got == 0 { + return Poll::Ready(Ok(())); + } + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "eof in compression frame length", + ))); + } + len_buf[*got..*got + filled.len()].copy_from_slice(filled); + *got += filled.len(); + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + let need = u32::from_le_bytes(*len_buf) as usize; + if need == 0 || need > MAX_FRAME_PAYLOAD { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("invalid compression frame length {need}"), + ))); + } + this.read.phase = ReadPhase::Body { + need, + buf: Vec::with_capacity(need), + }; + } + ReadPhase::Body { need, buf: body } => { + while body.len() < *need { + let mut tmp = [0u8; 8 * 1024]; + let want = (*need - body.len()).min(tmp.len()); + let mut rb = ReadBuf::new(&mut tmp[..want]); + match this.inner.as_mut().poll_read(cx, &mut rb) { + Poll::Ready(Ok(())) => { + let filled = rb.filled(); + if filled.is_empty() { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "eof in compression frame body", + ))); + } + body.extend_from_slice(filled); + } + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + } + let plain = Lz4Stream::::decode_frame(body)?; + this.read.phase = ReadPhase::Deliver { + data: plain, + off: 0, + }; + } + } + } + } +} + +impl AsyncWrite for Lz4Stream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let mut this = self.project(); + if this.write.finished { + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::WriteZero, + "compression stream already shut down", + ))); + } + + loop { + match poll_drain_pending(this.inner.as_mut(), this.write, cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + + if let Some(acked) = this.write.acked.take() { + this.write.clear_pending(); + return Poll::Ready(Ok(acked)); + } + + this.write.clear_pending(); + + if buf.is_empty() { + return Poll::Ready(Ok(0)); + } + + let take = buf.len().min(MAX_PLAIN_CHUNK); + this.write.encode_into_pending(&buf[..take])?; + this.write.acked = Some(take); + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut this = self.project(); + match poll_drain_pending(this.inner.as_mut(), this.write, cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + if this.write.acked.is_none() { + this.write.clear_pending(); + } + this.inner.as_mut().poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + match self.as_mut().poll_flush(cx) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(e)) => return Poll::Ready(Err(e)), + Poll::Pending => return Poll::Pending, + } + let mut this = self.project(); + this.write.finished = true; + this.inner.as_mut().poll_shutdown(cx) + } +} + +pub fn maybe_compress(stream: DynStream, algo: CompressionAlgo) -> DynStream { + match algo { + CompressionAlgo::None => stream, + CompressionAlgo::Lz4 => boxed_stream(Lz4Stream::new(stream)), + } +} diff --git a/core/src/config/client.rs b/core/src/config/client.rs index e01a6a7..2bb3dd8 100644 --- a/core/src/config/client.rs +++ b/core/src/config/client.rs @@ -205,12 +205,19 @@ pub struct TunnelTransportConfig { alias = "proxy_protocol_version" )] pub proxy_protocol_version: String, + + #[serde(default = "default_compression")] + pub compression: String, } fn default_bandwidth_limit_side() -> String { "client".into() } +fn default_compression() -> String { + "none".into() +} + fn default_auth_type() -> String { "token".into() } @@ -432,6 +439,8 @@ impl ClientConfig { t.transport.bandwidth_limit_side )); } + crate::compression::CompressionAlgo::parse(&t.transport.compression) + .map_err(|e| anyhow!("tunnel `{}` {e}", t.name))?; match proto.as_str() { "tcp" | "udp" => { if t.remote_port == 0 { diff --git a/core/src/lib.rs b/core/src/lib.rs index 4ac6306..cde9bf4 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -1,4 +1,5 @@ pub mod auth; +pub mod compression; pub mod config; pub mod io; pub mod limit; diff --git a/core/src/msg/types.rs b/core/src/msg/types.rs index 686c2b8..659c482 100644 --- a/core/src/msg/types.rs +++ b/core/src/msg/types.rs @@ -81,6 +81,9 @@ pub struct NewTunnel { #[serde(default)] pub bandwidth_limit_side: String, + + #[serde(default, skip_serializing_if = "String::is_empty")] + pub compression: String, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/server/src/control/session/register.rs b/server/src/control/session/register.rs index 6c0e192..95caa17 100644 --- a/server/src/control/session/register.rs +++ b/server/src/control/session/register.rs @@ -4,6 +4,8 @@ use crate::tunnel::{ format_local_addr, HttpTunnel, HttpsTunnel, RegisteredTunnel, TcpTunnel, UdpTunnel, }; use anyhow::{anyhow, Result}; +use orbien_core::compression::CompressionAlgo; +use orbien_core::limit::BandwidthLimiter; use orbien_core::msg::{self, CloseTunnel, Message, NewTunnel, NewTunnelResp}; use std::sync::Arc; @@ -13,6 +15,33 @@ impl Control { .new_tunnel(name, tunnel_type, &self.user, &self.session_id); } + fn tunnel_transport( + np: &NewTunnel, + ) -> Result<(Option>, CompressionAlgo)> { + let limiter = orbien_core::limit::limiter_if_side( + np.bandwidth, + &np.bandwidth_limit_side, + orbien_core::limit::BandwidthLimitSide::Server, + )?; + if let Some(ref l) = limiter { + tracing::info!( + tunnel = %np.tunnel_name, + bytes_per_sec = l.bytes_per_sec(), + mode = "server", + "bandwidth limit enabled" + ); + } + let compression = CompressionAlgo::parse(&np.compression)?; + if !compression.is_none() { + tracing::info!( + tunnel = %np.tunnel_name, + algo = compression.as_str(), + "data connection compression enabled" + ); + } + Ok((limiter, compression)) + } + pub(super) async fn handle_new_tunnel(self: &Arc, np: NewTunnel) -> Result<()> { let resp = match self.register_tunnel(&np).await { Ok(remote_addr) => NewTunnelResp { @@ -47,19 +76,7 @@ impl Control { return Err(anyhow!("invalid remote_port")); } - let limiter = orbien_core::limit::limiter_if_side( - np.bandwidth, - &np.bandwidth_limit_side, - orbien_core::limit::BandwidthLimitSide::Server, - )?; - if let Some(ref l) = limiter { - tracing::info!( - tunnel = %np.tunnel_name, - bytes_per_sec = l.bytes_per_sec(), - mode = "server", - "bandwidth limit enabled" - ); - } + let (limiter, compression) = Self::tunnel_transport(np)?; let bind_addr = self.cfg.proxy_addr.clone(); let remote_port = np.remote_port as u16; @@ -79,6 +96,7 @@ impl Control { remote_port, control, limiter, + compression, Arc::clone(&self.access), ) .await?; @@ -100,19 +118,7 @@ impl Control { .clone() .ok_or_else(|| anyhow!("http tunnel requires server httpGwPort > 0"))?; - let limiter = orbien_core::limit::limiter_if_side( - np.bandwidth, - &np.bandwidth_limit_side, - orbien_core::limit::BandwidthLimitSide::Server, - )?; - if let Some(ref l) = limiter { - tracing::info!( - tunnel = %np.tunnel_name, - bytes_per_sec = l.bytes_per_sec(), - mode = "server", - "bandwidth limit enabled" - ); - } + let (limiter, compression) = Self::tunnel_transport(np)?; let name = np.tunnel_name.clone(); { @@ -128,6 +134,7 @@ impl Control { Arc::clone(&gw), &self.cfg.root_domain, limiter, + compression, ) .await?; @@ -153,19 +160,7 @@ impl Control { .clone() .ok_or_else(|| anyhow!("https tunnel requires server httpsGwPort > 0"))?; - let limiter = orbien_core::limit::limiter_if_side( - np.bandwidth, - &np.bandwidth_limit_side, - orbien_core::limit::BandwidthLimitSide::Server, - )?; - if let Some(ref l) = limiter { - tracing::info!( - tunnel = %np.tunnel_name, - bytes_per_sec = l.bytes_per_sec(), - mode = "server", - "bandwidth limit enabled" - ); - } + let (limiter, compression) = Self::tunnel_transport(np)?; let name = np.tunnel_name.clone(); { @@ -181,6 +176,7 @@ impl Control { Arc::clone(&gw), &self.cfg.root_domain, limiter, + compression, ) .await?; @@ -205,19 +201,7 @@ impl Control { return Err(anyhow!("invalid remote_port")); } - let limiter = orbien_core::limit::limiter_if_side( - np.bandwidth, - &np.bandwidth_limit_side, - orbien_core::limit::BandwidthLimitSide::Server, - )?; - if let Some(ref l) = limiter { - tracing::info!( - tunnel = %np.tunnel_name, - bytes_per_sec = l.bytes_per_sec(), - mode = "server", - "bandwidth limit enabled" - ); - } + let (limiter, compression) = Self::tunnel_transport(np)?; let bind_addr = self.cfg.proxy_addr.clone(); let remote_port = np.remote_port as u16; @@ -238,6 +222,7 @@ impl Control { remote_port, control, limiter, + compression, packet_size, ) .await?; diff --git a/server/src/tunnel/gw.rs b/server/src/tunnel/gw.rs index 12f29ab..75f5600 100644 --- a/server/src/tunnel/gw.rs +++ b/server/src/tunnel/gw.rs @@ -13,6 +13,7 @@ pub struct HttpRoute { pub basic_auth_password: String, pub route_by_http_user: String, pub limiter: Option>, + pub compression: orbien_core::compression::CompressionAlgo, } type DomainIndex = HashMap>>; diff --git a/server/src/tunnel/http.rs b/server/src/tunnel/http.rs index c984bcb..37afa7f 100644 --- a/server/src/tunnel/http.rs +++ b/server/src/tunnel/http.rs @@ -7,7 +7,8 @@ use crate::control::Control; use crate::metrics::ServerMetrics; use anyhow::{anyhow, bail, Result}; use httparse::Status; -use orbien_core::limit::{maybe_limit, BandwidthLimiter}; +use orbien_core::compression::{wrap_data_conn, CompressionAlgo}; +use orbien_core::limit::BandwidthLimiter; use orbien_core::msg::NewTunnel; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -29,6 +30,7 @@ impl HttpTunnel { gw: Arc, sub_domain_host: &str, limiter: Option>, + compression: CompressionAlgo, ) -> Result { let domains = build_domains(&np.domains, sub_domain_host)?; let name = np.tunnel_name.clone(); @@ -53,6 +55,7 @@ impl HttpTunnel { basic_auth_password: basic_auth_password.clone(), route_by_http_user: route_by_http_user.clone(), limiter: limiter.clone(), + compression, }, ) .await?; @@ -214,7 +217,7 @@ async fn handle_http_ingress( ) .await?; - let mut data = maybe_limit(data, route.limiter.clone()); + let mut data = wrap_data_conn(data, route.limiter.clone(), route.compression); let head_len = raw.len() as u64; data.write_all(&raw).await?; tracing::debug!( diff --git a/server/src/tunnel/https.rs b/server/src/tunnel/https.rs index 7f9c706..d089615 100644 --- a/server/src/tunnel/https.rs +++ b/server/src/tunnel/https.rs @@ -3,7 +3,8 @@ use crate::access::{prepare_ingress, AccessPolicy}; use crate::control::Control; use crate::metrics; use anyhow::{anyhow, Result}; -use orbien_core::limit::{maybe_limit, BandwidthLimiter}; +use orbien_core::compression::{wrap_data_conn, CompressionAlgo}; +use orbien_core::limit::BandwidthLimiter; use orbien_core::msg::NewTunnel; use orbien_core::tls::{peek_client_hello_sni, PrefixedStream}; use std::collections::HashMap; @@ -17,6 +18,7 @@ pub struct HttpsRoute { pub tunnel_name: String, pub control: Weak, pub limiter: Option>, + pub compression: CompressionAlgo, } pub struct HttpsGw { @@ -74,6 +76,7 @@ impl HttpsTunnel { gw: Arc, sub_domain_host: &str, limiter: Option>, + compression: CompressionAlgo, ) -> Result { let domains = build_domains(&np.domains, sub_domain_host)?; let name = np.tunnel_name.clone(); @@ -87,6 +90,7 @@ impl HttpsTunnel { tunnel_name: name.clone(), control: Arc::downgrade(&control), limiter: limiter.clone(), + compression, }, ) .await?; @@ -192,7 +196,7 @@ async fn handle_https_ingress( ) .await?; - let data = maybe_limit(data, route.limiter.clone()); + let data = wrap_data_conn(data, route.limiter.clone(), route.compression); let user = PrefixedStream::new(prefix, ingress.stream); tracing::debug!( diff --git a/server/src/tunnel/tcp.rs b/server/src/tunnel/tcp.rs index aaf62f9..8264ec6 100644 --- a/server/src/tunnel/tcp.rs +++ b/server/src/tunnel/tcp.rs @@ -2,7 +2,8 @@ use crate::access::{prepare_ingress, AccessPolicy}; use crate::control::Control; use crate::metrics; use anyhow::Result; -use orbien_core::limit::{maybe_limit, BandwidthLimiter}; +use orbien_core::compression::{wrap_data_conn, CompressionAlgo}; +use orbien_core::limit::BandwidthLimiter; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use tokio::net::TcpListener; @@ -24,6 +25,7 @@ impl TcpTunnel { remote_port: u16, control: Arc, limiter: Option>, + compression: CompressionAlgo, access: Arc, ) -> Result { let addr = format!("{bind_addr}:{remote_port}"); @@ -63,6 +65,7 @@ impl TcpTunnel { stream, peer, lim, + compression, access, ) .await @@ -117,6 +120,7 @@ async fn handle_ingress( stream: tokio::net::TcpStream, peer: std::net::SocketAddr, limiter: Option>, + compression: CompressionAlgo, access: Arc, ) -> Result<()> { let ingress = prepare_ingress(stream, peer, &access).await?; @@ -135,7 +139,7 @@ async fn handle_ingress( ) .await?; - let data = maybe_limit(data, limiter); + let data = wrap_data_conn(data, limiter, compression); tracing::debug!( tunnel = %tunnel_name, diff --git a/server/src/tunnel/udp.rs b/server/src/tunnel/udp.rs index 4102f95..f4a3f68 100644 --- a/server/src/tunnel/udp.rs +++ b/server/src/tunnel/udp.rs @@ -1,7 +1,8 @@ use crate::control::Control; use crate::metrics::{MemMetrics, ServerMetrics}; use anyhow::Result; -use orbien_core::limit::{maybe_limit, BandwidthLimiter}; +use orbien_core::compression::{wrap_data_conn, CompressionAlgo}; +use orbien_core::limit::BandwidthLimiter; use orbien_core::msg::{self, Message, UdpPacket}; use orbien_core::udp::{forward_user_conn, CHANNEL_CAP, SERVER_DATA_READ_DEADLINE}; use std::sync::atomic::{AtomicBool, Ordering}; @@ -30,6 +31,7 @@ impl UdpTunnel { remote_port: u16, control: Arc, limiter: Option>, + compression: CompressionAlgo, packet_size: usize, ) -> Result { let addr = format!("{bind_addr}:{remote_port}"); @@ -66,6 +68,7 @@ impl UdpTunnel { tunnel_name, control, limiter, + compression, send_rx, read_tx, closed_flag, @@ -116,6 +119,7 @@ async fn data_conn_loop( tunnel_name: String, control: std::sync::Weak, limiter: Option>, + compression: CompressionAlgo, mut send_rx: mpsc::Receiver, read_tx: mpsc::Sender, closed: Arc, @@ -156,7 +160,7 @@ async fn data_conn_loop( } }; - let data = maybe_limit(data, limiter.clone()); + let data = wrap_data_conn(data, limiter.clone(), compression); let (reader, mut writer) = tokio::io::split(data); tracing::info!( tunnel = %tunnel_name,