788 lines
24 KiB
Rust
788 lines
24 KiB
Rust
use crate::mut_qstring::MutQString;
|
||
use crate::tun::{BorrowedTunIo, RewindStream};
|
||
use crate::{error, info, warn};
|
||
use bytes::Bytes;
|
||
use http::Uri;
|
||
use http_body_util::{BodyExt, Full};
|
||
use hyper::body::Incoming;
|
||
use hyper_util::rt::TokioIo;
|
||
use ipstack::IpStackStream;
|
||
use regex::Regex;
|
||
use rustls::RootCertStore;
|
||
use std::io;
|
||
use std::net::SocketAddr;
|
||
use std::os::fd::RawFd;
|
||
use std::sync::{Arc, Mutex as StdMutex, OnceLock};
|
||
use std::time::Duration;
|
||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||
use tokio::net::UdpSocket;
|
||
use tokio::task::JoinSet;
|
||
use tokio_rustls::StartHandshake;
|
||
use tokio_util::sync::CancellationToken;
|
||
use tun_rs::BorrowedAsyncDevice;
|
||
use crate::acceptor::build_tls_acceptor;
|
||
|
||
static ROOT_STORE: OnceLock<Arc<RootCertStore>> = OnceLock::new();
|
||
|
||
const BLOCKED_HOSTS: &[&str] = &["submit.backtrace.io"];
|
||
|
||
fn is_blocked_host(host: &str) -> bool {
|
||
let host = host.trim_end_matches('.');
|
||
BLOCKED_HOSTS
|
||
.iter()
|
||
.any(|blocked| host.eq_ignore_ascii_case(blocked))
|
||
}
|
||
|
||
const SNIFF_TIMEOUT: Duration = Duration::from_secs(5);
|
||
const TEARDOWN_DRAIN_TIMEOUT: Duration = Duration::from_secs(3);
|
||
const TEARDOWN_SETTLE: Duration = Duration::from_millis(50);
|
||
const TEARDOWN_FIN_TIMEOUT: Duration = Duration::from_millis(300);
|
||
const TEARDOWN_GRACE: Duration = Duration::from_secs(4);
|
||
|
||
fn root_store() -> Arc<RootCertStore> {
|
||
ROOT_STORE
|
||
.get_or_init(|| {
|
||
let mut root_store = RootCertStore::empty();
|
||
let native_certs = rustls_native_certs::load_native_certs();
|
||
for err in native_certs.errors {
|
||
warn!("加载证书警告: {err}");
|
||
}
|
||
for cert in native_certs.certs {
|
||
let _ = root_store.add(cert);
|
||
}
|
||
|
||
load_system_ca_dirs(&mut root_store);
|
||
|
||
if root_store.is_empty() {
|
||
root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||
}
|
||
Arc::new(root_store)
|
||
})
|
||
.clone()
|
||
}
|
||
|
||
fn load_system_ca_dirs(store: &mut RootCertStore) {
|
||
use std::io::BufReader;
|
||
|
||
const DIRS: &[&str] = &[
|
||
"/system/etc/security/cacerts", // Android 13-
|
||
"/apex/com.android.conscrypt/cacerts", // Android 14+
|
||
];
|
||
|
||
for dir in DIRS {
|
||
let Ok(entries) = std::fs::read_dir(dir) else {
|
||
continue;
|
||
};
|
||
for entry in entries.flatten() {
|
||
let Ok(file) = std::fs::File::open(entry.path()) else {
|
||
continue;
|
||
};
|
||
for cert in rustls_pemfile::certs(&mut BufReader::new(file)).flatten() {
|
||
let _ = store.add(cert);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn sniff_client_hello<S>(
|
||
io: &mut RewindStream<S>,
|
||
) -> io::Result<Option<(Option<String>, rustls::server::Accepted)>>
|
||
where
|
||
S: AsyncRead + AsyncWrite + Unpin,
|
||
{
|
||
let mut acceptor = rustls::server::Acceptor::default();
|
||
let mut buf = vec![0u8; 16 * 1024];
|
||
|
||
loop {
|
||
let n = match tokio::time::timeout(SNIFF_TIMEOUT, io.read(&mut buf)).await {
|
||
Ok(Ok(n)) => n,
|
||
Ok(Err(e)) => return Err(e),
|
||
Err(_) => return Ok(None), // 对端迟迟不发数据，当作非 TLS
|
||
};
|
||
if n == 0 {
|
||
return Ok(None);
|
||
}
|
||
|
||
if acceptor.read_tls(&mut &buf[..n]).is_err() {
|
||
return Ok(None);
|
||
}
|
||
match acceptor.accept() {
|
||
Ok(Some(accepted)) => {
|
||
let sni = accepted.client_hello().server_name().map(|s| s.to_string());
|
||
return Ok(Some((sni, accepted)));
|
||
}
|
||
Ok(None) => continue, // 还没凑齐
|
||
Err(_) => return Ok(None), // 不是合法 TLS
|
||
}
|
||
}
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
async fn handle_tls_connection<S>(
|
||
tcp: S,
|
||
dst: SocketAddr,
|
||
upstream: Upstream,
|
||
target_hosts: Vec<String>,
|
||
waiters: Waiters,
|
||
server_config: Arc<rustls::ServerConfig>,
|
||
cancel: CancellationToken,
|
||
) -> anyhow::Result<()>
|
||
where
|
||
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
||
{
|
||
let mut io = RewindStream::new(tcp);
|
||
let sniffed = match sniff_client_hello(&mut io).await {
|
||
Ok(v) => v,
|
||
Err(e) => {
|
||
warn!("client hello sniff error on {dst:?}: {e:?}");
|
||
return Ok(());
|
||
}
|
||
};
|
||
|
||
let Some((sni, accepted)) = sniffed else {
|
||
io.rewind();
|
||
io.stop_recording();
|
||
return passthrough_cancel_wrapped(io, dst, None, upstream, &cancel).await;
|
||
};
|
||
|
||
if let Some(host) = sni.as_deref()
|
||
&& is_blocked_host(host)
|
||
{
|
||
warn!("blocked host {host} ({dst:?}), connection rejected");
|
||
return Ok(());
|
||
}
|
||
|
||
let host_matched = sni
|
||
.as_deref()
|
||
.map(|s| target_hosts.iter().any(|h| s.eq_ignore_ascii_case(h)))
|
||
.unwrap_or(false);
|
||
|
||
if !host_matched {
|
||
io.rewind();
|
||
io.stop_recording();
|
||
let r = passthrough_cancel_wrapped(io, dst, sni.as_deref(), upstream, &cancel).await;
|
||
if r.is_err() {
|
||
warn!("tcp_passthrough error on {sni:?} {dst:?}");
|
||
}
|
||
return r;
|
||
}
|
||
|
||
io.stop_recording();
|
||
let client_tls = StartHandshake::from_parts(accepted, io)
|
||
.into_stream(server_config)
|
||
.await?;
|
||
|
||
let svc_host = sni.clone().unwrap_or_default();
|
||
let svc = hyper::service::service_fn(move |req: hyper::Request<Incoming>| {
|
||
let host = svc_host.clone();
|
||
let waiters = waiters.clone();
|
||
let upstream = upstream.clone();
|
||
async move {
|
||
Ok::<_, std::convert::Infallible>(
|
||
forward_request(req, dst, host, waiters, upstream).await,
|
||
)
|
||
}
|
||
});
|
||
|
||
let conn =
|
||
hyper::server::conn::http1::Builder::new().serve_connection(TokioIo::new(client_tls), svc);
|
||
tokio::pin!(conn);
|
||
|
||
let mut graceful = false;
|
||
loop {
|
||
tokio::select! {
|
||
res = conn.as_mut() => {
|
||
match res {
|
||
Ok(()) => {}
|
||
Err(e) => warn!("downstream http1 connection error: {e:?}"),
|
||
}
|
||
break;
|
||
}
|
||
_ = cancel.cancelled(), if !graceful => {
|
||
info!("teardown: graceful shutdown of downstream connection");
|
||
conn.as_mut().graceful_shutdown();
|
||
graceful = true;
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
async fn passthrough_cancel_wrapped<IO>(
|
||
io: IO,
|
||
dst: SocketAddr,
|
||
sni: Option<&str>,
|
||
upstream: Upstream,
|
||
cancel: &CancellationToken,
|
||
) -> anyhow::Result<()>
|
||
where
|
||
IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
||
{
|
||
tokio::select! {
|
||
_ = cancel.cancelled() => Ok(()),
|
||
r = tcp_passthrough(io, dst, sni, upstream) => r,
|
||
}
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
async fn forward_request(
|
||
mut req: hyper::Request<Incoming>,
|
||
dst: SocketAddr,
|
||
url_host: String,
|
||
waiters: Waiters,
|
||
upstream: Upstream,
|
||
) -> hyper::Response<Full<Bytes>> {
|
||
if let Some(rewritten) = force_full_mysekai(req.uri()) {
|
||
warn!("rewrite mysekai flag: {} -> {}", req.uri(), rewritten);
|
||
*req.uri_mut() = rewritten;
|
||
}
|
||
|
||
let uri = req.uri().clone();
|
||
let url = build_url(&url_host, &uri);
|
||
let path_and_query = uri
|
||
.path_and_query()
|
||
.map(|p| p.as_str().to_string())
|
||
.unwrap_or_else(|| "/".to_string());
|
||
|
||
let waiter = take_matching_waiter(&waiters, &path_and_query);
|
||
if let Some(w) = &waiter {
|
||
info!("step '{}' matched", w.id);
|
||
}
|
||
let want_body = waiter.is_some();
|
||
|
||
match forward_request_inner(req, dst, &url_host, upstream, want_body).await {
|
||
Ok((response, saved_body)) => {
|
||
if let (Some(w), Some(body)) = (waiter, saved_body) {
|
||
save_and_notify(w, url, body);
|
||
}
|
||
response
|
||
}
|
||
Err(e) => {
|
||
error!("forward to {dst} failed: {e:?}");
|
||
if let Some(w) = waiter {
|
||
crate::notify_step_result(&w.id, None);
|
||
}
|
||
hyper::Response::builder()
|
||
.status(hyper::StatusCode::BAD_GATEWAY)
|
||
.header("content-type", "text/plain; charset=utf-8")
|
||
.body(Full::new(Bytes::from(format!("mitm forward error: {e}"))))
|
||
.unwrap()
|
||
}
|
||
}
|
||
}
|
||
|
||
fn save_and_notify(waiter: Waiter, url: String, body: Bytes) {
|
||
tokio::spawn(async move {
|
||
let path = std::path::Path::new(&waiter.save_path);
|
||
if !path.is_absolute() {
|
||
error!("save_path 不是绝对路径（{}）", waiter.save_path);
|
||
}
|
||
if let Some(parent) = path.parent() {
|
||
let _ = tokio::fs::create_dir_all(parent).await;
|
||
}
|
||
match tokio::fs::write(path, &body).await {
|
||
Ok(()) => {
|
||
info!("step '{}' saved ({} bytes)", waiter.id, body.len());
|
||
crate::notify_step_result(&waiter.id, Some(&url));
|
||
}
|
||
Err(e) => {
|
||
error!("write {} failed: {e:?}", waiter.save_path);
|
||
crate::notify_step_result(&waiter.id, None);
|
||
}
|
||
}
|
||
});
|
||
}
|
||
|
||
async fn forward_request_inner(
|
||
req: hyper::Request<Incoming>,
|
||
dst: SocketAddr,
|
||
url_host: &str,
|
||
upstream: Upstream,
|
||
want_body: bool,
|
||
) -> anyhow::Result<(hyper::Response<Full<Bytes>>, Option<Bytes>)> {
|
||
let (mut parts, body) = req.into_parts();
|
||
let body_bytes = body.collect().await?.to_bytes();
|
||
|
||
// 拒绝长链接
|
||
for h in [
|
||
"connection",
|
||
"keep-alive",
|
||
"proxy-connection",
|
||
"transfer-encoding",
|
||
"upgrade",
|
||
"te",
|
||
"trailer",
|
||
"content-length",
|
||
] {
|
||
parts.headers.remove(h);
|
||
}
|
||
|
||
let mut client_config = rustls::ClientConfig::builder()
|
||
.with_root_certificates(root_store())
|
||
.with_no_client_auth();
|
||
client_config.alpn_protocols = vec![b"http/1.1".to_vec()];
|
||
let connector = tokio_rustls::TlsConnector::from(Arc::new(client_config));
|
||
let server_name = rustls::pki_types::ServerName::try_from(url_host.to_string())?;
|
||
let tcp_to_server = upstream.connect(dst, Some(url_host)).await?;
|
||
let server_tls = connector.connect(server_name, tcp_to_server).await?;
|
||
|
||
let (mut sender, conn) =
|
||
hyper::client::conn::http1::handshake::<_, Full<Bytes>>(TokioIo::new(server_tls)).await?;
|
||
tokio::spawn(async move {
|
||
if let Err(e) = conn.await {
|
||
warn!("upstream http1 connection error: {e:?}");
|
||
}
|
||
});
|
||
|
||
let upstream_req = hyper::Request::from_parts(parts, Full::new(body_bytes));
|
||
let resp = sender.send_request(upstream_req).await?;
|
||
|
||
let (mut res_parts, res_body) = resp.into_parts();
|
||
let res_bytes = res_body.collect().await?.to_bytes();
|
||
for h in [
|
||
"connection",
|
||
"keep-alive",
|
||
"proxy-connection",
|
||
"transfer-encoding",
|
||
"upgrade",
|
||
"te",
|
||
"trailer",
|
||
] {
|
||
res_parts.headers.remove(h);
|
||
}
|
||
|
||
let saved = if want_body {
|
||
Some(res_bytes.clone())
|
||
} else {
|
||
None
|
||
};
|
||
Ok((
|
||
hyper::Response::from_parts(res_parts, Full::new(res_bytes)),
|
||
saved,
|
||
))
|
||
}
|
||
|
||
fn build_url(host: &str, uri: &Uri) -> String {
|
||
if uri.scheme().is_some() && uri.authority().is_some() {
|
||
return uri.to_string();
|
||
}
|
||
let pq = uri.path_and_query().map(|p| p.as_str()).unwrap_or("/");
|
||
format!("https://{host}{pq}")
|
||
}
|
||
|
||
fn force_full_mysekai(uri: &Uri) -> Option<Uri> {
|
||
const KEY: &str = "isForceAllReloadOnlyMysekai";
|
||
let mut query = qstring::QString::from(uri.query()?);
|
||
if query.get(KEY) != Some("False") {
|
||
return None;
|
||
}
|
||
query.set(KEY, "True");
|
||
Uri::try_from(format!("{}?{}", uri.path(), query).as_str()).ok()
|
||
}
|
||
|
||
#[derive(serde::Deserialize, Debug, Clone)]
|
||
pub struct StepSpec {
|
||
pub id: String,
|
||
#[serde(default)]
|
||
pub hosts: Vec<String>,
|
||
pub path: String,
|
||
#[serde(rename = "savePath")]
|
||
pub save_path: String,
|
||
/// 持续运行：命中后不摘除 waiter，同类型的包可以反复抓
|
||
#[serde(default)]
|
||
pub repeat: bool,
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
pub struct Waiter {
|
||
pub id: String,
|
||
pub path: Regex,
|
||
pub save_path: String,
|
||
pub repeat: bool,
|
||
}
|
||
|
||
pub type Waiters = Arc<StdMutex<Vec<Waiter>>>;
|
||
|
||
fn take_matching_waiter(waiters: &StdMutex<Vec<Waiter>>, path_and_query: &str) -> Option<Waiter> {
|
||
let mut list = waiters.lock().unwrap();
|
||
let idx = list.iter().position(|w| w.path.is_match(path_and_query))?;
|
||
if list[idx].repeat {
|
||
// 持续运行：保留 waiter，让后续同类型的请求继续命中
|
||
Some(list[idx].clone())
|
||
} else {
|
||
Some(list.remove(idx))
|
||
}
|
||
}
|
||
|
||
#[derive(Clone, Debug)]
|
||
pub struct ProxyConfig {
|
||
pub tls: bool,
|
||
pub host: String,
|
||
pub port: u16,
|
||
pub auth: Option<String>,
|
||
}
|
||
|
||
impl ProxyConfig {
|
||
pub fn parse(input: &str) -> anyhow::Result<Option<Self>> {
|
||
let raw = input.trim();
|
||
if raw.is_empty() {
|
||
return Ok(None);
|
||
}
|
||
let url = url::Url::parse(raw)
|
||
.map_err(|e| anyhow::anyhow!("代理地址不是合法 URL: {raw} ({e})"))?;
|
||
let tls = match url.scheme() {
|
||
"http" => false,
|
||
"https" => true,
|
||
other => anyhow::bail!("代理只支持 http/https，当前是 {other}://"),
|
||
};
|
||
let host = url
|
||
.host_str()
|
||
.filter(|h| !h.is_empty())
|
||
.ok_or_else(|| anyhow::anyhow!("代理地址缺少 host: {raw}"))?
|
||
.to_string();
|
||
let port = url
|
||
.port_or_known_default()
|
||
.ok_or_else(|| anyhow::anyhow!("代理地址缺少端口: {raw}"))?;
|
||
let auth = if url.username().is_empty() {
|
||
None
|
||
} else {
|
||
let pass = url.password().unwrap_or("");
|
||
Some(base64_encode(
|
||
format!("{}:{}", url.username(), pass).as_bytes(),
|
||
))
|
||
};
|
||
Ok(Some(Self {
|
||
tls,
|
||
host,
|
||
port,
|
||
auth,
|
||
}))
|
||
}
|
||
}
|
||
|
||
fn base64_encode(data: &[u8]) -> String {
|
||
use base64::Engine as _;
|
||
base64::engine::general_purpose::STANDARD.encode(data)
|
||
}
|
||
|
||
type UpstreamIo = Box<dyn AsyncReadWrite + Send + Unpin>;
|
||
|
||
trait AsyncReadWrite: AsyncRead + AsyncWrite {}
|
||
impl<T: AsyncRead + AsyncWrite + ?Sized> AsyncReadWrite for T {}
|
||
|
||
#[derive(Clone)]
|
||
struct Upstream {
|
||
proxy: Option<Arc<ProxyConfig>>,
|
||
}
|
||
|
||
impl Upstream {
|
||
fn new(proxy: Option<ProxyConfig>) -> Self {
|
||
Self {
|
||
proxy: proxy.map(Arc::new),
|
||
}
|
||
}
|
||
|
||
async fn connect(&self, dst: SocketAddr, host: Option<&str>) -> anyhow::Result<UpstreamIo> {
|
||
match &self.proxy {
|
||
None => {
|
||
let tcp = tokio::net::TcpStream::connect(dst).await?;
|
||
tcp.set_nodelay(true).ok();
|
||
Ok(Box::new(tcp))
|
||
}
|
||
Some(proxy) => connect_via_proxy(proxy, dst, host).await,
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn connect_via_proxy(
|
||
proxy: &ProxyConfig,
|
||
dst: SocketAddr,
|
||
host: Option<&str>,
|
||
) -> anyhow::Result<UpstreamIo> {
|
||
let tcp = tokio::net::TcpStream::connect((proxy.host.as_str(), proxy.port))
|
||
.await
|
||
.map_err(|e| anyhow::anyhow!("连接代理 {}:{} 失败: {e}", proxy.host, proxy.port))?;
|
||
tcp.set_nodelay(true).ok();
|
||
|
||
let mut io: UpstreamIo = if proxy.tls {
|
||
let mut cfg = rustls::ClientConfig::builder()
|
||
.with_root_certificates(root_store())
|
||
.with_no_client_auth();
|
||
cfg.alpn_protocols = vec![b"http/1.1".to_vec()];
|
||
let connector = tokio_rustls::TlsConnector::from(Arc::new(cfg));
|
||
let name = rustls::pki_types::ServerName::try_from(proxy.host.clone())?;
|
||
Box::new(
|
||
connector
|
||
.connect(name, tcp)
|
||
.await
|
||
.map_err(|e| anyhow::anyhow!("与代理建立 TLS 失败: {e}"))?,
|
||
)
|
||
} else {
|
||
Box::new(tcp)
|
||
};
|
||
|
||
let authority = match host {
|
||
Some(h) if !h.is_empty() => format!("{h}:{}", dst.port()),
|
||
_ => dst.to_string(),
|
||
};
|
||
let mut req = format!("CONNECT {authority} HTTP/1.1\r\nHost: {authority}\r\n");
|
||
if let Some(auth) = &proxy.auth {
|
||
req.push_str(&format!("Proxy-Authorization: Basic {auth}\r\n"));
|
||
}
|
||
req.push_str("\r\n");
|
||
io.write_all(req.as_bytes()).await?;
|
||
io.flush().await?;
|
||
|
||
let mut head = Vec::with_capacity(256);
|
||
let mut byte = [0u8; 1];
|
||
loop {
|
||
let n = io.read(&mut byte).await?;
|
||
if n == 0 {
|
||
anyhow::bail!("代理在 CONNECT 阶段直接断开");
|
||
}
|
||
head.push(byte[0]);
|
||
if head.ends_with(b"\r\n\r\n") {
|
||
break;
|
||
}
|
||
if head.len() > 16 * 1024 {
|
||
anyhow::bail!("代理 CONNECT 响应头过长");
|
||
}
|
||
}
|
||
let head_str = String::from_utf8_lossy(&head);
|
||
let status_line = head_str.lines().next().unwrap_or("").trim().to_string();
|
||
let ok = status_line
|
||
.split_whitespace()
|
||
.nth(1)
|
||
.map(|code| code == "200")
|
||
.unwrap_or(false);
|
||
if !ok {
|
||
anyhow::bail!("代理 CONNECT 失败: {status_line}");
|
||
}
|
||
Ok(io)
|
||
}
|
||
|
||
pub async fn run_session(
|
||
tun_fd: i32,
|
||
target_hosts: Vec<String>,
|
||
proxy: Option<ProxyConfig>,
|
||
waiters: Waiters,
|
||
cancel: CancellationToken,
|
||
) -> anyhow::Result<()> {
|
||
let device = unsafe { BorrowedAsyncDevice::borrow_raw(tun_fd as RawFd)? };
|
||
|
||
let mut ipstack_config = ipstack::IpStackConfig::default();
|
||
ipstack_config.mtu(1500)?;
|
||
let device_ref: &BorrowedAsyncDevice<'static> = unsafe { std::mem::transmute(&device) };
|
||
let tun_io = BorrowedTunIo::new(device_ref);
|
||
let mut ip_stack = ipstack::IpStack::new(ipstack_config, tun_io);
|
||
let mut bridge_tasks = JoinSet::new();
|
||
|
||
let tls_acceptor = build_tls_acceptor(target_hosts.clone());
|
||
let upstream = Upstream::new(proxy);
|
||
|
||
loop {
|
||
let token = cancel.clone();
|
||
tokio::select! {
|
||
_ = cancel.cancelled() => {
|
||
break;
|
||
}
|
||
_ = bridge_tasks.join_next(), if !bridge_tasks.is_empty() => {}
|
||
|
||
res = ip_stack.accept() => {
|
||
let stream = match res {
|
||
Ok(s) => s,
|
||
Err(e) => {
|
||
error!("ip_stack accept error (device might be dropped): {e:?}");
|
||
break;
|
||
}
|
||
};
|
||
|
||
match stream {
|
||
IpStackStream::Udp(udp) => {
|
||
let dst = udp.peer_addr();
|
||
if dst.port() == 443 {
|
||
// drop quic roughly
|
||
drop(udp);
|
||
} else {
|
||
bridge_tasks.spawn(async move {
|
||
if let Err(e) = udp_passthrough(udp, dst, token).await {
|
||
error!("udp passthrough error: {e:?}");
|
||
}
|
||
});
|
||
}
|
||
}
|
||
IpStackStream::Tcp(mut tcp) => {
|
||
let dst = tcp.peer_addr();
|
||
if dst.port() == 443 {
|
||
let (app_io, mut tun_io) = tokio::io::duplex(256 * 1024);
|
||
|
||
let pump_token = cancel.clone();
|
||
bridge_tasks.spawn(async move {
|
||
tokio::select! {
|
||
_ = pump_token.cancelled() => {
|
||
let _ = tokio::time::timeout(
|
||
TEARDOWN_DRAIN_TIMEOUT,
|
||
tokio::io::copy(&mut tun_io, &mut tcp),
|
||
)
|
||
.await;
|
||
}
|
||
_ = tokio::io::copy_bidirectional(&mut tun_io, &mut tcp) => {}
|
||
}
|
||
tokio::time::sleep(TEARDOWN_SETTLE).await;
|
||
let _ = tokio::time::timeout(
|
||
TEARDOWN_FIN_TIMEOUT,
|
||
tcp.shutdown(),
|
||
)
|
||
.await;
|
||
});
|
||
|
||
let hosts = target_hosts.clone();
|
||
let server_config = tls_acceptor.clone();
|
||
let token = cancel.clone();
|
||
let waiters = waiters.clone();
|
||
let upstream = upstream.clone();
|
||
|
||
tokio::spawn(async move {
|
||
match handle_tls_connection(
|
||
app_io,
|
||
dst,
|
||
upstream,
|
||
hosts,
|
||
waiters,
|
||
server_config,
|
||
token,
|
||
).await {
|
||
Ok(()) => {}
|
||
Err(e) => error!("handle tls conn error: {e:?}"),
|
||
}
|
||
});
|
||
} else {
|
||
let pump_token = cancel.clone();
|
||
let upstream = upstream.clone();
|
||
bridge_tasks.spawn(async move {
|
||
let _ =
|
||
tcp_passthrough_cancellable(tcp, dst, None, upstream, pump_token)
|
||
.await;
|
||
});
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
let _ = tokio::time::timeout(TEARDOWN_GRACE, async {
|
||
while bridge_tasks.join_next().await.is_some() {}
|
||
})
|
||
.await;
|
||
bridge_tasks.abort_all();
|
||
while bridge_tasks.join_next().await.is_some() {}
|
||
|
||
info!("session stopped");
|
||
Ok(())
|
||
}
|
||
|
||
async fn tcp_passthrough<IO>(
|
||
mut client: IO,
|
||
dst: SocketAddr,
|
||
sni: Option<&str>,
|
||
upstream: Upstream,
|
||
) -> anyhow::Result<()>
|
||
where
|
||
IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
||
{
|
||
let mut server = upstream.connect(dst, sni).await?;
|
||
tokio::io::copy_bidirectional(&mut client, &mut server).await?;
|
||
Ok(())
|
||
}
|
||
|
||
async fn tcp_passthrough_cancellable<IO>(
|
||
mut client: IO,
|
||
dst: SocketAddr,
|
||
sni: Option<&str>,
|
||
upstream: Upstream,
|
||
cancel: CancellationToken,
|
||
) -> anyhow::Result<()>
|
||
where
|
||
IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
||
{
|
||
let mut server = upstream.connect(dst, sni).await?;
|
||
tokio::select! {
|
||
_ = cancel.cancelled() => {
|
||
let _ = tokio::time::timeout(
|
||
TEARDOWN_DRAIN_TIMEOUT,
|
||
tokio::io::copy(&mut client, &mut server),
|
||
)
|
||
.await;
|
||
}
|
||
res = tokio::io::copy_bidirectional(&mut client, &mut server) => {
|
||
res?;
|
||
}
|
||
}
|
||
tokio::time::sleep(TEARDOWN_SETTLE).await;
|
||
let _ = tokio::time::timeout(TEARDOWN_FIN_TIMEOUT, client.shutdown()).await;
|
||
let _ = tokio::time::timeout(TEARDOWN_FIN_TIMEOUT, server.shutdown()).await;
|
||
Ok(())
|
||
}
|
||
async fn udp_passthrough<IO>(
|
||
client: IO,
|
||
dst: SocketAddr,
|
||
cancel: CancellationToken,
|
||
) -> anyhow::Result<()>
|
||
where
|
||
IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
||
{
|
||
let bind_addr: SocketAddr = if dst.is_ipv4() {
|
||
"0.0.0.0:0".parse()?
|
||
} else {
|
||
"[::]:0".parse()?
|
||
};
|
||
|
||
let socket = UdpSocket::bind(bind_addr).await?;
|
||
|
||
socket.connect(dst).await?;
|
||
let socket = Arc::new(socket);
|
||
|
||
let (mut cr, mut cw) = tokio::io::split(client);
|
||
|
||
let mut client_buf = vec![0u8; 65535];
|
||
let mut server_buf = vec![0u8; 65535];
|
||
|
||
let idle_timeout = Duration::from_secs(12);
|
||
|
||
loop {
|
||
tokio::select! {
|
||
_ = cancel.cancelled() => break,
|
||
|
||
// c2s
|
||
res = tokio::time::timeout(idle_timeout, cr.read(&mut client_buf)) => {
|
||
match res {
|
||
Ok(Ok(0)) | Err(_) => break,
|
||
Ok(Ok(n)) => {
|
||
if socket.send(&client_buf[..n]).await.is_err() {
|
||
break;
|
||
}
|
||
}
|
||
Ok(Err(_)) => break,
|
||
}
|
||
}
|
||
|
||
// s2c
|
||
res = tokio::time::timeout(idle_timeout, socket.recv(&mut server_buf)) => {
|
||
match res {
|
||
Ok(Ok(n)) if n > 0 => {
|
||
if cw.write_all(&server_buf[..n]).await.is_err() {
|
||
break;
|
||
}
|
||
}
|
||
_ => break,
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(())
|
||
}
|