tensamin/omikron
Watch
6
0
Fork
You've already forked omikron
 0
Code Issues Pull requests 3 Projects Releases Packages Wiki Activity Actions

[WIP] WTransport

Browse source
This commit is contained in:
Alex Emmet 2026-03-04 17:12:50 +01:00
commit b2e6e903f7
4 changed files with 79 additions and 110 deletions
Download patch file Download diff file Expand all files Collapse all files

15
src/main.rs
View file

@ -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();
}

67
src/omega/omega_connection.rs
View file

@ -1,28 +1,18 @@
use crate::{
data::user::UserStatus,
get_private_key, log, log_cv_in, log_cv_out, log_err, log_in, log_out,
rho::{
connection::GeneralConnection,
rho_manager::{self, RHO_CONNECTIONS, connection_count},
},
get_private_key, log, log_cv_in, log_cv_out, log_err, log_in,
rho::rho_manager::{self, RHO_CONNECTIONS, connection_count},
util::{
crypto_helper::{decrypt_b64, secret_key_to_base64},
file_util::{load_file_buf, load_file_vec},
file_util::load_file_vec,
logger::PrintType,
},
};
use dashmap::DashMap;
use epsilon_core::{CommunicationType, CommunicationValue, DataTypes, DataValue, rand_u32};
use epsilon_native::{Receiver, Sender}; // Your existing types
use futures::prelude::*;
use once_cell::sync::Lazy;
use quinn::{ClientConfig, Endpoint};
use rustls::{
ClientConfig as RustlsClientConfig,
crypto::{CryptoProvider, aws_lc_rs},
pki_types::ServerName,
};
use std::{collections::HashMap, env, net::SocketAddr, sync::Arc, time::Duration};
use std::{collections::HashMap, env, sync::Arc, time::Duration};
use tokio::{
sync::{Mutex, RwLock, mpsc, watch},
task::JoinHandle,
@ -34,7 +24,7 @@ use uuid::Uuid;
// Configuration
// ============================================================================
const OMEGA_HOST_DEFAULT: &str = "188.114.97.0";
const OMEGA_HOST_DEFAULT: &str = "tensamin.net";
const OMEGA_PORT_DEFAULT: u16 = 9187;
const RECONNECT_DELAY: Duration = Duration::from_secs(5);
const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(300);
@ -93,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>>>,
@ -113,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 {
@ -121,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())),
@ -211,43 +211,26 @@ impl OmegaConnection {
async fn connect_once(self: Arc<Self>) -> Result<(), String> {
*self.state.write().await = ConnectionState::Connecting;
let addr_str = format!("{}:{}", self.host, self.port);
let addrs: Vec<SocketAddr> = tokio::net::lookup_host(&addr_str)
.await
.map_err(|e| format!("DNS lookup failed for {}: {}", addr_str, e))?
.collect();
if addrs.is_empty() {
return Err(format!("No addresses found for {}", addr_str));
}
// Prefer IPv4 if available, otherwise use first available (IPv6 or IPv4)
let remote_addr = addrs
.iter()
.find(|a| a.is_ipv4())
.copied()
.unwrap_or_else(|| addrs[0]);
// Build WebTransport URL
let url = format!("https://{}:{}", self.host, self.port);
log_in!(
0,
PrintType::Omega,
"Connecting to {} ({})...",
self.host,
remote_addr
"Connecting to {}...", // WebTransport uses URLs, not socket addresses
url
);
// Connect using epsilon_native wrapper
let (sender, receiver) = epsilon_native::client::connect(remote_addr)
// Connect using new epsilon_native API with certificate verification
let (sender, receiver) = epsilon_native::connect(&url, self.server_cert.clone())
.await
.map_err(|e| format!("Connection failed: {}", e))?;
log_in!(
0,
PrintType::Omega,
"QUIC connection established to {} (via {})",
addr_str,
remote_addr
"WebTransport connection established to {}", // Updated protocol name
url
);
// Store sender
@ -571,7 +554,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())

69
src/rho/server.rs
View file

@ -1,55 +1,26 @@
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))?;
let tls_cfg = load_tls().expect("TLS config failed");
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
);
let server_crypto = quinn::crypto::rustls::QuicServerConfig::try_from(tls_cfg)
.expect("Failed to convert to QuicServerConfig");
while let Some((sender, receiver)) = host.next().await {
tokio::spawn(async move {
let conn = GeneralConnection::new(sender, receiver);
conn.handle().await;
});
}
let server_cfg = ServerConfig::with_crypto(Arc::new(server_crypto));
let mut host: Host = epsilon_native::host(port, server_cfg).await.unwrap();
tokio::spawn(async move {
while let Some((sender, receiver)) = host.next().await {
tokio::spawn(async move {
GeneralConnection::new(sender, receiver).handle().await;
});
}
});
}
fn load_tls() -> Option<CryptoConfig> {
let _ = aws_lc_rs::default_provider().install_default();
let mut cert_pem = load_file_buf("certs", "cert.pem").ok()?;
let cert_chain = rustls_pemfile::certs(&mut cert_pem)
.collect::<Result<Vec<_>, _>>()
.ok()?;
let key_pem = load_file_vec("certs", "key.pem").ok()?;
let key_der = rustls_pemfile::private_key(&mut &*key_pem).ok()??;
let cfg = CryptoConfig::builder()
.with_no_client_auth()
.with_single_cert(cert_chain, key_der)
.ok()?;
Some(cfg)
Ok(())
}