refactor(tests): reformat code and update dependency management
- Reorganized import statements in adapter tests for better readability - Replaced or_insert_with(Vec::new) with or_default() in test closures - Updated Cargo.lock with new dependency versions and checksums - Added TLS features to tonic dependency configuration - Included sqlx, chrono, and uuid dependencies with specific features - Added jsonwebtoken and arc-swap as project dependencies - Reformatted assertion statements to comply with line length limits - Adjusted base64 import order in engine codec module - Updated protobuf include statement formatting
This commit is contained in:
zhenyi
+86
-71
@@ -1,11 +1,12 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use wtransport::{Connection, Endpoint, ServerConfig, Identity};
|
||||
use wtransport::{Connection, Endpoint, Identity, ServerConfig};
|
||||
|
||||
use crate::engine::codec;
|
||||
use crate::engine::packet::{Packet, PacketType};
|
||||
use crate::engine::server::EngineConfig;
|
||||
use crate::engine::session::{SessionState, SessionStore, TransportType};
|
||||
use crate::{ImksError, ImksResult};
|
||||
|
||||
pub async fn run_webtransport_server(
|
||||
port: u16,
|
||||
@@ -14,15 +15,18 @@ pub async fn run_webtransport_server(
|
||||
store: SessionStore,
|
||||
config: EngineConfig,
|
||||
on_message: Arc<dyn Fn(String, Packet) + Send + Sync>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let identity = Identity::load_pemfiles(cert_path, key_path).await?;
|
||||
) -> ImksResult<()> {
|
||||
let identity = Identity::load_pemfiles(cert_path, key_path)
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
|
||||
let server_config = ServerConfig::builder()
|
||||
.with_bind_default(port)
|
||||
.with_identity(identity)
|
||||
.build();
|
||||
|
||||
let server = Endpoint::server(server_config)?;
|
||||
let server =
|
||||
Endpoint::server(server_config).map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
|
||||
tracing::info!("WebTransport server listening on UDP port {}", port);
|
||||
|
||||
@@ -49,9 +53,14 @@ async fn handle_webtransport_session(
|
||||
store: SessionStore,
|
||||
config: EngineConfig,
|
||||
on_message: Arc<dyn Fn(String, Packet) + Send + Sync>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let request = incoming.await?;
|
||||
let connection = request.accept().await?;
|
||||
) -> ImksResult<()> {
|
||||
let request = incoming
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
let connection = request
|
||||
.accept()
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
|
||||
let sid = crate::engine::session::generate_sid();
|
||||
let mut rx = store.create(sid.clone(), TransportType::WebTransport);
|
||||
@@ -81,65 +90,57 @@ async fn handle_webtransport_session(
|
||||
// Reuse buffer across recv iterations instead of allocating 65KB each time
|
||||
let recv_handle = tokio::spawn(async move {
|
||||
let mut buf = vec![0u8; 65536];
|
||||
loop {
|
||||
match connection_recv.accept_bi().await {
|
||||
Ok((mut send, mut recv)) => {
|
||||
// Reset buffer length for the next read without deallocating
|
||||
buf.resize(65536, 0);
|
||||
match recv.read(&mut buf).await {
|
||||
Ok(Some(n)) => {
|
||||
if n > max_payload {
|
||||
tracing::warn!(
|
||||
"WebTransport payload too large ({}) for session {}",
|
||||
n,
|
||||
sid_clone
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if let Ok(packet) = codec::decode_packet_ws(&buf[..n]) {
|
||||
match packet.packet_type {
|
||||
PacketType::Ping => {
|
||||
let pong = Packet::pong("");
|
||||
if send_wt_packet_on_stream(&mut send, &pong)
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
PacketType::Pong => {
|
||||
if let Some(s) = store_clone.get(&sid_clone) {
|
||||
let mut s = s.write().await;
|
||||
s.update_ping();
|
||||
}
|
||||
}
|
||||
PacketType::Message => {
|
||||
let on_msg = on_message_clone.clone();
|
||||
let sid = sid_clone.clone();
|
||||
tokio::spawn(async move {
|
||||
on_msg(sid, packet);
|
||||
});
|
||||
}
|
||||
PacketType::Close => {
|
||||
if let Some(s) = store_clone.get(&sid_clone) {
|
||||
let mut s = s.write().await;
|
||||
s.set_state(SessionState::Closed);
|
||||
}
|
||||
store_clone.remove(&sid_clone);
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
while let Ok((mut send, mut recv)) = connection_recv.accept_bi().await {
|
||||
// Reset buffer length for the next read without deallocating
|
||||
buf.resize(65536, 0);
|
||||
match recv.read(&mut buf).await {
|
||||
Ok(Some(n)) => {
|
||||
if n > max_payload {
|
||||
tracing::warn!(
|
||||
"WebTransport payload too large ({}) for session {}",
|
||||
n,
|
||||
sid_clone
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if let Ok(packet) = codec::decode_packet_ws(&buf[..n]) {
|
||||
match packet.packet_type {
|
||||
PacketType::Ping => {
|
||||
let pong = Packet::pong("");
|
||||
if send_wt_packet_on_stream(&mut send, &pong).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
PacketType::Pong => {
|
||||
if let Some(s) = store_clone.get(&sid_clone) {
|
||||
let mut s = s.write().await;
|
||||
s.update_ping();
|
||||
}
|
||||
}
|
||||
PacketType::Message => {
|
||||
let on_msg = on_message_clone.clone();
|
||||
let sid = sid_clone.clone();
|
||||
tokio::spawn(async move {
|
||||
on_msg(sid, packet);
|
||||
});
|
||||
}
|
||||
PacketType::Close => {
|
||||
if let Some(s) = store_clone.get(&sid_clone) {
|
||||
let mut s = s.write().await;
|
||||
s.set_state(SessionState::Closed);
|
||||
}
|
||||
store_clone.remove(&sid_clone);
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
|
||||
Ok::<(), ImksError>(())
|
||||
});
|
||||
|
||||
let connection_send = connection.clone();
|
||||
@@ -191,18 +192,26 @@ async fn handle_webtransport_session(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_wt_packet(
|
||||
connection: &Connection,
|
||||
packet: &Packet,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let (mut send, _recv) = connection.open_bi().await?.await?;
|
||||
async fn send_wt_packet(connection: &Connection, packet: &Packet) -> ImksResult<()> {
|
||||
let (mut send, _recv) = connection
|
||||
.open_bi()
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
let encoded = codec::encode_packet_binary_ws(packet);
|
||||
let is_binary = matches!(packet.data, crate::engine::packet::PacketData::Binary(_));
|
||||
let header = codec::encode_webtransport_header(encoded.len(), is_binary);
|
||||
|
||||
send.write_all(&header).await?;
|
||||
send.write_all(&encoded).await?;
|
||||
send.finish().await?;
|
||||
send.write_all(&header)
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
send.write_all(&encoded)
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
send.finish()
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -210,14 +219,20 @@ async fn send_wt_packet(
|
||||
async fn send_wt_packet_on_stream(
|
||||
send: &mut wtransport::SendStream,
|
||||
packet: &Packet,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
) -> ImksResult<()> {
|
||||
let encoded = codec::encode_packet_binary_ws(packet);
|
||||
let is_binary = matches!(packet.data, crate::engine::packet::PacketData::Binary(_));
|
||||
let header = codec::encode_webtransport_header(encoded.len(), is_binary);
|
||||
|
||||
send.write_all(&header).await?;
|
||||
send.write_all(&encoded).await?;
|
||||
send.finish().await?;
|
||||
send.write_all(&header)
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
send.write_all(&encoded)
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
send.finish()
|
||||
.await
|
||||
.map_err(|e| ImksError::WebTransport(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user