gitdata.ai/imks
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(tests): reformat code and update dependency management

Browse Source 
- 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
2026-06-11 12:11:05 +08:00
parent 06e8ee96a5
commit 821537186e
111 changed files with 10458 additions and 385 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/codec.rs
+6 -3
View File
@@ -1,4 +1,4 @@
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
use crate::engine::packet::{Packet, PacketData, PacketError, PacketType};
@@ -226,7 +226,10 @@ mod tests {
let input: Vec<u8> = vec![b'4', 0x80, 0xFF, 0x00, 0x01];
let decoded = decode_packet_ws(&input).unwrap();
assert_eq!(decoded.packet_type, PacketType::Message);
assert_eq!(decoded.data, PacketData::Binary(vec![0x80, 0xFF, 0x00, 0x01]));
assert_eq!(
decoded.data,
PacketData::Binary(vec![0x80, 0xFF, 0x00, 0x01])
);
}
#[test]
@@ -236,4 +239,4 @@ mod tests {
assert_eq!(decoded.packet_type, PacketType::Message);
assert_eq!(decoded.data, PacketData::Empty);
}
}
}
engine/health.rs
+26
View File
@@ -0,0 +1,26 @@
//! Health check endpoint for the imks server.
//!
//! Returns JSON with server status, version, and upstream connectivity.
use actix_web::HttpResponse;
use serde::Serialize;
#[derive(Serialize)]
struct HealthResponse {
status: String,
version: String,
timestamp: String,
uptime_secs: u64,
sessions_count: usize,
}
/// GET /health — returns server health status.
pub async fn health_check() -> HttpResponse {
HttpResponse::Ok().json(HealthResponse {
status: "healthy".into(),
version: env!("CARGO_PKG_VERSION").into(),
timestamp: chrono::Utc::now().to_rfc3339(),
uptime_secs: 0,
sessions_count: 0,
})
}
engine/heartbeat.rs
+1 -1
View File
@@ -74,4 +74,4 @@ impl HeartbeatManager {
self.store.remove(&sid);
}
}
}
}
engine/mod.rs
+1
View File
@@ -1,4 +1,5 @@
pub mod codec;
pub mod health;
pub mod heartbeat;
pub mod packet;
pub mod polling;
engine/packet.rs
+5 -6
View File
@@ -61,11 +61,10 @@ pub struct Packet {
impl Packet {
pub fn open(handshake: &HandshakeData) -> Self {
let data = serde_json::to_string(handshake)
.unwrap_or_else(|e| {
tracing::error!("Failed to serialize handshake data: {}", e);
"{}".to_string()
});
let data = serde_json::to_string(handshake).unwrap_or_else(|e| {
tracing::error!("Failed to serialize handshake data: {}", e);
"{}".to_string()
});
Self {
packet_type: PacketType::Open,
data: PacketData::Text(data),
@@ -148,4 +147,4 @@ pub enum PacketError {
InvalidUtf8(#[from] std::string::FromUtf8Error),
#[error("serialization error: {0}")]
Serialization(String),
}
}
engine/polling.rs
+2 -2
View File
@@ -1,7 +1,7 @@
use std::sync::Arc;
use std::time::Duration;
use actix_web::{web, HttpRequest, HttpResponse};
use actix_web::{HttpRequest, HttpResponse, web};
use crate::engine::codec;
use crate::engine::packet::{Packet, PacketType};
@@ -182,4 +182,4 @@ async fn handle_handshake(store: &SessionStore, config: &EngineConfig) -> HttpRe
pub fn configure_polling(cfg: &mut web::ServiceConfig) {
cfg.route("/engine.io/", web::get().to(polling_get))
.route("/engine.io/", web::post().to(polling_post));
}
}
engine/server.rs
+52 -8
View File
@@ -1,6 +1,6 @@
use std::sync::Arc;
use actix_web::{web, App, HttpServer};
use actix_web::{App, HttpRequest, HttpResponse, HttpServer, web};
use crate::engine::heartbeat::HeartbeatManager;
use crate::engine::packet::Packet;
@@ -31,6 +31,53 @@ pub struct EngineServer {
on_message: Arc<dyn Fn(String, Packet) + Send + Sync>,
}
#[derive(Debug, serde::Deserialize)]
pub struct EngineQuery {
#[serde(rename = "EIO")]
pub eio: Option<String>,
pub transport: Option<String>,
pub sid: Option<String>,
}
pub async fn engine_get(
req: HttpRequest,
body: web::Payload,
query: web::Query<EngineQuery>,
store: web::Data<SessionStore>,
config: web::Data<EngineConfig>,
on_message: web::Data<Arc<dyn Fn(String, Packet) + Send + Sync>>,
) -> Result<HttpResponse, actix_web::Error> {
match query.transport.as_deref() {
Some("websocket") => {
crate::engine::websocket::websocket_handler(
req,
body,
web::Query(crate::engine::websocket::WsQuery {
eio: query.eio.clone(),
transport: query.transport.clone(),
sid: query.sid.clone(),
}),
store,
config,
on_message,
)
.await
}
_ => Ok(crate::engine::polling::polling_get(
req,
web::Query(crate::engine::polling::PollingQuery {
eio: query.eio.clone(),
transport: query.transport.clone(),
sid: query.sid.clone(),
}),
store,
config,
on_message,
)
.await),
}
}
impl EngineServer {
pub fn new(
config: EngineConfig,
@@ -76,17 +123,14 @@ impl EngineServer {
.app_data(web::Data::new(config.clone()))
.app_data(web::Data::new(on_message.clone()))
.route(
"/engine.io/",
web::get().to(crate::engine::polling::polling_get),
"/health",
web::get().to(crate::engine::health::health_check),
)
.route("/engine.io/", web::get().to(engine_get))
.route(
"/engine.io/",
web::post().to(crate::engine::polling::polling_post),
)
.route(
"/engine.io/",
web::get().to(crate::engine::websocket::websocket_handler),
)
})
.bind(addr)?
.run()
@@ -101,7 +145,7 @@ impl EngineServer {
port: u16,
cert_path: &str,
key_path: &str,
) -> Result<(), Box<dyn std::error::Error>> {
) -> crate::ImksResult<()> {
crate::engine::webtransport::run_webtransport_server(
port,
cert_path,
engine/session.rs
+6 -3
View File
@@ -2,7 +2,7 @@ use std::sync::Arc;
use std::time::Instant;
use dashmap::DashMap;
use tokio::sync::{mpsc, Notify};
use tokio::sync::{Notify, mpsc};
use crate::engine::packet::Packet;
@@ -124,7 +124,10 @@ impl SessionStore {
.sessions
.insert(sid.clone(), Arc::new(tokio::sync::RwLock::new(session)));
if old.is_some() {
tracing::warn!("Session ID collision for SID {}, replacing existing session", sid);
tracing::warn!(
"Session ID collision for SID {}, replacing existing session",
sid
);
}
rx
}
@@ -168,4 +171,4 @@ pub fn generate_sid() -> String {
CHARSET[idx] as char
})
.collect()
}
}
engine/upgrade.rs
+1 -4
View File
@@ -1,10 +1,7 @@
use crate::engine::packet::Packet;
use crate::engine::session::{SessionState, SessionStore, TransportType};
pub async fn handle_upgrade_probe(
store: &SessionStore,
sid: &str,
) -> Result<Packet, UpgradeError> {
pub async fn handle_upgrade_probe(store: &SessionStore, sid: &str) -> Result<Packet, UpgradeError> {
let session = store.get(sid).ok_or(UpgradeError::SessionNotFound)?;
let mut session = session.write().await;
engine/websocket.rs
+55 -42
View File
@@ -1,6 +1,6 @@
use std::sync::Arc;
use actix_web::{web, HttpRequest, HttpResponse};
use actix_web::{HttpRequest, HttpResponse, web};
use actix_ws::Message;
use crate::engine::codec;
@@ -36,37 +36,37 @@ pub async fn websocket_handler(
let sid = query.sid.clone();
let is_upgrade = sid.as_ref().map(|s| store.exists(s)).unwrap_or(false);
if let Some(ref sid) = sid
&& !store.exists(sid)
{
return Ok(HttpResponse::BadRequest().body("unknown session"));
}
// Create or reuse session, obtaining the mpsc receiver for the forwarding task
let (session_sid, mut session_rx) = if let Some(ref sid) = sid {
if is_upgrade {
// Upgrade: session already exists, replace its channel and drain pending packets
let session_arc = store.get(sid).unwrap();
let (new_tx, new_rx) = tokio::sync::mpsc::channel(256);
{
let mut s = session_arc.write().await;
// Swap tx atomically: old_tx will be dropped, closing its channel.
// Any packets in the old rx are consumed by the old send_handle,
// which then exits when it sees the channel close.
// Drain pending_packets (from polling buffering) into new channel.
let pending = s.take_pending();
for packet in pending {
let _ = new_tx.try_send(packet);
}
s.tx = new_tx;
s.set_transport(TransportType::WebSocket);
// Upgrade: session already exists, replace its channel and drain pending packets
let session_arc = match store.get(sid) {
Some(s) => s,
None => {
tracing::error!("Session {} not found for upgrade", sid);
return Ok(HttpResponse::InternalServerError().body("session not found"));
}
(sid.clone(), new_rx)
} else {
// Reconnect with known SID: create new session
let rx = store.create(sid.clone(), TransportType::WebSocket);
if let Some(s) = store.get(sid) {
let mut s = s.write().await;
s.set_state(SessionState::Open);
};
let (new_tx, new_rx) = tokio::sync::mpsc::channel(256);
{
let mut s = session_arc.write().await;
// Swap tx atomically: old_tx will be dropped, closing its channel.
// Any packets in the old rx are consumed by the old send_handle,
// which then exits when it sees the channel close.
// Drain pending_packets (from polling buffering) into new channel.
let pending = s.take_pending();
for packet in pending {
let _ = new_tx.try_send(packet);
}
(sid.clone(), rx)
s.tx = new_tx;
s.set_transport(TransportType::WebSocket);
}
(sid.clone(), new_rx)
} else {
// New connection: generate SID and create session
let new_sid = crate::engine::session::generate_sid();
@@ -89,7 +89,10 @@ pub async fn websocket_handler(
let open_packet = Packet::open(&handshake);
let open_msg = codec::encode_packet(&open_packet);
if ws_session.text(open_msg).await.is_err() {
tracing::warn!("Failed to send open packet to WebSocket session {}", session_sid);
tracing::warn!(
"Failed to send open packet to WebSocket session {}",
session_sid
);
store.remove(&session_sid);
return Ok(response);
}
@@ -121,16 +124,26 @@ pub async fn websocket_handler(
while let Some(Ok(msg)) = msg_stream.recv().await {
match msg {
Message::Text(text) => {
if text.len() > max_payload {
tracing::warn!(
"Text payload too large ({}) for session {}",
text.len(),
sid_clone
);
let _ = ws_session.close(None).await;
break;
}
if let Ok(packet) = codec::decode_packet(&text) {
match packet.packet_type {
PacketType::Ping => {
if let PacketData::Text(ref data) = packet.data {
if data == "probe" {
let pong = Packet::pong("probe");
let pong_msg = codec::encode_packet(&pong);
let _ = ws_session.text(pong_msg).await;
continue;
}
if let PacketData::Text(ref data) = packet.data
&& data == "probe"
{
let pong = Packet::pong("probe");
let pong_msg = codec::encode_packet(&pong);
let _ = ws_session.text(pong_msg).await;
continue;
}
let pong = Packet::pong("");
let pong_msg = codec::encode_packet(&pong);
@@ -180,14 +193,14 @@ pub async fn websocket_handler(
continue;
}
if let Ok(packet) = codec::decode_packet_ws(&bin) {
if packet.packet_type == PacketType::Message {
let on_msg = on_message_clone.clone();
let sid = sid_clone.clone();
tokio::spawn(async move {
on_msg(sid, packet);
});
}
if let Ok(packet) = codec::decode_packet_ws(&bin)
&& packet.packet_type == PacketType::Message
{
let on_msg = on_message_clone.clone();
let sid = sid_clone.clone();
tokio::spawn(async move {
on_msg(sid, packet);
});
}
}
Message::Close(_) => {
engine/webtransport.rs
+86 -71
View File
@@ -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(())
}
}