mercuryworkshop-mirror/epoxy-tls
1
0
Fork 0
mirror of https://github.com/MercuryWorkshop/epoxy-tls.git synced 2025-05-12 14:00:01 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

add speed limits

Browse source
This commit is contained in:
Toshit Chawda 2024-11-22 21:41:03 -08:00
parent 90936adb2c
commit 35431c3a41
No known key found for this signature in database
GPG key ID: 91480ED99E2B3D9D
4 changed files with 49 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
server/Cargo.toml
View file

@ -5,6 +5,7 @@ edition = "2021"
[dependencies]
anyhow = "1.0.86"
async-speed-limit = { version = "0.4.2", optional = true }
async-trait = { version = "0.1.81", optional = true }
bytes = "1.7.1"
cfg-if = "1.0.0"
@ -47,6 +48,7 @@ yaml = ["dep:serde_yaml"]
toml = ["dep:toml"]
twisp = ["dep:pty-process", "dep:libc", "dep:async-trait", "dep:shell-words"]
speed-limit = ["dep:async-speed-limit"]
[build-dependencies]
vergen-git2 = { version = "1.0.0", features = ["rustc"] }

12
server/src/config.rs
View file

@ -181,6 +181,13 @@ pub struct StreamConfig {
/// Buffer size of reads from TCP sockets.
pub buffer_size: usize,
#[cfg(feature = "speed-limit")]
/// Remote stream read limit in bytes/second.
pub read_limit: f64,
#[cfg(feature = "speed-limit")]
/// Remote stream write limit in bytes/second.
pub write_limit: f64,
/// Whether or not to allow Wisp clients to create UDP streams.
pub allow_udp: bool,
/// Whether or not to enable nonstandard legacy wsproxy UDP streams.
@ -407,6 +414,11 @@ impl Default for StreamConfig {
tcp_nodelay: false,
buffer_size: 16384,
#[cfg(feature = "speed-limit")]
read_limit: f64::INFINITY,
#[cfg(feature = "speed-limit")]
write_limit: f64::INFINITY,
allow_udp: true,
allow_wsproxy_udp: false,
#[cfg(feature = "twisp")]

16
server/src/handle/wisp/mod.rs
View file

@ -34,6 +34,11 @@ async fn copy_read_fast(
muxrx: MuxStreamAsyncRead,
mut tcptx: OwnedWriteHalf,
) -> std::io::Result<()> {
#[cfg(feature = "speed-limit")]
let limiter = async_speed_limit::Limiter::builder(CONFIG.stream.write_limit)
.refill(Duration::from_secs(1))
.clock(async_speed_limit::clock::StandardClock)
.build();
let mut muxrx = muxrx.compat();
loop {
let buf = muxrx.fill_buf().await?;
@ -42,6 +47,9 @@ async fn copy_read_fast(
return Ok(());
}
#[cfg(feature = "speed-limit")]
limiter.consume(buf.len()).await;
let i = tcptx.write(buf).await?;
if i == 0 {
return Err(std::io::ErrorKind::WriteZero.into());
@ -52,6 +60,11 @@ async fn copy_read_fast(
}
async fn copy_write_fast(muxtx: MuxStreamWrite, tcprx: OwnedReadHalf) -> anyhow::Result<()> {
#[cfg(feature = "speed-limit")]
let limiter = async_speed_limit::Limiter::builder(CONFIG.stream.read_limit)
.refill(Duration::from_secs(1))
.clock(async_speed_limit::clock::StandardClock)
.build();
let mut tcprx = BufReader::with_capacity(CONFIG.stream.buffer_size, tcprx);
loop {
let buf = tcprx.fill_buf().await?;
@ -61,6 +74,9 @@ async fn copy_write_fast(muxtx: MuxStreamWrite, tcprx: OwnedReadHalf) -> anyhow:
return Ok(());
}
#[cfg(feature = "speed-limit")]
limiter.consume(buf.len()).await;
muxtx.write(&buf).await?;
tcprx.consume(len);
}