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

34
Cargo.lock generated
View file

@ -1090,7 +1090,7 @@ dependencies = [
[[package]]
name = "epsilon-core"
version = "0.1.0"
source = "git+https://github.com/Tensamin/Epsilon.git#4bdd32b44a77ac980cede20a1d401004f52e4a29"
source = "git+https://github.com/Tensamin/Epsilon.git#1cdb6a4dab3cc1ef1e1ff36b6b58bcbcfa65330a"
dependencies = [
"byteorder",
"quinn",
@ -1102,7 +1102,7 @@ dependencies = [
[[package]]
name = "epsilon-native"
version = "0.1.0"
source = "git+https://github.com/Tensamin/Epsilon.git#4bdd32b44a77ac980cede20a1d401004f52e4a29"
source = "git+https://github.com/Tensamin/Epsilon.git#1cdb6a4dab3cc1ef1e1ff36b6b58bcbcfa65330a"
dependencies = [
"anyhow",
"async-trait",
@ -1115,6 +1115,7 @@ dependencies = [
"thiserror 2.0.18",
"tokio 1.50.0",
"wtransport",
"wtransport-proto",
]
[[package]]
@ -2725,9 +2726,9 @@ dependencies = [
[[package]]
name = "proc-macro-crate"
version = "3.4.0"
version = "3.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "219cb19e96be00ab2e37d6e299658a0cfa83e52429179969b0f0121b4ac46983"
checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f"
dependencies = [
"toml_edit",
]
@ -2854,9 +2855,9 @@ dependencies = [
[[package]]
name = "quote"
version = "1.0.44"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"proc-macro2",
]
@ -2964,6 +2965,7 @@ version = "0.14.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "10b99e0098aa4082912d4c649628623db6aba77335e4f4569ff5083a6448b32e"
dependencies = [
"aws-lc-rs",
"pem",
"ring 0.17.14",
"rustls-pki-types",
@ -4025,7 +4027,7 @@ dependencies = [
"indexmap",
"serde_core",
"serde_spanned",
"toml_datetime",
"toml_datetime 0.7.5+spec-1.1.0",
"toml_parser",
"toml_writer",
"winnow",
@ -4041,13 +4043,22 @@ dependencies = [
]
[[package]]
name = "toml_edit"
version = "0.23.10+spec-1.0.0"
name = "toml_datetime"
version = "1.0.0+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "84c8b9f757e028cee9fa244aea147aab2a9ec09d5325a9b01e0a49730c2b5269"
checksum = "32c2555c699578a4f59f0cc68e5116c8d7cabbd45e1409b989d4be085b53f13e"
dependencies = [
"serde_core",
]
[[package]]
name = "toml_edit"
version = "0.25.3+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a0a07913e63758bc95142d9863a5a45173b71515e68b690cad70cf99c3255ce1"
dependencies = [
"indexmap",
"toml_datetime",
"toml_datetime 1.0.0+spec-1.1.0",
"toml_parser",
"winnow",
]
@ -5169,6 +5180,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d43b0f71ce057da06bc0851b23ee24f3f86190b07203dd8f567d0b706a185202"
dependencies = [
"asn1-rs",
"aws-lc-rs",
"data-encoding",
"der-parser",
"lazy_static",

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