TLS transport and games framework on client and server

1 files changed, 115 insertions, 21 deletions

diff --git a/src/bin/client/connection.rs b/src/bin/client/connection.rs
index b07bf01..c5d506d 100644
--- a/src/bin/client/connection.rs
+++ b/src/bin/client/connection.rs
@@ -1,46 +1,76 @@
-use std::io::{BufRead, BufReader, Write};
+use crate::database;
+use native_tls::TlsConnector;
+use rusqlite::Connection;
+use std::io::{Read, Write};
 use std::net::TcpStream;
 use std::sync::mpsc;
+use std::sync::Arc;
+use std::sync::Mutex;
 use std::thread;
 use std::time::{Duration, Instant};
-use rusqlite::Connection;
-use crate::database;
 
-pub fn connect(addr: &str, conn: &Connection) -> Result<(TcpStream, Option<String>), Box<dyn std::error::Error>> {
+static UNSAFE_MODE: Mutex<bool> = Mutex::new(false);
+
+pub fn set_unsafe_mode(unsafe_mode: bool) {
+    let mut mode = UNSAFE_MODE.lock().unwrap();
+    *mode = unsafe_mode;
+}
+
+pub fn connect(
+    addr: &str,
+    conn: &Connection,
+) -> Result<(TlsStream, Option<String>), Box<dyn std::error::Error>> {
     let (tx, rx) = mpsc::channel();
     let addr_str = addr.to_string();
-    thread::spawn(move || {
-        match TcpStream::connect(&addr_str) {
-            Ok(s) => tx.send(Ok(s)).unwrap(),
-            Err(e) => tx.send(Err(e)).unwrap(),
-        }
+    let server_name = addr_str
+        .split(':')
+        .next()
+        .unwrap_or("127.0.0.1")
+        .to_string();
+    thread::spawn(move || match TcpStream::connect(&addr_str) {
+        Ok(s) => tx.send(Ok(s)).unwrap(),
+        Err(e) => tx.send(Err(e)).unwrap(),
     });
     match rx.recv_timeout(Duration::from_secs(3)) {
         Ok(Ok(s)) => {
-            let mut ping_writer = s.try_clone().unwrap();
-            let mut ping_reader = BufReader::new(s.try_clone().unwrap());
+            let unsafe_mode = *UNSAFE_MODE.lock().unwrap();
+
+            let connector = if unsafe_mode {
+                TlsConnector::builder()
+                    .danger_accept_invalid_certs(true)
+                    .danger_accept_invalid_hostnames(true)
+                    .build()?
+            } else {
+                TlsConnector::builder().build()?
+            };
+
+            let tls_stream = connector.connect(&server_name, s)?;
+            let mut tls = TlsStream::new(tls_stream);
+
             let start = Instant::now();
-            writeln!(ping_writer, "PING").unwrap();
+            tls.write_line("PING")?;
             let mut response = String::new();
-            ping_reader.read_line(&mut response).unwrap();
+            tls.read_line(&mut response)?;
             let elapsed = start.elapsed();
-            println!("\n\rConnected to {} with ping delay: {:?}", addr, elapsed);
-            
-            // Check for credentials and auto-auth
+            println!(
+                "\n\rConnected to {} with ping delay: {:?} (TLS)",
+                addr, elapsed
+            );
+
             if let Some((username, password)) = database::get_credentials(conn, addr) {
                 let auth_msg = format!("AUTH {} {}", username, password);
-                writeln!(ping_writer, "{}", auth_msg).unwrap();
+                tls.write_line(&auth_msg)?;
                 let mut auth_response = String::new();
-                ping_reader.read_line(&mut auth_response).unwrap();
+                tls.read_line(&mut auth_response)?;
                 if auth_response.trim().starts_with("OK") {
                     println!("Authenticated as {}", username);
-                    Ok((s, Some(username)))
+                    Ok((tls, Some(username)))
                 } else {
                     println!("Authentication failed: {}", auth_response.trim());
-                    Ok((s, None))
+                    Ok((tls, None))
                 }
             } else {
-                Ok((s, None))
+                Ok((tls, None))
             }
         }
         Ok(Err(e)) => Err(Box::new(e)),
@@ -48,3 +78,67 @@ pub fn connect(addr: &str, conn: &Connection) -> Result<(TcpStream, Option<Strin
         Err(_) => Err("Connection attempt failed".into()),
     }
 }
+
+pub struct TlsStream {
+    stream: Arc<Mutex<native_tls::TlsStream<TcpStream>>>,
+}
+
+unsafe impl Send for TlsStream {}
+unsafe impl Sync for TlsStream {}
+
+impl TlsStream {
+    pub fn new(stream: native_tls::TlsStream<TcpStream>) -> Self {
+        Self {
+            stream: Arc::new(Mutex::new(stream)),
+        }
+    }
+
+    pub fn read_line(&mut self, buf: &mut String) -> std::io::Result<usize> {
+        let mut byte_buf = [0u8; 1024];
+        let mut stream = self.stream.lock().unwrap();
+        loop {
+            let n = stream.read(&mut byte_buf)?;
+            if n == 0 {
+                return Ok(0);
+            }
+            let chunk = std::str::from_utf8(&byte_buf[..n]).unwrap_or("");
+            buf.push_str(chunk);
+            if buf.contains('\n') {
+                return Ok(n);
+            }
+        }
+    }
+
+    pub fn write_line(&mut self, msg: &str) -> std::io::Result<()> {
+        let mut stream = self.stream.lock().unwrap();
+        stream.write_all(msg.as_bytes())?;
+        stream.write_all(b"\n")?;
+        stream.flush()?;
+        Ok(())
+    }
+
+    pub fn try_clone(&self) -> std::io::Result<TlsStream> {
+        Ok(TlsStream {
+            stream: Arc::clone(&self.stream),
+        })
+    }
+}
+
+impl Read for TlsStream {
+    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
+        let mut stream = self.stream.lock().unwrap();
+        stream.read(buf)
+    }
+}
+
+impl Write for TlsStream {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        let mut stream = self.stream.lock().unwrap();
+        stream.write(buf)
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        let mut stream = self.stream.lock().unwrap();
+        stream.flush()
+    }
+}