mirror of https://github.com/ospab/ostp.git
106 lines
3.6 KiB
Rust
106 lines
3.6 KiB
Rust
use std::sync::{Arc, RwLock};
|
|
use tokio::net::TcpStream;
|
|
use anyhow::Result;
|
|
use crate::outbound::{OutboundConfig, connect_target};
|
|
use crate::dns::DnsServer;
|
|
|
|
#[derive(Clone)]
|
|
pub struct Router {
|
|
pub outbound_cfg: Arc<RwLock<Option<OutboundConfig>>>,
|
|
pub dns_server: Arc<DnsServer>,
|
|
pub debug: bool,
|
|
}
|
|
|
|
impl Router {
|
|
pub fn new(outbound_cfg: Option<OutboundConfig>, dns_server: Arc<DnsServer>, debug: bool) -> Self {
|
|
Self {
|
|
outbound_cfg: Arc::new(RwLock::new(outbound_cfg)),
|
|
dns_server,
|
|
debug,
|
|
}
|
|
}
|
|
|
|
/// TCP Target Routing
|
|
pub async fn route_tcp(&self, target: &str) -> Result<TcpStream> {
|
|
let cfg = {
|
|
let lock = self.outbound_cfg.read().unwrap();
|
|
lock.clone()
|
|
};
|
|
connect_target(target, cfg.as_ref(), self.debug).await
|
|
}
|
|
|
|
/// UDP Target Routing
|
|
pub async fn route_udp(&self, target: &str, server_udp: std::sync::Arc<tokio::net::UdpSocket>) -> Result<crate::outbound::UdpProxySocket> {
|
|
let cfg = {
|
|
let lock = self.outbound_cfg.read().unwrap();
|
|
lock.clone()
|
|
};
|
|
crate::outbound::connect_udp_target(target, cfg.as_ref(), self.debug, server_udp).await
|
|
}
|
|
|
|
/// Establish a UDP session router that can dynamically route packets
|
|
pub async fn route_udp_associate(&self, server_udp: std::sync::Arc<tokio::net::UdpSocket>) -> UdpSessionRouter {
|
|
let cfg = {
|
|
let lock = self.outbound_cfg.read().unwrap();
|
|
lock.clone()
|
|
};
|
|
|
|
let mut proxy = None;
|
|
if let Some(ref c) = cfg {
|
|
if c.enabled && c.protocol == "socks5" {
|
|
let proxy_addr = format!("{}:{}", c.address, c.port);
|
|
if let Ok(p) = crate::outbound::connect_udp_via_socks5(&proxy_addr, server_udp.clone()).await {
|
|
proxy = Some(Arc::new(p));
|
|
} else if self.debug {
|
|
tracing::warn!("Failed to establish SOCKS5 UDP Associate");
|
|
}
|
|
}
|
|
}
|
|
|
|
UdpSessionRouter {
|
|
direct: server_udp,
|
|
proxy,
|
|
cfg,
|
|
debug: self.debug,
|
|
}
|
|
}
|
|
|
|
/// Unified DNS Routing and Resolution (AdBlock / Custom Domains / DoH)
|
|
pub async fn route_dns(&self, client_ip: std::net::IpAddr, payload: &[u8]) -> Option<Vec<u8>> {
|
|
self.dns_server.resolve(payload, client_ip).await
|
|
}
|
|
}
|
|
|
|
pub struct UdpSessionRouter {
|
|
direct: Arc<tokio::net::UdpSocket>,
|
|
proxy: Option<Arc<crate::outbound::UdpProxySocket>>,
|
|
cfg: Option<OutboundConfig>,
|
|
debug: bool,
|
|
}
|
|
|
|
impl UdpSessionRouter {
|
|
pub async fn send_to(&self, data: &[u8], target: &str) -> Result<usize> {
|
|
if let Some(cfg) = &self.cfg {
|
|
if cfg.enabled {
|
|
let action = crate::outbound::select_outbound_action(target, "udp", cfg, self.debug).await;
|
|
if action == crate::outbound::OutboundAction::Block {
|
|
return Err(anyhow::anyhow!("blocked by outbound udp rule: {}", target));
|
|
}
|
|
if action == crate::outbound::OutboundAction::Proxy {
|
|
if let Some(p) = &self.proxy {
|
|
return p.send_to(data, target).await;
|
|
} else {
|
|
return Err(anyhow::anyhow!("UDP Proxy not available for proxy action (possibly proxy doesn't support UDP)"));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
self.direct.send_to(data, target).await.map_err(Into::into)
|
|
}
|
|
|
|
pub fn get_proxy_sock(&self) -> Option<Arc<crate::outbound::UdpProxySocket>> {
|
|
self.proxy.clone()
|
|
}
|
|
}
|
|
|