ChrisIsSoFun

TCP UDP server forwarder

Dec 6th, 2021
957
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. use std::collections::HashMap;
  2. use std::fs::File;
  3. use std::io::{Read, stdin, stdout};
  4. use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
  5. use std::ops::Deref;
  6. use std::str::FromStr;
  7. use std::sync::Arc;
  8. use std::time::SystemTime;
  9. use async_std::io::{BufReader, Cursor};
  10. use async_std::net::{TcpListener, TcpStream, UdpSocket};
  11. use async_std::sync::RwLock;
  12. use async_std::task;
  13. use chrono::{DateTime, Local};
  14. use futures::executor::block_on;
  15. use futures::{AsyncReadExt, AsyncWriteExt, join};
  16. use futures::io::BufWriter;
  17.  
  18.  
  19. #[async_std::main]
  20. async fn main() {
  21.     //println!("Yo");
  22.     main_async().await;
  23. }
  24. async fn main_async() -> std::io::Result<()> {
  25.     log("MAIN", "Loading config...");
  26.     let config = parse_config();
  27.  
  28.     match config {
  29.         Ok(conf) => {
  30.  
  31.             log("MAIN", "Server is running...");
  32.             if conf.tcp && conf.udp {
  33.                 let tcp = host_tcp(conf.clone());
  34.                 let udp = host_udp(conf.clone());
  35.  
  36.                 join!(tcp, udp);
  37.             }
  38.             else if conf.tcp {
  39.                 host_tcp(conf).await;
  40.             }
  41.             else if conf.udp {
  42.                 host_udp(conf).await;
  43.             }
  44.         }
  45.         Err(e) => {
  46.             log("MAIN", e);
  47.         }
  48.     }
  49.  
  50.     Ok(())
  51. }
  52.  
  53. async fn host_tcp(data : ServerData) {
  54.     let t = TcpListener::bind(data.addr1).await.unwrap();
  55.  
  56.  
  57.     while let Ok(mut c) = t.accept().await {
  58.         let mut split = c.0.split();
  59.  
  60.         let mut read = split.0;
  61.         let mut write = split.1;
  62.  
  63.         let mut data_clone = data.clone();
  64.  
  65.         log("TCP CONNECTIONS", format!("Client of address {} connected", c.1.to_string()).as_str());
  66.  
  67.         task::spawn(async move {
  68.             let mut outer_client = TcpStream::connect(data_clone.addr2).await.unwrap();
  69.             let mut split2 = outer_client.split();
  70.             let mut read2 = split2.0;
  71.             let mut write2 = split2.1;
  72.  
  73.             task::spawn(async move {
  74.                 let mut b = [0u8; 8192];
  75.  
  76.                 while let Ok(size) = read2.read(&mut b).await {
  77.                     write.write_all(&b[0..size]).await;
  78.                     //b.clear();
  79.                 }
  80.  
  81.                 write.close().await;
  82.                 log("TCP CLOSES", format!("Internal client of address {} disconnected", c.1.to_string()).as_str());
  83.             });
  84.             task::spawn(async move {
  85.                 let mut b = [0u8; 8192];
  86.  
  87.                 while let Ok(size) = read.read(&mut b).await {
  88.                     write2.write_all(&b[0..size]).await;
  89.                     //b.clear();
  90.                 }
  91.  
  92.                 write2.close().await;
  93.                 log("TCP CLOSES", format!("External client of address {} disconnected", c.1.to_string()).as_str())
  94.             })
  95.         });
  96.     }
  97. }
  98.  
  99. async fn host_udp(data : ServerData) {
  100.     let address_map : Arc<RwLock<HashMap<SocketAddr, Arc<UdpSocket>>>> = Arc::new(RwLock::new(HashMap::new()));
  101.     let u = Arc::new((UdpSocket::bind(data.addr1).await).unwrap());
  102.  
  103.     let mut b = vec![0; 8192];
  104.  
  105.     while let Ok((n, peer)) = u.recv_from(&mut b).await {
  106.         let mut c = Vec::<u8>::with_capacity(n);
  107.         c.extend_from_slice(&b[0..n]);
  108.  
  109.         let outer_server = u.clone();
  110.         let adr_map = address_map.clone();
  111.  
  112.         {
  113.             let x = adr_map.read().await;
  114.             let op = x.get(&peer);
  115.             match op {
  116.                 Some(client) => {
  117.                     client.send(&c).await;
  118.                     continue;
  119.                 },
  120.                 _ => {}
  121.             }
  122.         }
  123.  
  124.         log("UDP CONNECTIONS", format!("New UDP IP and Port identified: {}", peer).as_str());
  125.  
  126.         let x = Arc::new(UdpSocket::bind(SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0)).await.unwrap());
  127.         x.connect(data.addr2).await;
  128.  
  129.         {
  130.             let mut w = adr_map.write().await;
  131.             w.insert(peer.clone(), x.clone());
  132.         }
  133.  
  134.         let pclone = peer.clone();
  135.         task::spawn(async move {
  136.             let mut new_msg = vec![0; 8192];
  137.             while let Ok(n) = x.recv(&mut new_msg).await {
  138.                 outer_server.send_to(&new_msg[0..n], pclone).await;
  139.             }
  140.         });
  141.     }
  142. }
  143.  
  144. pub struct ServerData {
  145.     pub addr1 : SocketAddr, // The current binding ip
  146.     pub addr2 : SocketAddr, // The proxy server ip
  147.     pub tcp : bool,
  148.     pub udp : bool
  149. }
  150.  
  151. impl Clone for ServerData {
  152.     fn clone(&self) -> Self {
  153.         return ServerData {
  154.             addr1: self.addr1.clone(),
  155.             addr2: self.addr2.clone(),
  156.             tcp: self.tcp,
  157.             udp: self.udp
  158.         }
  159.     }
  160.  
  161. }
  162.  
  163. pub fn parse_config() -> Result<ServerData, &'static str>{
  164.    let r_file = File::open("config.json");
  165.  
  166.    let mut file;
  167.  
  168.    match r_file {
  169.        Ok(f) => {
  170.            file = f;
  171.        }
  172.        Err(_) => {
  173.            return Err("The config file does not exist");
  174.        }
  175.    }
  176.  
  177.    let mut contents = String::new();
  178.  
  179.    file.read_to_string(&mut contents);
  180.  
  181.    let result = json::parse(contents.as_str());
  182.  
  183.    match result {
  184.        Ok(_) => {},
  185.        Err(_) => {
  186.            return Err("The config file is invalid");
  187.        }
  188.    }
  189.  
  190.    let mut result = result.unwrap();
  191.  
  192.    let ip1 = result["ip1"].as_str();
  193.    let ip2 = result["ip2"].as_str();
  194.    let port1 = result["port1"].as_u16();
  195.    let port2 = result["port2"].as_u16();
  196.    let tcp = result["tcp"].as_bool();
  197.    let udp = result["udp"].as_bool();
  198.  
  199.    if ip1.is_none() || ip2.is_none() || port1.is_none() || port2.is_none() || tcp.is_none() || udp.is_none() {
  200.        return Err("The file has invalid values");
  201.    }
  202.  
  203.    let tcp = tcp.unwrap();
  204.    let udp = udp.unwrap();
  205.  
  206.    let addr1 = SocketAddr::new(IpAddr::V4(Ipv4Addr::from_str(ip1.unwrap()).unwrap()), port1.unwrap());
  207.    let addr2 = SocketAddr::new(IpAddr::V6(Ipv6Addr::from_str(ip2.unwrap()).unwrap()), port2.unwrap());
  208.  
  209.    return Ok(ServerData {
  210.        addr1,
  211.        addr2,
  212.        tcp,
  213.        udp
  214.    })
  215. }
  216.  
  217. pub fn log(from : &str, msg : &str) {
  218.    let time = format_time(SystemTime::now(), "%d/%m/%Y %T");
  219.    println!("{}", format!("[{} at {}] {}", time, from, msg).as_str());
  220. }
  221. fn format_time(dt: SystemTime, format: &str) -> String {
  222.    let datetime : DateTime<Local> = dt.into();
  223.    format!("{}", datetime.format("%d/%m/%Y %T"))
  224. }
RAW Paste Data