mirror of
https://github.com/edera-dev/krata.git
synced 2025-08-04 13:41:31 +00:00
krata: restructure packages for cleanliness
This commit is contained in:
Alex Zenla
276
crates/network/src/proxynat/icmp.rs
Normal file
276
crates/network/src/proxynat/icmp.rs
Normal file
@ -0,0 +1,276 @@
|
||||
use std::{
|
||||
net::{IpAddr, Ipv4Addr, Ipv6Addr},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use bytes::{BufMut, BytesMut};
|
||||
use etherparse::{
|
||||
IcmpEchoHeader, Icmpv4Header, Icmpv4Type, Icmpv6Header, Icmpv6Type, IpNumber, Ipv4Slice,
|
||||
Ipv6Slice, NetSlice, PacketBuilder, SlicedPacket,
|
||||
};
|
||||
use log::{debug, trace, warn};
|
||||
use smoltcp::wire::IpAddress;
|
||||
use tokio::{
|
||||
select,
|
||||
sync::mpsc::{Receiver, Sender},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
icmp::{IcmpClient, IcmpProtocol, IcmpReply},
|
||||
nat::handler::{NatHandler, NatHandlerContext},
|
||||
};
|
||||
|
||||
const ICMP_PING_TIMEOUT_SECS: u64 = 20;
|
||||
const ICMP_TIMEOUT_SECS: u64 = 30;
|
||||
|
||||
pub struct ProxyIcmpHandler {
|
||||
rx_sender: Sender<BytesMut>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl NatHandler for ProxyIcmpHandler {
|
||||
async fn receive(&self, data: &[u8]) -> Result<bool> {
|
||||
if self.rx_sender.is_closed() {
|
||||
Ok(true)
|
||||
} else {
|
||||
self.rx_sender.try_send(data.into())?;
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum ProxyIcmpSelect {
|
||||
Internal(BytesMut),
|
||||
Close,
|
||||
}
|
||||
|
||||
impl ProxyIcmpHandler {
|
||||
pub fn new(rx_sender: Sender<BytesMut>) -> Self {
|
||||
ProxyIcmpHandler { rx_sender }
|
||||
}
|
||||
|
||||
pub async fn spawn(
|
||||
&mut self,
|
||||
context: NatHandlerContext,
|
||||
rx_receiver: Receiver<BytesMut>,
|
||||
) -> Result<()> {
|
||||
let client = IcmpClient::new(match context.key.external_ip.addr {
|
||||
IpAddress::Ipv4(_) => IcmpProtocol::Icmpv4,
|
||||
IpAddress::Ipv6(_) => IcmpProtocol::Icmpv6,
|
||||
})?;
|
||||
tokio::spawn(async move {
|
||||
if let Err(error) = ProxyIcmpHandler::process(client, rx_receiver, context).await {
|
||||
warn!("processing of icmp proxy failed: {}", error);
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process(
|
||||
client: IcmpClient,
|
||||
mut rx_receiver: Receiver<BytesMut>,
|
||||
context: NatHandlerContext,
|
||||
) -> Result<()> {
|
||||
loop {
|
||||
let deadline = tokio::time::sleep(Duration::from_secs(ICMP_TIMEOUT_SECS));
|
||||
let selection = select! {
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyIcmpSelect::Internal(data)
|
||||
} else {
|
||||
ProxyIcmpSelect::Close
|
||||
},
|
||||
_ = deadline => ProxyIcmpSelect::Close,
|
||||
};
|
||||
|
||||
match selection {
|
||||
ProxyIcmpSelect::Internal(data) => {
|
||||
let packet = SlicedPacket::from_ethernet(&data)?;
|
||||
let Some(ref net) = packet.net else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match net {
|
||||
NetSlice::Ipv4(ipv4) => {
|
||||
ProxyIcmpHandler::process_ipv4(&context, ipv4, &client).await?
|
||||
}
|
||||
|
||||
NetSlice::Ipv6(ipv6) => {
|
||||
ProxyIcmpHandler::process_ipv6(&context, ipv6, &client).await?
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ProxyIcmpSelect::Close => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
context.reclaim().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process_ipv4(
|
||||
context: &NatHandlerContext,
|
||||
ipv4: &Ipv4Slice<'_>,
|
||||
client: &IcmpClient,
|
||||
) -> Result<()> {
|
||||
if ipv4.header().protocol() != IpNumber::ICMP {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let (header, payload) = Icmpv4Header::from_slice(ipv4.payload().payload)?;
|
||||
if let Icmpv4Type::EchoRequest(echo) = header.icmp_type {
|
||||
let IpAddr::V4(external_ipv4) = context.key.external_ip.addr.into() else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let context = context.clone();
|
||||
let client = client.clone();
|
||||
let payload = payload.to_vec();
|
||||
tokio::task::spawn(async move {
|
||||
if let Err(error) = ProxyIcmpHandler::process_echo_ipv4(
|
||||
context,
|
||||
client,
|
||||
external_ipv4,
|
||||
echo,
|
||||
payload,
|
||||
)
|
||||
.await
|
||||
{
|
||||
trace!("icmp4 echo failed: {}", error);
|
||||
}
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process_ipv6(
|
||||
context: &NatHandlerContext,
|
||||
ipv6: &Ipv6Slice<'_>,
|
||||
client: &IcmpClient,
|
||||
) -> Result<()> {
|
||||
if ipv6.header().next_header() != IpNumber::IPV6_ICMP {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let (header, payload) = Icmpv6Header::from_slice(ipv6.payload().payload)?;
|
||||
if let Icmpv6Type::EchoRequest(echo) = header.icmp_type {
|
||||
let IpAddr::V6(external_ipv6) = context.key.external_ip.addr.into() else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let context = context.clone();
|
||||
let client = client.clone();
|
||||
let payload = payload.to_vec();
|
||||
tokio::task::spawn(async move {
|
||||
if let Err(error) = ProxyIcmpHandler::process_echo_ipv6(
|
||||
context,
|
||||
client,
|
||||
external_ipv6,
|
||||
echo,
|
||||
payload,
|
||||
)
|
||||
.await
|
||||
{
|
||||
trace!("icmp6 echo failed: {}", error);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process_echo_ipv4(
|
||||
context: NatHandlerContext,
|
||||
client: IcmpClient,
|
||||
external_ipv4: Ipv4Addr,
|
||||
echo: IcmpEchoHeader,
|
||||
payload: Vec<u8>,
|
||||
) -> Result<()> {
|
||||
let reply = client
|
||||
.ping4(
|
||||
external_ipv4,
|
||||
echo.id,
|
||||
echo.seq,
|
||||
&payload,
|
||||
Duration::from_secs(ICMP_PING_TIMEOUT_SECS),
|
||||
)
|
||||
.await?;
|
||||
let Some(IcmpReply::Icmpv4 {
|
||||
header: _,
|
||||
echo,
|
||||
payload,
|
||||
}) = reply
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let packet = PacketBuilder::ethernet2(context.key.local_mac.0, context.key.client_mac.0);
|
||||
let packet = match (context.key.external_ip.addr, context.key.client_ip.addr) {
|
||||
(IpAddress::Ipv4(external_addr), IpAddress::Ipv4(client_addr)) => {
|
||||
packet.ipv4(external_addr.0, client_addr.0, 20)
|
||||
}
|
||||
_ => {
|
||||
return Err(anyhow!("IP endpoint mismatch"));
|
||||
}
|
||||
};
|
||||
let packet = packet.icmpv4_echo_reply(echo.id, echo.seq);
|
||||
let buffer = BytesMut::with_capacity(packet.size(payload.len()));
|
||||
let mut writer = buffer.writer();
|
||||
packet.write(&mut writer, &payload)?;
|
||||
let buffer = writer.into_inner();
|
||||
if let Err(error) = context.try_transmit(buffer) {
|
||||
debug!("failed to transmit icmp packet: {}", error);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process_echo_ipv6(
|
||||
context: NatHandlerContext,
|
||||
client: IcmpClient,
|
||||
external_ipv6: Ipv6Addr,
|
||||
echo: IcmpEchoHeader,
|
||||
payload: Vec<u8>,
|
||||
) -> Result<()> {
|
||||
let reply = client
|
||||
.ping6(
|
||||
external_ipv6,
|
||||
echo.id,
|
||||
echo.seq,
|
||||
&payload,
|
||||
Duration::from_secs(ICMP_PING_TIMEOUT_SECS),
|
||||
)
|
||||
.await?;
|
||||
let Some(IcmpReply::Icmpv6 {
|
||||
header: _,
|
||||
echo,
|
||||
payload,
|
||||
}) = reply
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let packet = PacketBuilder::ethernet2(context.key.local_mac.0, context.key.client_mac.0);
|
||||
let packet = match (context.key.external_ip.addr, context.key.client_ip.addr) {
|
||||
(IpAddress::Ipv6(external_addr), IpAddress::Ipv6(client_addr)) => {
|
||||
packet.ipv6(external_addr.0, client_addr.0, 20)
|
||||
}
|
||||
_ => {
|
||||
return Err(anyhow!("IP endpoint mismatch"));
|
||||
}
|
||||
};
|
||||
let packet = packet.icmpv6_echo_reply(echo.id, echo.seq);
|
||||
let buffer = BytesMut::with_capacity(packet.size(payload.len()));
|
||||
let mut writer = buffer.writer();
|
||||
packet.write(&mut writer, &payload)?;
|
||||
let buffer = writer.into_inner();
|
||||
if let Err(error) = context.try_transmit(buffer) {
|
||||
debug!("failed to transmit icmp packet: {}", error);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
77
crates/network/src/proxynat/mod.rs
Normal file
77
crates/network/src/proxynat/mod.rs
Normal file
@ -0,0 +1,77 @@
|
||||
use async_trait::async_trait;
|
||||
|
||||
use bytes::BytesMut;
|
||||
use log::warn;
|
||||
|
||||
use tokio::sync::mpsc::channel;
|
||||
|
||||
use crate::proxynat::udp::ProxyUdpHandler;
|
||||
|
||||
use crate::nat::handler::{NatHandler, NatHandlerContext, NatHandlerFactory};
|
||||
use crate::nat::key::NatKeyProtocol;
|
||||
|
||||
use self::icmp::ProxyIcmpHandler;
|
||||
use self::tcp::ProxyTcpHandler;
|
||||
|
||||
mod icmp;
|
||||
mod tcp;
|
||||
mod udp;
|
||||
|
||||
const RX_CHANNEL_QUEUE_LEN: usize = 1000;
|
||||
|
||||
pub struct ProxyNatHandlerFactory {}
|
||||
|
||||
impl Default for ProxyNatHandlerFactory {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProxyNatHandlerFactory {
|
||||
pub fn new() -> Self {
|
||||
Self {}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl NatHandlerFactory for ProxyNatHandlerFactory {
|
||||
async fn nat(&self, context: NatHandlerContext) -> Option<Box<dyn NatHandler>> {
|
||||
match context.key.protocol {
|
||||
NatKeyProtocol::Udp => {
|
||||
let (rx_sender, rx_receiver) = channel::<BytesMut>(RX_CHANNEL_QUEUE_LEN);
|
||||
let mut handler = ProxyUdpHandler::new(rx_sender);
|
||||
|
||||
if let Err(error) = handler.spawn(context, rx_receiver).await {
|
||||
warn!("unable to spawn udp proxy handler: {}", error);
|
||||
None
|
||||
} else {
|
||||
Some(Box::new(handler))
|
||||
}
|
||||
}
|
||||
|
||||
NatKeyProtocol::Icmp => {
|
||||
let (rx_sender, rx_receiver) = channel::<BytesMut>(RX_CHANNEL_QUEUE_LEN);
|
||||
let mut handler = ProxyIcmpHandler::new(rx_sender);
|
||||
|
||||
if let Err(error) = handler.spawn(context, rx_receiver).await {
|
||||
warn!("unable to spawn icmp proxy handler: {}", error);
|
||||
None
|
||||
} else {
|
||||
Some(Box::new(handler))
|
||||
}
|
||||
}
|
||||
|
||||
NatKeyProtocol::Tcp => {
|
||||
let (rx_sender, rx_receiver) = channel::<BytesMut>(RX_CHANNEL_QUEUE_LEN);
|
||||
let mut handler = ProxyTcpHandler::new(rx_sender);
|
||||
|
||||
if let Err(error) = handler.spawn(context, rx_receiver).await {
|
||||
warn!("unable to spawn tcp proxy handler: {}", error);
|
||||
None
|
||||
} else {
|
||||
Some(Box::new(handler))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
466
crates/network/src/proxynat/tcp.rs
Normal file
466
crates/network/src/proxynat/tcp.rs
Normal file
@ -0,0 +1,466 @@
|
||||
use std::{
|
||||
net::{IpAddr, SocketAddr},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use anyhow::Result;
|
||||
use async_trait::async_trait;
|
||||
use bytes::BytesMut;
|
||||
use etherparse::{EtherType, Ethernet2Header};
|
||||
use log::{debug, warn};
|
||||
use smoltcp::{
|
||||
iface::{Config, Interface, SocketSet, SocketStorage},
|
||||
phy::Medium,
|
||||
socket::tcp::{self, SocketBuffer, State},
|
||||
time::Instant,
|
||||
wire::{HardwareAddress, IpAddress, IpCidr},
|
||||
};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpStream,
|
||||
select,
|
||||
sync::mpsc::channel,
|
||||
};
|
||||
use tokio::{sync::mpsc::Receiver, sync::mpsc::Sender};
|
||||
|
||||
use crate::{
|
||||
chandev::ChannelDevice,
|
||||
nat::handler::{NatHandler, NatHandlerContext},
|
||||
};
|
||||
|
||||
const TCP_BUFFER_SIZE: usize = 65535;
|
||||
const TCP_IP_BUFFER_QUEUE_LEN: usize = 3000;
|
||||
const TCP_ACCEPT_TIMEOUT_SECS: u64 = 120;
|
||||
const TCP_DANGLE_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
pub struct ProxyTcpHandler {
|
||||
rx_sender: Sender<BytesMut>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl NatHandler for ProxyTcpHandler {
|
||||
async fn receive(&self, data: &[u8]) -> Result<bool> {
|
||||
if self.rx_sender.is_closed() {
|
||||
Ok(false)
|
||||
} else {
|
||||
self.rx_sender.try_send(data.into())?;
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum ProxyTcpAcceptSelect {
|
||||
Internal(BytesMut),
|
||||
TxIpPacket(BytesMut),
|
||||
TimePassed,
|
||||
DoNothing,
|
||||
Close,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum ProxyTcpDataSelect {
|
||||
ExternalRecv(usize),
|
||||
ExternalSent(usize),
|
||||
InternalRecv(BytesMut),
|
||||
TxIpPacket(BytesMut),
|
||||
TimePassed,
|
||||
DoNothing,
|
||||
Close,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum ProxyTcpFinishSelect {
|
||||
InternalRecv(BytesMut),
|
||||
TxIpPacket(BytesMut),
|
||||
Close,
|
||||
}
|
||||
|
||||
impl ProxyTcpHandler {
|
||||
pub fn new(rx_sender: Sender<BytesMut>) -> Self {
|
||||
ProxyTcpHandler { rx_sender }
|
||||
}
|
||||
|
||||
pub async fn spawn(
|
||||
&mut self,
|
||||
context: NatHandlerContext,
|
||||
rx_receiver: Receiver<BytesMut>,
|
||||
) -> Result<()> {
|
||||
let external_addr = match context.key.external_ip.addr {
|
||||
IpAddress::Ipv4(addr) => {
|
||||
SocketAddr::new(IpAddr::V4(addr.0.into()), context.key.external_ip.port)
|
||||
}
|
||||
IpAddress::Ipv6(addr) => {
|
||||
SocketAddr::new(IpAddr::V6(addr.0.into()), context.key.external_ip.port)
|
||||
}
|
||||
};
|
||||
|
||||
let socket = TcpStream::connect(external_addr).await?;
|
||||
tokio::spawn(async move {
|
||||
if let Err(error) = ProxyTcpHandler::process(context, socket, rx_receiver).await {
|
||||
warn!("processing of tcp proxy failed: {}", error);
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process(
|
||||
context: NatHandlerContext,
|
||||
mut external_socket: TcpStream,
|
||||
mut rx_receiver: Receiver<BytesMut>,
|
||||
) -> Result<()> {
|
||||
let (ip_sender, mut ip_receiver) = channel::<BytesMut>(TCP_IP_BUFFER_QUEUE_LEN);
|
||||
let mut external_buffer = vec![0u8; TCP_BUFFER_SIZE];
|
||||
|
||||
let mut device = ChannelDevice::new(
|
||||
context.mtu - Ethernet2Header::LEN,
|
||||
Medium::Ip,
|
||||
ip_sender.clone(),
|
||||
);
|
||||
let config = Config::new(HardwareAddress::Ip);
|
||||
|
||||
let tcp_rx_buffer = SocketBuffer::new(vec![0; TCP_BUFFER_SIZE]);
|
||||
let tcp_tx_buffer = SocketBuffer::new(vec![0; TCP_BUFFER_SIZE]);
|
||||
let internal_socket = tcp::Socket::new(tcp_rx_buffer, tcp_tx_buffer);
|
||||
let mut iface = Interface::new(config, &mut device, Instant::now());
|
||||
|
||||
iface.update_ip_addrs(|addrs| {
|
||||
let _ = addrs.push(IpCidr::new(context.key.external_ip.addr, 0));
|
||||
});
|
||||
|
||||
let mut sockets = SocketSet::new([SocketStorage::EMPTY]);
|
||||
let internal_socket_handle = sockets.add(internal_socket);
|
||||
let (mut external_r, mut external_w) = external_socket.split();
|
||||
|
||||
{
|
||||
let socket = sockets.get_mut::<tcp::Socket>(internal_socket_handle);
|
||||
socket.connect(
|
||||
iface.context(),
|
||||
context.key.client_ip,
|
||||
context.key.external_ip,
|
||||
)?;
|
||||
}
|
||||
|
||||
iface.poll(Instant::now(), &mut device, &mut sockets);
|
||||
|
||||
let mut sleeper: Option<tokio::time::Sleep> = None;
|
||||
loop {
|
||||
let socket = sockets.get_mut::<tcp::Socket>(internal_socket_handle);
|
||||
if socket.is_active() && socket.state() != State::SynSent {
|
||||
break;
|
||||
}
|
||||
|
||||
if socket.state() == State::Closed {
|
||||
break;
|
||||
}
|
||||
|
||||
let deadline = tokio::time::sleep(Duration::from_secs(TCP_ACCEPT_TIMEOUT_SECS));
|
||||
let selection = if let Some(sleep) = sleeper.take() {
|
||||
select! {
|
||||
biased;
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpAcceptSelect::Internal(data)
|
||||
} else {
|
||||
ProxyTcpAcceptSelect::Close
|
||||
},
|
||||
x = ip_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpAcceptSelect::TxIpPacket(data)
|
||||
} else {
|
||||
ProxyTcpAcceptSelect::Close
|
||||
},
|
||||
_ = sleep => ProxyTcpAcceptSelect::TimePassed,
|
||||
_ = deadline => ProxyTcpAcceptSelect::Close,
|
||||
}
|
||||
} else {
|
||||
select! {
|
||||
biased;
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpAcceptSelect::Internal(data)
|
||||
} else {
|
||||
ProxyTcpAcceptSelect::Close
|
||||
},
|
||||
x = ip_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpAcceptSelect::TxIpPacket(data)
|
||||
} else {
|
||||
ProxyTcpAcceptSelect::Close
|
||||
},
|
||||
_ = std::future::ready(()) => ProxyTcpAcceptSelect::DoNothing,
|
||||
_ = deadline => ProxyTcpAcceptSelect::Close,
|
||||
}
|
||||
};
|
||||
match selection {
|
||||
ProxyTcpAcceptSelect::TimePassed => {
|
||||
iface.poll(Instant::now(), &mut device, &mut sockets);
|
||||
}
|
||||
|
||||
ProxyTcpAcceptSelect::DoNothing => {
|
||||
sleeper = Some(tokio::time::sleep(Duration::from_micros(100)));
|
||||
}
|
||||
|
||||
ProxyTcpAcceptSelect::Internal(data) => {
|
||||
let (_, payload) = Ethernet2Header::from_slice(&data)?;
|
||||
device.rx = Some(payload.into());
|
||||
iface.poll(Instant::now(), &mut device, &mut sockets);
|
||||
}
|
||||
|
||||
ProxyTcpAcceptSelect::TxIpPacket(payload) => {
|
||||
let mut buffer = BytesMut::with_capacity(Ethernet2Header::LEN + payload.len());
|
||||
let header = Ethernet2Header {
|
||||
source: context.key.local_mac.0,
|
||||
destination: context.key.client_mac.0,
|
||||
ether_type: match context.key.external_ip.addr {
|
||||
IpAddress::Ipv4(_) => EtherType::IPV4,
|
||||
IpAddress::Ipv6(_) => EtherType::IPV6,
|
||||
},
|
||||
};
|
||||
buffer.extend_from_slice(&header.to_bytes());
|
||||
buffer.extend_from_slice(&payload);
|
||||
if let Err(error) = context.try_transmit(buffer) {
|
||||
debug!("failed to transmit tcp packet: {}", error);
|
||||
}
|
||||
}
|
||||
|
||||
ProxyTcpAcceptSelect::Close => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let accepted = if sockets
|
||||
.get_mut::<tcp::Socket>(internal_socket_handle)
|
||||
.is_active()
|
||||
{
|
||||
true
|
||||
} else {
|
||||
debug!("failed to accept tcp connection from client");
|
||||
false
|
||||
};
|
||||
|
||||
let mut already_shutdown = false;
|
||||
let mut sleeper: Option<tokio::time::Sleep> = None;
|
||||
loop {
|
||||
if !accepted {
|
||||
break;
|
||||
}
|
||||
|
||||
let socket = sockets.get_mut::<tcp::Socket>(internal_socket_handle);
|
||||
|
||||
match socket.state() {
|
||||
State::Closed
|
||||
| State::Listen
|
||||
| State::Closing
|
||||
| State::LastAck
|
||||
| State::TimeWait => {
|
||||
break;
|
||||
}
|
||||
State::FinWait1
|
||||
| State::SynSent
|
||||
| State::CloseWait
|
||||
| State::FinWait2
|
||||
| State::SynReceived
|
||||
| State::Established => {}
|
||||
}
|
||||
|
||||
let bytes_to_client = if socket.can_send() {
|
||||
socket.send_capacity() - socket.send_queue()
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
let (bytes_to_external, do_shutdown) = if socket.may_recv() {
|
||||
if let Ok(data) = socket.peek(TCP_BUFFER_SIZE) {
|
||||
if data.is_empty() {
|
||||
(None, false)
|
||||
} else {
|
||||
(Some(data), false)
|
||||
}
|
||||
} else {
|
||||
(None, false)
|
||||
}
|
||||
} else if !already_shutdown && matches!(socket.state(), State::CloseWait) {
|
||||
(None, true)
|
||||
} else {
|
||||
(None, false)
|
||||
};
|
||||
let selection = if let Some(sleep) = sleeper.take() {
|
||||
if !do_shutdown {
|
||||
select! {
|
||||
biased;
|
||||
x = ip_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpDataSelect::TxIpPacket(data)
|
||||
} else {
|
||||
ProxyTcpDataSelect::Close
|
||||
},
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpDataSelect::InternalRecv(data)
|
||||
} else {
|
||||
ProxyTcpDataSelect::Close
|
||||
},
|
||||
x = external_w.write(bytes_to_external.unwrap_or(b"")), if bytes_to_external.is_some() => ProxyTcpDataSelect::ExternalSent(x?),
|
||||
x = external_r.read(&mut external_buffer[..bytes_to_client]), if bytes_to_client > 0 => ProxyTcpDataSelect::ExternalRecv(x?),
|
||||
_ = sleep => ProxyTcpDataSelect::TimePassed,
|
||||
}
|
||||
} else {
|
||||
select! {
|
||||
biased;
|
||||
x = ip_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpDataSelect::TxIpPacket(data)
|
||||
} else {
|
||||
ProxyTcpDataSelect::Close
|
||||
},
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpDataSelect::InternalRecv(data)
|
||||
} else {
|
||||
ProxyTcpDataSelect::Close
|
||||
},
|
||||
_ = external_w.shutdown() => ProxyTcpDataSelect::ExternalSent(0),
|
||||
x = external_r.read(&mut external_buffer[..bytes_to_client]), if bytes_to_client > 0 => ProxyTcpDataSelect::ExternalRecv(x?),
|
||||
_ = sleep => ProxyTcpDataSelect::TimePassed,
|
||||
}
|
||||
}
|
||||
} else if !do_shutdown {
|
||||
select! {
|
||||
biased;
|
||||
x = ip_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpDataSelect::TxIpPacket(data)
|
||||
} else {
|
||||
ProxyTcpDataSelect::Close
|
||||
},
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpDataSelect::InternalRecv(data)
|
||||
} else {
|
||||
ProxyTcpDataSelect::Close
|
||||
},
|
||||
x = external_w.write(bytes_to_external.unwrap_or(b"")), if bytes_to_external.is_some() => ProxyTcpDataSelect::ExternalSent(x?),
|
||||
x = external_r.read(&mut external_buffer[..bytes_to_client]), if bytes_to_client > 0 => ProxyTcpDataSelect::ExternalRecv(x?),
|
||||
_ = std::future::ready(()) => ProxyTcpDataSelect::DoNothing,
|
||||
}
|
||||
} else {
|
||||
select! {
|
||||
biased;
|
||||
x = ip_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpDataSelect::TxIpPacket(data)
|
||||
} else {
|
||||
ProxyTcpDataSelect::Close
|
||||
},
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpDataSelect::InternalRecv(data)
|
||||
} else {
|
||||
ProxyTcpDataSelect::Close
|
||||
},
|
||||
_ = external_w.shutdown() => ProxyTcpDataSelect::ExternalSent(0),
|
||||
x = external_r.read(&mut external_buffer[..bytes_to_client]), if bytes_to_client > 0 => ProxyTcpDataSelect::ExternalRecv(x?),
|
||||
_ = std::future::ready(()) => ProxyTcpDataSelect::DoNothing,
|
||||
}
|
||||
};
|
||||
match selection {
|
||||
ProxyTcpDataSelect::ExternalRecv(size) => {
|
||||
if size == 0 {
|
||||
socket.close();
|
||||
} else {
|
||||
socket.send_slice(&external_buffer[..size])?;
|
||||
}
|
||||
}
|
||||
|
||||
ProxyTcpDataSelect::ExternalSent(size) => {
|
||||
if size == 0 {
|
||||
already_shutdown = true;
|
||||
} else {
|
||||
socket.recv(|_| (size, ()))?;
|
||||
}
|
||||
}
|
||||
|
||||
ProxyTcpDataSelect::InternalRecv(data) => {
|
||||
let (_, payload) = Ethernet2Header::from_slice(&data)?;
|
||||
device.rx = Some(payload.into());
|
||||
iface.poll(Instant::now(), &mut device, &mut sockets);
|
||||
}
|
||||
|
||||
ProxyTcpDataSelect::TxIpPacket(payload) => {
|
||||
let mut buffer = BytesMut::with_capacity(Ethernet2Header::LEN + payload.len());
|
||||
let header = Ethernet2Header {
|
||||
source: context.key.local_mac.0,
|
||||
destination: context.key.client_mac.0,
|
||||
ether_type: match context.key.external_ip.addr {
|
||||
IpAddress::Ipv4(_) => EtherType::IPV4,
|
||||
IpAddress::Ipv6(_) => EtherType::IPV6,
|
||||
},
|
||||
};
|
||||
buffer.extend_from_slice(&header.to_bytes());
|
||||
buffer.extend_from_slice(&payload);
|
||||
if let Err(error) = context.try_transmit(buffer) {
|
||||
debug!("failed to transmit tcp packet: {}", error);
|
||||
}
|
||||
}
|
||||
|
||||
ProxyTcpDataSelect::TimePassed => {
|
||||
iface.poll(Instant::now(), &mut device, &mut sockets);
|
||||
}
|
||||
|
||||
ProxyTcpDataSelect::DoNothing => {
|
||||
sleeper = Some(tokio::time::sleep(Duration::from_micros(100)));
|
||||
}
|
||||
|
||||
ProxyTcpDataSelect::Close => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = external_socket.shutdown().await;
|
||||
drop(external_socket);
|
||||
|
||||
loop {
|
||||
let deadline = tokio::time::sleep(Duration::from_secs(TCP_DANGLE_TIMEOUT_SECS));
|
||||
tokio::pin!(deadline);
|
||||
|
||||
let selection = select! {
|
||||
biased;
|
||||
x = ip_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpFinishSelect::TxIpPacket(data)
|
||||
} else {
|
||||
ProxyTcpFinishSelect::Close
|
||||
},
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyTcpFinishSelect::InternalRecv(data)
|
||||
} else {
|
||||
ProxyTcpFinishSelect::Close
|
||||
},
|
||||
_ = deadline => ProxyTcpFinishSelect::Close,
|
||||
};
|
||||
|
||||
match selection {
|
||||
ProxyTcpFinishSelect::InternalRecv(data) => {
|
||||
let (_, payload) = Ethernet2Header::from_slice(&data)?;
|
||||
device.rx = Some(payload.into());
|
||||
iface.poll(Instant::now(), &mut device, &mut sockets);
|
||||
}
|
||||
|
||||
ProxyTcpFinishSelect::TxIpPacket(payload) => {
|
||||
let mut buffer = BytesMut::with_capacity(Ethernet2Header::LEN + payload.len());
|
||||
let header = Ethernet2Header {
|
||||
source: context.key.local_mac.0,
|
||||
destination: context.key.client_mac.0,
|
||||
ether_type: match context.key.external_ip.addr {
|
||||
IpAddress::Ipv4(_) => EtherType::IPV4,
|
||||
IpAddress::Ipv6(_) => EtherType::IPV6,
|
||||
},
|
||||
};
|
||||
buffer.extend_from_slice(&header.to_bytes());
|
||||
buffer.extend_from_slice(&payload);
|
||||
if let Err(error) = context.try_transmit(buffer) {
|
||||
debug!("failed to transmit tcp packet: {}", error);
|
||||
}
|
||||
}
|
||||
|
||||
ProxyTcpFinishSelect::Close => {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
context.reclaim().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
142
crates/network/src/proxynat/udp.rs
Normal file
142
crates/network/src/proxynat/udp.rs
Normal file
@ -0,0 +1,142 @@
|
||||
use std::{
|
||||
net::{IpAddr, SocketAddr},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use bytes::{BufMut, BytesMut};
|
||||
use etherparse::{PacketBuilder, SlicedPacket, UdpSlice};
|
||||
use log::{debug, warn};
|
||||
use smoltcp::wire::IpAddress;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
select,
|
||||
};
|
||||
use tokio::{sync::mpsc::Receiver, sync::mpsc::Sender};
|
||||
use udp_stream::UdpStream;
|
||||
|
||||
use crate::nat::handler::{NatHandler, NatHandlerContext};
|
||||
|
||||
const UDP_TIMEOUT_SECS: u64 = 60;
|
||||
|
||||
pub struct ProxyUdpHandler {
|
||||
rx_sender: Sender<BytesMut>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl NatHandler for ProxyUdpHandler {
|
||||
async fn receive(&self, data: &[u8]) -> Result<bool> {
|
||||
if self.rx_sender.is_closed() {
|
||||
Ok(true)
|
||||
} else {
|
||||
self.rx_sender.try_send(data.into())?;
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum ProxyUdpSelect {
|
||||
External(usize),
|
||||
Internal(BytesMut),
|
||||
Close,
|
||||
}
|
||||
|
||||
impl ProxyUdpHandler {
|
||||
pub fn new(rx_sender: Sender<BytesMut>) -> Self {
|
||||
ProxyUdpHandler { rx_sender }
|
||||
}
|
||||
|
||||
pub async fn spawn(
|
||||
&mut self,
|
||||
context: NatHandlerContext,
|
||||
rx_receiver: Receiver<BytesMut>,
|
||||
) -> Result<()> {
|
||||
let external_addr = match context.key.external_ip.addr {
|
||||
IpAddress::Ipv4(addr) => {
|
||||
SocketAddr::new(IpAddr::V4(addr.0.into()), context.key.external_ip.port)
|
||||
}
|
||||
IpAddress::Ipv6(addr) => {
|
||||
SocketAddr::new(IpAddr::V6(addr.0.into()), context.key.external_ip.port)
|
||||
}
|
||||
};
|
||||
|
||||
let socket = UdpStream::connect(external_addr).await?;
|
||||
tokio::spawn(async move {
|
||||
if let Err(error) = ProxyUdpHandler::process(context, socket, rx_receiver).await {
|
||||
warn!("processing of udp proxy failed: {}", error);
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn process(
|
||||
context: NatHandlerContext,
|
||||
mut socket: UdpStream,
|
||||
mut rx_receiver: Receiver<BytesMut>,
|
||||
) -> Result<()> {
|
||||
let mut external_buffer = vec![0u8; 2048];
|
||||
|
||||
loop {
|
||||
let deadline = tokio::time::sleep(Duration::from_secs(UDP_TIMEOUT_SECS));
|
||||
let selection = select! {
|
||||
x = rx_receiver.recv() => if let Some(data) = x {
|
||||
ProxyUdpSelect::Internal(data)
|
||||
} else {
|
||||
ProxyUdpSelect::Close
|
||||
},
|
||||
x = socket.read(&mut external_buffer) => ProxyUdpSelect::External(x?),
|
||||
_ = deadline => ProxyUdpSelect::Close,
|
||||
};
|
||||
|
||||
match selection {
|
||||
ProxyUdpSelect::External(size) => {
|
||||
let data = &external_buffer[0..size];
|
||||
let packet =
|
||||
PacketBuilder::ethernet2(context.key.local_mac.0, context.key.client_mac.0);
|
||||
let packet = match (context.key.external_ip.addr, context.key.client_ip.addr) {
|
||||
(IpAddress::Ipv4(external_addr), IpAddress::Ipv4(client_addr)) => {
|
||||
packet.ipv4(external_addr.0, client_addr.0, 20)
|
||||
}
|
||||
(IpAddress::Ipv6(external_addr), IpAddress::Ipv6(client_addr)) => {
|
||||
packet.ipv6(external_addr.0, client_addr.0, 20)
|
||||
}
|
||||
_ => {
|
||||
return Err(anyhow!("IP endpoint mismatch"));
|
||||
}
|
||||
};
|
||||
let packet =
|
||||
packet.udp(context.key.external_ip.port, context.key.client_ip.port);
|
||||
let buffer = BytesMut::with_capacity(packet.size(data.len()));
|
||||
let mut writer = buffer.writer();
|
||||
packet.write(&mut writer, data)?;
|
||||
let buffer = writer.into_inner();
|
||||
if let Err(error) = context.try_transmit(buffer) {
|
||||
debug!("failed to transmit udp packet: {}", error);
|
||||
}
|
||||
}
|
||||
ProxyUdpSelect::Internal(data) => {
|
||||
let packet = SlicedPacket::from_ethernet(&data)?;
|
||||
let Some(ref net) = packet.net else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let Some(ip) = net.ip_payload_ref() else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let udp = UdpSlice::from_slice(ip.payload)?;
|
||||
socket.write_all(udp.payload()).await?;
|
||||
}
|
||||
ProxyUdpSelect::Close => {
|
||||
drop(socket);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
context.reclaim().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user