lil udp support

src/client.rs

@@ -59,7 +59,7 @@ impl ServerCertVerifier for NoCertVerify {
     }
 }
 
-async fn open_connection(
+pub async fn open_connection(
     host: SocketAddr,
 ) -> Result<(Endpoint, Connection), Box<dyn Error>> {
     let mut client_crypto = rustls::ClientConfig::builder()
@@ -81,35 +81,7 @@ async fn open_connection(
     Ok((endpoint, conn))
 }
 
-async fn open_request(
+pub async fn close_connection(
-    conn: &mut Connection,
-    remote: SocketAddr,
-    password: &str
-) -> Result<(SendStream, RecvStream), Box<dyn Error>> {
-    let (mut send, recv) = conn
-        .open_bi()
-        .await?;
-
-    let request = format!(
-        "GET /index.html\r\nHost: {}\r\nAuthentication: {}\r\n\r\n",
-        remote,
-        bcrypt::hash(format!("{}{password}", conn.stable_id()), DEFAULT_COST)?
-    );
-    send.write_all(request.as_bytes()).await?;
-
-    Ok((send, recv))
-}
-
-async fn close_request(
-    mut send: SendStream,
-    mut recv: RecvStream
-) -> Result<(), Box<dyn Error>> {
-    send.finish()?;
-    recv.stop(0u32.into())?;
-    Ok(())
-}
-
-async fn close_connection(
     endpoint: Endpoint,
     conn: Connection,
 ) -> Result<(), Box<dyn Error>> {
@@ -117,3 +89,33 @@ async fn close_connection(
     endpoint.wait_idle().await;
     Ok(())
 }
+
+pub async fn open_request(
+    conn: &mut Connection,
+    remote: SocketAddr,
+    password: &str,
+    udp: bool
+) -> Result<(SendStream, RecvStream), Box<dyn Error>> {
+    let (mut send, recv) = conn
+        .open_bi()
+        .await?;
+
+    let request = format!(
+        "GET /index.html\r\nHost: {}\r\nAuthentication: {}\r\nUDP: {}\r\n\r\n",
+        remote,
+        bcrypt::hash(format!("{}{password}", conn.stable_id()), DEFAULT_COST)?,
+        if udp { "1" } else { "0" }
+    );
+    send.write_all(request.as_bytes()).await?;
+
+    Ok((send, recv))
+}
+
+pub async fn close_request(
+    mut send: SendStream,
+    mut recv: RecvStream
+) -> Result<(), Box<dyn Error>> {
+    send.finish()?;
+    recv.stop(0u32.into())?;
+    Ok(())
+}

src/main.rs

@@ -1,4 +1,5 @@
 use clap::Parser;
+use server::run_server;
 
 mod client;
 mod server;
@@ -25,7 +26,10 @@ async fn main() {
     let args = Args::parse();
 
     if let Some(host) = args.bind {
-        todo!()
+        run_server(
+            host.parse().expect("error parsing host"),
+            &args.password
+        ).await.expect("error running server");
     } else if let Some(host) = args.connect {
         todo!()
     } else {

src/server.rs

@@ -1,7 +1,7 @@
 use std::{error::Error, net::{IpAddr, SocketAddr}, str, sync::Arc};
 use quinn::crypto::rustls::QuicServerConfig;
 use rustls::pki_types::PrivatePkcs8KeyDer;
-use tokio::{io::{AsyncReadExt, AsyncWriteExt}, net::TcpStream};
+use tokio::{io::{AsyncReadExt, AsyncWriteExt}, net::{TcpStream, UdpSocket}};
 
 pub async fn run_server(host: SocketAddr, password: &str) -> Result<(), Box<dyn Error>> {
     let cert = rcgen::generate_simple_self_signed(vec![
@@ -96,6 +96,7 @@ async fn handle_request(
     let mut value = vec![];
     
     let mut remote = String::new();
+    let mut udp = false;
     
     for n in head_data {
         stack = match (stack, n) {
@@ -106,8 +107,10 @@ async fn handle_request(
                 } else if key == b"Authentication" {
                     let passhash = String::from_utf8(value.clone())?;
                     if !bcrypt::verify(password.clone(), &format!("{stable_id}{passhash}"))? {
-                        return Err("bad passhash error!!! not nice env!!".into())
+                        return Err("bad passhash error!!! not nice env!!".into());
                     }
+                } else if key == b"UDP" {
+                    udp = value == b"1";
                 }
                 
                 is_key = true;
@@ -140,37 +143,85 @@ async fn handle_request(
     }
     
     if stack != 4 {
-        return Err("bad request very bad".into())
+        return Err("bad request very bad".into());
     }
 
     let remote: SocketAddr = remote.parse()?;
 
     if is_local_address(&remote) {
-        return Err("backdoor attack!!! absolutely not good!!!!!!".into())
+        return Err("backdoor attack!!! absolutely not good!!!!!!".into());
     }
-    
-    let stream = TcpStream::connect(remote).await?;
-    let (mut remote_recv, mut remote_send) = stream.into_split();
 
-    remote_send.write_all(&mut body_data).await?;
+    if udp {
+        let local_addr: SocketAddr = if remote.is_ipv4() {
+            "0.0.0.0:0"
+        } else {
+            "[::]:0"
+        }
+        .parse()?;
+        
+        let stream = UdpSocket::bind(local_addr).await?;
+        stream.connect(&remote).await?;
+        let stream = Arc::new(stream);
 
-    tokio::spawn(async move  {
+        stream.send(&body_data).await?;
-        loop {
+
-            let mut buf = [0; 1024];
+        tokio::spawn({
-            let Ok(len) = remote_recv.read(&mut buf).await else { break; };
+            let stream = stream.clone();
-            if len == 0 { break; };
+            async move  {
-            let Ok(_) = send.write_all(&buf[..len]).await else { break; };
+                loop {
-        }
+                    let mut buf = [0; 1024];
-    });
+                    let Ok(len) = stream.recv(&mut buf).await else { break; };
+                    if len == 0 { break; };
+                    let Ok(_) = send.write_all(&buf[..len]).await else { break; };
+                }
+            }
+        });
+
+        tokio::spawn({
+            let stream = stream.clone();
+            async move  {
+                loop {
+                    let mut buf = [0; 1024];
+                    let Ok(Some(mut len)) = recv.read(&mut buf).await else { break; };
+                    if len == 0 { break; };
+                    let mut sent_all = 0;
+                    loop {
+                        let Ok(sent) = stream.send(&buf[sent_all..len]).await else { return; };
+                        if sent < len {
+                            len -= sent;
+                            sent_all += sent;
+                        } else {
+                            break;
+                        }
+                    }
+                }
+            }
+        });
+    } else {
+        let stream = TcpStream::connect(remote).await?;
+        let (mut remote_recv, mut remote_send) = stream.into_split();
+
+        remote_send.write_all(&body_data).await?;
+
+        tokio::spawn(async move  {
+            loop {
+                let mut buf = [0; 1024];
+                let Ok(len) = remote_recv.read(&mut buf).await else { break; };
+                if len == 0 { break; };
+                let Ok(_) = send.write_all(&buf[..len]).await else { break; };
+            }
+        });
     
-    tokio::spawn(async move {
+        tokio::spawn(async move {
-        loop {
+            loop {
-            let mut buf = [0; 1024];
+                let mut buf = [0; 1024];
-            let Ok(Some(len)) = recv.read(&mut buf).await else { break; };
+                let Ok(Some(len)) = recv.read(&mut buf).await else { break; };
-            if len == 0 { break; };
+                if len == 0 { break; };
-            let Ok(_) = remote_send.write_all(&buf[..len]).await else { break; };
+                let Ok(_) = remote_send.write_all(&buf[..len]).await else { break; };
-        }
+            }
-    });
+        });
+    }
 
     Ok(())
 }