Merge branch 'main' of github.com:Tensamin/Omikron

src/main.rs

@@ -9,7 +9,7 @@ use std::env;
 
 use dotenv::dotenv;
 use once_cell::sync::Lazy;
-use rustls::crypto::aws_lc_rs;
+use rustls::crypto::aws_lc_rs::default_provider;
 
 use crate::{
     calls::call_util::garbage_collect_calls,
@@ -32,15 +32,18 @@ pub fn get_public_key() -> x448::PublicKey {
 
 #[tokio::main]
 async fn main() {
-    aws_lc_rs::default_provider().install_default();
+    if let Err(_) = default_provider().install_default() {
+        println!("Error loading Provider");
+        return;
+    }
     dotenv().ok();
     startup();
 
-    start(959).await;
-
-    garbage_collect_calls();
-
     get_omega_connection();
+    tokio::spawn(async move {
+        let _ = start(959).await;
+    });
+    garbage_collect_calls();
 
     tokio::signal::ctrl_c().await.unwrap();
 }

src/omega/omega_connection.rs

@@ -83,13 +83,13 @@ impl ConnectionState {
 // ============================================================================
 // Omega Connection (Client-side with auto-reconnect)
 // ============================================================================
-
 pub struct OmegaConnection {
     state: Arc<RwLock<ConnectionState>>,
     sender: Arc<RwLock<Option<Arc<Sender>>>>,
     connection_loop_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
     host: String,
     port: u16,
+    server_cert: Vec<u8>,
     last_ping: Arc<Mutex<i64>>,
     heartbeat_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
     message_send_times: Arc<Mutex<HashMap<Uuid, Instant>>>,
@@ -103,6 +103,15 @@ impl OmegaConnection {
     }
 
     pub fn with_host(host: &str, port: u16) -> Self {
+        // Load server certificate from default location
+        let server_cert =
+            load_file_vec("certs", "cert.pem").expect("Failed to load server certificate");
+
+        Self::with_host_and_cert(host, port, server_cert)
+    }
+
+    // New constructor that accepts certificate directly
+    pub fn with_host_and_cert(host: &str, port: u16, server_cert: Vec<u8>) -> Self {
         let (shutdown_tx, _) = watch::channel(false);
 
         OmegaConnection {
@@ -111,6 +120,7 @@ impl OmegaConnection {
             connection_loop_handle: Arc::new(Mutex::new(None)),
             host: host.to_string(),
             port,
+            server_cert, // Store certificate for connection
             last_ping: Arc::new(Mutex::new(-1)),
             heartbeat_handle: Arc::new(Mutex::new(None)),
             message_send_times: Arc::new(Mutex::new(HashMap::new())),
@@ -535,7 +545,7 @@ impl OmegaConnection {
 
         match tokio::time::timeout(timeout, rx.recv()).await {
             Ok(Some(response_cv)) => Ok(response_cv),
-            Ok(None) => Err("Channel closed".to_string()),
+            Ok(_) => Err("Channel closed".to_string()),
             Err(_) => {
                 WAITING_TASKS.remove(&msg_id);
                 Err("Request timed out".to_string())

src/rho/server.rs

@@ -1,23 +1,13 @@
-use crate::{
-    log,
-    rho::connection::GeneralConnection,
-    util::file_util::{load_file_buf, load_file_vec},
-};
+use crate::{log, rho::connection::GeneralConnection, util::file_util::load_file_vec};
 use epsilon_native::Host;
-use quinn::ServerConfig;
-use rustls::{
-    ServerConfig as CryptoConfig,
-    crypto::{CryptoProvider, aws_lc_rs},
-    pki_types::{
-        CertificateDer, PrivateKeyDer,
-        pem::{PemObject, SectionKind},
-    },
-};
-use std::sync::Arc;
 
-pub async fn start(port: u16) {
-    let _ = aws_lc_rs::default_provider().install_default();
+pub async fn start(port: u16) -> Result<(), Box<dyn std::error::Error>> {
+    let cert_pem = load_file_vec("certs", "cert.pem")
+        .map_err(|e| format!("Failed to load certificate: {}", e))?;
+    let key_pem = load_file_vec("certs", "key.pem")
+        .map_err(|e| format!("Failed to load private key: {}", e))?;
 
+<<<<<<< HEAD
     let key_pem = load_file_vec("certs", "key.pem").unwrap();
     let cert_pem = load_file_vec("certs", "cert.pem").unwrap();
 
@@ -29,4 +19,22 @@ pub async fn start(port: u16) {
             });
         }
     });
+=======
+    let mut host: Host = epsilon_native::host(port, cert_pem, key_pem).await?;
+    log!(
+        0,
+        crate::util::logger::PrintType::Omikron,
+        "Webtransport Server listening on port {}",
+        port
+    );
+
+    while let Some((sender, receiver)) = host.next().await {
+        tokio::spawn(async move {
+            let conn = GeneralConnection::new(sender, receiver);
+            conn.handle().await;
+        });
+    }
+
+    Ok(())
+>>>>>>> b2e6e903f789a81dd3629c21c08f03bbb430280b
 }