mirror of https://github.com/telemt/telemt.git
123 lines
3.8 KiB
Rust
123 lines
3.8 KiB
Rust
use super::*;
|
|
use crate::config::{UpstreamConfig, UpstreamType};
|
|
use std::sync::Arc;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};
|
|
use tokio::net::TcpListener;
|
|
use tokio::time::Duration;
|
|
|
|
fn new_upstream_manager(stats: Arc<Stats>) -> Arc<UpstreamManager> {
|
|
Arc::new(UpstreamManager::new(
|
|
vec![UpstreamConfig {
|
|
upstream_type: UpstreamType::Direct {
|
|
interface: None,
|
|
bind_addresses: None,
|
|
},
|
|
weight: 1,
|
|
enabled: true,
|
|
scopes: String::new(),
|
|
selected_scope: String::new(),
|
|
}],
|
|
1,
|
|
1,
|
|
1,
|
|
1,
|
|
false,
|
|
stats,
|
|
))
|
|
}
|
|
|
|
async fn run_probe_capture(
|
|
body_sent: usize,
|
|
tls_len: u16,
|
|
enable_shape_hardening: bool,
|
|
floor: usize,
|
|
cap: usize,
|
|
) -> Vec<u8> {
|
|
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let backend_addr = listener.local_addr().unwrap();
|
|
|
|
let mut cfg = ProxyConfig::default();
|
|
cfg.general.beobachten = false;
|
|
cfg.censorship.mask = true;
|
|
cfg.censorship.mask_unix_sock = None;
|
|
cfg.censorship.mask_host = Some("127.0.0.1".to_string());
|
|
cfg.censorship.mask_port = backend_addr.port();
|
|
cfg.censorship.mask_shape_hardening = enable_shape_hardening;
|
|
cfg.censorship.mask_shape_bucket_floor_bytes = floor;
|
|
cfg.censorship.mask_shape_bucket_cap_bytes = cap;
|
|
|
|
let accept_task = tokio::spawn(async move {
|
|
let (mut stream, _) = listener.accept().await.unwrap();
|
|
let mut got = Vec::new();
|
|
let _ = tokio::time::timeout(Duration::from_secs(2), stream.read_to_end(&mut got)).await;
|
|
got
|
|
});
|
|
|
|
let (server_side, mut client_side) = duplex(65536);
|
|
let handler = tokio::spawn(handle_client_stream(
|
|
server_side,
|
|
"198.51.100.188:56888".parse().unwrap(),
|
|
Arc::new(cfg),
|
|
Arc::new(Stats::new()),
|
|
new_upstream_manager(Arc::new(Stats::new())),
|
|
Arc::new(ReplayChecker::new(128, Duration::from_secs(60))),
|
|
Arc::new(BufferPool::new()),
|
|
Arc::new(SecureRandom::new()),
|
|
None,
|
|
Arc::new(RouteRuntimeController::new(RelayRouteMode::Direct)),
|
|
None,
|
|
Arc::new(UserIpTracker::new()),
|
|
Arc::new(BeobachtenStore::new()),
|
|
false,
|
|
));
|
|
|
|
let mut probe = vec![0u8; 5 + body_sent];
|
|
probe[0] = 0x16;
|
|
probe[1] = 0x03;
|
|
probe[2] = 0x01;
|
|
probe[3..5].copy_from_slice(&tls_len.to_be_bytes());
|
|
probe[5..].fill(0x66);
|
|
|
|
client_side.write_all(&probe).await.unwrap();
|
|
client_side.shutdown().await.unwrap();
|
|
|
|
let result = tokio::time::timeout(Duration::from_secs(4), handler)
|
|
.await
|
|
.unwrap()
|
|
.unwrap();
|
|
assert!(result.is_ok());
|
|
|
|
tokio::time::timeout(Duration::from_secs(4), accept_task)
|
|
.await
|
|
.unwrap()
|
|
.unwrap()
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn shape_hardening_disabled_keeps_original_probe_length() {
|
|
let got = run_probe_capture(17, 600, false, 512, 4096).await;
|
|
assert_eq!(got.len(), 22);
|
|
assert_eq!(&got[..5], &[0x16, 0x03, 0x01, 0x02, 0x58]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn shape_hardening_enabled_pads_small_probe_to_floor_bucket() {
|
|
let got = run_probe_capture(17, 600, true, 512, 4096).await;
|
|
assert_eq!(got.len(), 512);
|
|
assert_eq!(&got[..5], &[0x16, 0x03, 0x01, 0x02, 0x58]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn shape_hardening_enabled_pads_mid_probe_to_next_bucket() {
|
|
let got = run_probe_capture(511, 600, true, 512, 4096).await;
|
|
assert_eq!(got.len(), 1024);
|
|
assert_eq!(&got[..5], &[0x16, 0x03, 0x01, 0x02, 0x58]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn shape_hardening_respects_cap_and_avoids_padding_above_cap() {
|
|
let got = run_probe_capture(5000, 7000, true, 512, 4096).await;
|
|
assert_eq!(got.len(), 5005);
|
|
assert_eq!(&got[..5], &[0x16, 0x03, 0x01, 0x1b, 0x58]);
|
|
}
|