untested server

README.md

@@ -14,10 +14,10 @@ On server:
 ```bash
 git clone https://git.meex.lol/MeexReay/unrknize
 cd unrknize/unrknize-server
-cargo run -- somepassword.example.com 127.0.0.1:443
-# certificates can be self-signed, or i would even say that they SHOULD be self-signed
-# use nginx or smth if you already have https port used
-# somepassword.example.com can be changed on anything else, for example just your ip address
+cargo run -- somepassword.example.com 127.0.0.1:9090 127.0.0.1:443
+# 127.0.0.1:9090 is the server where you want client to connect, but RKN detects it by protocol handshake
+# use nginx or smth if you already have https port (443) used or use any other port (less believable)
+# somepassword.example.com can be changed on anything else, for example just your ip address, its used for SNI
 ```
 
 On client:

unrknize-server/src/main.rs

@@ -1,33 +1,71 @@
-use std::env;
+use std::io::{Read, Write};
+use std::{env, thread};
 use std::error::Error;
-use std::net::{TcpListener, TcpStream};
-use std::sync::Arc;
+use std::net::{Shutdown, TcpListener, TcpStream};
+use std::sync::{Arc, Mutex};
 
 use rcgen::generate_simple_self_signed;
 use rustls::crypto::aws_lc_rs::sign::any_supported_type;
 use rustls::pki_types::pem::PemObject;
-use rustls::pki_types::{CertificateDer, PrivateKeyDer};
+use rustls::pki_types::PrivateKeyDer;
 use rustls::server::ResolvesServerCertUsingSni;
 use rustls::sign::CertifiedKey;
-use rustls::{ServerConfig, ServerConnection, StreamOwned};
+use rustls::{ServerConfig, ServerConnection, Stream};
 use threadpool::ThreadPool;
 
-fn accept_client(stream: TcpStream, config: Arc<ServerConfig>, sni_domain: String) -> Result<(), Box<dyn Error>> {
+fn accept_client(stream: &mut TcpStream, config: Arc<ServerConfig>, sni_domain: String, host: String) -> Result<(), Box<dyn Error>> {
+    let remote_stream = Arc::new(TcpStream::connect(host)?);
+
     let connection = ServerConnection::new(config)?;
-
-    let mut stream = StreamOwned::new(connection, stream);
-
-    while stream.conn.is_handshaking() {
-        stream.conn.complete_io(&mut stream.sock)?;
-    }
-
-    if stream.conn.server_name() != Some(&sni_domain) {
-        return Err("unexpected sni domain".into());
-    }
-
+    let connection = Arc::new(Mutex::new(connection));
     
+    {
+        let mut conn = connection.lock().unwrap();
+        let mut tls_stream = Stream::new(&mut *conn, stream);
+        while tls_stream.conn.is_handshaking() {
+            tls_stream.conn.complete_io(&mut tls_stream.sock)?;
+        }
 
-    Ok(())
+        if tls_stream.conn.server_name() != Some(&sni_domain) {
+            return Err("unexpected sni domain".into());
+        }
+    }
+
+    thread::spawn({
+        let mut stream = stream.try_clone()?;
+        let connection = connection.clone();
+        let remote_stream = remote_stream.clone();
+
+        move || {
+            loop {
+                let mut buffer = vec![0; 4096];
+                let n = (&*remote_stream).read(&mut buffer).unwrap();
+                if n != 0 {
+                    buffer.truncate(n);
+
+                    let mut conn = connection.lock().unwrap();
+                    let mut tls_stream = Stream::new(&mut *conn, &mut stream);
+                    tls_stream.write_all(&buffer[..n]).unwrap();
+                    tls_stream.flush().unwrap();
+                }
+            }
+        }
+    });
+
+    let mut stream = stream.try_clone()?;
+    let connection = connection.clone();
+    
+    loop {
+        let mut conn = connection.lock().unwrap();
+        let mut tls_stream = Stream::new(&mut *conn, &mut stream);
+
+        let mut buffer = vec![0; 4096];
+        let n = tls_stream.read(&mut buffer)?;
+        if n != 0 {
+            buffer.truncate(n);
+            (&*remote_stream).write(&buffer)?;
+        }
+    }
 }
 
 fn main() -> Result<(), Box<dyn Error>> {
@@ -41,6 +79,9 @@ fn main() -> Result<(), Box<dyn Error>> {
     let host = args
         .next()
         .expect("missing host argument");
+    let local_host = args
+        .next()
+        .expect("missing local host argument");
 
     let certified_key = generate_simple_self_signed(vec![sni_domain.to_string()]).unwrap();
     let certified_key = CertifiedKey::new(
@@ -59,19 +100,22 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     let threadpool = ThreadPool::new(10);
 
-    let listener = TcpListener::bind(host).unwrap();
+    let listener = TcpListener::bind(local_host).unwrap();
 
     for stream in listener.incoming() {
-        let stream = stream.expect("listener got broken");
+        let mut stream = stream.expect("listener got broken");
 
         threadpool.execute({
             let config = config.clone();
             let sni_domain = sni_domain.clone();
+            let host = host.clone();
 
             move || {
-                if let Err(e) = accept_client(stream, config, sni_domain) {
+                if let Err(e) = accept_client(&mut stream, config, sni_domain, host) {
                     println!("error connection: {e:?}");
                 }
+
+                let _ = stream.shutdown(Shutdown::Both);
             }
         });
     }