mirrors/krata
1
0
Fork 0
mirror of https://github.com/edera-dev/krata.git synced 2025-11-03 23:29:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

network: split nat into separate mods

Browse Source
This commit is contained in:
Alex Zenla
2024-02-13 18:01:52 +00:00
parent 521ee93349
commit 21707daa98
11 changed files with 250 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
network/src/backend.rs
View File

@ -1,15 +1,13 @@
use crate::autonet::NetworkMetadata;
use crate::chandev::ChannelDevice;
use crate::nat::NatRouter;
use crate::pkt::RecvPacket;
use crate::nat::Nat;
use crate::proxynat::ProxyNatHandlerFactory;
use crate::raw_socket::{AsyncRawSocketChannel, RawSocketHandle, RawSocketProtocol};
use crate::vbridge::{BridgeJoinHandle, VirtualBridge};
use anyhow::{anyhow, Result};
use bytes::BytesMut;
use etherparse::SlicedPacket;
use futures::TryStreamExt;
use log::{debug, info, trace, warn};
use log::{info, trace, warn};
use smoltcp::iface::{Config, Interface, SocketSet};
use smoltcp::phy::Medium;
use smoltcp::time::Instant;
@ -30,7 +28,6 @@ pub struct NetworkBackend {
enum NetworkStackSelect {
Receive(Option<BytesMut>),
Send(Option<BytesMut>),
Reclaim,
}
struct NetworkStack<'a> {
@ -39,7 +36,7 @@ struct NetworkStack<'a> {
udev: ChannelDevice,
interface: Interface,
sockets: SocketSet<'a>,
router: NatRouter,
nat: Nat,
bridge: BridgeJoinHandle,
}
@ -50,7 +47,6 @@ impl NetworkStack<'_> {
x = self.bridge.from_bridge_receiver.recv() => NetworkStackSelect::Send(x),
x = self.bridge.from_broadcast_receiver.recv() => NetworkStackSelect::Send(x.ok()),
x = self.tx.recv() => NetworkStackSelect::Send(x),
_ = self.router.process_reclaim() => NetworkStackSelect::Reclaim,
};
match what {
@ -59,16 +55,13 @@ impl NetworkStack<'_> {
trace!("failed to send guest packet to bridge: {}", error);
}
if let Ok(slice) = SlicedPacket::from_ethernet(&packet) {
let packet = RecvPacket::new(&packet, &slice)?;
if let Err(error) = self.router.process(&packet).await {
debug!("router failed to process packet: {}", error);
}
self.udev.rx = Some(packet.raw.into());
self.interface
.poll(Instant::now(), &mut self.udev, &mut self.sockets);
if let Err(error) = self.nat.receive_sender.try_send(packet.clone()) {
trace!("failed to send guest packet to nat: {}", error);
}
self.udev.rx = Some(packet);
self.interface
.poll(Instant::now(), &mut self.udev, &mut self.sockets);
}
NetworkStackSelect::Send(Some(packet)) => {
@ -80,8 +73,6 @@ impl NetworkStack<'_> {
NetworkStackSelect::Receive(None) | NetworkStackSelect::Send(None) => {
return Ok(false);
}
NetworkStackSelect::Reclaim => {}
}
Ok(true)
@ -134,7 +125,7 @@ impl NetworkBackend {
let (tx_sender, tx_receiver) = channel::<BytesMut>(TX_CHANNEL_BUFFER_LEN);
let mut udev = ChannelDevice::new(mtu, Medium::Ethernet, tx_sender.clone());
let mac = self.metadata.gateway.mac;
let nat = NatRouter::new(mtu, proxy, mac, addresses.clone(), tx_sender.clone());
let nat = Nat::new(mtu, proxy, mac, addresses.clone(), tx_sender.clone())?;
let hardware_addr = HardwareAddress::Ethernet(mac);
let config = Config::new(hardware_addr);
let mut iface = Interface::new(config, &mut udev, Instant::now());
@ -145,14 +136,14 @@ impl NetworkBackend {
});
let sockets = SocketSet::new(vec![]);
let handle = self.bridge.join(self.metadata.guest.mac).await?;
let kdev = AsyncRawSocketChannel::new(kdev)?;
let kdev = AsyncRawSocketChannel::new(mtu, kdev)?;
Ok(NetworkStack {
tx: tx_receiver,
kdev,
udev,
interface: iface,
sockets,
router: nat,
nat,
bridge: handle,
})
}

36
network/src/nat/handler.rs Normal file
View File

@ -0,0 +1,36 @@
use anyhow::Result;
use async_trait::async_trait;
use bytes::BytesMut;
use tokio::sync::mpsc::Sender;
use super::key::NatKey;
#[derive(Debug, Clone)]
pub struct NatHandlerContext {
pub mtu: usize,
pub key: NatKey,
pub transmit_sender: Sender<BytesMut>,
pub reclaim_sender: Sender<NatKey>,
}
impl NatHandlerContext {
pub fn try_transmit(&self, buffer: BytesMut) -> Result<()> {
self.transmit_sender.try_send(buffer)?;
Ok(())
}
pub async fn reclaim(&self) -> Result<()> {
self.reclaim_sender.try_send(self.key)?;
Ok(())
}
}
#[async_trait]
pub trait NatHandler: Send {
async fn receive(&self, packet: &[u8]) -> Result<bool>;
}
#[async_trait]
pub trait NatHandlerFactory: Send {
async fn nat(&self, context: NatHandlerContext) -> Option<Box<dyn NatHandler>>;
}

29
network/src/nat/key.rs Normal file
View File

@ -0,0 +1,29 @@
use std::fmt::Display;
use smoltcp::wire::{EthernetAddress, IpEndpoint};
#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
pub enum NatKeyProtocol {
Tcp,
Udp,
Icmp,
}
#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
pub struct NatKey {
pub protocol: NatKeyProtocol,
pub client_mac: EthernetAddress,
pub local_mac: EthernetAddress,
pub client_ip: IpEndpoint,
pub external_ip: IpEndpoint,
}
impl Display for NatKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{} -> {} {:?} {} -> {}",
self.client_mac, self.local_mac, self.protocol, self.client_ip, self.external_ip
)
}
}

42
network/src/nat/mod.rs Normal file
View File

@ -0,0 +1,42 @@
use anyhow::Result;
use tokio::sync::mpsc::Sender;
use self::handler::NatHandlerFactory;
use self::processor::NatProcessor;
use bytes::BytesMut;
use smoltcp::wire::EthernetAddress;
use smoltcp::wire::IpCidr;
use tokio::task::JoinHandle;
pub mod handler;
pub mod key;
pub mod processor;
pub mod table;
pub struct Nat {
pub receive_sender: Sender<BytesMut>,
task: JoinHandle<()>,
}
impl Nat {
pub fn new(
mtu: usize,
factory: Box<dyn NatHandlerFactory>,
local_mac: EthernetAddress,
local_cidrs: Vec<IpCidr>,
transmit_sender: Sender<BytesMut>,
) -> Result<Self> {
let (receive_sender, task) =
NatProcessor::launch(mtu, factory, local_mac, local_cidrs, transmit_sender)?;
Ok(Self {
receive_sender,
task,
})
}
}
impl Drop for Nat {
fn drop(&mut self) {
self.task.abort();
}
}

203
network/src/nat.rs → network/src/nat/processor.rs
View File

@ -1,7 +1,6 @@
use crate::pkt::RecvPacket;
use crate::pkt::RecvPacketIp;
use anyhow::Result;
use async_trait::async_trait;
use bytes::BytesMut;
use etherparse::Icmpv4Header;
use etherparse::Icmpv4Type;
@ -11,126 +10,110 @@ use etherparse::IpNumber;
use etherparse::IpPayloadSlice;
use etherparse::Ipv4Slice;
use etherparse::Ipv6Slice;
use etherparse::SlicedPacket;
use etherparse::TcpHeaderSlice;
use etherparse::UdpHeaderSlice;
use log::warn;
use log::{debug, trace};
use smoltcp::wire::EthernetAddress;
use smoltcp::wire::IpAddress;
use smoltcp::wire::IpCidr;
use smoltcp::wire::IpEndpoint;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::fmt::Display;
use tokio::select;
use tokio::sync::mpsc::channel;
use tokio::sync::mpsc::Receiver;
use tokio::sync::mpsc::Sender;
use tokio::task::JoinHandle;
use super::handler::NatHandler;
use super::handler::NatHandlerContext;
use super::handler::NatHandlerFactory;
use super::key::NatKey;
use super::key::NatKeyProtocol;
use super::table::NatTable;
const RECLAIM_CHANNEL_QUEUE_LEN: usize = 10;
const RECEIVE_CHANNEL_QUEUE_LEN: usize = 30;
#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
pub enum NatKeyProtocol {
Tcp,
Udp,
Icmp,
}
#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
pub struct NatKey {
pub protocol: NatKeyProtocol,
pub client_mac: EthernetAddress,
pub local_mac: EthernetAddress,
pub client_ip: IpEndpoint,
pub external_ip: IpEndpoint,
}
impl Display for NatKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{} -> {} {:?} {} -> {}",
self.client_mac, self.local_mac, self.protocol, self.client_ip, self.external_ip
)
}
}
#[derive(Debug, Clone)]
pub struct NatHandlerContext {
pub mtu: usize,
pub key: NatKey,
tx_sender: Sender<BytesMut>,
reclaim_sender: Sender<NatKey>,
}
impl NatHandlerContext {
pub fn try_send(&self, buffer: BytesMut) -> Result<()> {
self.tx_sender.try_send(buffer)?;
Ok(())
}
pub async fn reclaim(&self) -> Result<()> {
self.reclaim_sender.try_send(self.key)?;
Ok(())
}
}
#[async_trait]
pub trait NatHandler: Send {
async fn receive(&self, packet: &[u8]) -> Result<bool>;
}
#[async_trait]
pub trait NatHandlerFactory: Send {
async fn nat(&self, context: NatHandlerContext) -> Option<Box<dyn NatHandler>>;
}
pub struct NatTable {
inner: HashMap<NatKey, Box<dyn NatHandler>>,
}
impl Default for NatTable {
fn default() -> Self {
Self::new()
}
}
impl NatTable {
pub fn new() -> Self {
Self {
inner: HashMap::new(),
}
}
}
pub struct NatRouter {
pub struct NatProcessor {
mtu: usize,
local_mac: EthernetAddress,
local_cidrs: Vec<IpCidr>,
factory: Box<dyn NatHandlerFactory>,
table: NatTable,
tx_sender: Sender<BytesMut>,
factory: Box<dyn NatHandlerFactory>,
transmit_sender: Sender<BytesMut>,
reclaim_sender: Sender<NatKey>,
reclaim_receiver: Receiver<NatKey>,
receive_receiver: Receiver<BytesMut>,
}
impl NatRouter {
pub fn new(
enum NatProcessorSelect {
Reclaim(Option<NatKey>),
ReceivedPacket(Option<BytesMut>),
}
impl NatProcessor {
pub fn launch(
mtu: usize,
factory: Box<dyn NatHandlerFactory>,
local_mac: EthernetAddress,
local_cidrs: Vec<IpCidr>,
tx_sender: Sender<BytesMut>,
) -> Self {
transmit_sender: Sender<BytesMut>,
) -> Result<(Sender<BytesMut>, JoinHandle<()>)> {
let (reclaim_sender, reclaim_receiver) = channel(RECLAIM_CHANNEL_QUEUE_LEN);
Self {
let (receive_sender, receive_receiver) = channel(RECEIVE_CHANNEL_QUEUE_LEN);
let mut processor = Self {
mtu,
local_mac,
local_cidrs,
factory,
table: NatTable::new(),
tx_sender,
transmit_sender,
reclaim_sender,
receive_receiver,
reclaim_receiver,
};
let handle = tokio::task::spawn(async move {
if let Err(error) = processor.process().await {
warn!("nat processing failed: {}", error);
}
});
Ok((receive_sender, handle))
}
pub async fn process(&mut self) -> Result<()> {
loop {
let selection = select! {
x = self.reclaim_receiver.recv() => NatProcessorSelect::Reclaim(x),
x = self.receive_receiver.recv() => NatProcessorSelect::ReceivedPacket(x),
};
match selection {
NatProcessorSelect::Reclaim(Some(key)) => {
if self.table.inner.remove(&key).is_some() {
debug!("reclaimed nat key: {}", key);
}
}
NatProcessorSelect::ReceivedPacket(Some(packet)) => {
if let Ok(slice) = SlicedPacket::from_ethernet(&packet) {
let Ok(packet) = RecvPacket::new(&packet, &slice) else {
continue;
};
self.process_packet(&packet).await?;
}
}
NatProcessorSelect::ReceivedPacket(None) | NatProcessorSelect::Reclaim(None) => {
break
}
}
}
Ok(())
}
pub async fn process_reclaim(&mut self) -> Result<Option<NatKey>> {
@ -146,7 +129,7 @@ impl NatRouter {
})
}
pub async fn process<'a>(&mut self, packet: &RecvPacket<'a>) -> Result<()> {
pub async fn process_packet<'a>(&mut self, packet: &RecvPacket<'a>) -> Result<()> {
let Some(ether) = packet.ether else {
return Ok(());
};
@ -180,7 +163,7 @@ impl NatRouter {
let context = NatHandlerContext {
mtu: self.mtu,
key,
tx_sender: self.tx_sender.clone(),
transmit_sender: self.transmit_sender.clone(),
reclaim_sender: self.reclaim_sender.clone(),
};
let handler: Option<&mut Box<dyn NatHandler>> = match self.table.inner.entry(key) {
@ -251,28 +234,6 @@ impl NatRouter {
})
}
pub fn extract_key_tcp<'a>(
&mut self,
packet: &RecvPacket<'a>,
source_addr: IpAddress,
dest_addr: IpAddress,
payload: &IpPayloadSlice<'a>,
) -> Result<Option<NatKey>> {
let Some(ether) = packet.ether else {
return Ok(None);
};
let header = TcpHeaderSlice::from_slice(payload.payload)?;
let source = IpEndpoint::new(source_addr, header.source_port());
let dest = IpEndpoint::new(dest_addr, header.destination_port());
Ok(Some(NatKey {
protocol: NatKeyProtocol::Tcp,
client_mac: EthernetAddress(ether.source()),
local_mac: EthernetAddress(ether.destination()),
client_ip: source,
external_ip: dest,
}))
}
pub fn extract_key_udp<'a>(
&mut self,
packet: &RecvPacket<'a>,
@ -344,4 +305,26 @@ impl NatRouter {
external_ip: dest,
}))
}
pub fn extract_key_tcp<'a>(
&mut self,
packet: &RecvPacket<'a>,
source_addr: IpAddress,
dest_addr: IpAddress,
payload: &IpPayloadSlice<'a>,
) -> Result<Option<NatKey>> {
let Some(ether) = packet.ether else {
return Ok(None);
};
let header = TcpHeaderSlice::from_slice(payload.payload)?;
let source = IpEndpoint::new(source_addr, header.source_port());
let dest = IpEndpoint::new(dest_addr, header.destination_port());
Ok(Some(NatKey {
protocol: NatKeyProtocol::Tcp,
client_mac: EthernetAddress(ether.source()),
local_mac: EthernetAddress(ether.destination()),
client_ip: source,
external_ip: dest,
}))
}
}

21
network/src/nat/table.rs Normal file
View File

@ -0,0 +1,21 @@
use std::collections::HashMap;
use super::{handler::NatHandler, key::NatKey};
pub struct NatTable {
pub inner: HashMap<NatKey, Box<dyn NatHandler>>,
}
impl Default for NatTable {
fn default() -> Self {
Self::new()
}
}
impl NatTable {
pub fn new() -> Self {
Self {
inner: HashMap::new(),
}
}
}

6
network/src/proxynat/icmp.rs
View File

@ -19,7 +19,7 @@ use tokio::{
use crate::{
icmp::{IcmpClient, IcmpProtocol, IcmpReply},
nat::{NatHandler, NatHandlerContext},
nat::handler::{NatHandler, NatHandlerContext},
};
const ICMP_PING_TIMEOUT_SECS: u64 = 20;
@ -223,7 +223,7 @@ impl ProxyIcmpHandler {
let mut writer = buffer.writer();
packet.write(&mut writer, &payload)?;
let buffer = writer.into_inner();
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_transmit(buffer) {
debug!("failed to transmit icmp packet: {}", error);
}
Ok(())
@ -268,7 +268,7 @@ impl ProxyIcmpHandler {
let mut writer = buffer.writer();
packet.write(&mut writer, &payload)?;
let buffer = writer.into_inner();
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_transmit(buffer) {
debug!("failed to transmit icmp packet: {}", error);
}
Ok(())

4
network/src/proxynat/mod.rs
View File

@ -5,10 +5,10 @@ use log::warn;
use tokio::sync::mpsc::channel;
use crate::nat::NatHandlerContext;
use crate::proxynat::udp::ProxyUdpHandler;
use crate::nat::{NatHandler, NatHandlerFactory, NatKeyProtocol};
use crate::nat::handler::{NatHandler, NatHandlerContext, NatHandlerFactory};
use crate::nat::key::NatKeyProtocol;
use self::icmp::ProxyIcmpHandler;
use self::tcp::ProxyTcpHandler;

8
network/src/proxynat/tcp.rs
View File

@ -25,7 +25,7 @@ use tokio::{sync::mpsc::Receiver, sync::mpsc::Sender};
use crate::{
chandev::ChannelDevice,
nat::{NatHandler, NatHandlerContext},
nat::handler::{NatHandler, NatHandlerContext},
};
const TCP_BUFFER_SIZE: usize = 65535;
@ -216,7 +216,7 @@ impl ProxyTcpHandler {
};
buffer.extend_from_slice(&header.to_bytes());
buffer.extend_from_slice(&payload);
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_transmit(buffer) {
debug!("failed to transmit tcp packet: {}", error);
}
}
@ -389,7 +389,7 @@ impl ProxyTcpHandler {
};
buffer.extend_from_slice(&header.to_bytes());
buffer.extend_from_slice(&payload);
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_transmit(buffer) {
debug!("failed to transmit tcp packet: {}", error);
}
}
@ -449,7 +449,7 @@ impl ProxyTcpHandler {
};
buffer.extend_from_slice(&header.to_bytes());
buffer.extend_from_slice(&payload);
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_transmit(buffer) {
debug!("failed to transmit tcp packet: {}", error);
}
}

4
network/src/proxynat/udp.rs
View File

@ -16,7 +16,7 @@ use tokio::{
use tokio::{sync::mpsc::Receiver, sync::mpsc::Sender};
use udp_stream::UdpStream;
use crate::nat::{NatHandler, NatHandlerContext};
use crate::nat::handler::{NatHandler, NatHandlerContext};
const UDP_TIMEOUT_SECS: u64 = 60;
@ -111,7 +111,7 @@ impl ProxyUdpHandler {
let mut writer = buffer.writer();
packet.write(&mut writer, data)?;
let buffer = writer.into_inner();
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_transmit(buffer) {
debug!("failed to transmit udp packet: {}", error);
}
}

10
network/src/raw_socket.rs
View File

@ -204,10 +204,10 @@ enum AsyncRawSocketChannelSelect {
}
impl AsyncRawSocketChannel {
pub fn new(socket: RawSocketHandle) -> Result<AsyncRawSocketChannel> {
pub fn new(mtu: usize, socket: RawSocketHandle) -> Result<AsyncRawSocketChannel> {
let (transmit_sender, transmit_receiver) = channel(RAW_SOCKET_TRANSMIT_QUEUE_LEN);
let (receive_sender, receive_receiver) = channel(RAW_SOCKET_RECEIVE_QUEUE_LEN);
let task = AsyncRawSocketChannel::launch(socket, transmit_receiver, receive_sender)?;
let task = AsyncRawSocketChannel::launch(mtu, socket, transmit_receiver, receive_sender)?;
Ok(AsyncRawSocketChannel {
sender: transmit_sender,
receiver: receive_receiver,
@ -216,13 +216,14 @@ impl AsyncRawSocketChannel {
}
fn launch(
mtu: usize,
socket: RawSocketHandle,
transmit_receiver: Receiver<BytesMut>,
receive_sender: Sender<BytesMut>,
) -> Result<JoinHandle<()>> {
Ok(tokio::task::spawn(async move {
if let Err(error) =
AsyncRawSocketChannel::process(socket, transmit_receiver, receive_sender).await
AsyncRawSocketChannel::process(mtu, socket, transmit_receiver, receive_sender).await
{
warn!("failed to process raw socket: {}", error);
}
@ -230,6 +231,7 @@ impl AsyncRawSocketChannel {
}
async fn process(
mtu: usize,
socket: RawSocketHandle,
mut transmit_receiver: Receiver<BytesMut>,
receive_sender: Sender<BytesMut>,
@ -237,6 +239,7 @@ impl AsyncRawSocketChannel {
let socket = unsafe { std::net::UdpSocket::from_raw_fd(socket.into_raw_fd()) };
let socket = UdpSocket::from_std(socket)?;
let mut buffer = vec![0; mtu];
loop {
let selection = select! {
x = transmit_receiver.recv() => AsyncRawSocketChannelSelect::TransmitPacket(x),
@ -245,7 +248,6 @@ impl AsyncRawSocketChannel {
match selection {
AsyncRawSocketChannelSelect::Readable(_) => {
let mut buffer = vec![0; 1500];
match socket.try_recv(&mut buffer) {
Ok(len) => {
if len == 0 {