mirror of https://github.com/by-sonic/tglock.git
307 lines
9.6 KiB
Rust
307 lines
9.6 KiB
Rust
use std::net::Ipv4Addr;
|
||
use std::sync::Arc;
|
||
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
||
use tokio::net::{TcpListener, TcpStream};
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
use tokio_tungstenite::tungstenite;
|
||
use tungstenite::client::IntoClientRequest;
|
||
|
||
pub struct ProxyStats {
|
||
pub running: AtomicBool,
|
||
pub active_conn: AtomicU32,
|
||
pub total_conn: AtomicU32,
|
||
pub ws_active: AtomicU32,
|
||
}
|
||
|
||
impl ProxyStats {
|
||
pub fn new() -> Arc<Self> {
|
||
Arc::new(Self {
|
||
running: AtomicBool::new(false),
|
||
active_conn: AtomicU32::new(0),
|
||
total_conn: AtomicU32::new(0),
|
||
ws_active: AtomicU32::new(0),
|
||
})
|
||
}
|
||
}
|
||
|
||
pub async fn run_proxy(port: u16, stats: Arc<ProxyStats>) -> Result<(), String> {
|
||
let addr = format!("127.0.0.1:{}", port);
|
||
let listener = TcpListener::bind(&addr)
|
||
.await
|
||
.map_err(|e| format!("Не удалось занять порт {}: {}", port, e))?;
|
||
|
||
stats.running.store(true, Ordering::SeqCst);
|
||
|
||
loop {
|
||
if !stats.running.load(Ordering::SeqCst) {
|
||
break;
|
||
}
|
||
tokio::select! {
|
||
result = listener.accept() => {
|
||
if let Ok((stream, _)) = result {
|
||
let st = stats.clone();
|
||
st.active_conn.fetch_add(1, Ordering::Relaxed);
|
||
st.total_conn.fetch_add(1, Ordering::Relaxed);
|
||
tokio::spawn(async move {
|
||
let _ = handle_socks5(stream, &st).await;
|
||
st.active_conn.fetch_sub(1, Ordering::Relaxed);
|
||
});
|
||
}
|
||
}
|
||
_ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {}
|
||
}
|
||
}
|
||
|
||
stats.running.store(false, Ordering::SeqCst);
|
||
Ok(())
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// DC extraction from obfuscated2 init packet (same method as tg-ws-proxy)
|
||
// ---------------------------------------------------------------------------
|
||
|
||
fn extract_dc_from_init(init: &[u8; 64]) -> Option<u8> {
|
||
use aes::Aes256;
|
||
use cipher::{KeyIvInit, StreamCipher};
|
||
type Aes256Ctr = ctr::Ctr128BE<Aes256>;
|
||
|
||
let key = &init[8..40];
|
||
let iv = &init[40..56];
|
||
|
||
let mut dec = [0u8; 64];
|
||
dec.copy_from_slice(init);
|
||
|
||
let mut cipher = Aes256Ctr::new(key.into(), iv.into());
|
||
cipher.apply_keystream(&mut dec);
|
||
|
||
let dc_id = i32::from_le_bytes([dec[60], dec[61], dec[62], dec[63]]);
|
||
let dc = dc_id.unsigned_abs() as u8;
|
||
if (1..=5).contains(&dc) {
|
||
Some(dc)
|
||
} else {
|
||
None
|
||
}
|
||
}
|
||
|
||
fn dc_from_ip(ip: Ipv4Addr) -> Option<u8> {
|
||
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,
|
||
}),
|
||
(91, 105) | (185, 76) => Some(2),
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
fn is_telegram_ip(addr: &str) -> bool {
|
||
addr.parse::<Ipv4Addr>()
|
||
.ok()
|
||
.and_then(dc_from_ip)
|
||
.is_some()
|
||
}
|
||
|
||
/// Endpoint format used by the proven tg-ws-proxy project
|
||
fn ws_url(dc: u8) -> String {
|
||
format!("wss://kws{}.web.telegram.org/apiws", dc)
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// SOCKS5 handler
|
||
// ---------------------------------------------------------------------------
|
||
|
||
async fn handle_socks5(
|
||
mut stream: TcpStream,
|
||
stats: &ProxyStats,
|
||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||
stream.set_nodelay(true)?;
|
||
|
||
// --- auth negotiation ---
|
||
let mut buf = [0u8; 258];
|
||
let n = stream.read(&mut buf).await?;
|
||
if n < 2 || buf[0] != 0x05 {
|
||
return Err("Not SOCKS5".into());
|
||
}
|
||
stream.write_all(&[0x05, 0x00]).await?;
|
||
|
||
// --- CONNECT request ---
|
||
let n = stream.read(&mut buf).await?;
|
||
if n < 7 || buf[0] != 0x05 || buf[1] != 0x01 {
|
||
stream.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
|
||
return Err("Bad CONNECT".into());
|
||
}
|
||
|
||
let (dest_addr, dest_port) = parse_dest(&buf[3..n])?;
|
||
let is_tg = is_telegram_ip(&dest_addr);
|
||
|
||
// SOCKS5 success (we handle the connection ourselves)
|
||
stream
|
||
.write_all(&[0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, 0x04, 0x38])
|
||
.await?;
|
||
|
||
if is_tg {
|
||
// Read the first 64 bytes — obfuscated2 init packet
|
||
let mut init = [0u8; 64];
|
||
stream.read_exact(&mut init).await?;
|
||
|
||
// Extract DC from init packet (primary), fall back to IP-based
|
||
let dc = extract_dc_from_init(&init).unwrap_or_else(|| {
|
||
dest_addr
|
||
.parse::<Ipv4Addr>()
|
||
.ok()
|
||
.and_then(dc_from_ip)
|
||
.unwrap_or(2)
|
||
});
|
||
|
||
stats.ws_active.fetch_add(1, Ordering::Relaxed);
|
||
|
||
// Try WebSocket tunnel; fall back to direct TCP on failure
|
||
let ws_result = relay_via_ws(stream, dc, &init).await;
|
||
|
||
stats.ws_active.fetch_sub(1, Ordering::Relaxed);
|
||
|
||
if let Err(e) = ws_result {
|
||
return Err(format!("DC{} tunnel: {}", dc, e).into());
|
||
}
|
||
} else {
|
||
// Non-Telegram — direct TCP passthrough
|
||
let target = format!("{}:{}", dest_addr, dest_port);
|
||
match TcpStream::connect(&target).await {
|
||
Ok(remote) => {
|
||
let _ = remote.set_nodelay(true);
|
||
relay_tcp(stream, remote).await;
|
||
}
|
||
Err(e) => {
|
||
return Err(format!("TCP connect {}: {}", target, e).into());
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
fn parse_dest(data: &[u8]) -> Result<(String, u16), Box<dyn std::error::Error + Send + Sync>> {
|
||
match data[0] {
|
||
0x01 => {
|
||
if data.len() < 7 { return Err("short".into()); }
|
||
let ip = format!("{}.{}.{}.{}", data[1], data[2], data[3], data[4]);
|
||
let port = u16::from_be_bytes([data[5], data[6]]);
|
||
Ok((ip, port))
|
||
}
|
||
0x03 => {
|
||
let len = data[1] as usize;
|
||
if data.len() < 2 + len + 2 { return Err("short".into()); }
|
||
let domain = std::str::from_utf8(&data[2..2 + len])?.to_string();
|
||
let port = u16::from_be_bytes([data[2 + len], data[3 + len]]);
|
||
Ok((domain, port))
|
||
}
|
||
0x04 => {
|
||
if data.len() < 19 { return Err("short".into()); }
|
||
let port = u16::from_be_bytes([data[17], data[18]]);
|
||
let mut segs = [0u16; 8];
|
||
for i in 0..8 {
|
||
segs[i] = u16::from_be_bytes([data[1 + i * 2], data[2 + i * 2]]);
|
||
}
|
||
let ip = std::net::Ipv6Addr::new(
|
||
segs[0], segs[1], segs[2], segs[3], segs[4], segs[5], segs[6], segs[7],
|
||
);
|
||
Ok((ip.to_string(), port))
|
||
}
|
||
_ => Err("unknown addr type".into()),
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// WebSocket tunnel — reads init first, then relays
|
||
// ---------------------------------------------------------------------------
|
||
|
||
async fn relay_via_ws(
|
||
tcp_stream: TcpStream,
|
||
dc: u8,
|
||
init: &[u8; 64],
|
||
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
||
use futures_util::{SinkExt, StreamExt};
|
||
|
||
let url = ws_url(dc);
|
||
let mut request = url.as_str().into_client_request()?;
|
||
|
||
request
|
||
.headers_mut()
|
||
.insert("Sec-WebSocket-Protocol", "binary".parse()?);
|
||
|
||
let connector = tokio_tungstenite::Connector::NativeTls(
|
||
native_tls::TlsConnector::new().map_err(|e| format!("TLS: {}", e))?,
|
||
);
|
||
|
||
let (mut ws, _resp) = tokio_tungstenite::connect_async_tls_with_config(
|
||
request, None, false, Some(connector),
|
||
)
|
||
.await?;
|
||
|
||
let (mut tcp_rx, mut tcp_tx) = tokio::io::split(tcp_stream);
|
||
|
||
// Send the buffered 64-byte init as the first WebSocket message
|
||
ws.send(tungstenite::Message::Binary(init.to_vec())).await?;
|
||
|
||
// Single loop: handles TCP→WS, WS→TCP, and Ping/Pong in one place.
|
||
// This ensures Pong replies are sent immediately so the server
|
||
// doesn't kill the connection after a timeout.
|
||
let mut buf = vec![0u8; 32768];
|
||
|
||
loop {
|
||
tokio::select! {
|
||
biased;
|
||
|
||
ws_msg = ws.next() => {
|
||
match ws_msg {
|
||
Some(Ok(tungstenite::Message::Binary(data))) => {
|
||
if tcp_tx.write_all(data.as_ref()).await.is_err() {
|
||
break;
|
||
}
|
||
}
|
||
Some(Ok(tungstenite::Message::Ping(payload))) => {
|
||
let _ = ws.send(tungstenite::Message::Pong(payload)).await;
|
||
}
|
||
Some(Ok(tungstenite::Message::Close(_))) | None => break,
|
||
Some(Err(_)) => break,
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
n = tcp_rx.read(&mut buf) => {
|
||
match n {
|
||
Ok(0) | Err(_) => break,
|
||
Ok(n) => {
|
||
let msg = tungstenite::Message::Binary(buf[..n].to_vec());
|
||
if ws.send(msg).await.is_err() {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
let _ = ws.close(None).await;
|
||
Ok(())
|
||
}
|
||
|
||
async fn relay_tcp(client: TcpStream, remote: TcpStream) {
|
||
let (mut cr, mut cw) = tokio::io::split(client);
|
||
let (mut rr, mut rw) = tokio::io::split(remote);
|
||
tokio::select! {
|
||
_ = tokio::io::copy(&mut cr, &mut rw) => {}
|
||
_ = tokio::io::copy(&mut rr, &mut cw) => {}
|
||
}
|
||
}
|