mirror of
https://github.com/telemt/telemt.git
synced 2026-04-18 11:04:09 +03:00
Implement idle timeout for masking relay and add corresponding tests
- Introduced `copy_with_idle_timeout` function to handle reading and writing with an idle timeout. - Updated the proxy masking logic to use the new idle timeout function. - Added tests to verify that idle relays are closed by the idle timeout before the global relay timeout. - Ensured that connect refusal paths respect the masking budget and that responses followed by silence are cut off by the idle timeout. - Added tests for adversarial scenarios where clients may attempt to drip-feed data beyond the idle timeout.
This commit is contained in:
@@ -24,8 +24,36 @@ const MASK_TIMEOUT: Duration = Duration::from_millis(50);
|
||||
const MASK_RELAY_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
#[cfg(test)]
|
||||
const MASK_RELAY_TIMEOUT: Duration = Duration::from_millis(200);
|
||||
#[cfg(not(test))]
|
||||
const MASK_RELAY_IDLE_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
#[cfg(test)]
|
||||
const MASK_RELAY_IDLE_TIMEOUT: Duration = Duration::from_millis(100);
|
||||
const MASK_BUFFER_SIZE: usize = 8192;
|
||||
|
||||
async fn copy_with_idle_timeout<R, W>(reader: &mut R, writer: &mut W)
|
||||
where
|
||||
R: AsyncRead + Unpin,
|
||||
W: AsyncWrite + Unpin,
|
||||
{
|
||||
let mut buf = vec![0u8; MASK_BUFFER_SIZE];
|
||||
loop {
|
||||
let read_res = timeout(MASK_RELAY_IDLE_TIMEOUT, reader.read(&mut buf)).await;
|
||||
let n = match read_res {
|
||||
Ok(Ok(n)) => n,
|
||||
Ok(Err(_)) | Err(_) => break,
|
||||
};
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
|
||||
let write_res = timeout(MASK_RELAY_IDLE_TIMEOUT, writer.write_all(&buf[..n])).await;
|
||||
match write_res {
|
||||
Ok(Ok(())) => {}
|
||||
Ok(Err(_)) | Err(_) => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn write_proxy_header_with_timeout<W>(mask_write: &mut W, header: &[u8]) -> bool
|
||||
where
|
||||
W: AsyncWrite + Unpin,
|
||||
@@ -264,11 +292,11 @@ where
|
||||
|
||||
let _ = tokio::join!(
|
||||
async {
|
||||
let _ = tokio::io::copy(&mut reader, &mut mask_write).await;
|
||||
copy_with_idle_timeout(&mut reader, &mut mask_write).await;
|
||||
let _ = mask_write.shutdown().await;
|
||||
},
|
||||
async {
|
||||
let _ = tokio::io::copy(&mut mask_read, &mut writer).await;
|
||||
copy_with_idle_timeout(&mut mask_read, &mut writer).await;
|
||||
let _ = writer.shutdown().await;
|
||||
}
|
||||
);
|
||||
|
||||
Reference in New Issue
Block a user