mirrors/krata
1
0
Fork 0
mirror of https://github.com/edera-dev/krata.git synced 2025-08-04 21:51:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

network: utilize bytes crate

Browse Source
This commit is contained in:
Alex Zenla
2024-02-12 17:01:47 +00:00
parent ddeab7610d
commit 341907a536
10 changed files with 86 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -5,6 +5,7 @@ use std::{
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use bytes::BytesMut;
use etherparse::{
IcmpEchoHeader, Icmpv4Header, Icmpv4Type, Icmpv6Header, Icmpv6Type, IpNumber, Ipv4Slice,
Ipv6Slice, NetSlice, PacketBuilder, SlicedPacket,
@ -25,7 +26,7 @@ const ICMP_PING_TIMEOUT_SECS: u64 = 20;
const ICMP_TIMEOUT_SECS: u64 = 30;
pub struct ProxyIcmpHandler {
rx_sender: Sender<Vec<u8>>,
rx_sender: Sender<BytesMut>,
}
#[async_trait]
@ -34,26 +35,26 @@ impl NatHandler for ProxyIcmpHandler {
if self.rx_sender.is_closed() {
Ok(true)
} else {
self.rx_sender.try_send(data.to_vec())?;
self.rx_sender.try_send(data.into())?;
Ok(true)
}
}
}
enum ProxyIcmpSelect {
Internal(Vec<u8>),
Internal(BytesMut),
Close,
}
impl ProxyIcmpHandler {
pub fn new(rx_sender: Sender<Vec<u8>>) -> Self {
pub fn new(rx_sender: Sender<BytesMut>) -> Self {
ProxyIcmpHandler { rx_sender }
}
pub async fn spawn(
&mut self,
context: NatHandlerContext,
rx_receiver: Receiver<Vec<u8>>,
rx_receiver: Receiver<BytesMut>,
) -> Result<()> {
let client = IcmpClient::new(match context.key.external_ip.addr {
IpAddress::Ipv4(_) => IcmpProtocol::Icmpv4,
@ -69,7 +70,7 @@ impl ProxyIcmpHandler {
async fn process(
client: IcmpClient,
mut rx_receiver: Receiver<Vec<u8>>,
mut rx_receiver: Receiver<BytesMut>,
context: NatHandlerContext,
) -> Result<()> {
loop {
@ -222,7 +223,7 @@ impl ProxyIcmpHandler {
let packet = packet.icmpv4_echo_reply(echo.id, echo.seq);
let mut buffer: Vec<u8> = Vec::new();
packet.write(&mut buffer, &payload)?;
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_send(buffer.as_slice().into()) {
debug!("failed to transmit icmp packet: {}", error);
}
Ok(())
@ -265,7 +266,7 @@ impl ProxyIcmpHandler {
let packet = packet.icmpv6_echo_reply(echo.id, echo.seq);
let mut buffer: Vec<u8> = Vec::new();
packet.write(&mut buffer, &payload)?;
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_send(buffer.as_slice().into()) {
debug!("failed to transmit icmp packet: {}", error);
}
Ok(())

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

@ -1,5 +1,6 @@
use async_trait::async_trait;
use bytes::BytesMut;
use log::warn;
use tokio::sync::mpsc::channel;
@ -16,6 +17,8 @@ mod icmp;
mod tcp;
mod udp;
const RX_CHANNEL_BOUND: usize = 300;
pub struct ProxyNatHandlerFactory {}
impl Default for ProxyNatHandlerFactory {
@ -35,7 +38,7 @@ 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::<Vec<u8>>(4);
let (rx_sender, rx_receiver) = channel::<BytesMut>(RX_CHANNEL_BOUND);
let mut handler = ProxyUdpHandler::new(rx_sender);
if let Err(error) = handler.spawn(context, rx_receiver).await {
@ -47,7 +50,7 @@ impl NatHandlerFactory for ProxyNatHandlerFactory {
}
NatKeyProtocol::Icmp => {
let (rx_sender, rx_receiver) = channel::<Vec<u8>>(300);
let (rx_sender, rx_receiver) = channel::<BytesMut>(RX_CHANNEL_BOUND);
let mut handler = ProxyIcmpHandler::new(rx_sender);
if let Err(error) = handler.spawn(context, rx_receiver).await {
@ -59,7 +62,7 @@ impl NatHandlerFactory for ProxyNatHandlerFactory {
}
NatKeyProtocol::Tcp => {
let (rx_sender, rx_receiver) = channel::<Vec<u8>>(300);
let (rx_sender, rx_receiver) = channel::<BytesMut>(RX_CHANNEL_BOUND);
let mut handler = ProxyTcpHandler::new(rx_sender);
if let Err(error) = handler.spawn(context, rx_receiver).await {

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

@ -5,6 +5,7 @@ use std::{
use anyhow::Result;
use async_trait::async_trait;
use bytes::BytesMut;
use etherparse::{EtherType, Ethernet2Header};
use log::{debug, warn};
use smoltcp::{
@ -32,7 +33,7 @@ const TCP_ACCEPT_TIMEOUT_SECS: u64 = 120;
const TCP_DANGLE_TIMEOUT_SECS: u64 = 10;
pub struct ProxyTcpHandler {
rx_sender: Sender<Vec<u8>>,
rx_sender: Sender<BytesMut>,
}
#[async_trait]
@ -41,7 +42,7 @@ impl NatHandler for ProxyTcpHandler {
if self.rx_sender.is_closed() {
Ok(false)
} else {
self.rx_sender.try_send(data.to_vec())?;
self.rx_sender.try_send(data.into())?;
Ok(true)
}
}
@ -49,8 +50,8 @@ impl NatHandler for ProxyTcpHandler {
#[derive(Debug)]
enum ProxyTcpAcceptSelect {
Internal(Vec<u8>),
TxIpPacket(Vec<u8>),
Internal(BytesMut),
TxIpPacket(BytesMut),
TimePassed,
DoNothing,
Close,
@ -60,8 +61,8 @@ enum ProxyTcpAcceptSelect {
enum ProxyTcpDataSelect {
ExternalRecv(usize),
ExternalSent(usize),
InternalRecv(Vec<u8>),
TxIpPacket(Vec<u8>),
InternalRecv(BytesMut),
TxIpPacket(BytesMut),
TimePassed,
DoNothing,
Close,
@ -69,20 +70,20 @@ enum ProxyTcpDataSelect {
#[derive(Debug)]
enum ProxyTcpFinishSelect {
InternalRecv(Vec<u8>),
TxIpPacket(Vec<u8>),
InternalRecv(BytesMut),
TxIpPacket(BytesMut),
Close,
}
impl ProxyTcpHandler {
pub fn new(rx_sender: Sender<Vec<u8>>) -> Self {
pub fn new(rx_sender: Sender<BytesMut>) -> Self {
ProxyTcpHandler { rx_sender }
}
pub async fn spawn(
&mut self,
context: NatHandlerContext,
rx_receiver: Receiver<Vec<u8>>,
rx_receiver: Receiver<BytesMut>,
) -> Result<()> {
let external_addr = match context.key.external_ip.addr {
IpAddress::Ipv4(addr) => {
@ -105,9 +106,9 @@ impl ProxyTcpHandler {
async fn process(
context: NatHandlerContext,
mut external_socket: TcpStream,
mut rx_receiver: Receiver<Vec<u8>>,
mut rx_receiver: Receiver<BytesMut>,
) -> Result<()> {
let (ip_sender, mut ip_receiver) = channel::<Vec<u8>>(300);
let (ip_sender, mut ip_receiver) = channel::<BytesMut>(300);
let mut external_buffer = vec![0u8; TCP_BUFFER_SIZE];
let mut device = ChannelDevice::new(
@ -197,7 +198,7 @@ impl ProxyTcpHandler {
ProxyTcpAcceptSelect::Internal(data) => {
let (_, payload) = Ethernet2Header::from_slice(&data)?;
device.rx = Some(payload.to_vec());
device.rx = Some(payload.into());
iface.poll(Instant::now(), &mut device, &mut sockets);
}
@ -213,7 +214,7 @@ impl ProxyTcpHandler {
};
header.write(&mut buffer)?;
buffer.extend_from_slice(&payload);
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_send(buffer.as_slice().into()) {
debug!("failed to transmit tcp packet: {}", error);
}
}
@ -370,7 +371,7 @@ impl ProxyTcpHandler {
ProxyTcpDataSelect::InternalRecv(data) => {
let (_, payload) = Ethernet2Header::from_slice(&data)?;
device.rx = Some(payload.to_vec());
device.rx = Some(payload.into());
iface.poll(Instant::now(), &mut device, &mut sockets);
}
@ -386,7 +387,7 @@ impl ProxyTcpHandler {
};
header.write(&mut buffer)?;
buffer.extend_from_slice(&payload);
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_send(buffer.as_slice().into()) {
debug!("failed to transmit tcp packet: {}", error);
}
}
@ -430,7 +431,7 @@ impl ProxyTcpHandler {
match selection {
ProxyTcpFinishSelect::InternalRecv(data) => {
let (_, payload) = Ethernet2Header::from_slice(&data)?;
device.rx = Some(payload.to_vec());
device.rx = Some(payload.into());
iface.poll(Instant::now(), &mut device, &mut sockets);
}
@ -446,7 +447,7 @@ impl ProxyTcpHandler {
};
header.write(&mut buffer)?;
buffer.extend_from_slice(&payload);
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_send(buffer.as_slice().into()) {
debug!("failed to transmit tcp packet: {}", error);
}
}

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

@ -5,6 +5,7 @@ use std::{
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use bytes::BytesMut;
use etherparse::{PacketBuilder, SlicedPacket, UdpSlice};
use log::{debug, warn};
use smoltcp::wire::IpAddress;
@ -20,7 +21,7 @@ use crate::nat::{NatHandler, NatHandlerContext};
const UDP_TIMEOUT_SECS: u64 = 60;
pub struct ProxyUdpHandler {
rx_sender: Sender<Vec<u8>>,
rx_sender: Sender<BytesMut>,
}
#[async_trait]
@ -29,7 +30,7 @@ impl NatHandler for ProxyUdpHandler {
if self.rx_sender.is_closed() {
Ok(true)
} else {
self.rx_sender.try_send(data.to_vec())?;
self.rx_sender.try_send(data.into())?;
Ok(true)
}
}
@ -37,19 +38,19 @@ impl NatHandler for ProxyUdpHandler {
enum ProxyUdpSelect {
External(usize),
Internal(Vec<u8>),
Internal(BytesMut),
Close,
}
impl ProxyUdpHandler {
pub fn new(rx_sender: Sender<Vec<u8>>) -> Self {
pub fn new(rx_sender: Sender<BytesMut>) -> Self {
ProxyUdpHandler { rx_sender }
}
pub async fn spawn(
&mut self,
context: NatHandlerContext,
rx_receiver: Receiver<Vec<u8>>,
rx_receiver: Receiver<BytesMut>,
) -> Result<()> {
let external_addr = match context.key.external_ip.addr {
IpAddress::Ipv4(addr) => {
@ -72,7 +73,7 @@ impl ProxyUdpHandler {
async fn process(
context: NatHandlerContext,
mut socket: UdpStream,
mut rx_receiver: Receiver<Vec<u8>>,
mut rx_receiver: Receiver<BytesMut>,
) -> Result<()> {
let mut external_buffer = vec![0u8; 2048];
@ -108,7 +109,7 @@ impl ProxyUdpHandler {
packet.udp(context.key.external_ip.port, context.key.client_ip.port);
let mut buffer: Vec<u8> = Vec::new();
packet.write(&mut buffer, data)?;
if let Err(error) = context.try_send(buffer) {
if let Err(error) = context.try_send(buffer.as_slice().into()) {
debug!("failed to transmit udp packet: {}", error);
}
}