diff --git a/src/executables/backend_server_insecure.rs b/src/executables/backend_server_insecure.rs index 78addba..3483ad9 100644 --- a/src/executables/backend_server_insecure.rs +++ b/src/executables/backend_server_insecure.rs @@ -1,4 +1,4 @@ -use rust_nex::reggie::{RemoteController, UnitPacketRead}; +use rust_nex::reggie::{RemoteController, UnitPacketRead, WebStreamSocket}; use log::{error, info}; use once_cell::sync::Lazy; use rustls::client::danger::HandshakeSignatureValid; @@ -65,7 +65,15 @@ async fn main() { let listen = TcpListener::bind(SocketAddrV4::new(*OWN_IP_PRIVATE, *SERVER_PORT)).await.unwrap(); + + while let Ok((stream, addr)) = listen.accept().await { + let Ok(websocket) = tokio_tungstenite::accept_async(stream).await else { + continue; + }; + + let stream = WebStreamSocket::new(websocket); + let mut stream = match acceptor.accept(stream).await { Ok(v) => v, Err(e) => { diff --git a/src/executables/backend_server_secure.rs b/src/executables/backend_server_secure.rs index 8ba76e1..4163026 100644 --- a/src/executables/backend_server_secure.rs +++ b/src/executables/backend_server_secure.rs @@ -1,6 +1,6 @@ use std::io::Cursor; use rust_nex::rmc::structures::RmcSerialize; -use rust_nex::reggie::{RemoteController, UnitPacketRead}; +use rust_nex::reggie::{RemoteController, UnitPacketRead, WebStreamSocket}; use std::net::SocketAddrV4; use std::sync::Arc; use std::sync::atomic::AtomicU32; @@ -50,6 +50,12 @@ async fn main() { MatchmakeManager::initialize_garbage_collect_thread(weak_mmm).await; while let Ok((stream, addr)) = listen.accept().await { + let Ok(websocket) = tokio_tungstenite::accept_async(stream).await else { + continue; + }; + + let stream = WebStreamSocket::new(websocket); + let mut stream = match acceptor.accept(stream).await { Ok(v) => v, Err(e) => { diff --git a/src/executables/proxy_insecure.rs b/src/executables/proxy_insecure.rs index b0c55dd..9742618 100644 --- a/src/executables/proxy_insecure.rs +++ b/src/executables/proxy_insecure.rs @@ -1,5 +1,5 @@ -use rust_nex::reggie::LocalProxy; +use rust_nex::reggie::{tls_connect_to, LocalProxy}; use std::env; use std::ffi::CStr; use std::io::{Read, Write}; @@ -107,8 +107,11 @@ async fn main() { return; } - let mut stream - = establish_tls_connection_to(&dest, &dest).await; + let Ok(mut stream) + = tls_connect_to(&dest).await else { + error!("failed to connect"); + return; + }; if let Err(e) = stream.send_buffer(&ConnectionInitData{ prudpsock_addr: conn.socket_addr, diff --git a/src/executables/proxy_secure.rs b/src/executables/proxy_secure.rs index e484a72..30bcf23 100644 --- a/src/executables/proxy_secure.rs +++ b/src/executables/proxy_secure.rs @@ -13,7 +13,7 @@ use rust_nex::prudp::packet::VirtualPort; use rust_nex::prudp::router::Router; use rust_nex::prudp::secure::Secure; use rust_nex::prudp::unsecure::Unsecure; -use rust_nex::reggie::{establish_tls_connection_to, ProxyManagement, RemoteController}; +use rust_nex::reggie::{establish_tls_connection_to, tls_connect_to, ProxyManagement, RemoteController}; use rust_nex::rmc::response::ErrorCode; use rust_nex::rnex_proxy_common::ConnectionInitData; use rust_nex::reggie::ServerCluster::Auth; @@ -92,8 +92,11 @@ async fn main() { return; } - let mut stream - = establish_tls_connection_to(&dest, &dest).await; + let Ok(mut stream) + = tls_connect_to(&dest).await else { + error!("failed to connect"); + return; + }; if let Err(e) = stream.send_buffer(&ConnectionInitData{ prudpsock_addr: conn.socket_addr, diff --git a/src/reggie.rs b/src/reggie.rs index 53bb5db..c9a35e2 100644 --- a/src/reggie.rs +++ b/src/reggie.rs @@ -16,7 +16,7 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio::net::{TcpListener, TcpStream}; use tokio_rustls::{TlsAcceptor, TlsConnector}; use tokio_rustls::client::TlsStream; -use tokio_tungstenite::{connect_async, WebSocketStream}; +use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; use tokio_tungstenite::tungstenite::Message; use webpki::anchor_from_trusted_cert; use rust_nex::common::setup; @@ -307,16 +307,22 @@ pub enum ConnectError{ DataSendError(#[from] io::Error), } -pub async fn rmc_connect_to(url: &str, init_data: U, create_func: F) -> Result, ConnectError> - where - F: FnOnce(RmcConnection) -> Arc{ +pub async fn tls_connect_to(url: &str) -> Result>>, ConnectError>{ let (stream, _)= connect_async(format!("ws://{}/", url)).await?; let webstreamsocket = WebStreamSocket::new(stream); let connector = get_configured_tls_connector().await; - let mut connection = connector.connect(ServerName::try_from(url.to_string()).unwrap(), webstreamsocket).await.unwrap(); + let connection = connector.connect(ServerName::try_from(url.to_string()).unwrap(), webstreamsocket).await?; + + Ok(connection) +} + +pub async fn rmc_connect_to(url: &str, init_data: U, create_func: F) -> Result, ConnectError> + where + F: FnOnce(RmcConnection) -> Arc{ + let mut connection = tls_connect_to(url).await?; connection.send_buffer(&init_data.to_data()).await?;