astelm/tglock
1
0
Fork 0
mirror of https://github.com/by-sonic/tglock.git synced 2026-08-03 06:45:51 +03:00
Code Issues Packages Projects Releases Wiki Activity

feat: rebuild TGLock with adaptive transport and Tauri UI

Browse Source
This commit is contained in:
Никита Митусов
2026-07-29 15:22:54 +03:00
parent 39a7772151
commit c25bea1d92
43 changed files with 9074 additions and 512 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main.rs
+233 -349
View File
@@ -1,384 +1,246 @@
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
mod mtproto;
mod proxy;
mod transport;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use tauri::{Manager, State};
use eframe::egui;
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
struct Settings {
lan_mode: bool,
port: u16,
worker_domain: String,
}
// -- Colors (GitHub Dark inspired) ------------------------------------------
impl Default for Settings {
fn default() -> Self {
Self {
lan_mode: false,
port: proxy::DEFAULT_PORT,
worker_domain: String::new(),
}
}
}
const BG: egui::Color32 = egui::Color32::from_rgb(13, 17, 23);
const SURFACE: egui::Color32 = egui::Color32::from_rgb(22, 27, 34);
const BORDER: egui::Color32 = egui::Color32::from_rgb(48, 54, 61);
const ACCENT: egui::Color32 = egui::Color32::from_rgb(88, 166, 255);
const GREEN: egui::Color32 = egui::Color32::from_rgb(63, 185, 80);
const RED: egui::Color32 = egui::Color32::from_rgb(248, 81, 73);
const TEXT: egui::Color32 = egui::Color32::from_rgb(230, 237, 243);
const TEXT2: egui::Color32 = egui::Color32::from_rgb(139, 148, 158);
const AD_BG: egui::Color32 = egui::Color32::from_rgb(17, 21, 28);
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct LogLine {
timestamp: String,
message: String,
error: bool,
}
fn main() -> eframe::Result<()> {
eframe::run_native(
"TGLock",
eframe::NativeOptions {
viewport: egui::ViewportBuilder::default()
.with_inner_size([520.0, 620.0])
.with_min_inner_size([420.0, 500.0])
.with_title("TGLock"),
..Default::default()
},
Box::new(|cc| {
apply_theme(&cc.egui_ctx);
Ok(Box::new(App::new()))
}),
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct StatusSnapshot {
running: bool,
active_connections: u32,
tunnels: u32,
data_center: Option<u16>,
route: String,
failures: u32,
uptime_seconds: u64,
port: u16,
logs: Vec<LogLine>,
}
struct AppState {
stats: Arc<proxy::Stats>,
settings: Mutex<Settings>,
active_port: Mutex<u16>,
started_at: Mutex<Option<Instant>>,
logs: Arc<Mutex<Vec<LogLine>>>,
settings_path: PathBuf,
}
impl AppState {
fn new(settings_path: PathBuf) -> Self {
let settings = std::fs::read(&settings_path)
.ok()
.and_then(|contents| serde_json::from_slice(&contents).ok())
.unwrap_or_default();
Self {
stats: proxy::Stats::new(),
settings: Mutex::new(settings),
active_port: Mutex::new(proxy::DEFAULT_PORT),
started_at: Mutex::new(None),
logs: Arc::new(Mutex::new(Vec::new())),
settings_path,
}
}
fn log(&self, message: impl Into<String>, error: bool) {
let mut logs = self.logs.lock().unwrap();
logs.push(LogLine {
timestamp: current_time(),
message: message.into(),
error,
});
if logs.len() > 100 {
logs.remove(0);
}
}
fn snapshot(&self) -> StatusSnapshot {
let data_center = self.stats.last_dc.load(Ordering::Relaxed);
let route = match self.stats.last_route.load(Ordering::Relaxed) {
1 => "Telegram WebSocket",
2 => "Cloudflare Worker",
_ => "Автоматический маршрут",
};
StatusSnapshot {
running: self.stats.running.load(Ordering::SeqCst),
active_connections: self.stats.active.load(Ordering::Relaxed),
tunnels: self.stats.ws.load(Ordering::Relaxed),
data_center: (data_center > 0).then_some(data_center),
route: route.to_owned(),
failures: self.stats.ws_failures.load(Ordering::Relaxed),
uptime_seconds: self
.started_at
.lock()
.unwrap()
.map_or(0, |started| started.elapsed().as_secs()),
port: *self.active_port.lock().unwrap(),
logs: self.logs.lock().unwrap().clone(),
}
}
fn persist_settings(&self, settings: &Settings) -> Result<(), String> {
if let Some(parent) = self.settings_path.parent() {
std::fs::create_dir_all(parent)
.map_err(|error| format!("Не удалось создать папку настроек: {error}"))?;
}
let contents = serde_json::to_vec_pretty(settings)
.map_err(|error| format!("Не удалось сохранить настройки: {error}"))?;
std::fs::write(&self.settings_path, contents)
.map_err(|error| format!("Не удалось сохранить настройки: {error}"))
}
}
fn current_time() -> String {
let seconds = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
format!(
"{:02}:{:02}:{:02}",
(seconds / 3600) % 24,
(seconds / 60) % 60,
seconds % 60
)
}
fn apply_theme(ctx: &egui::Context) {
let mut v = egui::Visuals::dark();
v.panel_fill = BG;
v.window_fill = SURFACE;
v.extreme_bg_color = BG;
v.faint_bg_color = SURFACE;
v.override_text_color = Some(TEXT);
v.widgets.noninteractive.bg_fill = SURFACE;
v.widgets.noninteractive.fg_stroke = egui::Stroke::new(1.0, TEXT2);
v.widgets.noninteractive.bg_stroke = egui::Stroke::new(1.0, BORDER);
v.widgets.inactive.bg_fill = egui::Color32::from_rgb(33, 38, 45);
v.widgets.inactive.fg_stroke = egui::Stroke::new(1.0, TEXT);
v.widgets.inactive.bg_stroke = egui::Stroke::new(1.0, BORDER);
v.widgets.hovered.bg_fill = egui::Color32::from_rgb(48, 54, 61);
v.widgets.hovered.fg_stroke = egui::Stroke::new(1.0, TEXT);
v.widgets.active.bg_fill = ACCENT;
v.widgets.active.fg_stroke = egui::Stroke::new(1.0, BG);
ctx.set_visuals(v);
#[tauri::command]
fn get_status(state: State<'_, AppState>) -> StatusSnapshot {
state.snapshot()
}
// -- Log --------------------------------------------------------------------
#[derive(Clone)]
struct LogLine {
ts: String,
msg: String,
err: bool,
#[tauri::command]
fn get_settings(state: State<'_, AppState>) -> Settings {
state.settings.lock().unwrap().clone()
}
fn now_ts() -> String {
let s = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
format!("{:02}:{:02}:{:02}", (s / 3600) % 24, (s / 60) % 60, s % 60)
#[tauri::command]
fn save_settings(settings: Settings, state: State<'_, AppState>) -> Result<Settings, String> {
if state.stats.running.load(Ordering::SeqCst) {
return Err("Сначала выключите защиту".into());
}
if settings.port == 0 {
return Err("Порт должен быть от 1 до 65535".into());
}
state.persist_settings(&settings)?;
*state.settings.lock().unwrap() = settings.clone();
state.log("Настройки сохранены", false);
Ok(settings)
}
fn log(log: &Arc<Mutex<Vec<LogLine>>>, msg: &str, err: bool) {
log.lock().unwrap().push(LogLine {
ts: now_ts(),
msg: msg.into(),
err,
});
}
// -- App --------------------------------------------------------------------
struct App {
stats: Arc<proxy::Stats>,
log: Arc<Mutex<Vec<LogLine>>>,
started_at: Option<Instant>,
lan_mode: bool,
port_str: String,
active_port: u16,
}
impl App {
fn new() -> Self {
Self {
stats: proxy::Stats::new(),
log: Arc::new(Mutex::new(Vec::new())),
started_at: None,
lan_mode: false,
port_str: proxy::DEFAULT_PORT.to_string(),
active_port: proxy::DEFAULT_PORT,
}
#[tauri::command]
fn start_proxy(state: State<'_, AppState>) -> Result<StatusSnapshot, String> {
if state.stats.running.load(Ordering::SeqCst) {
return Ok(state.snapshot());
}
fn running(&self) -> bool {
self.stats.running.load(Ordering::SeqCst)
}
let settings = state.settings.lock().unwrap().clone();
state.stats.set_worker_domain(&settings.worker_domain);
*state.active_port.lock().unwrap() = settings.port;
*state.started_at.lock().unwrap() = Some(Instant::now());
state.log("Запускаю защищённый маршрут…", false);
fn start(&mut self) {
if self.running() { return; }
let port: u16 = match self.port_str.trim().parse() {
Ok(p) if p > 0 => p,
_ => {
log(&self.log, "Неверный порт", true);
let stats = state.stats.clone();
let logs = state.logs.clone();
let lan_mode = settings.lan_mode;
let port = settings.port;
std::thread::spawn(move || {
let runtime = match tokio::runtime::Runtime::new() {
Ok(runtime) => runtime,
Err(error) => {
push_log(&logs, format!("Не удалось запустить сервис: {error}"), true);
return;
}
};
self.active_port = port;
self.started_at = Some(Instant::now());
let stats = self.stats.clone();
let lg = self.log.clone();
let lan = self.lan_mode;
log(&lg, "Запускаю прокси...", false);
std::thread::spawn(move || {
let rt = tokio::runtime::Runtime::new().unwrap();
let r = rt.block_on(proxy::run(stats, lan, port));
if let Err(e) = r {
log(&lg, &format!("Ошибка: {}", e), true);
}
});
std::thread::sleep(std::time::Duration::from_millis(250));
if self.running() {
let addr = if lan { "0.0.0.0" } else { "127.0.0.1" };
log(&self.log, &format!("SOCKS5 на {}:{}", addr, port), false);
if lan {
log(&self.log, "LAN-режим: другие устройства могут подключаться", false);
}
if let Err(error) = runtime.block_on(proxy::run(stats, lan_mode, port)) {
push_log(&logs, format!("Ошибка подключения: {error}"), true);
}
});
std::thread::sleep(std::time::Duration::from_millis(220));
if !state.stats.running.load(Ordering::SeqCst) {
*state.started_at.lock().unwrap() = None;
return Err(state
.logs
.lock()
.unwrap()
.last()
.map(|line| line.message.clone())
.unwrap_or_else(|| "Не удалось запустить прокси".into()));
}
fn stop(&mut self) {
self.stats.running.store(false, Ordering::SeqCst);
self.started_at = None;
log(&self.log, "Остановлен", false);
}
fn uptime_str(&self) -> String {
match self.started_at {
Some(t) => {
let s = t.elapsed().as_secs();
format!("{:02}:{:02}:{:02}", s / 3600, (s / 60) % 60, s % 60)
}
None => "--:--:--".into(),
}
}
state.log(
format!(
"Прокси запущен на {}:{}",
if settings.lan_mode {
"0.0.0.0"
} else {
"127.0.0.1"
},
settings.port
),
false,
);
let host = if settings.lan_mode {
local_ip().unwrap_or_else(|| "127.0.0.1".into())
} else {
"127.0.0.1".into()
};
let _ = open::that(format!(
"tg://proxy?server={host}&port={}&secret={}",
settings.port,
state.stats.telegram_secret()
));
state.log("Открываю подключение в Telegram…", false);
Ok(state.snapshot())
}
impl eframe::App for App {
fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
ctx.request_repaint_after(std::time::Duration::from_millis(300));
let on = self.running();
let active = self.stats.active.load(Ordering::Relaxed);
let total = self.stats.total.load(Ordering::Relaxed);
let ws = self.stats.ws.load(Ordering::Relaxed);
let dc = self.stats.last_dc.load(Ordering::Relaxed);
// === Ad bar (top) ===
egui::TopBottomPanel::top("ad").show(ctx, |ui| {
egui::Frame::new()
.fill(AD_BG)
.inner_margin(egui::Margin::symmetric(12, 6))
.show(ui, |ui| {
ui.horizontal(|ui| {
ui.colored_label(ACCENT, egui::RichText::new("RoseVPN").size(12.0).strong());
ui.colored_label(TEXT2, egui::RichText::new("Обход для всех приложений").size(11.0));
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
if ui.add(egui::Button::new(
egui::RichText::new("@rosevpnru_bot").size(11.0).strong().color(ACCENT)
).frame(false)).clicked() {
let _ = open::that("https://t.me/rosevpnru_bot");
}
});
});
});
});
// === Log (bottom) ===
egui::TopBottomPanel::bottom("log")
.min_height(120.0)
.show(ctx, |ui| {
ui.add_space(4.0);
ui.colored_label(TEXT2, egui::RichText::new("LOG").size(11.0));
ui.separator();
egui::ScrollArea::vertical()
.auto_shrink([false, false])
.stick_to_bottom(true)
.show(ui, |ui| {
for e in self.log.lock().unwrap().iter() {
let c = if e.err { RED } else { TEXT2 };
ui.colored_label(c, egui::RichText::new(
format!("{} {}", e.ts, e.msg)
).size(11.5).monospace());
}
});
});
// === Stats bar ===
egui::TopBottomPanel::bottom("stats").show(ctx, |ui| {
egui::Frame::new()
.fill(SURFACE)
.inner_margin(egui::Margin::symmetric(16, 8))
.show(ui, |ui| {
ui.horizontal(|ui| {
stat(ui, "Соединения", &active.to_string());
ui.add_space(20.0);
stat(ui, "WS-туннели", &ws.to_string());
ui.add_space(20.0);
stat(ui, "DC", &if dc > 0 { dc.to_string() } else { "—".into() });
ui.add_space(20.0);
stat(ui, "Всего", &total.to_string());
ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
stat(ui, "Аптайм", &self.uptime_str());
});
});
});
});
// === Main ===
egui::CentralPanel::default().show(ctx, |ui| {
ui.vertical_centered(|ui| {
ui.add_space(30.0);
// Title
ui.colored_label(TEXT, egui::RichText::new("TGLock").size(32.0).strong());
ui.add_space(4.0);
ui.colored_label(TEXT2, egui::RichText::new("WebSocket-туннель для Telegram").size(13.0));
ui.add_space(24.0);
// Status indicator
let (dot_color, status_text) = if on {
(GREEN, "Подключено")
} else {
(egui::Color32::from_rgb(80, 80, 80), "Отключено")
};
ui.horizontal(|ui| {
let center = ui.available_width() / 2.0 - 50.0;
ui.add_space(center);
let (r, _) = ui.allocate_exact_size(egui::vec2(10.0, 10.0), egui::Sense::hover());
ui.painter().circle_filled(r.center(), 5.0, dot_color);
ui.colored_label(
if on { GREEN } else { TEXT2 },
egui::RichText::new(status_text).size(14.0).strong(),
);
});
ui.add_space(20.0);
// Options (only when stopped)
if !on {
ui.horizontal(|ui| {
let center = ui.available_width() / 2.0 - 130.0;
ui.add_space(center);
ui.colored_label(TEXT2, egui::RichText::new("Порт:").size(12.0));
let port_edit = egui::TextEdit::singleline(&mut self.port_str)
.desired_width(55.0)
.font(egui::TextStyle::Monospace);
ui.add(port_edit);
ui.add_space(12.0);
ui.checkbox(&mut self.lan_mode, "");
ui.colored_label(TEXT2, egui::RichText::new("LAN").size(12.0));
ui.colored_label(
egui::Color32::from_rgb(80, 85, 95),
egui::RichText::new("(0.0.0.0)").size(10.5),
);
});
ui.add_space(8.0);
}
// Big button
if !on {
let btn = ui.add_sized(
[260.0, 52.0],
egui::Button::new(
egui::RichText::new("ПОДКЛЮЧИТЬ").size(18.0).strong().color(BG)
).fill(ACCENT).corner_radius(8.0),
);
if btn.clicked() {
self.start();
}
} else {
let btn = ui.add_sized(
[260.0, 52.0],
egui::Button::new(
egui::RichText::new("ОТКЛЮЧИТЬ").size(18.0).strong().color(TEXT)
).fill(egui::Color32::from_rgb(40, 45, 52)).corner_radius(8.0),
);
if btn.clicked() {
self.stop();
}
}
ui.add_space(24.0);
// Setup section
egui::Frame::new()
.fill(SURFACE)
.corner_radius(8.0)
.inner_margin(16.0)
.show(ui, |ui| {
ui.set_width(360.0);
ui.colored_label(TEXT, egui::RichText::new("Настройка Telegram").size(14.0).strong());
ui.add_space(6.0);
let server_addr = if self.lan_mode && on {
local_ip().unwrap_or_else(|| "127.0.0.1".into())
} else {
"127.0.0.1".into()
};
let display_port = if on { self.active_port } else {
self.port_str.trim().parse().unwrap_or(proxy::DEFAULT_PORT)
};
if on {
if ui.add(egui::Button::new(
egui::RichText::new("Настроить автоматически").size(13.0).color(ACCENT)
).frame(false)).clicked() {
let _ = open::that(format!("tg://socks?server={}&port={}", server_addr, display_port));
log(&self.log, "Открываю настройку Telegram...", false);
}
ui.add_space(4.0);
}
ui.colored_label(TEXT2, egui::RichText::new("Настройки → Продвинутые → Тип соединения → SOCKS5").size(11.5));
ui.add_space(4.0);
egui::Grid::new("cfg").num_columns(2).spacing([12.0, 3.0]).show(ui, |ui| {
ui.colored_label(TEXT2, "Сервер");
ui.monospace(&server_addr);
ui.end_row();
ui.colored_label(TEXT2, "Порт");
ui.monospace(format!("{}", display_port));
ui.end_row();
});
});
ui.add_space(16.0);
// How it works (compact)
ui.colored_label(TEXT2, egui::RichText::new(
"Трафик Telegram → SOCKS5 → WSS → web.telegram.org → DC"
).size(11.0));
ui.colored_label(TEXT2, egui::RichText::new(
"Провайдер видит обычный HTTPS. Остальной трафик не затрагивается."
).size(11.0));
});
});
}
#[tauri::command]
fn stop_proxy(state: State<'_, AppState>) -> StatusSnapshot {
state.stats.stop();
*state.started_at.lock().unwrap() = None;
state.log("Защита выключена", false);
state.snapshot()
}
fn stat(ui: &mut egui::Ui, label: &str, value: &str) {
ui.vertical(|ui| {
ui.colored_label(TEXT2, egui::RichText::new(label).size(10.0));
ui.colored_label(TEXT, egui::RichText::new(value).size(13.0).strong().monospace());
fn push_log(logs: &Arc<Mutex<Vec<LogLine>>>, message: String, error: bool) {
logs.lock().unwrap().push(LogLine {
timestamp: current_time(),
message,
error,
});
}
@@ -387,3 +249,25 @@ fn local_ip() -> Option<String> {
socket.connect("8.8.8.8:80").ok()?;
Some(socket.local_addr().ok()?.ip().to_string())
}
fn main() {
tauri::Builder::default()
.setup(|app| {
let settings_path = app
.path()
.app_config_dir()
.map_err(|error| error.to_string())?
.join("settings.json");
app.manage(AppState::new(settings_path));
Ok(())
})
.invoke_handler(tauri::generate_handler![
get_status,
get_settings,
save_settings,
start_proxy,
stop_proxy
])
.run(tauri::generate_context!())
.expect("failed to run TGLock");
}
src/mtproto.rs
+318
View File
@@ -0,0 +1,318 @@
use aes::Aes256;
use cipher::{KeyIvInit, StreamCipher};
use rand::{rngs::OsRng, RngCore};
use sha2::{Digest, Sha256};
#[cfg(not(test))]
use std::path::PathBuf;
type AesCtr = ctr::Ctr128BE<Aes256>;
const INIT_LEN: usize = 64;
const KEY_START: usize = 8;
const KEY_END: usize = 40;
const IV_END: usize = 56;
const TAG_START: usize = 56;
const DC_START: usize = 60;
const ABRIDGED: [u8; 4] = [0xef; 4];
const INTERMEDIATE: [u8; 4] = [0xee; 4];
const PADDED_INTERMEDIATE: [u8; 4] = [0xdd; 4];
pub struct ClientInit {
pub dc: u16,
pub media: bool,
pub relay_init: [u8; INIT_LEN],
pub crypto: CryptoContext,
}
pub struct CryptoContext {
client_decrypt: AesCtr,
client_encrypt: AesCtr,
telegram_encrypt: AesCtr,
telegram_decrypt: AesCtr,
}
impl CryptoContext {
pub fn client_to_telegram(&mut self, data: &mut [u8]) {
self.client_decrypt.apply_keystream(data);
self.telegram_encrypt.apply_keystream(data);
}
pub fn telegram_to_client(&mut self, data: &mut [u8]) {
self.telegram_decrypt.apply_keystream(data);
self.client_encrypt.apply_keystream(data);
}
}
pub fn generate_secret() -> [u8; 16] {
let mut secret = [0; 16];
OsRng.fill_bytes(&mut secret);
secret
}
#[cfg(not(test))]
pub fn load_or_create_secret() -> [u8; 16] {
let Some(path) = secret_path() else {
return generate_secret();
};
if let Ok(value) = std::fs::read_to_string(&path) {
if let Some(secret) = parse_secret_hex(value.trim()) {
return secret;
}
}
let secret = generate_secret();
if let Some(parent) = path.parent() {
let _ = std::fs::create_dir_all(parent);
}
write_secret_file(&path, &secret_hex(&secret));
secret
}
#[cfg(not(test))]
fn secret_path() -> Option<PathBuf> {
#[cfg(target_os = "windows")]
{
std::env::var_os("APPDATA")
.map(PathBuf::from)
.map(|path| path.join("TGLock").join("secret"))
}
#[cfg(target_os = "macos")]
{
std::env::var_os("HOME")
.map(PathBuf::from)
.map(|path| path.join("Library/Application Support/TGLock/secret"))
}
#[cfg(all(unix, not(target_os = "macos")))]
{
if let Some(path) = std::env::var_os("XDG_CONFIG_HOME") {
return Some(PathBuf::from(path).join("tglock").join("secret"));
}
std::env::var_os("HOME")
.map(PathBuf::from)
.map(|path| path.join(".config/tglock/secret"))
}
}
#[cfg(all(not(test), unix))]
fn write_secret_file(path: &std::path::Path, value: &str) {
use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
if let Ok(mut file) = std::fs::OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.mode(0o600)
.open(path)
{
let _ = file.write_all(value.as_bytes());
}
}
#[cfg(all(not(test), not(unix)))]
fn write_secret_file(path: &std::path::Path, value: &str) {
let _ = std::fs::write(path, value);
}
pub fn secret_hex(secret: &[u8; 16]) -> String {
let mut output = String::with_capacity(32);
for byte in secret {
use std::fmt::Write;
let _ = write!(output, "{:02x}", byte);
}
output
}
fn parse_secret_hex(value: &str) -> Option<[u8; 16]> {
if value.len() != 32 {
return None;
}
let mut secret = [0; 16];
for (index, byte) in secret.iter_mut().enumerate() {
*byte = u8::from_str_radix(&value[index * 2..index * 2 + 2], 16).ok()?;
}
Some(secret)
}
pub fn telegram_secret(secret: &[u8; 16]) -> String {
format!("dd{}", secret_hex(secret))
}
pub fn parse_client_init(init: &[u8; INIT_LEN], secret: &[u8; 16]) -> Option<ClientInit> {
let client_dec_key = secret_key(&init[KEY_START..KEY_END], secret);
let client_dec_iv: [u8; 16] = init[KEY_END..IV_END].try_into().ok()?;
let mut client_decrypt = AesCtr::new((&client_dec_key).into(), (&client_dec_iv).into());
let mut decrypted = *init;
client_decrypt.apply_keystream(&mut decrypted);
let protocol_tag: [u8; 4] = decrypted[TAG_START..DC_START].try_into().ok()?;
if !matches!(protocol_tag, ABRIDGED | INTERMEDIATE | PADDED_INTERMEDIATE) {
return None;
}
let dc_index = i16::from_le_bytes([decrypted[DC_START], decrypted[DC_START + 1]]);
let dc = dc_index.unsigned_abs();
if !matches!(dc, 1..=5 | 203) {
return None;
}
let relay_init = generate_relay_init(protocol_tag, dc_index);
let crypto = build_crypto_context(init, secret, &relay_init)?;
Some(ClientInit {
dc,
media: dc_index < 0,
relay_init,
crypto,
})
}
fn secret_key(prekey: &[u8], secret: &[u8; 16]) -> [u8; 32] {
let mut hash = Sha256::new();
hash.update(prekey);
hash.update(secret);
hash.finalize().into()
}
fn generate_relay_init(protocol_tag: [u8; 4], dc_index: i16) -> [u8; INIT_LEN] {
loop {
let mut init = [0; INIT_LEN];
OsRng.fill_bytes(&mut init);
if is_reserved_init(&init) {
continue;
}
let key: [u8; 32] = init[KEY_START..KEY_END].try_into().unwrap();
let iv: [u8; 16] = init[KEY_END..IV_END].try_into().unwrap();
let mut cipher = AesCtr::new((&key).into(), (&iv).into());
let mut encrypted = init;
cipher.apply_keystream(&mut encrypted);
let mut tail = [0; 8];
tail[..4].copy_from_slice(&protocol_tag);
tail[4..6].copy_from_slice(&dc_index.to_le_bytes());
OsRng.fill_bytes(&mut tail[6..]);
for index in 0..8 {
init[TAG_START + index] ^= tail[index] ^ encrypted[TAG_START + index];
}
return init;
}
}
fn is_reserved_init(init: &[u8; INIT_LEN]) -> bool {
init[0] == 0xef
|| &init[..4] == b"HEAD"
|| &init[..4] == b"POST"
|| &init[..4] == b"GET "
|| init[..4] == [0xee; 4]
|| init[..4] == [0xdd; 4]
|| init[..4] == [0x16, 0x03, 0x01, 0x02]
|| init[4..8] == [0; 4]
}
fn build_crypto_context(
client_init: &[u8; INIT_LEN],
secret: &[u8; 16],
relay_init: &[u8; INIT_LEN],
) -> Option<CryptoContext> {
let client_dec_key = secret_key(&client_init[KEY_START..KEY_END], secret);
let client_dec_iv: [u8; 16] = client_init[KEY_END..IV_END].try_into().ok()?;
let mut client_decrypt = AesCtr::new((&client_dec_key).into(), (&client_dec_iv).into());
client_decrypt.apply_keystream(&mut [0; INIT_LEN]);
let reversed_client: Vec<_> = client_init[KEY_START..IV_END]
.iter()
.rev()
.copied()
.collect();
let client_enc_key = secret_key(&reversed_client[..32], secret);
let client_enc_iv: [u8; 16] = reversed_client[32..].try_into().ok()?;
let client_encrypt = AesCtr::new((&client_enc_key).into(), (&client_enc_iv).into());
let relay_enc_key: [u8; 32] = relay_init[KEY_START..KEY_END].try_into().ok()?;
let relay_enc_iv: [u8; 16] = relay_init[KEY_END..IV_END].try_into().ok()?;
let mut telegram_encrypt = AesCtr::new((&relay_enc_key).into(), (&relay_enc_iv).into());
telegram_encrypt.apply_keystream(&mut [0; INIT_LEN]);
let reversed_relay: Vec<_> = relay_init[KEY_START..IV_END]
.iter()
.rev()
.copied()
.collect();
let relay_dec_key: [u8; 32] = reversed_relay[..32].try_into().ok()?;
let relay_dec_iv: [u8; 16] = reversed_relay[32..].try_into().ok()?;
let telegram_decrypt = AesCtr::new((&relay_dec_key).into(), (&relay_dec_iv).into());
Some(CryptoContext {
client_decrypt,
client_encrypt,
telegram_encrypt,
telegram_decrypt,
})
}
#[cfg(test)]
pub(crate) fn test_client_init(secret: &[u8; 16], dc_index: i16) -> [u8; INIT_LEN] {
tests::generate_client_init(secret, PADDED_INTERMEDIATE, dc_index)
}
#[cfg(test)]
mod tests {
use super::*;
pub(super) fn generate_client_init(
secret: &[u8; 16],
protocol_tag: [u8; 4],
dc_index: i16,
) -> [u8; INIT_LEN] {
let mut init = generate_relay_init(protocol_tag, dc_index);
let key = secret_key(&init[KEY_START..KEY_END], secret);
let iv: [u8; 16] = init[KEY_END..IV_END].try_into().unwrap();
let mut cipher = AesCtr::new((&key).into(), (&iv).into());
let mut encrypted = init;
cipher.apply_keystream(&mut encrypted);
let mut tail = [0; 8];
tail[..4].copy_from_slice(&protocol_tag);
tail[4..6].copy_from_slice(&dc_index.to_le_bytes());
tail[6..].copy_from_slice(&[17, 23]);
for index in 0..8 {
init[TAG_START + index] ^= tail[index] ^ encrypted[TAG_START + index];
}
init
}
#[test]
fn parses_secret_protected_media_init() {
let secret = [42; 16];
let init = generate_client_init(&secret, PADDED_INTERMEDIATE, -4);
let parsed = parse_client_init(&init, &secret).unwrap();
assert_eq!(parsed.dc, 4);
assert!(parsed.media);
}
#[test]
fn rejects_wrong_secret() {
let init = generate_client_init(&[42; 16], INTERMEDIATE, 2);
assert!(parse_client_init(&init, &[7; 16]).is_none());
}
#[test]
fn telegram_link_secret_has_padded_intermediate_prefix() {
assert_eq!(
telegram_secret(&[0xab; 16]),
"ddabababababababababababababababab"
);
}
#[test]
fn parses_persisted_secret() {
assert_eq!(
parse_secret_hex("00112233445566778899aabbccddeeff"),
Some([
0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd,
0xee, 0xff,
])
);
assert_eq!(parse_secret_hex("not-a-secret"), None);
}
}
src/proxy.rs
+350 -79
View File
@@ -1,20 +1,26 @@
use std::net::Ipv4Addr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU8, Ordering};
use std::sync::atomic::{AtomicBool, AtomicU16, AtomicU32, AtomicU8, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio_tungstenite::tungstenite;
use tungstenite::client::IntoClientRequest;
pub const DEFAULT_PORT: u16 = 1080;
const IO_TIMEOUT: Duration = Duration::from_secs(10);
pub struct Stats {
pub running: AtomicBool,
pub active: AtomicU32,
pub total: AtomicU32,
pub ws: AtomicU32,
pub last_dc: AtomicU8,
pub last_dc: AtomicU16,
pub ws_failures: AtomicU32,
/// See `transport::RouteKind::ui_code`.
pub last_route: AtomicU8,
transport: crate::transport::TransportEngine,
secret: [u8; 16],
shutdown: Mutex<Option<tokio::sync::watch::Sender<bool>>>,
}
impl Stats {
@@ -24,9 +30,43 @@ impl Stats {
active: AtomicU32::new(0),
total: AtomicU32::new(0),
ws: AtomicU32::new(0),
last_dc: AtomicU8::new(0),
last_dc: AtomicU16::new(0),
ws_failures: AtomicU32::new(0),
last_route: AtomicU8::new(0),
transport: crate::transport::TransportEngine::new(),
secret: initial_secret(),
shutdown: Mutex::new(None),
})
}
pub fn telegram_secret(&self) -> String {
crate::mtproto::telegram_secret(&self.secret)
}
pub fn set_worker_domain(&self, domain: &str) {
let domains = domain
.split([',', ';', ' '])
.filter(|value| !value.trim().is_empty())
.map(str::to_owned)
.collect::<Vec<_>>();
self.transport.set_worker_domains(&domains);
}
pub fn stop(&self) {
if let Some(sender) = self.shutdown.lock().unwrap().take() {
let _ = sender.send(true);
}
}
}
#[cfg(not(test))]
fn initial_secret() -> [u8; 16] {
crate::mtproto::load_or_create_secret()
}
#[cfg(test)]
fn initial_secret() -> [u8; 16] {
crate::mtproto::generate_secret()
}
pub async fn run(stats: Arc<Stats>, lan: bool, port: u16) -> Result<(), String> {
@@ -37,25 +77,38 @@ pub async fn run(stats: Arc<Stats>, lan: bool, port: u16) -> Result<(), String>
.map_err(|e| format!("Port {} busy: {}", port, e))?;
stats.running.store(true, Ordering::SeqCst);
let (shutdown_tx, mut shutdown_rx) = tokio::sync::watch::channel(false);
*stats.shutdown.lock().unwrap() = Some(shutdown_tx);
let mut tasks = tokio::task::JoinSet::new();
loop {
if !stats.running.load(Ordering::SeqCst) {
break;
}
tokio::select! {
Ok((stream, _)) = listener.accept() => {
let s = stats.clone();
s.active.fetch_add(1, Ordering::Relaxed);
s.total.fetch_add(1, Ordering::Relaxed);
tokio::spawn(async move {
let _ = handle(stream, &s).await;
s.active.fetch_sub(1, Ordering::Relaxed);
});
result = listener.accept() => {
match result {
Ok((stream, _)) => {
let s = stats.clone();
s.active.fetch_add(1, Ordering::Relaxed);
s.total.fetch_add(1, Ordering::Relaxed);
tasks.spawn(async move {
let _ = handle(stream, &s, !lan).await;
s.active.fetch_sub(1, Ordering::Relaxed);
});
}
Err(error) => {
stats.running.store(false, Ordering::SeqCst);
return Err(format!("Accept failed: {}", error));
}
}
}
_ = tokio::time::sleep(Duration::from_millis(150)) => {}
_ = shutdown_rx.changed() => break,
Some(_) = tasks.join_next(), if !tasks.is_empty() => {}
}
}
tasks.abort_all();
while tasks.join_next().await.is_some() {}
stats.active.store(0, Ordering::Relaxed);
stats.ws.store(0, Ordering::Relaxed);
stats.running.store(false, Ordering::SeqCst);
Ok(())
}
@@ -63,79 +116,179 @@ pub async fn run(stats: Arc<Stats>, lan: bool, port: u16) -> Result<(), String>
// -- SOCKS5 -----------------------------------------------------------------
async fn handle(
s: TcpStream,
stats: &Stats,
allow_direct: bool,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let mut first = [0; 1];
let peeked = tokio::time::timeout(IO_TIMEOUT, s.peek(&mut first))
.await
.map_err(|_| "client protocol detection timeout")??;
if peeked == 0 {
return Ok(());
}
if first[0] == 0x05 {
handle_socks5(s, stats, allow_direct).await
} else {
handle_mtproto(s, stats).await
}
}
async fn handle_socks5(
mut s: TcpStream,
stats: &Stats,
allow_direct: bool,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
s.set_nodelay(true)?;
let mut buf = [0u8; 258];
let n = s.read(&mut buf).await?;
if n < 2 || buf[0] != 0x05 {
return Err("not socks5".into());
}
s.write_all(&[0x05, 0x00]).await?;
let n = s.read(&mut buf).await?;
if n < 7 || buf[0] != 0x05 || buf[1] != 0x01 {
s.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
return Err("bad connect".into());
}
let (addr, port) = parse_addr(&buf[3..n])?;
let (addr, port) = tokio::time::timeout(IO_TIMEOUT, read_socks5_request(&mut s))
.await
.map_err(|_| "SOCKS5 handshake timeout")??;
let tg = addr.parse::<Ipv4Addr>().ok().and_then(dc_from_ip).is_some();
if !tg && !allow_direct {
write_socks_reply(&mut s, 0x02).await?;
return Err("LAN mode only permits Telegram destinations".into());
}
// success reply
s.write_all(&[0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, 0x04, 0x38]).await?;
s.write_all(&[0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, 0x04, 0x38])
.await?;
if tg {
// Read 64-byte obfuscated2 init → extract real DC
let mut init = [0u8; 64];
s.read_exact(&mut init).await?;
let dc = dc_from_init(&init).unwrap_or_else(|| {
addr.parse::<Ipv4Addr>().ok().and_then(dc_from_ip).unwrap_or(2)
let (dc, media) = dc_from_init(&init).unwrap_or_else(|| {
addr.parse::<Ipv4Addr>()
.ok()
.and_then(dc_from_ip)
.map_or((2, false), |dc| (dc, false))
});
stats.last_dc.store(dc, Ordering::Relaxed);
stats.ws.fetch_add(1, Ordering::Relaxed);
let r = ws_tunnel(s, dc, &init).await;
let r = ws_tunnel(s, dc, media, &init, None, stats).await;
stats.ws.fetch_sub(1, Ordering::Relaxed);
if r.is_err() {
stats.ws_failures.fetch_add(1, Ordering::Relaxed);
}
r?;
} else {
let remote = TcpStream::connect(format!("{}:{}", addr, port)).await?;
let remote = tokio::time::timeout(IO_TIMEOUT, TcpStream::connect((addr.as_str(), port)))
.await
.map_err(|_| "direct connection timeout")??;
let _ = remote.set_nodelay(true);
tcp_relay(s, remote).await;
}
Ok(())
}
fn parse_addr(d: &[u8]) -> Result<(String, u16), Box<dyn std::error::Error + Send + Sync>> {
match d[0] {
0x01 if d.len() >= 7 => {
Ok((format!("{}.{}.{}.{}", d[1], d[2], d[3], d[4]),
u16::from_be_bytes([d[5], d[6]])))
async fn handle_mtproto(
mut stream: TcpStream,
stats: &Stats,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
stream.set_nodelay(true)?;
let mut init = [0; 64];
tokio::time::timeout(IO_TIMEOUT, stream.read_exact(&mut init))
.await
.map_err(|_| "MTProto init timeout")??;
let parsed = crate::mtproto::parse_client_init(&init, &stats.secret)
.ok_or("invalid MTProto init or secret")?;
stats.last_dc.store(parsed.dc, Ordering::Relaxed);
stats.ws.fetch_add(1, Ordering::Relaxed);
let result = ws_tunnel(
stream,
parsed.dc,
parsed.media,
&parsed.relay_init,
Some(parsed.crypto),
stats,
)
.await;
stats.ws.fetch_sub(1, Ordering::Relaxed);
if result.is_err() {
stats.ws_failures.fetch_add(1, Ordering::Relaxed);
}
result
}
async fn read_socks5_request<S>(
stream: &mut S,
) -> Result<(String, u16), Box<dyn std::error::Error + Send + Sync>>
where
S: AsyncRead + AsyncWrite + Unpin,
{
let version = stream.read_u8().await?;
let method_count = stream.read_u8().await? as usize;
if version != 0x05 || method_count == 0 {
return Err("invalid SOCKS5 greeting".into());
}
let mut methods = vec![0; method_count];
stream.read_exact(&mut methods).await?;
if !methods.contains(&0x00) {
stream.write_all(&[0x05, 0xff]).await?;
return Err("SOCKS5 client does not support no-auth mode".into());
}
stream.write_all(&[0x05, 0x00]).await?;
let version = stream.read_u8().await?;
let command = stream.read_u8().await?;
let reserved = stream.read_u8().await?;
let address_type = stream.read_u8().await?;
if version != 0x05 || reserved != 0 {
return Err("invalid SOCKS5 request".into());
}
if command != 0x01 {
write_socks_reply(stream, 0x07).await?;
return Err("unsupported SOCKS5 command".into());
}
let address = match address_type {
0x01 => {
let mut octets = [0; 4];
stream.read_exact(&mut octets).await?;
Ipv4Addr::from(octets).to_string()
}
0x03 => {
let l = d[1] as usize;
if d.len() < 2 + l + 2 { return Err("short".into()); }
Ok((std::str::from_utf8(&d[2..2 + l])?.into(),
u16::from_be_bytes([d[2 + l], d[3 + l]])))
let length = stream.read_u8().await? as usize;
if length == 0 {
return Err("empty SOCKS5 domain".into());
}
let mut domain = vec![0; length];
stream.read_exact(&mut domain).await?;
String::from_utf8(domain)?
}
0x04 if d.len() >= 19 => {
let mut seg = [0u16; 8];
for i in 0..8 { seg[i] = u16::from_be_bytes([d[1 + i * 2], d[2 + i * 2]]); }
let ip = std::net::Ipv6Addr::new(seg[0],seg[1],seg[2],seg[3],seg[4],seg[5],seg[6],seg[7]);
Ok((ip.to_string(), u16::from_be_bytes([d[17], d[18]])))
0x04 => {
let mut octets = [0; 16];
stream.read_exact(&mut octets).await?;
std::net::Ipv6Addr::from(octets).to_string()
}
_ => Err("bad addr".into()),
}
_ => {
write_socks_reply(stream, 0x08).await?;
return Err("unsupported SOCKS5 address type".into());
}
};
let port = stream.read_u16().await?;
Ok((address, port))
}
async fn write_socks_reply<S>(stream: &mut S, status: u8) -> std::io::Result<()>
where
S: AsyncWrite + Unpin,
{
stream
.write_all(&[0x05, status, 0x00, 0x01, 127, 0, 0, 1, 0, 0])
.await
}
// -- DC detection -----------------------------------------------------------
fn dc_from_init(init: &[u8; 64]) -> Option<u8> {
fn dc_from_init(init: &[u8; 64]) -> Option<(u16, bool)> {
use aes::Aes256;
use cipher::{KeyIvInit, StreamCipher};
type Ctr = ctr::Ctr128BE<Aes256>;
@@ -144,47 +297,49 @@ fn dc_from_init(init: &[u8; 64]) -> Option<u8> {
let mut c = Ctr::new(init[8..40].into(), init[40..56].into());
c.apply_keystream(&mut dec);
let id = i32::from_le_bytes([dec[60], dec[61], dec[62], dec[63]]);
let dc = id.unsigned_abs() as u8;
(1..=5).contains(&dc).then_some(dc)
let id = i16::from_le_bytes([dec[60], dec[61]]);
let dc = id.unsigned_abs();
matches!(dc, 1..=5 | 203).then_some((dc, id < 0))
}
fn dc_from_ip(ip: Ipv4Addr) -> Option<u8> {
fn dc_from_ip(ip: Ipv4Addr) -> Option<u16> {
let o = ip.octets();
match (o[0], o[1]) {
(149, 154) => Some(match o[2] { 160..=163 => 1, 164..=167 => 2, 168..=171 => 3, 172..=175 => 1, _ => 2 }),
(91, 108) => Some(match o[2] { 56..=59 => 5, 8..=11 => 3, 12..=15 => 4, _ => 2 }),
(149, 154) => Some(match o[2] {
160..=163 => 1,
164..=167 => 2,
168..=171 => 3,
172..=175 => 1,
_ => 2,
}),
(91, 108) => Some(match o[2] {
56..=59 => 5,
8..=11 => 3,
12..=15 => 4,
_ => 2,
}),
(91, 105) if o[2] == 192 => Some(203),
(91, 105) | (185, 76) => Some(2),
_ => None,
}
}
fn ws_url(dc: u8) -> String {
format!("wss://kws{}.web.telegram.org/apiws", dc)
}
// -- WebSocket tunnel -------------------------------------------------------
async fn ws_tunnel(
tcp: TcpStream,
dc: u8,
dc: u16,
media: bool,
init: &[u8; 64],
mut crypto: Option<crate::mtproto::CryptoContext>,
stats: &Stats,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
use futures_util::{SinkExt, StreamExt};
let mut req = ws_url(dc).as_str().into_client_request()?;
req.headers_mut().insert("Sec-WebSocket-Protocol", "binary".parse()?);
let tls = native_tls::TlsConnector::new().map_err(|e| format!("tls: {}", e))?;
let connector = tokio_tungstenite::Connector::NativeTls(tls);
let (mut ws, _) = tokio::time::timeout(
Duration::from_secs(10),
tokio_tungstenite::connect_async_tls_with_config(req, None, false, Some(connector)),
)
.await
.map_err(|_| "WS connect timeout")?
.map_err(|e| format!("WS: {}", e))?;
let (mut ws, connected) = stats.transport.connect(dc, media).await?;
stats
.last_route
.store(connected.route.kind.ui_code(), Ordering::Relaxed);
let (mut tcp_r, mut tcp_w) = tokio::io::split(tcp);
@@ -198,7 +353,10 @@ async fn ws_tunnel(
biased;
msg = ws.next() => match msg {
Some(Ok(tungstenite::Message::Binary(data))) => {
Some(Ok(tungstenite::Message::Binary(mut data))) => {
if let Some(crypto) = &mut crypto {
crypto.telegram_to_client(data.as_mut());
}
tcp_w.write_all(data.as_ref()).await?;
tcp_w.flush().await?;
}
@@ -213,6 +371,9 @@ async fn ws_tunnel(
n = tcp_r.read(&mut buf) => match n {
Ok(0) | Err(_) => break,
Ok(n) => {
if let Some(crypto) = &mut crypto {
crypto.client_to_telegram(&mut buf[..n]);
}
ws.send(tungstenite::Message::Binary(buf[..n].to_vec())).await?;
}
},
@@ -231,3 +392,113 @@ async fn tcp_relay(a: TcpStream, b: TcpStream) {
_ = tokio::io::copy(&mut br, &mut aw) => {}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
#[tokio::test]
async fn parses_fragmented_ipv4_socks5_handshake() {
let (mut client, mut server) = tokio::io::duplex(64);
let server_task =
tokio::spawn(async move { read_socks5_request(&mut server).await.unwrap() });
for part in [&[0x05][..], &[0x01, 0x00]] {
client.write_all(part).await.unwrap();
tokio::task::yield_now().await;
}
let mut method_reply = [0; 2];
client.read_exact(&mut method_reply).await.unwrap();
assert_eq!(method_reply, [0x05, 0x00]);
for byte in [0x05, 0x01, 0x00, 0x01, 149, 154, 167, 50, 0x01, 0xbb] {
client.write_all(&[byte]).await.unwrap();
tokio::task::yield_now().await;
}
assert_eq!(
server_task.await.unwrap(),
("149.154.167.50".to_owned(), 443)
);
}
#[tokio::test]
async fn rejects_auth_only_clients() {
let (mut client, mut server) = tokio::io::duplex(16);
let server_task =
tokio::spawn(async move { read_socks5_request(&mut server).await.is_err() });
client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
let mut reply = [0; 2];
client.read_exact(&mut reply).await.unwrap();
assert_eq!(reply, [0x05, 0xff]);
assert!(server_task.await.unwrap());
}
#[test]
fn maps_known_telegram_networks_to_dc() {
assert_eq!(dc_from_ip("149.154.160.1".parse().unwrap()), Some(1));
assert_eq!(dc_from_ip("149.154.167.255".parse().unwrap()), Some(2));
assert_eq!(dc_from_ip("91.108.58.1".parse().unwrap()), Some(5));
assert_eq!(dc_from_ip("1.1.1.1".parse().unwrap()), None);
}
#[tokio::test]
async fn stops_listener_and_active_tasks_cleanly() {
let reservation = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let port = reservation.local_addr().unwrap().port();
drop(reservation);
let stats = Stats::new();
let server_stats = stats.clone();
let server = tokio::spawn(async move { run(server_stats, false, port).await });
tokio::time::timeout(Duration::from_secs(2), async {
while !stats.running.load(Ordering::SeqCst) {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
stats.stop();
assert!(tokio::time::timeout(Duration::from_secs(2), server)
.await
.unwrap()
.unwrap()
.is_ok());
assert!(!stats.running.load(Ordering::SeqCst));
}
#[tokio::test]
#[ignore = "requires live Telegram network access"]
async fn accepts_mtproto_and_builds_live_media_tunnel() {
let reservation = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
let port = reservation.local_addr().unwrap().port();
drop(reservation);
let stats = Stats::new();
let server_stats = stats.clone();
let server = tokio::spawn(async move { run(server_stats, false, port).await });
while !stats.running.load(Ordering::SeqCst) {
tokio::task::yield_now().await;
}
let mut client = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
let init = crate::mtproto::test_client_init(&stats.secret, -4);
client.write_all(&init).await.unwrap();
tokio::time::timeout(Duration::from_secs(10), async {
while stats.last_route.load(Ordering::Relaxed) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
assert_eq!(stats.last_dc.load(Ordering::Relaxed), 4);
assert_eq!(stats.ws_failures.load(Ordering::Relaxed), 0);
stats.stop();
server.await.unwrap().unwrap();
}
}
src/transport.rs
+386
View File
@@ -0,0 +1,386 @@
use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
const CONNECT_TIMEOUT: Duration = Duration::from_secs(4);
const FAILURE_BACKOFF_INITIAL: Duration = Duration::from_secs(30);
const FAILURE_BACKOFF_MAX: Duration = Duration::from_secs(30 * 60);
pub type TelegramWebSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum RouteKind {
TelegramIp,
AlternateTelegramIp,
SystemDns,
CloudflareWorker,
}
impl RouteKind {
pub fn ui_code(self) -> u8 {
match self {
Self::TelegramIp => 1,
Self::AlternateTelegramIp => 2,
Self::SystemDns => 3,
Self::CloudflareWorker => 4,
}
}
}
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct Route {
pub connect_host: String,
pub websocket_host: String,
pub path: String,
pub kind: RouteKind,
}
#[derive(Clone, Debug)]
pub struct ConnectedRoute {
pub route: Route,
}
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
struct DcKey {
dc: u16,
media: bool,
}
#[derive(Clone, Copy, Debug)]
struct RouteHealth {
failures: u32,
retry_at: Instant,
}
#[derive(Default)]
struct HealthState {
routes: HashMap<Route, RouteHealth>,
preferred: HashMap<DcKey, Route>,
}
#[derive(Default)]
pub struct TransportEngine {
health: Mutex<HealthState>,
worker_domains: Mutex<Vec<String>>,
}
impl TransportEngine {
pub fn new() -> Self {
Self::default()
}
pub fn set_worker_domains(&self, domains: &[String]) {
let mut normalized = Vec::new();
for domain in domains {
let domain = domain.trim().to_ascii_lowercase();
if valid_domain(&domain) && !normalized.contains(&domain) {
normalized.push(domain);
}
}
*self.worker_domains.lock().unwrap() = normalized;
}
pub async fn connect(
&self,
dc: u16,
media: bool,
) -> Result<(TelegramWebSocket, ConnectedRoute), String> {
let key = DcKey { dc, media };
let candidates = self.ordered_candidates(key);
let mut errors = Vec::new();
for route in candidates {
match connect_route(&route).await {
Ok(websocket) => {
self.record_success(key, &route);
return Ok((websocket, ConnectedRoute { route }));
}
Err(error) => {
self.record_failure(&route);
errors.push(format!(
"{} via {}: {}",
route.websocket_host, route.connect_host, error
));
}
}
}
Err(format!(
"all Telegram routes for DC{}{} failed: {}",
dc,
if media { " media" } else { "" },
errors.join("; ")
))
}
fn ordered_candidates(&self, key: DcKey) -> Vec<Route> {
let now = Instant::now();
let all_routes = self.routes_for_key(key);
let health = self.health.lock().unwrap();
let preferred = health.preferred.get(&key);
let mut candidates: Vec<_> = all_routes
.iter()
.filter(|route| {
health
.routes
.get(route)
.is_none_or(|route_health| route_health.retry_at <= now)
})
.cloned()
.collect();
candidates.sort_by_key(|route| {
let preferred_rank = u8::from(preferred != Some(route));
let kind_rank = match route.kind {
RouteKind::TelegramIp => 0,
RouteKind::AlternateTelegramIp => 1,
RouteKind::SystemDns => 2,
RouteKind::CloudflareWorker => 3,
};
(preferred_rank, kind_rank)
});
// If every route is cooling down, retry the one that becomes available first.
if candidates.is_empty() {
if let Some((route, _)) = health
.routes
.iter()
.filter(|(route, _)| all_routes.contains(route))
.min_by_key(|(_, route_health)| route_health.retry_at)
{
candidates.push(route.clone());
}
}
candidates
}
fn routes_for_key(&self, key: DcKey) -> Vec<Route> {
let mut routes = routes_for_dc(key.dc, key.media);
let Some(destination) = telegram_ips(key.dc).first() else {
return routes;
};
for domain in self.worker_domains.lock().unwrap().iter() {
routes.push(Route {
connect_host: domain.clone(),
websocket_host: domain.clone(),
path: format!("/apiws?dst={}&dc={}", destination, key.dc),
kind: RouteKind::CloudflareWorker,
});
}
routes
}
fn record_success(&self, key: DcKey, route: &Route) {
let mut health = self.health.lock().unwrap();
health.routes.remove(route);
health.preferred.insert(key, route.clone());
}
fn record_failure(&self, route: &Route) {
let mut health = self.health.lock().unwrap();
let failures = health
.routes
.get(route)
.map_or(1, |route_health| route_health.failures.saturating_add(1));
let exponent = failures.saturating_sub(1).min(6);
let delay = FAILURE_BACKOFF_INITIAL
.saturating_mul(2_u32.pow(exponent))
.min(FAILURE_BACKOFF_MAX);
health.routes.insert(
route.clone(),
RouteHealth {
failures,
retry_at: Instant::now() + delay,
},
);
}
}
fn canonical_dc(dc: u16) -> u16 {
if dc == 203 {
2
} else {
dc
}
}
fn telegram_ips(dc: u16) -> &'static [&'static str] {
match dc {
1 => &["149.154.175.50"],
2 => &["149.154.167.51", "149.154.167.220"],
3 => &["149.154.175.100"],
4 => &["149.154.167.91", "149.154.167.220"],
5 => &["149.154.171.5"],
203 => &["91.105.192.100"],
_ => &[],
}
}
pub fn routes_for_dc(dc: u16, media: bool) -> Vec<Route> {
let websocket_dc = canonical_dc(dc);
let primary = format!("kws{}.web.telegram.org", websocket_dc);
let secondary = format!("kws{}-1.web.telegram.org", websocket_dc);
let websocket_hosts = if media {
[secondary, primary]
} else {
[primary, secondary]
};
let ips = telegram_ips(dc);
let mut routes = Vec::new();
for websocket_host in &websocket_hosts {
for (index, ip) in ips.iter().enumerate() {
routes.push(Route {
connect_host: (*ip).to_owned(),
websocket_host: websocket_host.clone(),
path: "/apiws".to_owned(),
kind: if index == 0 {
RouteKind::TelegramIp
} else {
RouteKind::AlternateTelegramIp
},
});
}
routes.push(Route {
connect_host: websocket_host.clone(),
websocket_host: websocket_host.clone(),
path: "/apiws".to_owned(),
kind: RouteKind::SystemDns,
});
}
routes
}
async fn connect_route(route: &Route) -> Result<TelegramWebSocket, String> {
let tcp = tokio::time::timeout(
CONNECT_TIMEOUT,
TcpStream::connect((route.connect_host.as_str(), 443)),
)
.await
.map_err(|_| "TCP connect timeout".to_owned())?
.map_err(|error| format!("TCP connect: {}", error))?;
tcp.set_nodelay(true)
.map_err(|error| format!("TCP_NODELAY: {}", error))?;
let url = format!("wss://{}{}", route.websocket_host, route.path);
let mut request = url
.as_str()
.into_client_request()
.map_err(|error| format!("WebSocket request: {}", error))?;
request.headers_mut().insert(
"Sec-WebSocket-Protocol",
"binary"
.parse()
.map_err(|error| format!("WebSocket protocol header: {}", error))?,
);
// The URI host remains the real Telegram hostname even when the TCP socket
// is opened to a pinned IP. Native TLS therefore validates Telegram's
// certificate and sends the correct SNI.
let tls = native_tls::TlsConnector::new().map_err(|error| format!("TLS setup: {}", error))?;
let connector = tokio_tungstenite::Connector::NativeTls(tls);
tokio::time::timeout(
CONNECT_TIMEOUT,
tokio_tungstenite::client_async_tls_with_config(request, tcp, None, Some(connector)),
)
.await
.map_err(|_| "TLS/WebSocket timeout".to_owned())?
.map(|(websocket, _)| websocket)
.map_err(|error| format!("TLS/WebSocket handshake: {}", error))
}
fn valid_domain(domain: &str) -> bool {
if domain.is_empty() || domain.len() > 253 || !domain.contains('.') {
return false;
}
domain.split('.').all(|label| {
!label.is_empty()
&& label.len() <= 63
&& !label.starts_with('-')
&& !label.ends_with('-')
&& label
.chars()
.all(|character| character.is_ascii_alphanumeric() || character == '-')
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn media_prefers_dash_one_websocket_host() {
let routes = routes_for_dc(2, true);
assert_eq!(routes[0].websocket_host, "kws2-1.web.telegram.org");
assert_eq!(routes[0].connect_host, "149.154.167.51");
}
#[test]
fn dc203_uses_dc2_websocket_and_its_own_ip() {
let routes = routes_for_dc(203, false);
assert_eq!(routes[0].websocket_host, "kws2.web.telegram.org");
assert_eq!(routes[0].connect_host, "91.105.192.100");
}
#[test]
fn successful_route_becomes_preferred() {
let engine = TransportEngine::new();
let key = DcKey {
dc: 2,
media: false,
};
let preferred = routes_for_dc(2, false)[2].clone();
engine.record_success(key, &preferred);
assert_eq!(engine.ordered_candidates(key)[0], preferred);
}
#[test]
fn failed_route_enters_cooldown() {
let engine = TransportEngine::new();
let key = DcKey {
dc: 2,
media: false,
};
let failed = routes_for_dc(2, false)[0].clone();
engine.record_failure(&failed);
assert!(!engine.ordered_candidates(key).contains(&failed));
}
#[test]
fn worker_is_explicit_and_validated() {
let engine = TransportEngine::new();
engine.set_worker_domains(&[
"Example.User.Workers.dev".to_owned(),
"https://invalid.example/path".to_owned(),
]);
let routes = engine.routes_for_key(DcKey { dc: 4, media: true });
let worker = routes
.iter()
.find(|route| route.kind == RouteKind::CloudflareWorker)
.unwrap();
assert_eq!(worker.websocket_host, "example.user.workers.dev");
assert_eq!(worker.path, "/apiws?dst=149.154.167.91&dc=4");
assert_eq!(
routes
.iter()
.filter(|route| route.kind == RouteKind::CloudflareWorker)
.count(),
1
);
}
#[tokio::test]
#[ignore = "requires live Telegram network access"]
async fn connects_to_all_production_data_centers() {
let engine = TransportEngine::new();
for dc in [1, 2, 3, 4, 5, 203] {
let (mut websocket, connected) = engine.connect(dc, false).await.unwrap();
assert!(!connected.route.websocket_host.is_empty());
websocket.close(None).await.unwrap();
}
}
}