mirror of
				https://github.com/edera-dev/krata.git
				synced 2025-11-03 23:29:39 +00:00 
			
		
		
		
	network: split out proxynat code
This commit is contained in:
		@ -24,7 +24,6 @@ use tokio::sync::mpsc::Sender;
 | 
				
			|||||||
pub enum NatKeyProtocol {
 | 
					pub enum NatKeyProtocol {
 | 
				
			||||||
    Tcp,
 | 
					    Tcp,
 | 
				
			||||||
    Udp,
 | 
					    Udp,
 | 
				
			||||||
    Ping,
 | 
					 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
 | 
					#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
 | 
				
			||||||
 | 
				
			|||||||
							
								
								
									
										49
									
								
								network/src/proxynat/mod.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										49
									
								
								network/src/proxynat/mod.rs
									
									
									
									
									
										Normal file
									
								
							@ -0,0 +1,49 @@
 | 
				
			|||||||
 | 
					mod udp;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					use async_trait::async_trait;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					use log::{debug, warn};
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					use tokio::sync::mpsc::channel;
 | 
				
			||||||
 | 
					use tokio::sync::mpsc::Sender;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					use crate::proxynat::udp::ProxyUdpHandler;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					use crate::nat::{NatHandler, NatHandlerFactory, NatKey, NatKeyProtocol};
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					pub struct ProxyNatHandlerFactory {}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					impl ProxyNatHandlerFactory {
 | 
				
			||||||
 | 
					    pub fn new() -> Self {
 | 
				
			||||||
 | 
					        Self {}
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					#[async_trait]
 | 
				
			||||||
 | 
					impl NatHandlerFactory for ProxyNatHandlerFactory {
 | 
				
			||||||
 | 
					    async fn nat(&self, key: NatKey, sender: Sender<Vec<u8>>) -> Option<Box<dyn NatHandler>> {
 | 
				
			||||||
 | 
					        debug!("creating proxy nat entry for key: {}", key);
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					        match key.protocol {
 | 
				
			||||||
 | 
					            NatKeyProtocol::Udp => {
 | 
				
			||||||
 | 
					                let (rx_sender, rx_receiver) = channel::<Vec<u8>>(4);
 | 
				
			||||||
 | 
					                let mut handler = ProxyUdpHandler::new(key, rx_sender);
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					                if let Err(error) = handler.spawn(rx_receiver, sender.clone()).await {
 | 
				
			||||||
 | 
					                    warn!("unable to spawn udp proxy handler: {}", error);
 | 
				
			||||||
 | 
					                    None
 | 
				
			||||||
 | 
					                } else {
 | 
				
			||||||
 | 
					                    Some(Box::new(handler))
 | 
				
			||||||
 | 
					                }
 | 
				
			||||||
 | 
					            }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					            _ => None,
 | 
				
			||||||
 | 
					        }
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					pub enum ProxyNatSelect {
 | 
				
			||||||
 | 
					    External(usize),
 | 
				
			||||||
 | 
					    Internal(Vec<u8>),
 | 
				
			||||||
 | 
					    Closed,
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
@ -11,49 +11,19 @@ use smoltcp::{
 | 
				
			|||||||
use tokio::{
 | 
					use tokio::{
 | 
				
			||||||
    io::{AsyncReadExt, AsyncWriteExt},
 | 
					    io::{AsyncReadExt, AsyncWriteExt},
 | 
				
			||||||
    select,
 | 
					    select,
 | 
				
			||||||
    sync::mpsc::channel,
 | 
					 | 
				
			||||||
};
 | 
					};
 | 
				
			||||||
use tokio::{sync::mpsc::Receiver, sync::mpsc::Sender};
 | 
					use tokio::{sync::mpsc::Receiver, sync::mpsc::Sender};
 | 
				
			||||||
use udp_stream::UdpStream;
 | 
					use udp_stream::UdpStream;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
use crate::nat::{NatHandler, NatHandlerFactory, NatKey, NatKeyProtocol};
 | 
					use crate::nat::{NatHandler, NatKey};
 | 
				
			||||||
 | 
					
 | 
				
			||||||
pub struct ProxyNatHandlerFactory {}
 | 
					use super::ProxyNatSelect;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
struct ProxyUdpHandler {
 | 
					pub struct ProxyUdpHandler {
 | 
				
			||||||
    key: NatKey,
 | 
					    key: NatKey,
 | 
				
			||||||
    rx_sender: Sender<Vec<u8>>,
 | 
					    rx_sender: Sender<Vec<u8>>,
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
impl ProxyNatHandlerFactory {
 | 
					 | 
				
			||||||
    pub fn new() -> Self {
 | 
					 | 
				
			||||||
        Self {}
 | 
					 | 
				
			||||||
    }
 | 
					 | 
				
			||||||
}
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
#[async_trait]
 | 
					 | 
				
			||||||
impl NatHandlerFactory for ProxyNatHandlerFactory {
 | 
					 | 
				
			||||||
    async fn nat(&self, key: NatKey, sender: Sender<Vec<u8>>) -> Option<Box<dyn NatHandler>> {
 | 
					 | 
				
			||||||
        debug!("creating proxy nat entry for key: {}", key);
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
        match key.protocol {
 | 
					 | 
				
			||||||
            NatKeyProtocol::Udp => {
 | 
					 | 
				
			||||||
                let (rx_sender, rx_receiver) = channel::<Vec<u8>>(4);
 | 
					 | 
				
			||||||
                let mut handler = ProxyUdpHandler { key, rx_sender };
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
                if let Err(error) = handler.spawn(rx_receiver, sender.clone()).await {
 | 
					 | 
				
			||||||
                    warn!("unable to spawn udp proxy handler: {}", error);
 | 
					 | 
				
			||||||
                    None
 | 
					 | 
				
			||||||
                } else {
 | 
					 | 
				
			||||||
                    Some(Box::new(handler))
 | 
					 | 
				
			||||||
                }
 | 
					 | 
				
			||||||
            }
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
            _ => None,
 | 
					 | 
				
			||||||
        }
 | 
					 | 
				
			||||||
    }
 | 
					 | 
				
			||||||
}
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
#[async_trait]
 | 
					#[async_trait]
 | 
				
			||||||
impl NatHandler for ProxyUdpHandler {
 | 
					impl NatHandler for ProxyUdpHandler {
 | 
				
			||||||
    async fn receive(&self, data: &[u8]) -> Result<()> {
 | 
					    async fn receive(&self, data: &[u8]) -> Result<()> {
 | 
				
			||||||
@ -62,14 +32,12 @@ impl NatHandler for ProxyUdpHandler {
 | 
				
			|||||||
    }
 | 
					    }
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
enum ProxySelect {
 | 
					 | 
				
			||||||
    External(usize),
 | 
					 | 
				
			||||||
    Internal(Vec<u8>),
 | 
					 | 
				
			||||||
    Closed,
 | 
					 | 
				
			||||||
}
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
impl ProxyUdpHandler {
 | 
					impl ProxyUdpHandler {
 | 
				
			||||||
    async fn spawn(
 | 
					    pub fn new(key: NatKey, rx_sender: Sender<Vec<u8>>) -> Self {
 | 
				
			||||||
 | 
					        ProxyUdpHandler { key, rx_sender }
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    pub async fn spawn(
 | 
				
			||||||
        &mut self,
 | 
					        &mut self,
 | 
				
			||||||
        rx_receiver: Receiver<Vec<u8>>,
 | 
					        rx_receiver: Receiver<Vec<u8>>,
 | 
				
			||||||
        tx_sender: Sender<Vec<u8>>,
 | 
					        tx_sender: Sender<Vec<u8>>,
 | 
				
			||||||
@ -109,15 +77,15 @@ impl ProxyUdpHandler {
 | 
				
			|||||||
        loop {
 | 
					        loop {
 | 
				
			||||||
            let selection = select! {
 | 
					            let selection = select! {
 | 
				
			||||||
                x = rx_receiver.recv() => if let Some(data) = x {
 | 
					                x = rx_receiver.recv() => if let Some(data) = x {
 | 
				
			||||||
                    ProxySelect::Internal(data)
 | 
					                    ProxyNatSelect::Internal(data)
 | 
				
			||||||
                } else {
 | 
					                } else {
 | 
				
			||||||
                    ProxySelect::Closed
 | 
					                    ProxyNatSelect::Closed
 | 
				
			||||||
                },
 | 
					                },
 | 
				
			||||||
                x = socket.read(&mut external_buffer) => ProxySelect::External(x?),
 | 
					                x = socket.read(&mut external_buffer) => ProxyNatSelect::External(x?),
 | 
				
			||||||
            };
 | 
					            };
 | 
				
			||||||
 | 
					
 | 
				
			||||||
            match selection {
 | 
					            match selection {
 | 
				
			||||||
                ProxySelect::External(size) => {
 | 
					                ProxyNatSelect::External(size) => {
 | 
				
			||||||
                    let data = &external_buffer[0..size];
 | 
					                    let data = &external_buffer[0..size];
 | 
				
			||||||
                    let packet = PacketBuilder::ethernet2(key.local_mac.0, key.client_mac.0);
 | 
					                    let packet = PacketBuilder::ethernet2(key.local_mac.0, key.client_mac.0);
 | 
				
			||||||
                    let packet = match (key.external_ip.addr, key.client_ip.addr) {
 | 
					                    let packet = match (key.external_ip.addr, key.client_ip.addr) {
 | 
				
			||||||
@ -138,8 +106,7 @@ impl ProxyUdpHandler {
 | 
				
			|||||||
                        debug!("failed to transmit udp packet: {}", error);
 | 
					                        debug!("failed to transmit udp packet: {}", error);
 | 
				
			||||||
                    }
 | 
					                    }
 | 
				
			||||||
                }
 | 
					                }
 | 
				
			||||||
                ProxySelect::Internal(data) => {
 | 
					                ProxyNatSelect::Internal(data) => {
 | 
				
			||||||
                    debug!("udp socket to handle data: {:?}", data);
 | 
					 | 
				
			||||||
                    let packet = SlicedPacket::from_ethernet(&data)?;
 | 
					                    let packet = SlicedPacket::from_ethernet(&data)?;
 | 
				
			||||||
                    let Some(ref net) = packet.net else {
 | 
					                    let Some(ref net) = packet.net else {
 | 
				
			||||||
                        continue;
 | 
					                        continue;
 | 
				
			||||||
@ -150,10 +117,9 @@ impl ProxyUdpHandler {
 | 
				
			|||||||
                    };
 | 
					                    };
 | 
				
			||||||
 | 
					
 | 
				
			||||||
                    let udp = UdpSlice::from_slice(ip.payload)?;
 | 
					                    let udp = UdpSlice::from_slice(ip.payload)?;
 | 
				
			||||||
                    debug!("UDP from internal: {:?}", udp.payload());
 | 
					 | 
				
			||||||
                    socket.write_all(udp.payload()).await?;
 | 
					                    socket.write_all(udp.payload()).await?;
 | 
				
			||||||
                }
 | 
					                }
 | 
				
			||||||
                ProxySelect::Closed => warn!("UDP socket closed"),
 | 
					                ProxyNatSelect::Closed => warn!("UDP socket closed"),
 | 
				
			||||||
            }
 | 
					            }
 | 
				
			||||||
        }
 | 
					        }
 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
		Reference in New Issue
	
	Block a user