Bluemangoo/IndPjsk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 6 Wiki Activity

复用连接

Browse Source
This commit is contained in:
Bluemangoo
2026-10-28 13:51:20 +08:00
parent 0851713876
commit 5739f7f8e7
15 changed files with 581 additions and 376 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core_vpn/src/mitm.rs
+106 -90
View File
@@ -12,12 +12,11 @@ use rustls::RootCertStore;
use std::io;
use std::net::SocketAddr;
use std::os::fd::RawFd;
use std::sync::{Arc, OnceLock};
use std::sync::{Arc, Mutex as StdMutex, OnceLock};
use std::time::Duration;
use http::Uri;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::UdpSocket;
use tokio::sync::Mutex;
use tokio::task::JoinSet;
use tokio_rustls::StartHandshake;
use tokio_util::sync::CancellationToken;
@@ -115,11 +114,9 @@ async fn handle_tls_connection<S>(
dst: SocketAddr,
upstream: Upstream,
target_hosts: Vec<String>,
target_path: Regex,
save_path: String,
waiters: Waiters,
server_config: Arc<rustls::ServerConfig>,
cancel: CancellationToken,
result_tx: Arc<Mutex<Option<tokio::sync::oneshot::Sender<String>>>>,
) -> anyhow::Result<()>
where
S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
@@ -136,7 +133,6 @@ where
let Some((sni, accepted)) = sniffed else {
io.rewind();
io.stop_recording();
// 非 TLS：拿不到 SNI，走代理时只能按 IP CONNECT
return passthrough_cancel_wrapped(io, dst, None, upstream, &cancel).await;
};
@@ -154,7 +150,6 @@ where
}
return r;
}
info!("Encountered target host.");
io.stop_recording();
let client_tls = StartHandshake::from_parts(accepted, io)
@@ -162,16 +157,13 @@ where
.await?;
let svc_host = sni.clone().unwrap_or_default();
let svc_upstream = upstream;
let svc = hyper::service::service_fn(move |req: hyper::Request<Incoming>| {
let host = svc_host.clone();
let path = target_path.clone();
let save = save_path.clone();
let result_tx = result_tx.clone();
let upstream = svc_upstream.clone();
let waiters = waiters.clone();
let upstream = upstream.clone();
async move {
Ok::<_, std::convert::Infallible>(
forward_request(req, dst, host, path, save, result_tx, upstream).await,
forward_request(req, dst, host, waiters, upstream).await,
)
}
});
@@ -191,6 +183,7 @@ where
break;
}
_ = cancel.cancelled(), if !graceful => {
info!("teardown: graceful shutdown of downstream connection");
conn.as_mut().graceful_shutdown();
graceful = true;
}
@@ -218,38 +211,42 @@ where
#[allow(clippy::too_many_arguments)]
async fn forward_request(
req: hyper::Request<Incoming>,
mut req: hyper::Request<Incoming>,
dst: SocketAddr,
url_host: String,
target_path: Regex,
save_path: String,
result_tx: Arc<Mutex<Option<tokio::sync::oneshot::Sender<String>>>>,
waiters: Waiters,
upstream: Upstream,
) -> hyper::Response<Full<Bytes>> {
match forward_request_inner(req, dst, &url_host, &target_path, &upstream).await {
Ok((response, matched_body)) => {
if let Some((url, body)) = matched_body {
info!("Target matched! Downloading {} bytes", body.len());
let path = std::path::Path::new(&save_path);
if !path.is_absolute() {
error!("Failed to resolve path: {save_path}");
}
if let Some(parent) = path.parent() {
let _ = tokio::fs::create_dir_all(parent).await;
}
match tokio::fs::write(path, &body).await {
Ok(()) => {
if let Some(tx) = result_tx.lock().await.take() {
let _ = tx.send(url);
}
}
Err(e) => error!("write {save_path} failed: {e:?}"),
}
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")
@@ -259,28 +256,39 @@ async fn forward_request(
}
}
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,
target_path: &Regex,
upstream: &Upstream,
) -> anyhow::Result<(hyper::Response<Full<Bytes>>, Option<(String, Bytes)>)> {
upstream: Upstream,
want_body: bool,
) -> anyhow::Result<(hyper::Response<Full<Bytes>>, Option<Bytes>)> {
let (mut parts, body) = req.into_parts();
if let Some(rewritten) = force_full_mysekai(&parts.uri) {
warn!("rewrite mysekai flag: {} -> {}", parts.uri, rewritten);
parts.uri = rewritten;
}
let parts_uri = parts.uri.clone();
let path_and_query = parts_uri
.path_and_query()
.map(|p| p.as_str().to_string())
.unwrap_or_else(|| "/".to_string());
let body_bytes = body.collect().await?.to_bytes();
// 拒绝长链接，反正抓到就跑。
// 拒绝长链接
for h in [
"connection",
"keep-alive",
@@ -297,7 +305,6 @@ async fn forward_request_inner(
let mut client_config = rustls::ClientConfig::builder()
.with_root_certificates(root_store())
.with_no_client_auth();
// 问就是有服务器不接受 H2
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())?;
@@ -329,17 +336,14 @@ async fn forward_request_inner(
res_parts.headers.remove(h);
}
let is_target = target_path.is_match(&path_and_query);
let url = build_url(url_host, &parts_uri);
let save = if is_target {
Some((url, res_bytes.clone()))
let saved = if want_body {
Some(res_bytes.clone())
} else {
None
};
Ok((
hyper::Response::from_parts(res_parts, Full::new(res_bytes)),
save,
saved,
))
}
@@ -362,6 +366,30 @@ fn force_full_mysekai(uri: &Uri) -> Option<Uri> {
}
#[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,
}
pub struct Waiter {
pub id: String,
pub path: Regex,
pub save_path: String,
}
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))?;
Some(list.remove(idx))
}
#[derive(Clone, Debug)]
pub struct ProxyConfig {
pub tls: bool,
@@ -500,14 +528,13 @@ async fn connect_via_proxy(
Ok(io)
}
pub async fn run_capture_async(
pub async fn run_session(
tun_fd: i32,
target_hosts: Vec<String>,
target_path: Regex,
save_path: String,
proxy: Option<ProxyConfig>,
cancel_token: CancellationToken,
) -> anyhow::Result<Option<String>> {
waiters: Waiters,
cancel: CancellationToken,
) -> anyhow::Result<()> {
let device = unsafe { BorrowedAsyncDevice::borrow_raw(tun_fd as RawFd)? };
let mut ipstack_config = ipstack::IpStackConfig::default();
@@ -520,19 +547,11 @@ pub async fn run_capture_async(
let tls_acceptor = build_tls_acceptor();
let upstream = Upstream::new(proxy);
let (result_tx, mut result_rx) = tokio::sync::oneshot::channel::<String>();
let result_tx = Arc::new(Mutex::new(Some(result_tx)));
let result: anyhow::Result<Option<String>> = loop {
let token = cancel_token.clone();
loop {
let token = cancel.clone();
tokio::select! {
_ = cancel_token.cancelled() => {
let final_result = result_rx.try_recv().ok();
break Ok(final_result);
}
Ok(url) = &mut result_rx => {
cancel_token.cancel();
break Ok(Some(url));
_ = cancel.cancelled() => {
break;
}
_ = bridge_tasks.join_next(), if !bridge_tasks.is_empty() => {}
@@ -541,7 +560,7 @@ pub async fn run_capture_async(
Ok(s) => s,
Err(e) => {
error!("ip_stack accept error (device might be dropped): {e:?}");
break Err(anyhow::Error::from(e));
break;
}
};
@@ -551,7 +570,7 @@ pub async fn run_capture_async(
if dst.port() == 443 {
// drop quic roughly
drop(udp);
} else{
} else {
bridge_tasks.spawn(async move {
if let Err(e) = udp_passthrough(udp, dst, token).await {
error!("udp passthrough error: {e:?}");
@@ -564,7 +583,7 @@ pub async fn run_capture_async(
if dst.port() == 443 {
let (app_io, mut tun_io) = tokio::io::duplex(256 * 1024);
let pump_token = cancel_token.clone();
let pump_token = cancel.clone();
bridge_tasks.spawn(async move {
tokio::select! {
_ = pump_token.cancelled() => {
@@ -585,11 +604,9 @@ pub async fn run_capture_async(
});
let hosts = target_hosts.clone();
let path = target_path.clone();
let save = save_path.clone();
let server_config = tls_acceptor.clone();
let token = cancel_token.clone();
let result_tx = result_tx.clone();
let token = cancel.clone();
let waiters = waiters.clone();
let upstream = upstream.clone();
tokio::spawn(async move {
@@ -598,18 +615,16 @@ pub async fn run_capture_async(
dst,
upstream,
hosts,
path,
save,
waiters,
server_config,
token,
result_tx
).await {
Ok(()) => { }
Ok(()) => {}
Err(e) => error!("handle tls conn error: {e:?}"),
}
});
}else{
let pump_token = cancel_token.clone();
} else {
let pump_token = cancel.clone();
let upstream = upstream.clone();
bridge_tasks.spawn(async move {
let _ =
@@ -622,7 +637,7 @@ pub async fn run_capture_async(
}
}
}
};
}
let _ = tokio::time::timeout(TEARDOWN_GRACE, async {
while bridge_tasks.join_next().await.is_some() {}
@@ -631,7 +646,8 @@ pub async fn run_capture_async(
bridge_tasks.abort_all();
while bridge_tasks.join_next().await.is_some() {}
result
info!("session stopped");
Ok(())
}
async fn tcp_passthrough<IO>(