mirror of
https://github.com/zhom/donutbrowser.git
synced 2026-07-11 15:16:34 +02:00
zhom
2265 lines
76 KiB
Rust
2265 lines
76 KiB
Rust
use crate::proxy_storage::ProxyConfig;
|
|
use crate::traffic_stats::{get_traffic_tracker, init_traffic_tracker};
|
|
use http_body_util::{BodyExt, Full};
|
|
use hyper::body::Bytes;
|
|
use hyper::server::conn::http1;
|
|
use hyper::service::service_fn;
|
|
use hyper::{Method, Request, Response, StatusCode};
|
|
use hyper_util::rt::TokioIo;
|
|
use regex_lite::Regex;
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::convert::Infallible;
|
|
use std::io;
|
|
use std::net::SocketAddr;
|
|
use std::pin::Pin;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use std::sync::{Arc, Mutex, OnceLock};
|
|
use std::task::{Context, Poll};
|
|
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
|
|
use tokio::net::TcpStream;
|
|
|
|
/// Combined read+write trait for tunnel target streams, allowing
|
|
/// `handle_connect_from_buffer` to handle plain TCP, SOCKS, and
|
|
/// Shadowsocks through the same bidirectional-copy path.
|
|
pub(crate) trait AsyncStream: AsyncRead + AsyncWrite + Unpin + Send {}
|
|
impl<T: AsyncRead + AsyncWrite + Unpin + Send> AsyncStream for T {}
|
|
pub(crate) type BoxedAsyncStream = Box<dyn AsyncStream>;
|
|
use url::Url;
|
|
|
|
enum CompiledRule {
|
|
Regex(Regex),
|
|
Exact(String),
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct BypassMatcher {
|
|
rules: Arc<Vec<CompiledRule>>,
|
|
}
|
|
|
|
impl BypassMatcher {
|
|
pub fn new(rules: &[String]) -> Self {
|
|
let compiled = rules
|
|
.iter()
|
|
.map(|rule| match Regex::new(rule) {
|
|
Ok(re) => CompiledRule::Regex(re),
|
|
Err(_) => CompiledRule::Exact(rule.clone()),
|
|
})
|
|
.collect();
|
|
Self {
|
|
rules: Arc::new(compiled),
|
|
}
|
|
}
|
|
|
|
pub fn should_bypass(&self, host: &str) -> bool {
|
|
self.rules.iter().any(|rule| match rule {
|
|
CompiledRule::Regex(re) => re.is_match(host),
|
|
CompiledRule::Exact(exact) => host == exact,
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct BlocklistMatcher {
|
|
domains: Arc<HashSet<String>>,
|
|
}
|
|
|
|
impl Default for BlocklistMatcher {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl BlocklistMatcher {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
domains: Arc::new(HashSet::new()),
|
|
}
|
|
}
|
|
|
|
pub fn from_file(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
|
|
let content = std::fs::read_to_string(path)?;
|
|
let domains: HashSet<String> = content
|
|
.lines()
|
|
.filter(|line| !line.starts_with('#') && !line.trim().is_empty())
|
|
.map(|line| line.trim().to_lowercase())
|
|
.collect();
|
|
log::info!("[blocklist] Loaded {} domains from {}", domains.len(), path);
|
|
Ok(Self {
|
|
domains: Arc::new(domains),
|
|
})
|
|
}
|
|
|
|
pub fn is_blocked(&self, host: &str) -> bool {
|
|
if self.domains.is_empty() {
|
|
return false;
|
|
}
|
|
let host_lower = host.to_lowercase();
|
|
// Exact match
|
|
if self.domains.contains(host_lower.as_str()) {
|
|
return true;
|
|
}
|
|
// Suffix matching: check parent domains (like uBlock)
|
|
let mut start = 0;
|
|
while let Some(dot_pos) = host_lower[start..].find('.') {
|
|
start += dot_pos + 1;
|
|
if self.domains.contains(&host_lower[start..]) {
|
|
return true;
|
|
}
|
|
}
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Wrapper stream that counts bytes read and written
|
|
struct CountingStream<S> {
|
|
inner: S,
|
|
bytes_read: Arc<AtomicU64>,
|
|
bytes_written: Arc<AtomicU64>,
|
|
}
|
|
|
|
impl<S> CountingStream<S> {
|
|
fn new(inner: S) -> Self {
|
|
Self {
|
|
inner,
|
|
bytes_read: Arc::new(AtomicU64::new(0)),
|
|
bytes_written: Arc::new(AtomicU64::new(0)),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<S: AsyncRead + Unpin> AsyncRead for CountingStream<S> {
|
|
fn poll_read(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &mut ReadBuf<'_>,
|
|
) -> Poll<io::Result<()>> {
|
|
let filled_before = buf.filled().len();
|
|
let result = Pin::new(&mut self.inner).poll_read(cx, buf);
|
|
if let Poll::Ready(Ok(())) = &result {
|
|
let bytes_read = buf.filled().len() - filled_before;
|
|
if bytes_read > 0 {
|
|
self
|
|
.bytes_read
|
|
.fetch_add(bytes_read as u64, Ordering::Relaxed);
|
|
// Update global tracker - count as received (data coming into proxy)
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
tracker.add_bytes_received(bytes_read as u64);
|
|
}
|
|
}
|
|
}
|
|
result
|
|
}
|
|
}
|
|
|
|
impl<S: AsyncWrite + Unpin> AsyncWrite for CountingStream<S> {
|
|
fn poll_write(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &[u8],
|
|
) -> Poll<io::Result<usize>> {
|
|
let result = Pin::new(&mut self.inner).poll_write(cx, buf);
|
|
if let Poll::Ready(Ok(n)) = &result {
|
|
self.bytes_written.fetch_add(*n as u64, Ordering::Relaxed);
|
|
// Update global tracker - count as sent (data going out of proxy)
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
tracker.add_bytes_sent(*n as u64);
|
|
}
|
|
}
|
|
result
|
|
}
|
|
|
|
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
Pin::new(&mut self.inner).poll_flush(cx)
|
|
}
|
|
|
|
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
Pin::new(&mut self.inner).poll_shutdown(cx)
|
|
}
|
|
}
|
|
|
|
// Wrapper to prepend consumed bytes to a stream
|
|
struct PrependReader {
|
|
prepended: Vec<u8>,
|
|
prepended_pos: usize,
|
|
inner: TcpStream,
|
|
}
|
|
|
|
impl AsyncRead for PrependReader {
|
|
fn poll_read(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &mut ReadBuf<'_>,
|
|
) -> Poll<io::Result<()>> {
|
|
// First, read from prepended bytes if any
|
|
if self.prepended_pos < self.prepended.len() {
|
|
let available = self.prepended.len() - self.prepended_pos;
|
|
let to_copy = available.min(buf.remaining());
|
|
buf.put_slice(&self.prepended[self.prepended_pos..self.prepended_pos + to_copy]);
|
|
self.prepended_pos += to_copy;
|
|
return Poll::Ready(Ok(()));
|
|
}
|
|
|
|
// Then read from inner stream
|
|
Pin::new(&mut self.inner).poll_read(cx, buf)
|
|
}
|
|
}
|
|
|
|
impl AsyncWrite for PrependReader {
|
|
fn poll_write(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &[u8],
|
|
) -> Poll<io::Result<usize>> {
|
|
Pin::new(&mut self.inner).poll_write(cx, buf)
|
|
}
|
|
|
|
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
Pin::new(&mut self.inner).poll_flush(cx)
|
|
}
|
|
|
|
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
Pin::new(&mut self.inner).poll_shutdown(cx)
|
|
}
|
|
}
|
|
|
|
async fn handle_request(
|
|
req: Request<hyper::body::Incoming>,
|
|
upstream_url: Option<String>,
|
|
bypass_matcher: BypassMatcher,
|
|
blocklist_matcher: BlocklistMatcher,
|
|
) -> Result<Response<Full<Bytes>>, Infallible> {
|
|
// Handle CONNECT method for HTTPS tunneling
|
|
if req.method() == Method::CONNECT {
|
|
return handle_connect(req, upstream_url, bypass_matcher, blocklist_matcher).await;
|
|
}
|
|
|
|
// Handle regular HTTP requests
|
|
handle_http(req, upstream_url, bypass_matcher, blocklist_matcher).await
|
|
}
|
|
|
|
async fn handle_connect(
|
|
req: Request<hyper::body::Incoming>,
|
|
upstream_url: Option<String>,
|
|
bypass_matcher: BypassMatcher,
|
|
blocklist_matcher: BlocklistMatcher,
|
|
) -> Result<Response<Full<Bytes>>, Infallible> {
|
|
let authority = req.uri().authority().cloned();
|
|
|
|
if let Some(authority) = authority {
|
|
let target_addr = format!("{}", authority);
|
|
|
|
// Parse target host and port
|
|
let (target_host, target_port) = if let Some(colon_pos) = target_addr.find(':') {
|
|
let host = &target_addr[..colon_pos];
|
|
let port: u16 = target_addr[colon_pos + 1..].parse().unwrap_or(443);
|
|
(host, port)
|
|
} else {
|
|
(&target_addr[..], 443)
|
|
};
|
|
|
|
// Block if domain is in the DNS blocklist (before any connection)
|
|
if blocklist_matcher.is_blocked(target_host) {
|
|
log::debug!("[blocklist] Blocked CONNECT to {}", target_host);
|
|
let mut response = Response::new(Full::new(Bytes::from("Blocked by DNS blocklist")));
|
|
*response.status_mut() = StatusCode::FORBIDDEN;
|
|
return Ok(response);
|
|
}
|
|
|
|
// If no upstream proxy, or bypass rule matches, connect directly
|
|
if upstream_url.is_none()
|
|
|| upstream_url
|
|
.as_ref()
|
|
.map(|s| s == "DIRECT")
|
|
.unwrap_or(false)
|
|
|| bypass_matcher.should_bypass(target_host)
|
|
{
|
|
match TcpStream::connect(&target_addr).await {
|
|
Ok(_stream) => {
|
|
let mut response = Response::new(Full::new(Bytes::from("")));
|
|
*response.status_mut() = StatusCode::from_u16(200).unwrap();
|
|
return Ok(response);
|
|
}
|
|
Err(e) => {
|
|
log::error!("Failed to connect to {}: {}", target_addr, e);
|
|
let mut response =
|
|
Response::new(Full::new(Bytes::from(format!("Connection failed: {}", e))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Connect through upstream proxy
|
|
let upstream = match upstream_url.as_ref().and_then(|u| Url::parse(u).ok()) {
|
|
Some(url) => url,
|
|
None => {
|
|
let mut response = Response::new(Full::new(Bytes::from("Invalid upstream URL")));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
};
|
|
|
|
let scheme = upstream.scheme();
|
|
match scheme {
|
|
"http" | "https" => {
|
|
// Use manual CONNECT for HTTP/HTTPS proxies
|
|
match connect_via_http_proxy(&upstream, target_host, target_port).await {
|
|
Ok(_) => {
|
|
let mut response = Response::new(Full::new(Bytes::from("")));
|
|
*response.status_mut() = StatusCode::from_u16(200).unwrap();
|
|
Ok(response)
|
|
}
|
|
Err(e) => {
|
|
log::error!("HTTP proxy CONNECT failed: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!(
|
|
"Proxy connection failed: {}",
|
|
e
|
|
))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
Ok(response)
|
|
}
|
|
}
|
|
}
|
|
"socks4" | "socks5" => {
|
|
// Use async-socks5 for SOCKS proxies
|
|
let host = upstream.host_str().unwrap_or("127.0.0.1");
|
|
let port = upstream.port().unwrap_or(1080);
|
|
let socks_addr = format!("{}:{}", host, port);
|
|
|
|
let (username, password) = upstream_userpass(&upstream);
|
|
let auth = (!username.is_empty()).then_some((username.as_str(), password.as_str()));
|
|
|
|
match connect_via_socks(
|
|
&socks_addr,
|
|
target_host,
|
|
target_port,
|
|
scheme == "socks5",
|
|
auth,
|
|
)
|
|
.await
|
|
{
|
|
Ok(_stream) => {
|
|
let mut response = Response::new(Full::new(Bytes::from("")));
|
|
*response.status_mut() = StatusCode::from_u16(200).unwrap();
|
|
Ok(response)
|
|
}
|
|
Err(e) => {
|
|
log::error!("SOCKS connection failed: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!(
|
|
"SOCKS connection failed: {}",
|
|
e
|
|
))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
Ok(response)
|
|
}
|
|
}
|
|
}
|
|
_ => {
|
|
let mut response = Response::new(Full::new(Bytes::from("Unsupported upstream scheme")));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
Ok(response)
|
|
}
|
|
}
|
|
} else {
|
|
let mut response = Response::new(Full::new(Bytes::from("Bad Request")));
|
|
*response.status_mut() = StatusCode::BAD_REQUEST;
|
|
Ok(response)
|
|
}
|
|
}
|
|
|
|
async fn connect_via_http_proxy(
|
|
upstream: &Url,
|
|
target_host: &str,
|
|
target_port: u16,
|
|
) -> Result<TcpStream, Box<dyn std::error::Error>> {
|
|
let proxy_host = upstream.host_str().unwrap_or("127.0.0.1");
|
|
let proxy_port = upstream.port().unwrap_or(8080);
|
|
let mut stream = tokio::time::timeout(
|
|
UPSTREAM_DIAL_TIMEOUT,
|
|
TcpStream::connect((proxy_host, proxy_port)),
|
|
)
|
|
.await
|
|
.map_err(|_| format!("upstream proxy connect to {proxy_host}:{proxy_port} timed out"))??;
|
|
|
|
// Add proxy authentication if provided
|
|
let mut connect_req = format!(
|
|
"CONNECT {}:{} HTTP/1.1\r\nHost: {}:{}\r\n",
|
|
target_host, target_port, target_host, target_port
|
|
);
|
|
|
|
let (username, password) = upstream_userpass(upstream);
|
|
if !username.is_empty() {
|
|
use base64::{engine::general_purpose, Engine as _};
|
|
let auth = general_purpose::STANDARD.encode(format!("{}:{}", username, password));
|
|
connect_req.push_str(&format!("Proxy-Authorization: Basic {}\r\n", auth));
|
|
}
|
|
|
|
connect_req.push_str("\r\n");
|
|
|
|
stream.write_all(connect_req.as_bytes()).await?;
|
|
|
|
let mut buffer = [0u8; 4096];
|
|
let n = tokio::time::timeout(UPSTREAM_DIAL_TIMEOUT, stream.read(&mut buffer))
|
|
.await
|
|
.map_err(|_| "upstream proxy CONNECT response timed out")??;
|
|
let response = String::from_utf8_lossy(&buffer[..n]);
|
|
|
|
if response.starts_with("HTTP/1.1 200") || response.starts_with("HTTP/1.0 200") {
|
|
Ok(stream)
|
|
} else {
|
|
Err(format!("Upstream proxy CONNECT failed: {}", response).into())
|
|
}
|
|
}
|
|
|
|
/// Extract percent-decoded (username, password) from the upstream URL.
|
|
///
|
|
/// `url::Url::username()` / `Url::password()` return percent-encoded ASCII
|
|
/// strings per the WHATWG spec. `build_proxy_url` on the producer side
|
|
/// already percent-encodes the credentials with `urlencoding::encode`, so
|
|
/// we must decode here — otherwise the upstream SOCKS5 / HTTP CONNECT
|
|
/// receives `%40` instead of `@`, breaking RFC1929 user/password
|
|
/// authentication or HTTP Basic-Auth
|
|
fn upstream_userpass(upstream: &Url) -> (String, String) {
|
|
let username = urlencoding::decode(upstream.username())
|
|
.map(|cow| cow.into_owned())
|
|
.unwrap_or_default();
|
|
let password = urlencoding::decode(upstream.password().unwrap_or(""))
|
|
.map(|cow| cow.into_owned())
|
|
.unwrap_or_default();
|
|
(username, password)
|
|
}
|
|
|
|
/// Transparent AsyncRead/AsyncWrite wrapper that logs every read/write
|
|
/// byte of the SOCKS5 handshake. Used only during the handshake — the
|
|
/// inner stream is taken back via `into_inner` once the handshake
|
|
/// completes, so the tunnel phase pays no overhead
|
|
struct SocksHandshakeLogger<S> {
|
|
inner: S,
|
|
label: String,
|
|
}
|
|
|
|
impl<S> SocksHandshakeLogger<S> {
|
|
fn new(inner: S, label: String) -> Self {
|
|
Self { inner, label }
|
|
}
|
|
|
|
fn into_inner(self) -> S {
|
|
self.inner
|
|
}
|
|
}
|
|
|
|
impl<S: AsyncRead + Unpin> AsyncRead for SocksHandshakeLogger<S> {
|
|
fn poll_read(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &mut ReadBuf<'_>,
|
|
) -> Poll<io::Result<()>> {
|
|
let before = buf.filled().len();
|
|
let result = Pin::new(&mut self.inner).poll_read(cx, buf);
|
|
if let Poll::Ready(Ok(())) = &result {
|
|
let after = buf.filled().len();
|
|
if after > before {
|
|
let bytes = &buf.filled()[before..after];
|
|
log::trace!(
|
|
"[socks-handshake:{}] <- {} byte(s): {:02x?}",
|
|
self.label,
|
|
bytes.len(),
|
|
bytes
|
|
);
|
|
} else {
|
|
log::trace!("[socks-handshake:{}] <- EOF (peer closed)", self.label);
|
|
}
|
|
}
|
|
result
|
|
}
|
|
}
|
|
|
|
impl<S: AsyncWrite + Unpin> AsyncWrite for SocksHandshakeLogger<S> {
|
|
fn poll_write(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &[u8],
|
|
) -> Poll<io::Result<usize>> {
|
|
let result = Pin::new(&mut self.inner).poll_write(cx, buf);
|
|
if let Poll::Ready(Ok(n)) = &result {
|
|
log::trace!(
|
|
"[socks-handshake:{}] -> {} byte(s): {:02x?}",
|
|
self.label,
|
|
n,
|
|
&buf[..*n]
|
|
);
|
|
}
|
|
result
|
|
}
|
|
|
|
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
Pin::new(&mut self.inner).poll_flush(cx)
|
|
}
|
|
|
|
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
Pin::new(&mut self.inner).poll_shutdown(cx)
|
|
}
|
|
}
|
|
|
|
async fn connect_via_socks(
|
|
socks_addr: &str,
|
|
target_host: &str,
|
|
target_port: u16,
|
|
is_socks5: bool,
|
|
auth: Option<(&str, &str)>,
|
|
) -> Result<TcpStream, Box<dyn std::error::Error>> {
|
|
let stream = tokio::time::timeout(UPSTREAM_DIAL_TIMEOUT, TcpStream::connect(socks_addr))
|
|
.await
|
|
.map_err(|_| format!("SOCKS upstream connect to {socks_addr} timed out"))??;
|
|
|
|
if is_socks5 {
|
|
// SOCKS5 connection using async_socks5
|
|
use async_socks5::{connect, AddrKind, Auth};
|
|
|
|
let target = if let Ok(ip) = target_host.parse::<std::net::IpAddr>() {
|
|
AddrKind::Ip(std::net::SocketAddr::new(ip, target_port))
|
|
} else {
|
|
AddrKind::Domain(target_host.to_string(), target_port)
|
|
};
|
|
|
|
let auth_info: Option<Auth> = auth.map(|(user, pass)| Auth {
|
|
username: user.to_string(),
|
|
password: pass.to_string(),
|
|
});
|
|
|
|
let has_auth = auth_info.is_some();
|
|
log::trace!(
|
|
"[socks-handshake] dialing {} (target={}:{}, has_auth={})",
|
|
socks_addr,
|
|
target_host,
|
|
target_port,
|
|
has_auth
|
|
);
|
|
|
|
// Disable Nagle so the kernel doesn't further delay/coalesce the
|
|
// syscalls issued when BufStream flushes
|
|
let _ = stream.set_nodelay(true);
|
|
|
|
// BufStream wrapping is required: async_socks5 calls write_u8 for every
|
|
// single-byte SOCKS5 / RFC1929 field, and on a raw TcpStream each call
|
|
// becomes its own TCP segment. Some upstream SOCKS5 implementations
|
|
// treat such a "fragmented auth submission" as a misbehaving client
|
|
// and silently FIN instead of returning an RFC1929 status. BufStream
|
|
// coalesces those small writes into one syscall on flush — this is
|
|
// the usage pattern shown in the async_socks5 README
|
|
let label = format!("{socks_addr}->{target_host}:{target_port}");
|
|
let logged = SocksHandshakeLogger::new(stream, label);
|
|
let mut buffered = tokio::io::BufStream::new(logged);
|
|
let handshake = tokio::time::timeout(
|
|
UPSTREAM_DIAL_TIMEOUT,
|
|
connect(&mut buffered, target, auth_info),
|
|
)
|
|
.await;
|
|
// Unwrap the layered stream: BufStream → SocksHandshakeLogger → TcpStream
|
|
let stream = buffered.into_inner().into_inner();
|
|
match handshake {
|
|
Ok(Ok(_)) => {
|
|
log::trace!("[socks-handshake] handshake completed ok");
|
|
Ok(stream)
|
|
}
|
|
Ok(Err(e)) => {
|
|
log::trace!("[socks-handshake] handshake failed: {:?}", e);
|
|
Err(e.into())
|
|
}
|
|
Err(_) => {
|
|
log::trace!("[socks-handshake] handshake timed out");
|
|
Err("SOCKS5 upstream handshake timed out".into())
|
|
}
|
|
}
|
|
} else {
|
|
let mut stream = stream;
|
|
// SOCKS4 - simplified implementation
|
|
let ip: std::net::IpAddr = target_host.parse()?;
|
|
|
|
let mut request = vec![0x04, 0x01]; // SOCKS4, CONNECT
|
|
request.extend_from_slice(&target_port.to_be_bytes());
|
|
match ip {
|
|
std::net::IpAddr::V4(ipv4) => {
|
|
request.extend_from_slice(&ipv4.octets());
|
|
}
|
|
std::net::IpAddr::V6(_) => {
|
|
return Err("SOCKS4 does not support IPv6".into());
|
|
}
|
|
}
|
|
request.push(0); // NULL terminator for userid
|
|
|
|
stream.write_all(&request).await?;
|
|
|
|
let mut response = [0u8; 8];
|
|
stream.read_exact(&mut response).await?;
|
|
|
|
if response[1] != 0x5A {
|
|
return Err("SOCKS4 connection failed".into());
|
|
}
|
|
|
|
Ok(stream)
|
|
}
|
|
}
|
|
|
|
async fn handle_http_via_socks4(
|
|
req: Request<hyper::body::Incoming>,
|
|
upstream_url: &str,
|
|
) -> Result<Response<Full<Bytes>>, Infallible> {
|
|
// Extract domain for traffic tracking
|
|
let domain = req
|
|
.uri()
|
|
.host()
|
|
.map(|h| h.to_string())
|
|
.unwrap_or_else(|| "unknown".to_string());
|
|
|
|
// Parse upstream SOCKS4 proxy URL
|
|
let upstream = match Url::parse(upstream_url) {
|
|
Ok(url) => url,
|
|
Err(e) => {
|
|
log::error!("Failed to parse SOCKS4 proxy URL: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from("Invalid proxy URL")));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
};
|
|
|
|
let socks_host = upstream.host_str().unwrap_or("127.0.0.1");
|
|
let socks_port = upstream.port().unwrap_or(1080);
|
|
let socks_addr = format!("{}:{}", socks_host, socks_port);
|
|
|
|
// Parse target from request URI
|
|
let target_uri = req.uri();
|
|
let target_host = target_uri.host().unwrap_or("localhost");
|
|
let target_port = target_uri.port_u16().unwrap_or(80);
|
|
|
|
// Connect to SOCKS4 proxy
|
|
let mut socks_stream = match TcpStream::connect(&socks_addr).await {
|
|
Ok(stream) => stream,
|
|
Err(e) => {
|
|
log::error!("Failed to connect to SOCKS4 proxy {}: {}", socks_addr, e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!(
|
|
"Failed to connect to SOCKS4 proxy: {}",
|
|
e
|
|
))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
};
|
|
|
|
// Build a SOCKS4a CONNECT request. We deliberately do NOT resolve the target
|
|
// hostname locally: tokio::net::lookup_host would call the HOST resolver
|
|
// (getaddrinfo), leaking the destination domain to the host's DNS server and
|
|
// defeating the per-profile proxy. SOCKS4a has the PROXY resolve the name —
|
|
// send the sentinel IP 0.0.0.x (x != 0), then the NULL-terminated userid, then
|
|
// the NULL-terminated hostname. (Most SOCKS4 proxies support 4a; a legacy
|
|
// SOCKS4-only proxy without remote DNS cannot be used leak-free for plaintext
|
|
// HTTP — prefer SOCKS5 there.)
|
|
let mut socks_request = vec![0x04, 0x01]; // SOCKS4, CONNECT
|
|
socks_request.extend_from_slice(&target_port.to_be_bytes());
|
|
socks_request.extend_from_slice(&[0, 0, 0, 1]); // 0.0.0.1 => SOCKS4a remote-DNS marker
|
|
socks_request.push(0); // empty userid, NULL-terminated
|
|
socks_request.extend_from_slice(target_host.as_bytes()); // hostname for the proxy to resolve
|
|
socks_request.push(0); // NULL-terminated hostname
|
|
|
|
// Send SOCKS4 CONNECT request
|
|
if let Err(e) = socks_stream.write_all(&socks_request).await {
|
|
log::error!("Failed to send SOCKS4 CONNECT request: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!(
|
|
"Failed to send SOCKS4 request: {}",
|
|
e
|
|
))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
|
|
// Read SOCKS4 response
|
|
let mut socks_response = [0u8; 8];
|
|
if let Err(e) = socks_stream.read_exact(&mut socks_response).await {
|
|
log::error!("Failed to read SOCKS4 response: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!(
|
|
"Failed to read SOCKS4 response: {}",
|
|
e
|
|
))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
|
|
// Check SOCKS4 response (second byte should be 0x5A for success)
|
|
if socks_response[1] != 0x5A {
|
|
log::error!(
|
|
"SOCKS4 connection failed, response code: {}",
|
|
socks_response[1]
|
|
);
|
|
let mut response = Response::new(Full::new(Bytes::from("SOCKS4 connection failed")));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
|
|
// Now send the HTTP request through the SOCKS4 connection
|
|
// Build HTTP request line
|
|
let method = req.method().as_str();
|
|
let path = target_uri
|
|
.path_and_query()
|
|
.map(|pq| pq.as_str())
|
|
.unwrap_or("/");
|
|
let http_version = if req.version() == hyper::Version::HTTP_11 {
|
|
"HTTP/1.1"
|
|
} else {
|
|
"HTTP/1.0"
|
|
};
|
|
|
|
let mut http_request = format!("{} {} {}\r\n", method, path, http_version);
|
|
|
|
// Add Host header if not present
|
|
let mut has_host = false;
|
|
for (name, value) in req.headers().iter() {
|
|
if name.as_str().eq_ignore_ascii_case("host") {
|
|
has_host = true;
|
|
}
|
|
// Skip proxy-specific headers
|
|
if name.as_str().eq_ignore_ascii_case("proxy-authorization")
|
|
|| name.as_str().eq_ignore_ascii_case("proxy-connection")
|
|
|| name.as_str().eq_ignore_ascii_case("proxy-authenticate")
|
|
{
|
|
continue;
|
|
}
|
|
// Skip Content-Length and Transfer-Encoding - we'll add our own Content-Length
|
|
// based on the collected body size. Having both violates HTTP/1.1 (RFC 7230).
|
|
if name.as_str().eq_ignore_ascii_case("content-length")
|
|
|| name.as_str().eq_ignore_ascii_case("transfer-encoding")
|
|
{
|
|
continue;
|
|
}
|
|
if let Ok(val) = value.to_str() {
|
|
http_request.push_str(&format!("{}: {}\r\n", name.as_str(), val));
|
|
}
|
|
}
|
|
|
|
if !has_host {
|
|
http_request.push_str(&format!("Host: {}:{}\r\n", target_host, target_port));
|
|
}
|
|
|
|
// Get body
|
|
let body_bytes = match req.collect().await {
|
|
Ok(collected) => collected.to_bytes(),
|
|
Err(_) => Bytes::new(),
|
|
};
|
|
|
|
// Add Content-Length if there's a body
|
|
if !body_bytes.is_empty() {
|
|
http_request.push_str(&format!("Content-Length: {}\r\n", body_bytes.len()));
|
|
}
|
|
|
|
http_request.push_str("\r\n");
|
|
|
|
// Send HTTP request
|
|
if let Err(e) = socks_stream.write_all(http_request.as_bytes()).await {
|
|
log::error!("Failed to send HTTP request through SOCKS4: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!(
|
|
"Failed to send HTTP request: {}",
|
|
e
|
|
))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
|
|
// Send body if present
|
|
if !body_bytes.is_empty() {
|
|
if let Err(e) = socks_stream.write_all(&body_bytes).await {
|
|
log::error!("Failed to send HTTP body through SOCKS4: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!(
|
|
"Failed to send HTTP body: {}",
|
|
e
|
|
))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
}
|
|
|
|
// Read HTTP response
|
|
let mut response_buffer = Vec::with_capacity(8192);
|
|
let mut temp_buf = [0u8; 4096];
|
|
let mut content_length: Option<usize> = None;
|
|
let mut is_chunked = false;
|
|
|
|
// Read until we have complete headers
|
|
loop {
|
|
match socks_stream.read(&mut temp_buf).await {
|
|
Ok(0) => break, // Connection closed
|
|
Ok(n) => {
|
|
response_buffer.extend_from_slice(&temp_buf[..n]);
|
|
// Check for end of headers (\r\n\r\n)
|
|
if let Some(pos) = response_buffer.windows(4).position(|w| w == b"\r\n\r\n") {
|
|
// Parse headers
|
|
let headers_str = String::from_utf8_lossy(&response_buffer[..pos + 4]);
|
|
for line in headers_str.lines() {
|
|
let line_lower = line.to_lowercase();
|
|
if line_lower.starts_with("content-length:") {
|
|
if let Some(len_str) = line.split(':').nth(1) {
|
|
if let Ok(len) = len_str.trim().parse::<usize>() {
|
|
content_length = Some(len);
|
|
}
|
|
}
|
|
} else if line_lower.starts_with("transfer-encoding:") && line_lower.contains("chunked")
|
|
{
|
|
is_chunked = true;
|
|
}
|
|
}
|
|
// Read body if Content-Length is specified and we don't have it all
|
|
if let Some(cl) = content_length {
|
|
let body_start = pos + 4;
|
|
let body_received = response_buffer.len() - body_start;
|
|
if body_received < cl {
|
|
// Read remaining body (but don't use read_exact as connection might close)
|
|
let remaining = cl - body_received;
|
|
let mut read_so_far = 0;
|
|
while read_so_far < remaining {
|
|
match socks_stream.read(&mut temp_buf).await {
|
|
Ok(0) => break, // Connection closed
|
|
Ok(m) => {
|
|
let to_read = (remaining - read_so_far).min(m);
|
|
response_buffer.extend_from_slice(&temp_buf[..to_read]);
|
|
read_so_far += to_read;
|
|
if to_read < m {
|
|
// More data than needed, might be next response - stop here
|
|
break;
|
|
}
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
}
|
|
} else if !is_chunked {
|
|
// No Content-Length and not chunked - read until connection closes
|
|
// But limit to reasonable size to avoid memory issues
|
|
let max_body_size = 10 * 1024 * 1024; // 10MB max
|
|
while response_buffer.len() < max_body_size {
|
|
match socks_stream.read(&mut temp_buf).await {
|
|
Ok(0) => break, // Connection closed
|
|
Ok(n) => {
|
|
response_buffer.extend_from_slice(&temp_buf[..n]);
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
}
|
|
// Note: Chunked encoding is complex to parse manually, so we'll read what we can
|
|
// For full chunked support, we'd need a proper HTTP parser
|
|
break;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
log::error!("Error reading HTTP response from SOCKS4: {}", e);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Parse HTTP response
|
|
let response_str = String::from_utf8_lossy(&response_buffer);
|
|
let mut lines = response_str.lines();
|
|
let status_line = lines.next().unwrap_or("HTTP/1.1 500 Internal Server Error");
|
|
let status_parts: Vec<&str> = status_line.split_whitespace().collect();
|
|
let status_code = status_parts
|
|
.get(1)
|
|
.and_then(|s| s.parse::<u16>().ok())
|
|
.unwrap_or(500);
|
|
|
|
// Find header/body boundary
|
|
let header_end = response_buffer
|
|
.windows(4)
|
|
.position(|w| w == b"\r\n\r\n")
|
|
.map(|p| p + 4)
|
|
.unwrap_or(response_buffer.len());
|
|
|
|
let body = response_buffer[header_end..].to_vec();
|
|
|
|
// Record request in traffic tracker
|
|
let response_size = body.len() as u64;
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
tracker.record_request(&domain, body_bytes.len() as u64, response_size);
|
|
}
|
|
|
|
let mut hyper_response = Response::new(Full::new(Bytes::from(body)));
|
|
*hyper_response.status_mut() = StatusCode::from_u16(status_code).unwrap();
|
|
|
|
Ok(hyper_response)
|
|
}
|
|
|
|
/// Handle plain HTTP requests through a Shadowsocks upstream.
|
|
/// reqwest doesn't support SS natively, so we connect through the SS tunnel
|
|
/// manually and forward the HTTP request/response.
|
|
async fn handle_http_via_shadowsocks(
|
|
req: Request<hyper::body::Incoming>,
|
|
upstream: &Url,
|
|
) -> Result<Response<Full<Bytes>>, Infallible> {
|
|
let domain = req
|
|
.uri()
|
|
.host()
|
|
.map(|h| h.to_string())
|
|
.unwrap_or_else(|| "unknown".to_string());
|
|
let port = req.uri().port_u16().unwrap_or(80);
|
|
|
|
let ss_host = upstream.host_str().unwrap_or("127.0.0.1");
|
|
let ss_port = upstream.port().unwrap_or(8388);
|
|
let method_str = urlencoding::decode(upstream.username())
|
|
.unwrap_or_default()
|
|
.to_string();
|
|
let password = urlencoding::decode(upstream.password().unwrap_or(""))
|
|
.unwrap_or_default()
|
|
.to_string();
|
|
|
|
let cipher = match method_str.parse::<shadowsocks::crypto::CipherKind>() {
|
|
Ok(c) => c,
|
|
Err(_) => {
|
|
let mut resp = Response::new(Full::new(Bytes::from(format!(
|
|
"Bad SS cipher: {method_str}"
|
|
))));
|
|
*resp.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(resp);
|
|
}
|
|
};
|
|
|
|
let context = shadowsocks::context::Context::new_shared(shadowsocks::config::ServerType::Local);
|
|
let svr_cfg = match shadowsocks::config::ServerConfig::new(
|
|
shadowsocks::config::ServerAddr::from((ss_host.to_string(), ss_port)),
|
|
&password,
|
|
cipher,
|
|
) {
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
let mut resp = Response::new(Full::new(Bytes::from(format!("SS config error: {e}"))));
|
|
*resp.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(resp);
|
|
}
|
|
};
|
|
|
|
let target_addr = shadowsocks::relay::Address::DomainNameAddress(domain.clone(), port);
|
|
|
|
let mut stream = match shadowsocks::relay::tcprelay::proxy_stream::ProxyClientStream::connect(
|
|
context,
|
|
&svr_cfg,
|
|
target_addr,
|
|
)
|
|
.await
|
|
{
|
|
Ok(s) => s,
|
|
Err(e) => {
|
|
let mut resp = Response::new(Full::new(Bytes::from(format!("SS connect: {e}"))));
|
|
*resp.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(resp);
|
|
}
|
|
};
|
|
|
|
// Build and send the HTTP request through the SS tunnel
|
|
let path = req
|
|
.uri()
|
|
.path_and_query()
|
|
.map(|pq| pq.as_str())
|
|
.unwrap_or("/");
|
|
let method = req.method().as_str();
|
|
let mut raw_req = format!("{method} {path} HTTP/1.1\r\nHost: {domain}\r\nConnection: close\r\n");
|
|
for (name, value) in req.headers() {
|
|
if name != "host" && name != "connection" {
|
|
raw_req.push_str(&format!("{}: {}\r\n", name, value.to_str().unwrap_or("")));
|
|
}
|
|
}
|
|
raw_req.push_str("\r\n");
|
|
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
if let Err(e) = stream.write_all(raw_req.as_bytes()).await {
|
|
let mut resp = Response::new(Full::new(Bytes::from(format!("SS write: {e}"))));
|
|
*resp.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(resp);
|
|
}
|
|
|
|
let mut response_buf = Vec::new();
|
|
if let Err(e) = stream.read_to_end(&mut response_buf).await {
|
|
log::warn!("SS read error (may be partial): {e}");
|
|
}
|
|
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
tracker.record_request(&domain, raw_req.len() as u64, response_buf.len() as u64);
|
|
}
|
|
|
|
// Parse the raw HTTP response
|
|
let response_str = String::from_utf8_lossy(&response_buf);
|
|
let header_end = response_str.find("\r\n\r\n").unwrap_or(response_str.len());
|
|
let status_line = response_str
|
|
.lines()
|
|
.next()
|
|
.unwrap_or("HTTP/1.1 502 Bad Gateway");
|
|
let status_code: u16 = status_line
|
|
.split_whitespace()
|
|
.nth(1)
|
|
.and_then(|s| s.parse().ok())
|
|
.unwrap_or(502);
|
|
let body = if header_end + 4 < response_buf.len() {
|
|
&response_buf[header_end + 4..]
|
|
} else {
|
|
b""
|
|
};
|
|
|
|
let mut hyper_response = Response::new(Full::new(Bytes::from(body.to_vec())));
|
|
*hyper_response.status_mut() =
|
|
StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_GATEWAY);
|
|
|
|
Ok(hyper_response)
|
|
}
|
|
|
|
async fn handle_http(
|
|
req: Request<hyper::body::Incoming>,
|
|
upstream_url: Option<String>,
|
|
bypass_matcher: BypassMatcher,
|
|
blocklist_matcher: BlocklistMatcher,
|
|
) -> Result<Response<Full<Bytes>>, Infallible> {
|
|
// Extract domain for traffic tracking
|
|
let domain = req
|
|
.uri()
|
|
.host()
|
|
.map(|h| h.to_string())
|
|
.unwrap_or_else(|| "unknown".to_string());
|
|
|
|
// Block if domain is in the DNS blocklist (before any connection)
|
|
if blocklist_matcher.is_blocked(&domain) {
|
|
log::debug!("[blocklist] Blocked HTTP request to {}", domain);
|
|
let mut response = Response::new(Full::new(Bytes::from("Blocked by DNS blocklist")));
|
|
*response.status_mut() = StatusCode::FORBIDDEN;
|
|
return Ok(response);
|
|
}
|
|
|
|
log::trace!(
|
|
"Handling HTTP request: {} {} (host: {:?})",
|
|
req.method(),
|
|
req.uri(),
|
|
req.uri().host()
|
|
);
|
|
|
|
let should_bypass = bypass_matcher.should_bypass(&domain);
|
|
|
|
// Handle proxy types that reqwest doesn't support natively
|
|
if !should_bypass {
|
|
if let Some(ref upstream) = upstream_url {
|
|
if upstream != "DIRECT" {
|
|
if let Ok(url) = Url::parse(upstream) {
|
|
match url.scheme() {
|
|
"socks4" => {
|
|
return handle_http_via_socks4(req, upstream).await;
|
|
}
|
|
"ss" | "shadowsocks" => {
|
|
return handle_http_via_shadowsocks(req, &url).await;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Use reqwest for HTTP/HTTPS/SOCKS5 proxies
|
|
use reqwest::Client;
|
|
|
|
let client_builder = Client::builder();
|
|
let client = if should_bypass {
|
|
client_builder.build().unwrap_or_default()
|
|
} else if let Some(ref upstream) = upstream_url {
|
|
if upstream == "DIRECT" {
|
|
client_builder.build().unwrap_or_default()
|
|
} else {
|
|
// Build reqwest client with proxy
|
|
match build_reqwest_client_with_proxy(upstream) {
|
|
Ok(c) => c,
|
|
Err(e) => {
|
|
log::error!("Failed to create proxy client: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!(
|
|
"Proxy configuration error: {}",
|
|
e
|
|
))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
return Ok(response);
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
client_builder.build().unwrap_or_default()
|
|
};
|
|
|
|
// Convert hyper request to reqwest request
|
|
let uri = req.uri().to_string();
|
|
let method = req.method().clone();
|
|
let headers = req.headers().clone();
|
|
|
|
let mut request_builder = match method.as_str() {
|
|
"GET" => client.get(&uri),
|
|
"POST" => client.post(&uri),
|
|
"PUT" => client.put(&uri),
|
|
"DELETE" => client.delete(&uri),
|
|
"PATCH" => client.patch(&uri),
|
|
"HEAD" => client.head(&uri),
|
|
_ => {
|
|
let mut response = Response::new(Full::new(Bytes::from("Unsupported method")));
|
|
*response.status_mut() = StatusCode::METHOD_NOT_ALLOWED;
|
|
return Ok(response);
|
|
}
|
|
};
|
|
|
|
// Copy headers, but skip proxy-specific headers that shouldn't be forwarded
|
|
for (name, value) in headers.iter() {
|
|
// Skip proxy-specific headers - these are for the local proxy, not the upstream
|
|
if name.as_str().eq_ignore_ascii_case("proxy-authorization")
|
|
|| name.as_str().eq_ignore_ascii_case("proxy-connection")
|
|
|| name.as_str().eq_ignore_ascii_case("proxy-authenticate")
|
|
{
|
|
continue;
|
|
}
|
|
if let Ok(val) = value.to_str() {
|
|
request_builder = request_builder.header(name.as_str(), val);
|
|
}
|
|
}
|
|
|
|
// Get body
|
|
let body_bytes = match req.collect().await {
|
|
Ok(collected) => collected.to_bytes(),
|
|
Err(_) => Bytes::new(),
|
|
};
|
|
|
|
if !body_bytes.is_empty() {
|
|
request_builder = request_builder.body(body_bytes.to_vec());
|
|
}
|
|
|
|
// Execute request
|
|
match request_builder.send().await {
|
|
Ok(response) => {
|
|
let status = response.status();
|
|
let headers = response.headers().clone();
|
|
let body = response.bytes().await.unwrap_or_default();
|
|
|
|
// Record request in traffic tracker
|
|
let response_size = body.len() as u64;
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
tracker.record_request(&domain, body_bytes.len() as u64, response_size);
|
|
}
|
|
|
|
let mut hyper_response = Response::new(Full::new(body));
|
|
*hyper_response.status_mut() = StatusCode::from_u16(status.as_u16()).unwrap();
|
|
|
|
// Copy response headers
|
|
for (name, value) in headers.iter() {
|
|
if let Ok(val) = value.to_str() {
|
|
hyper_response
|
|
.headers_mut()
|
|
.insert(name, val.parse().unwrap());
|
|
}
|
|
}
|
|
|
|
Ok(hyper_response)
|
|
}
|
|
Err(e) => {
|
|
log::error!("Request failed: {}", e);
|
|
let mut response = Response::new(Full::new(Bytes::from(format!("Request failed: {}", e))));
|
|
*response.status_mut() = StatusCode::BAD_GATEWAY;
|
|
Ok(response)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn build_reqwest_client_with_proxy(
|
|
upstream_url: &str,
|
|
) -> Result<reqwest::Client, Box<dyn std::error::Error>> {
|
|
use reqwest::Proxy;
|
|
|
|
let client_builder = reqwest::Client::builder();
|
|
|
|
// Parse the upstream URL
|
|
let url = Url::parse(upstream_url)?;
|
|
let scheme = url.scheme();
|
|
|
|
let proxy = match scheme {
|
|
"http" | "https" => {
|
|
// For HTTP/HTTPS proxies, reqwest handles them directly
|
|
// Note: HTTPS proxy URLs still use HTTP CONNECT method, reqwest handles TLS automatically
|
|
Proxy::http(upstream_url)?
|
|
}
|
|
"socks5" => {
|
|
// Donut: force REMOTE (proxy-side) DNS for plaintext HTTP over a SOCKS5
|
|
// upstream. reqwest maps the bare `socks5` scheme to DnsResolve::Local,
|
|
// which resolves the destination hostname on the HOST (getaddrinfo) BEFORE
|
|
// connecting — leaking the destination domain to the host's DNS resolver
|
|
// and defeating the per-profile proxy. The `socks5h` scheme maps to
|
|
// DnsResolve::Proxy, so the proxy resolves the hostname and nothing leaks.
|
|
// (The CONNECT/HTTPS path already does remote DNS via connect_via_socks's
|
|
// AddrKind::Domain.)
|
|
let remote_dns_url = match upstream_url.strip_prefix("socks5://") {
|
|
Some(rest) => format!("socks5h://{rest}"),
|
|
None => upstream_url.to_string(),
|
|
};
|
|
Proxy::all(remote_dns_url)?
|
|
}
|
|
"socks4" => {
|
|
// SOCKS4 is handled manually in handle_http_via_socks4
|
|
// This should not be reached, but return error as fallback
|
|
return Err("SOCKS4 should be handled manually".into());
|
|
}
|
|
_ => {
|
|
return Err(format!("Unsupported proxy scheme: {}", scheme).into());
|
|
}
|
|
};
|
|
|
|
Ok(client_builder.proxy(proxy).build()?)
|
|
}
|
|
|
|
/// Handle a single proxy connection (used by both the proxy worker and in-process proxy checks).
|
|
pub async fn handle_proxy_connection(
|
|
mut stream: tokio::net::TcpStream,
|
|
upstream_url: Option<String>,
|
|
bypass_matcher: BypassMatcher,
|
|
blocklist_matcher: BlocklistMatcher,
|
|
) {
|
|
let _ = stream.set_nodelay(true);
|
|
|
|
if stream.readable().await.is_err() {
|
|
return;
|
|
}
|
|
|
|
let mut peek_buffer = [0u8; 16];
|
|
match stream.read(&mut peek_buffer).await {
|
|
Ok(0) => {}
|
|
Ok(n) => {
|
|
let request_start_upper = String::from_utf8_lossy(&peek_buffer[..n.min(7)]).to_uppercase();
|
|
let is_connect = request_start_upper.starts_with("CONNECT");
|
|
|
|
if is_connect {
|
|
let mut full_request = Vec::with_capacity(4096);
|
|
full_request.extend_from_slice(&peek_buffer[..n]);
|
|
|
|
let mut remaining = [0u8; 4096];
|
|
let mut total_read = n;
|
|
let max_reads = 100;
|
|
let mut reads = 0;
|
|
|
|
loop {
|
|
if reads >= max_reads {
|
|
break;
|
|
}
|
|
match stream.read(&mut remaining).await {
|
|
Ok(0) => {
|
|
if full_request.ends_with(b"\r\n\r\n")
|
|
|| full_request.ends_with(b"\n\n")
|
|
|| total_read > 0
|
|
{
|
|
break;
|
|
}
|
|
return;
|
|
}
|
|
Ok(m) => {
|
|
reads += 1;
|
|
total_read += m;
|
|
full_request.extend_from_slice(&remaining[..m]);
|
|
if full_request.ends_with(b"\r\n\r\n") || full_request.ends_with(b"\n\n") {
|
|
break;
|
|
}
|
|
}
|
|
Err(_) => {
|
|
if total_read > 0 {
|
|
break;
|
|
}
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Err(e) = handle_connect_from_buffer(
|
|
stream,
|
|
full_request,
|
|
upstream_url,
|
|
bypass_matcher,
|
|
blocklist_matcher,
|
|
)
|
|
.await
|
|
{
|
|
let msg = e.to_string();
|
|
if let Some(suppressed) = log_throttle(&msg) {
|
|
if suppressed > 0 {
|
|
log::warn!(
|
|
"CONNECT tunnel ended with error: {msg} ({suppressed} more suppressed in last 30s)"
|
|
);
|
|
} else {
|
|
log::warn!("CONNECT tunnel ended with error: {msg}");
|
|
}
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
// Non-CONNECT: prepend consumed bytes and pass to hyper
|
|
let prepended_bytes = peek_buffer[..n].to_vec();
|
|
let prepended_reader = PrependReader {
|
|
prepended: prepended_bytes,
|
|
prepended_pos: 0,
|
|
inner: stream,
|
|
};
|
|
let io = TokioIo::new(prepended_reader);
|
|
let service = service_fn(move |req| {
|
|
handle_request(
|
|
req,
|
|
upstream_url.clone(),
|
|
bypass_matcher.clone(),
|
|
blocklist_matcher.clone(),
|
|
)
|
|
});
|
|
|
|
let _ = http1::Builder::new().serve_connection(io, service).await;
|
|
}
|
|
Err(_) => {}
|
|
}
|
|
}
|
|
|
|
/// Render an upstream proxy URL for logging with any embedded credentials
|
|
/// stripped. `config.upstream_url` carries `scheme://user:pass@host:port`, and
|
|
/// these logs land in a world-readable file under the system temp dir, so the
|
|
/// userinfo must never be emitted.
|
|
fn redacted_upstream(upstream: &str) -> String {
|
|
if upstream.is_empty() {
|
|
return "none".to_string();
|
|
}
|
|
match Url::parse(upstream) {
|
|
Ok(u) => match (u.host_str(), u.port()) {
|
|
(Some(host), Some(port)) => format!("{}://{host}:{port}", u.scheme()),
|
|
(Some(host), None) => format!("{}://{host}", u.scheme()),
|
|
_ => "<redacted>".to_string(),
|
|
},
|
|
Err(_) => "<redacted>".to_string(),
|
|
}
|
|
}
|
|
|
|
pub async fn run_proxy_server(config: ProxyConfig) -> Result<(), Box<dyn std::error::Error>> {
|
|
log::info!(
|
|
"Proxy worker starting, looking for config id: {}",
|
|
config.id
|
|
);
|
|
|
|
// Load the config from disk to get the latest state
|
|
let config = match crate::proxy_storage::get_proxy_config(&config.id) {
|
|
Some(c) => c,
|
|
None => {
|
|
log::error!("Config not found for id: {}", config.id);
|
|
return Err("Config not found".into());
|
|
}
|
|
};
|
|
|
|
log::info!(
|
|
"Found config: id={}, port={:?}, upstream={}, profile_id={:?}",
|
|
config.id,
|
|
config.local_port,
|
|
redacted_upstream(&config.upstream_url),
|
|
config.profile_id
|
|
);
|
|
|
|
// Initialize traffic tracker with profile ID if available.
|
|
// This can be called multiple times to update the tracker.
|
|
init_traffic_tracker(config.id.clone(), config.profile_id.clone());
|
|
|
|
// Determine the bind address
|
|
let bind_addr = SocketAddr::from(([127, 0, 0, 1], config.local_port.unwrap_or(0)));
|
|
|
|
log::info!("Attempting to bind proxy server to {}", bind_addr);
|
|
|
|
// Bind to the port. Use SO_REUSEADDR so that a freshly-restarted worker
|
|
// can bind a port that the previous worker left in TIME_WAIT, and retry
|
|
// briefly to absorb transient races with the OS releasing the socket.
|
|
let listener = {
|
|
let mut attempts: u32 = 0;
|
|
loop {
|
|
let socket = tokio::net::TcpSocket::new_v4()?;
|
|
let _ = socket.set_reuseaddr(true);
|
|
match socket.bind(bind_addr) {
|
|
Ok(()) => match socket.listen(1024) {
|
|
Ok(l) => break l,
|
|
Err(e) if attempts < 5 => {
|
|
attempts += 1;
|
|
let delay = std::time::Duration::from_millis(200 * u64::from(attempts));
|
|
log::warn!(
|
|
"listen() on {} failed (attempt {}/5): {}, retrying in {}ms",
|
|
bind_addr,
|
|
attempts,
|
|
e,
|
|
delay.as_millis()
|
|
);
|
|
tokio::time::sleep(delay).await;
|
|
}
|
|
Err(e) => {
|
|
return Err(format!("Failed to listen on {bind_addr} after 5 attempts: {e}").into())
|
|
}
|
|
},
|
|
Err(e) if attempts < 5 => {
|
|
attempts += 1;
|
|
let delay = std::time::Duration::from_millis(200 * u64::from(attempts));
|
|
log::warn!(
|
|
"bind() on {} failed (attempt {}/5): {}, retrying in {}ms",
|
|
bind_addr,
|
|
attempts,
|
|
e,
|
|
delay.as_millis()
|
|
);
|
|
tokio::time::sleep(delay).await;
|
|
}
|
|
Err(e) => return Err(format!("Failed to bind {bind_addr} after 5 attempts: {e}").into()),
|
|
}
|
|
}
|
|
};
|
|
let actual_port = listener.local_addr()?.port();
|
|
|
|
log::info!("Successfully bound to port {}", actual_port);
|
|
|
|
// Protocol served to the browser: "socks5" (Wayfern) or "http" (default).
|
|
let local_protocol = config.local_protocol_or_default();
|
|
let serve_socks5 = local_protocol == "socks5";
|
|
|
|
// Update config with actual port and local_url (scheme matches the protocol
|
|
// we serve, so the parent's readiness check and any consumer see the truth)
|
|
let mut updated_config = config.clone();
|
|
updated_config.local_port = Some(actual_port);
|
|
updated_config.local_url = Some(format!(
|
|
"{}://127.0.0.1:{}",
|
|
if serve_socks5 { "socks5" } else { "http" },
|
|
actual_port
|
|
));
|
|
|
|
if !crate::proxy_storage::update_proxy_config(&updated_config) {
|
|
log::error!("Failed to update proxy config");
|
|
return Err("Failed to update proxy config".into());
|
|
}
|
|
|
|
let upstream_url = if updated_config.upstream_url == "DIRECT" {
|
|
None
|
|
} else {
|
|
Some(updated_config.upstream_url.clone())
|
|
};
|
|
|
|
log::info!(
|
|
"Proxy server listening on 127.0.0.1:{} (ready to accept connections)",
|
|
actual_port
|
|
);
|
|
log::info!("Proxy server entering accept loop - process should stay alive");
|
|
|
|
// Start a background task to write lightweight session snapshots for real-time updates
|
|
// These are much smaller than full stats and can be written frequently (~100 bytes every 2 seconds)
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
let tracker_clone = tracker.clone();
|
|
tokio::spawn(async move {
|
|
let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(2));
|
|
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
|
|
loop {
|
|
interval.tick().await;
|
|
// Write lightweight session snapshot (only current counters, ~100 bytes)
|
|
if let Err(e) = tracker_clone.write_session_snapshot() {
|
|
log::debug!("Failed to write session snapshot: {}", e);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
// Start a background task to periodically flush traffic stats to disk
|
|
// Use adaptive flush frequency: every 5 seconds when active, every 30 seconds when idle
|
|
tokio::spawn(async move {
|
|
let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(5));
|
|
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
|
let mut last_activity_time = std::time::Instant::now();
|
|
let mut last_flush_time = std::time::Instant::now();
|
|
let mut current_interval_secs = 5u64;
|
|
|
|
loop {
|
|
interval.tick().await;
|
|
// Catch panics so a poisoned lock or unexpected error inside
|
|
// flush_to_disk doesn't abort the flush task and leave stats
|
|
// unwritten for the lifetime of the worker. The captured state
|
|
// is all Copy or atomic-assignment, so AssertUnwindSafe is sound.
|
|
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
let (sent, recv, requests) = tracker.get_snapshot();
|
|
let current_bytes = sent + recv;
|
|
let time_since_activity = last_activity_time.elapsed();
|
|
let time_since_flush = last_flush_time.elapsed();
|
|
let has_traffic = current_bytes > 0 || requests > 0;
|
|
|
|
let desired_interval_secs =
|
|
if has_traffic || time_since_activity < std::time::Duration::from_secs(30) {
|
|
5u64
|
|
} else {
|
|
30u64
|
|
};
|
|
|
|
if desired_interval_secs != current_interval_secs {
|
|
current_interval_secs = desired_interval_secs;
|
|
interval =
|
|
tokio::time::interval(tokio::time::Duration::from_secs(desired_interval_secs));
|
|
}
|
|
|
|
let flush_interval = std::time::Duration::from_secs(desired_interval_secs);
|
|
let should_flush = time_since_flush >= flush_interval;
|
|
|
|
if should_flush {
|
|
match tracker.flush_to_disk() {
|
|
Ok(Some((sent, recv))) => {
|
|
last_flush_time = std::time::Instant::now();
|
|
if sent > 0 || recv > 0 {
|
|
last_activity_time = std::time::Instant::now();
|
|
}
|
|
}
|
|
Ok(None) => {
|
|
last_flush_time = std::time::Instant::now();
|
|
}
|
|
Err(e) => {
|
|
log::error!("Failed to flush traffic stats: {}", e);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}));
|
|
if let Err(panic) = result {
|
|
log::error!("Panic caught in proxy traffic flush task; continuing: {panic:?}");
|
|
}
|
|
}
|
|
});
|
|
|
|
// Self-reaping supervisor. The worker is a detached process that outlives the
|
|
// GUI, so it cannot rely on the GUI's in-memory death-monitor (which is lost
|
|
// when the GUI restarts). Once the GUI records the browser PID this worker
|
|
// serves, poll it and exit when that browser is gone — never while it is
|
|
// alive, and never before a PID is recorded (covers the launch window and
|
|
// pre-upgrade configs lacking the field). A 2-miss debounce avoids exiting on
|
|
// a transient sysinfo false-negative under load / sleep-wake.
|
|
//
|
|
// This runs on a DEDICATED OS THREAD, not a tokio task. If the worker's
|
|
// accept/dial path ever busy-loops (e.g. a client retry-storm against a
|
|
// failing upstream), it saturates the async runtime, and a tokio-based
|
|
// supervisor would never be scheduled — leaving the worker spinning forever
|
|
// even after its browser exits or its config is deleted (observed in the
|
|
// field as pegged-CPU orphans that survive config deletion). A real thread
|
|
// with a blocking sleep cannot be starved that way, so the worker always
|
|
// reaps itself. Every call here is synchronous and safe off the runtime.
|
|
{
|
|
let watch_id = config.id.clone();
|
|
std::thread::spawn(move || {
|
|
let mut consecutive_misses: u32 = 0;
|
|
loop {
|
|
std::thread::sleep(std::time::Duration::from_secs(15));
|
|
match crate::proxy_storage::get_proxy_config(&watch_id) {
|
|
Some(cfg) => match cfg.browser_pid {
|
|
Some(bpid) if bpid != 0 => {
|
|
if crate::proxy_storage::is_process_running(bpid) {
|
|
consecutive_misses = 0;
|
|
} else {
|
|
consecutive_misses += 1;
|
|
if consecutive_misses >= 2 {
|
|
log::info!("Browser PID {bpid} for config {watch_id} is gone; worker exiting");
|
|
crate::proxy_storage::delete_proxy_config(&watch_id);
|
|
std::process::exit(0);
|
|
}
|
|
}
|
|
}
|
|
// No browser PID recorded yet (launch window / old config): keep running.
|
|
_ => consecutive_misses = 0,
|
|
},
|
|
// Our own config was removed (e.g. GUI stopped us): nothing to serve.
|
|
None => {
|
|
log::info!("Proxy config {watch_id} was removed; worker exiting");
|
|
std::process::exit(0);
|
|
}
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
let bypass_matcher = BypassMatcher::new(&config.bypass_rules);
|
|
let blocklist_matcher = if let Some(ref path) = config.blocklist_file {
|
|
match BlocklistMatcher::from_file(path) {
|
|
Ok(m) => m,
|
|
Err(e) => {
|
|
log::error!("[blocklist] Failed to load from {}: {}", path, e);
|
|
BlocklistMatcher::new()
|
|
}
|
|
}
|
|
} else {
|
|
BlocklistMatcher::new()
|
|
};
|
|
|
|
// Bound concurrent connection handlers. A client retry-storm (e.g. a browser
|
|
// hammering CONNECT requests while DNS is failing) must not spawn unbounded
|
|
// tasks,
|
|
// each of which parks a Tokio blocking thread inside getaddrinfo — that is
|
|
// what exhausted the resolver pool and pegged the CPU on long-lived workers.
|
|
// A real browser never approaches this ceiling; waiting for a permit
|
|
// backpressures a storm instead of amplifying it.
|
|
let conn_semaphore = Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_CONNECTIONS));
|
|
|
|
// Keep the runtime alive with an infinite loop
|
|
// This ensures the process doesn't exit even if there are no active connections
|
|
loop {
|
|
match listener.accept().await {
|
|
Ok((stream, _peer_addr)) => {
|
|
// The semaphore is never closed, so acquire cannot fail.
|
|
let permit = conn_semaphore
|
|
.clone()
|
|
.acquire_owned()
|
|
.await
|
|
.expect("connection semaphore is never closed");
|
|
let upstream = upstream_url.clone();
|
|
let matcher = bypass_matcher.clone();
|
|
let blocker = blocklist_matcher.clone();
|
|
if serve_socks5 {
|
|
tokio::task::spawn(async move {
|
|
let _permit = permit;
|
|
crate::socks5_local::handle_socks5_connection(stream, upstream, matcher, blocker).await;
|
|
});
|
|
} else {
|
|
tokio::task::spawn(async move {
|
|
let _permit = permit;
|
|
handle_proxy_connection(stream, upstream, matcher, blocker).await;
|
|
});
|
|
}
|
|
}
|
|
Err(e) => {
|
|
log::error!("Error accepting connection: {:?}", e);
|
|
// Continue accepting connections even if one fails
|
|
// Add a small delay to avoid busy-waiting on errors
|
|
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_connect_from_buffer(
|
|
mut client_stream: TcpStream,
|
|
request_buffer: Vec<u8>,
|
|
upstream_url: Option<String>,
|
|
bypass_matcher: BypassMatcher,
|
|
blocklist_matcher: BlocklistMatcher,
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
|
// Parse the CONNECT request from the buffer
|
|
let request_str = String::from_utf8_lossy(&request_buffer);
|
|
let lines: Vec<&str> = request_str.lines().collect();
|
|
|
|
if lines.is_empty() {
|
|
let _ = client_stream
|
|
.write_all(b"HTTP/1.1 400 Bad Request\r\n\r\n")
|
|
.await;
|
|
return Err("Empty CONNECT request".into());
|
|
}
|
|
|
|
// Parse CONNECT request: "CONNECT host:port HTTP/1.1"
|
|
let parts: Vec<&str> = lines[0].split_whitespace().collect();
|
|
if parts.len() < 2 || parts[0] != "CONNECT" {
|
|
let _ = client_stream
|
|
.write_all(b"HTTP/1.1 400 Bad Request\r\n\r\n")
|
|
.await;
|
|
return Err("Invalid CONNECT request".into());
|
|
}
|
|
|
|
let target = parts[1];
|
|
let (target_host, target_port) = if let Some(colon_pos) = target.find(':') {
|
|
let host = &target[..colon_pos];
|
|
let port: u16 = target[colon_pos + 1..].parse().unwrap_or(443);
|
|
(host, port)
|
|
} else {
|
|
(target, 443)
|
|
};
|
|
|
|
// Block if domain is in the DNS blocklist (before any connection)
|
|
if blocklist_matcher.is_blocked(target_host) {
|
|
log::debug!("[blocklist] Blocked CONNECT tunnel to {}", target_host);
|
|
let _ = client_stream
|
|
.write_all(b"HTTP/1.1 403 Forbidden\r\nContent-Length: 24\r\n\r\nBlocked by DNS blocklist")
|
|
.await;
|
|
return Ok(());
|
|
}
|
|
|
|
// Record domain access in traffic tracker
|
|
let domain = target_host.to_string();
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
tracker.record_request(&domain, 0, 0);
|
|
}
|
|
|
|
log::debug!(
|
|
"CONNECT {}:{} (upstream={})",
|
|
target_host,
|
|
target_port,
|
|
upstream_url
|
|
.as_deref()
|
|
.map(redacted_upstream)
|
|
.unwrap_or_else(|| "DIRECT".to_string())
|
|
);
|
|
|
|
// Connect to target (directly or via upstream proxy).
|
|
let target_stream = connect_to_target_via_upstream(
|
|
target_host,
|
|
target_port,
|
|
upstream_url.as_deref(),
|
|
&bypass_matcher,
|
|
)
|
|
.await?;
|
|
|
|
// Send 200 Connection Established response to client
|
|
// CRITICAL: Must flush after writing to ensure response is sent before tunneling
|
|
client_stream
|
|
.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")
|
|
.await?;
|
|
client_stream.flush().await?;
|
|
|
|
log::trace!("Sent 200 Connection Established response, starting tunnel");
|
|
|
|
tunnel_streams(client_stream, target_stream, domain).await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Upper bound on concurrent connection handlers per worker. A real browser
|
|
/// never holds anywhere near this many simultaneous tunnels; the cap stops a
|
|
/// client retry-storm from spawning unbounded tasks (each of which parks a
|
|
/// Tokio blocking thread inside getaddrinfo).
|
|
const MAX_CONCURRENT_CONNECTIONS: usize = 512;
|
|
|
|
/// Connect timeout for the direct (no-upstream) dial path. Bounds a wedged
|
|
/// `getaddrinfo` so a broken resolver can't park a blocking thread for the
|
|
/// full OS timeout.
|
|
const DIRECT_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
|
|
|
|
/// Overall timeout for dialing an UPSTREAM proxy (TCP connect + CONNECT/SOCKS/SS
|
|
/// handshake). Without it, an upstream that accepts TCP but stalls before
|
|
/// replying hangs the worker task forever and holds a connection slot; under
|
|
/// load (e.g. two profiles sharing one proxy) the slots exhaust and the browser
|
|
/// sees `ERR_PROXY_CONNECTION_FAILED` until the profile is restarted. A
|
|
/// bounded dial fails fast and releases the slot.
|
|
const UPSTREAM_DIAL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
|
|
|
|
/// Per-host failure state (last failure instant, consecutive failure count) for
|
|
/// the direct dial path. Process-global — each worker is its own process.
|
|
fn direct_dial_failures() -> &'static Mutex<HashMap<String, (std::time::Instant, u32)>> {
|
|
static M: OnceLock<Mutex<HashMap<String, (std::time::Instant, u32)>>> = OnceLock::new();
|
|
M.get_or_init(|| Mutex::new(HashMap::new()))
|
|
}
|
|
|
|
/// If `host` is inside its failure backoff window, return the remaining time so
|
|
/// the caller can short-circuit without a fresh getaddrinfo/connect. Never
|
|
/// mutates state, so the window always expires and the path self-heals once
|
|
/// DNS recovers.
|
|
fn direct_backoff_remaining(host: &str) -> Option<std::time::Duration> {
|
|
let map = direct_dial_failures();
|
|
let guard = map.lock().unwrap();
|
|
let (last, fails) = guard.get(host).copied()?;
|
|
// Exponential window capped at 30s: 2, 4, 8, 16, 30, 30, ...
|
|
let window = std::time::Duration::from_secs((1u64 << fails.min(5)).min(30));
|
|
let elapsed = last.elapsed();
|
|
if elapsed < window {
|
|
Some(window - elapsed)
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// Record a direct-dial failure for `host`, growing its backoff window.
|
|
fn direct_backoff_record(host: &str) {
|
|
let map = direct_dial_failures();
|
|
let mut guard = map.lock().unwrap();
|
|
// Bound memory against a page that emits many distinct failing hosts.
|
|
if guard.len() > 2048 {
|
|
guard.retain(|_, (last, _)| last.elapsed() < std::time::Duration::from_secs(60));
|
|
}
|
|
let entry = guard
|
|
.entry(host.to_string())
|
|
.or_insert_with(|| (std::time::Instant::now(), 0));
|
|
entry.0 = std::time::Instant::now();
|
|
entry.1 = entry.1.saturating_add(1);
|
|
}
|
|
|
|
/// Clear `host`'s failure state after a successful dial.
|
|
fn direct_backoff_clear(host: &str) {
|
|
direct_dial_failures().lock().unwrap().remove(host);
|
|
}
|
|
|
|
/// Dial a target directly (no upstream) with a connect timeout and per-host
|
|
/// failure backoff. This is the server-side counterpart to the browser's
|
|
/// instant client-side retry: when a host's DNS/connect is failing (e.g. the
|
|
/// macOS resolver wedges after sleep/wake), repeated CONNECT requests
|
|
/// short-circuit
|
|
/// here instead of each spawning a fresh blocking getaddrinfo — which is what
|
|
/// let a retry-storm exhaust the blocking thread pool and peg the CPU.
|
|
async fn dial_direct(host: &str, port: u16) -> Result<TcpStream, Box<dyn std::error::Error>> {
|
|
if let Some(remaining) = direct_backoff_remaining(host) {
|
|
return Err(
|
|
format!(
|
|
"skipping direct dial to {host}: backing off ~{}s after repeated connect failures",
|
|
remaining.as_secs().max(1)
|
|
)
|
|
.into(),
|
|
);
|
|
}
|
|
match tokio::time::timeout(DIRECT_CONNECT_TIMEOUT, TcpStream::connect((host, port))).await {
|
|
Ok(Ok(stream)) => {
|
|
let _ = stream.set_nodelay(true);
|
|
direct_backoff_clear(host);
|
|
Ok(stream)
|
|
}
|
|
Ok(Err(e)) => {
|
|
direct_backoff_record(host);
|
|
Err(e.into())
|
|
}
|
|
Err(_) => {
|
|
direct_backoff_record(host);
|
|
Err(
|
|
format!(
|
|
"direct connect to {host}:{port} timed out after {}s",
|
|
DIRECT_CONNECT_TIMEOUT.as_secs()
|
|
)
|
|
.into(),
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Rate-limit a repetitive log line keyed by `key`: returns `Some(suppressed)`
|
|
/// when the caller should emit (first time or after a 30s window, with the
|
|
/// count dropped since the last emit), or `None` to skip. Stops a connect/DNS
|
|
/// storm from writing the same WARN millions of times (the line that grew
|
|
/// worker logs to 100MB).
|
|
pub(crate) fn log_throttle(key: &str) -> Option<u64> {
|
|
fn throttle_map() -> &'static Mutex<HashMap<String, (std::time::Instant, u64)>> {
|
|
static M: OnceLock<Mutex<HashMap<String, (std::time::Instant, u64)>>> = OnceLock::new();
|
|
M.get_or_init(|| Mutex::new(HashMap::new()))
|
|
}
|
|
let map = throttle_map();
|
|
let mut guard = map.lock().unwrap();
|
|
if guard.len() > 2048 {
|
|
guard.retain(|_, (last, _)| last.elapsed() < std::time::Duration::from_secs(60));
|
|
}
|
|
let now = std::time::Instant::now();
|
|
match guard.get_mut(key) {
|
|
Some((last, suppressed)) => {
|
|
if now.duration_since(*last) >= std::time::Duration::from_secs(30) {
|
|
let dropped = *suppressed;
|
|
*last = now;
|
|
*suppressed = 0;
|
|
Some(dropped)
|
|
} else {
|
|
*suppressed += 1;
|
|
None
|
|
}
|
|
}
|
|
None => {
|
|
guard.insert(key.to_string(), (now, 0));
|
|
Some(0)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Establish a stream to `target_host:target_port`, either directly or through
|
|
/// the configured upstream proxy. Shared by the HTTP CONNECT path and the
|
|
/// local SOCKS5 server so every upstream type (direct, HTTP/HTTPS CONNECT,
|
|
/// SOCKS4/5, Shadowsocks) is dialed in exactly one place. Returns a
|
|
/// `BoxedAsyncStream` so the caller can tunnel over any upstream uniformly.
|
|
pub(crate) async fn connect_to_target_via_upstream(
|
|
target_host: &str,
|
|
target_port: u16,
|
|
upstream_url: Option<&str>,
|
|
bypass_matcher: &BypassMatcher,
|
|
) -> Result<BoxedAsyncStream, Box<dyn std::error::Error>> {
|
|
let should_bypass = bypass_matcher.should_bypass(target_host);
|
|
// Helper: configure outbound TCP to match browser TCP fingerprint
|
|
let configure_tcp = |stream: &TcpStream| {
|
|
let _ = stream.set_nodelay(true);
|
|
};
|
|
let target_stream: BoxedAsyncStream = match upstream_url {
|
|
None | Some("DIRECT") => Box::new(dial_direct(target_host, target_port).await?),
|
|
_ if should_bypass => Box::new(dial_direct(target_host, target_port).await?),
|
|
Some(upstream_url_str) => {
|
|
let upstream = Url::parse(upstream_url_str)?;
|
|
let scheme = upstream.scheme();
|
|
|
|
match scheme {
|
|
"http" | "https" => {
|
|
let proxy_host = upstream.host_str().unwrap_or("127.0.0.1");
|
|
let proxy_port = upstream.port().unwrap_or(8080);
|
|
let mut proxy_stream = tokio::time::timeout(
|
|
UPSTREAM_DIAL_TIMEOUT,
|
|
TcpStream::connect((proxy_host, proxy_port)),
|
|
)
|
|
.await
|
|
.map_err(|_| {
|
|
format!("upstream proxy connect to {proxy_host}:{proxy_port} timed out")
|
|
})??;
|
|
configure_tcp(&proxy_stream);
|
|
|
|
let mut connect_req = format!(
|
|
"CONNECT {}:{} HTTP/1.1\r\nHost: {}:{}\r\n",
|
|
target_host, target_port, target_host, target_port
|
|
);
|
|
|
|
let (username, password) = upstream_userpass(&upstream);
|
|
if !username.is_empty() {
|
|
use base64::{engine::general_purpose, Engine as _};
|
|
let auth = general_purpose::STANDARD.encode(format!("{}:{}", username, password));
|
|
connect_req.push_str(&format!("Proxy-Authorization: Basic {}\r\n", auth));
|
|
}
|
|
|
|
connect_req.push_str("\r\n");
|
|
|
|
proxy_stream.write_all(connect_req.as_bytes()).await?;
|
|
|
|
let mut buffer = [0u8; 4096];
|
|
let n = tokio::time::timeout(UPSTREAM_DIAL_TIMEOUT, proxy_stream.read(&mut buffer))
|
|
.await
|
|
.map_err(|_| "upstream proxy CONNECT response timed out")??;
|
|
let response_full = String::from_utf8_lossy(&buffer[..n]).to_string();
|
|
let status_line = response_full.lines().next().unwrap_or("").to_string();
|
|
|
|
if !response_full.starts_with("HTTP/1.1 200")
|
|
&& !response_full.starts_with("HTTP/1.0 200")
|
|
{
|
|
log::warn!(
|
|
"Upstream CONNECT to {}:{} via {}:{} rejected: {}",
|
|
target_host,
|
|
target_port,
|
|
proxy_host,
|
|
proxy_port,
|
|
status_line
|
|
);
|
|
return Err(format!("Upstream proxy CONNECT failed: {response_full}").into());
|
|
}
|
|
|
|
// Detect the buffer-drop race where the upstream returned the
|
|
// 200 response coalesced with destination bytes — those bytes
|
|
// would otherwise be silently discarded and the browser would
|
|
// see a TLS stream missing its first record.
|
|
let header_end_in_buffer = response_full.find("\r\n\r\n").map(|i| i + 4);
|
|
if let Some(end) = header_end_in_buffer {
|
|
if end < n {
|
|
log::warn!(
|
|
"Upstream CONNECT response coalesced {} byte(s) of payload — these would be dropped without forwarding",
|
|
n - end
|
|
);
|
|
}
|
|
}
|
|
|
|
log::info!(
|
|
"Upstream CONNECT to {}:{} via {}:{} accepted ({})",
|
|
target_host,
|
|
target_port,
|
|
proxy_host,
|
|
proxy_port,
|
|
status_line
|
|
);
|
|
|
|
Box::new(proxy_stream)
|
|
}
|
|
"socks4" | "socks5" => {
|
|
let socks_host = upstream.host_str().unwrap_or("127.0.0.1");
|
|
let socks_port = upstream.port().unwrap_or(1080);
|
|
let socks_addr = format!("{}:{}", socks_host, socks_port);
|
|
|
|
let (username, password) = upstream_userpass(&upstream);
|
|
let auth = (!username.is_empty()).then_some((username.as_str(), password.as_str()));
|
|
|
|
let stream = connect_via_socks(
|
|
&socks_addr,
|
|
target_host,
|
|
target_port,
|
|
scheme == "socks5",
|
|
auth,
|
|
)
|
|
.await?;
|
|
Box::new(stream)
|
|
}
|
|
"ss" | "shadowsocks" => {
|
|
// Shadowsocks: URL format is ss://method:password@host:port
|
|
// where "method" is the cipher (e.g. aes-256-gcm, chacha20-ietf-poly1305)
|
|
// and "password" is the SS server password.
|
|
let ss_host = upstream.host_str().unwrap_or("127.0.0.1");
|
|
let ss_port = upstream.port().unwrap_or(8388);
|
|
|
|
// The "username" field carries the cipher method
|
|
let method_str = urlencoding::decode(upstream.username())
|
|
.unwrap_or_default()
|
|
.to_string();
|
|
let password = urlencoding::decode(upstream.password().unwrap_or(""))
|
|
.unwrap_or_default()
|
|
.to_string();
|
|
|
|
if method_str.is_empty() || password.is_empty() {
|
|
return Err(
|
|
"Shadowsocks requires method and password (URL: ss://method:password@host:port)"
|
|
.into(),
|
|
);
|
|
}
|
|
|
|
let cipher = method_str.parse::<shadowsocks::crypto::CipherKind>().map_err(|_| {
|
|
format!("Unsupported Shadowsocks cipher: {method_str}. Use e.g. aes-256-gcm, chacha20-ietf-poly1305, aes-128-gcm")
|
|
})?;
|
|
|
|
let context =
|
|
shadowsocks::context::Context::new_shared(shadowsocks::config::ServerType::Local);
|
|
let svr_cfg = shadowsocks::config::ServerConfig::new(
|
|
shadowsocks::config::ServerAddr::from((ss_host.to_string(), ss_port)),
|
|
&password,
|
|
cipher,
|
|
)
|
|
.map_err(|e| format!("Invalid Shadowsocks config: {e}"))?;
|
|
|
|
let target_addr =
|
|
shadowsocks::relay::Address::DomainNameAddress(target_host.to_string(), target_port);
|
|
|
|
let stream = tokio::time::timeout(
|
|
UPSTREAM_DIAL_TIMEOUT,
|
|
shadowsocks::relay::tcprelay::proxy_stream::ProxyClientStream::connect(
|
|
context,
|
|
&svr_cfg,
|
|
target_addr,
|
|
),
|
|
)
|
|
.await
|
|
.map_err(|_| "Shadowsocks connection timed out".to_string())?
|
|
.map_err(|e| format!("Shadowsocks connection failed: {e}"))?;
|
|
|
|
Box::new(stream)
|
|
}
|
|
_ => {
|
|
return Err(format!("Unsupported upstream proxy scheme: {}", scheme).into());
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
Ok(target_stream)
|
|
}
|
|
|
|
/// Bidirectionally relay `client_stream` <-> `target_stream` until either side
|
|
/// closes, counting bytes for traffic stats and attributing them to `domain`.
|
|
/// The caller is responsible for having already sent any protocol-specific
|
|
/// success reply (HTTP `200` or SOCKS5 reply) before calling this.
|
|
pub(crate) async fn tunnel_streams(
|
|
client_stream: TcpStream,
|
|
target_stream: BoxedAsyncStream,
|
|
domain: String,
|
|
) {
|
|
// Wrap streams to count bytes transferred
|
|
let counting_client = CountingStream::new(client_stream);
|
|
let counting_target = CountingStream::new(target_stream);
|
|
|
|
// Get references for final stats
|
|
let client_read_counter = counting_client.bytes_read.clone();
|
|
let client_write_counter = counting_client.bytes_written.clone();
|
|
let target_read_counter = counting_target.bytes_read.clone();
|
|
let target_write_counter = counting_target.bytes_written.clone();
|
|
|
|
// Split streams for bidirectional copying
|
|
let (mut client_read, mut client_write) = tokio::io::split(counting_client);
|
|
let (mut target_read, mut target_write) = tokio::io::split(counting_target);
|
|
|
|
log::trace!("Starting bidirectional tunnel");
|
|
|
|
// Spawn two tasks to forward data in both directions
|
|
let client_to_target = tokio::spawn(async move {
|
|
let result = tokio::io::copy(&mut client_read, &mut target_write).await;
|
|
match result {
|
|
Ok(bytes) => {
|
|
log::trace!("Tunneled {bytes} bytes from client->target");
|
|
}
|
|
Err(e) => {
|
|
log::debug!("Error forwarding client->target: {e:?}");
|
|
}
|
|
}
|
|
});
|
|
|
|
let target_to_client = tokio::spawn(async move {
|
|
let result = tokio::io::copy(&mut target_read, &mut client_write).await;
|
|
match result {
|
|
Ok(bytes) => {
|
|
log::trace!("Tunneled {bytes} bytes from target->client");
|
|
}
|
|
Err(e) => {
|
|
log::debug!("Error forwarding target->client: {e:?}");
|
|
}
|
|
}
|
|
});
|
|
|
|
// Wait for either direction to finish (connection closed)
|
|
tokio::select! {
|
|
_ = client_to_target => {
|
|
log::trace!("Client->target tunnel closed");
|
|
}
|
|
_ = target_to_client => {
|
|
log::trace!("Target->client tunnel closed");
|
|
}
|
|
}
|
|
|
|
// Log final byte counts and update domain stats
|
|
let final_sent =
|
|
client_read_counter.load(Ordering::Relaxed) + target_write_counter.load(Ordering::Relaxed);
|
|
let final_recv =
|
|
target_read_counter.load(Ordering::Relaxed) + client_write_counter.load(Ordering::Relaxed);
|
|
log::trace!("Tunnel closed - sent: {final_sent} bytes, received: {final_recv} bytes");
|
|
|
|
// Update domain-specific byte counts now that tunnel is complete
|
|
if let Some(tracker) = get_traffic_tracker() {
|
|
tracker.update_domain_bytes(&domain, final_sent, final_recv);
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::io::Write;
|
|
|
|
/// Build an upstream URL with `urlencoding::encode`-d user/pass,
|
|
/// mirroring what `proxy_manager::build_proxy_url` actually emits
|
|
fn parse_encoded_upstream(scheme: &str, user: &str, pass: &str) -> Url {
|
|
let s = format!(
|
|
"{}://{}:{}@127.0.0.1:1080",
|
|
scheme,
|
|
urlencoding::encode(user),
|
|
urlencoding::encode(pass),
|
|
);
|
|
Url::parse(&s).unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn upstream_userpass_handles_plain_ascii() {
|
|
let u = parse_encoded_upstream("socks5", "alice", "secret123");
|
|
assert_eq!(upstream_userpass(&u), ("alice".into(), "secret123".into()));
|
|
}
|
|
|
|
#[test]
|
|
fn upstream_userpass_decodes_special_chars() {
|
|
// These characters all get percent-encoded by build_proxy_url before
|
|
// landing in the URL, and must be decoded back to the original literal
|
|
// before being handed off to the upstream
|
|
let cases = [
|
|
("alice", "p@ssw0rd"),
|
|
("alice", "p:assw0rd"),
|
|
("alice", "p ass word"),
|
|
("alice", "abc/d+e=f"),
|
|
("alice", "100%off!"),
|
|
("alice", "测试密码"),
|
|
("u@name", "v@lue"),
|
|
];
|
|
for (user, pass) in cases {
|
|
let u = parse_encoded_upstream("socks5", user, pass);
|
|
assert_eq!(
|
|
upstream_userpass(&u),
|
|
(user.to_string(), pass.to_string()),
|
|
"decode failed: user={user:?} pass={pass:?}"
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn upstream_userpass_empty_when_no_credentials() {
|
|
let u = Url::parse("socks5://127.0.0.1:1080").unwrap();
|
|
assert_eq!(upstream_userpass(&u), (String::new(), String::new()));
|
|
}
|
|
|
|
#[test]
|
|
fn upstream_userpass_handles_username_only() {
|
|
let s = format!("socks5://{}@127.0.0.1:1080", urlencoding::encode("u@name"));
|
|
let u = Url::parse(&s).unwrap();
|
|
assert_eq!(upstream_userpass(&u), ("u@name".into(), String::new()));
|
|
}
|
|
|
|
#[test]
|
|
fn test_blocklist_exact_match() {
|
|
let mut matcher = BlocklistMatcher::new();
|
|
let mut domains = HashSet::new();
|
|
domains.insert("example.com".to_string());
|
|
domains.insert("tracker.net".to_string());
|
|
matcher.domains = Arc::new(domains);
|
|
|
|
assert!(matcher.is_blocked("example.com"));
|
|
assert!(matcher.is_blocked("tracker.net"));
|
|
assert!(!matcher.is_blocked("safe.com"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_blocklist_subdomain_match() {
|
|
let mut matcher = BlocklistMatcher::new();
|
|
let mut domains = HashSet::new();
|
|
domains.insert("example.com".to_string());
|
|
matcher.domains = Arc::new(domains);
|
|
|
|
assert!(matcher.is_blocked("foo.example.com"));
|
|
assert!(matcher.is_blocked("bar.baz.example.com"));
|
|
assert!(matcher.is_blocked("a.b.c.example.com"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_blocklist_no_false_positives() {
|
|
let mut matcher = BlocklistMatcher::new();
|
|
let mut domains = HashSet::new();
|
|
domains.insert("example.com".to_string());
|
|
matcher.domains = Arc::new(domains);
|
|
|
|
// "notexample.com" should NOT match "example.com"
|
|
assert!(!matcher.is_blocked("notexample.com"));
|
|
assert!(!matcher.is_blocked("myexample.com"));
|
|
// But subdomain should
|
|
assert!(matcher.is_blocked("sub.example.com"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_blocklist_empty_blocks_nothing() {
|
|
let matcher = BlocklistMatcher::new();
|
|
assert!(!matcher.is_blocked("anything.com"));
|
|
assert!(!matcher.is_blocked("example.com"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_blocklist_case_insensitive() {
|
|
let mut matcher = BlocklistMatcher::new();
|
|
let mut domains = HashSet::new();
|
|
domains.insert("example.com".to_string());
|
|
matcher.domains = Arc::new(domains);
|
|
|
|
assert!(matcher.is_blocked("EXAMPLE.COM"));
|
|
assert!(matcher.is_blocked("Example.Com"));
|
|
assert!(matcher.is_blocked("FOO.EXAMPLE.COM"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_blocklist_from_file() {
|
|
let mut tmpfile = tempfile::NamedTempFile::new().unwrap();
|
|
writeln!(tmpfile, "# This is a comment").unwrap();
|
|
writeln!(tmpfile).unwrap();
|
|
writeln!(tmpfile, "tracker.example.com").unwrap();
|
|
writeln!(tmpfile, "ads.network.com").unwrap();
|
|
writeln!(tmpfile, "# Another comment").unwrap();
|
|
writeln!(tmpfile, "malware.site").unwrap();
|
|
tmpfile.flush().unwrap();
|
|
|
|
let matcher = BlocklistMatcher::from_file(tmpfile.path().to_str().unwrap()).unwrap();
|
|
|
|
assert!(matcher.is_blocked("tracker.example.com"));
|
|
assert!(matcher.is_blocked("ads.network.com"));
|
|
assert!(matcher.is_blocked("malware.site"));
|
|
assert!(matcher.is_blocked("sub.malware.site"));
|
|
assert!(!matcher.is_blocked("safe.com"));
|
|
// Comments and empty lines should be skipped: 3 domains loaded
|
|
assert_eq!(matcher.domains.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_blocklist_comments_skipped() {
|
|
let mut tmpfile = tempfile::NamedTempFile::new().unwrap();
|
|
writeln!(tmpfile, "# Title: HaGeZi's Light DNS Blocklist").unwrap();
|
|
writeln!(tmpfile, "# Description: test").unwrap();
|
|
writeln!(tmpfile, "# Version: 2026.0330.0928.01").unwrap();
|
|
writeln!(tmpfile).unwrap();
|
|
writeln!(tmpfile, "domain1.com").unwrap();
|
|
writeln!(tmpfile, "domain2.com").unwrap();
|
|
tmpfile.flush().unwrap();
|
|
|
|
let matcher = BlocklistMatcher::from_file(tmpfile.path().to_str().unwrap()).unwrap();
|
|
assert_eq!(matcher.domains.len(), 2);
|
|
assert!(matcher.is_blocked("domain1.com"));
|
|
assert!(matcher.is_blocked("domain2.com"));
|
|
}
|
|
}
|