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
socket/adapter/local.rs
+39 -13
View File
@@ -5,7 +5,7 @@ use async_trait::async_trait;
use dashmap::DashMap;
use uuid::Uuid;
use crate::socket::adapter::{Adapter, AdapterError, BroadcastOptions, SocketInfo};
use crate::socket::adapter::{Adapter, AdapterError, BroadcastOptions, LocalSendFn, SocketInfo};
use crate::socket::packet::Packet;
pub struct LocalAdapter {
@@ -16,7 +16,7 @@ pub struct LocalAdapter {
pub socket_sids: Arc<DashMap<String, String>>,
/// socket_sid → namespace path
socket_namespace: Arc<DashMap<String, String>>,
send_fn: Arc<dyn Fn(&str, &Packet) -> Result<(), String> + Send + Sync>,
send_fn: LocalSendFn,
}
impl LocalAdapter {
@@ -68,7 +68,11 @@ impl LocalAdapter {
#[async_trait]
impl Adapter for LocalAdapter {
async fn broadcast(&self, packet: &Packet, opts: &BroadcastOptions) -> Result<(), AdapterError> {
async fn broadcast(
&self,
packet: &Packet,
opts: &BroadcastOptions,
) -> Result<(), AdapterError> {
let namespace = &packet.namespace;
let sids = self.collect_matching_sids(opts, namespace);
for sid in &sids {
@@ -87,9 +91,16 @@ impl Adapter for LocalAdapter {
Ok(())
}
async fn register(&self, socket_sid: &str, engine_sid: &str, ns: &str) -> Result<(), AdapterError> {
self.socket_sids.insert(socket_sid.to_string(), engine_sid.to_string());
self.socket_namespace.insert(socket_sid.to_string(), ns.to_string());
async fn register(
&self,
socket_sid: &str,
engine_sid: &str,
ns: &str,
) -> Result<(), AdapterError> {
self.socket_sids
.insert(socket_sid.to_string(), engine_sid.to_string());
self.socket_namespace
.insert(socket_sid.to_string(), ns.to_string());
Ok(())
}
@@ -99,8 +110,16 @@ impl Adapter for LocalAdapter {
async fn add(&self, sid: &str, room: &str, ns: &str) -> Result<(), AdapterError> {
let key = Self::room_key(ns, room);
self.rooms.entry(key).or_insert_with(HashSet::new).value_mut().insert(sid.to_string());
self.socket_rooms.entry(sid.to_string()).or_insert_with(HashSet::new).value_mut().insert(room.to_string());
self.rooms
.entry(key)
.or_default()
.value_mut()
.insert(sid.to_string());
self.socket_rooms
.entry(sid.to_string())
.or_default()
.value_mut()
.insert(room.to_string());
Ok(())
}
@@ -137,10 +156,14 @@ impl Adapter for LocalAdapter {
}
}
self.socket_sids.remove(sid);
self.socket_namespace.remove(sid);
Ok(())
}
async fn fetch_sockets(&self, opts: &BroadcastOptions) -> Result<Vec<SocketInfo>, AdapterError> {
async fn fetch_sockets(
&self,
opts: &BroadcastOptions,
) -> Result<Vec<SocketInfo>, AdapterError> {
// fetch_sockets needs namespace context; use an empty namespace to match all
// (this method is typically called for inspection, not delivery)
let sids: Vec<String> = if opts.rooms.is_empty() {
@@ -164,11 +187,13 @@ impl Adapter for LocalAdapter {
continue;
}
if self.socket_sids.contains_key(sid) {
let namespace = self.socket_namespace
let namespace = self
.socket_namespace
.get(sid)
.map(|r| r.value().clone())
.unwrap_or_default();
let rooms = self.socket_rooms
let rooms = self
.socket_rooms
.get(sid)
.map(|r| r.value().clone())
.unwrap_or_default();
@@ -183,7 +208,8 @@ impl Adapter for LocalAdapter {
}
async fn socket_rooms(&self, sid: &str) -> Result<HashSet<String>, AdapterError> {
Ok(self.socket_rooms
Ok(self
.socket_rooms
.get(sid)
.map(|r| r.value().clone())
.unwrap_or_default())
@@ -196,4 +222,4 @@ impl Adapter for LocalAdapter {
async fn close(&self) -> Result<(), AdapterError> {
Ok(())
}
}
}
socket/adapter/mod.rs
+18 -5
View File
@@ -1,14 +1,20 @@
pub mod local;
pub mod redis;
pub mod nats;
pub mod redis;
use std::collections::HashSet;
use std::sync::Arc;
use async_trait::async_trait;
use thiserror::Error;
use crate::socket::packet::Packet;
/// Alias for cross-node broadcast callback functions.
pub type LocalBroadcastFn = Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>;
/// Alias for local send-to-socket callback functions.
pub type LocalSendFn = Arc<dyn Fn(&str, &Packet) -> Result<(), String> + Send + Sync>;
#[derive(Error, Debug)]
pub enum AdapterError {
#[error("Redis error: {0}")]
@@ -72,11 +78,13 @@ pub enum BusMessage {
#[async_trait]
pub trait Adapter: Send + Sync + 'static {
async fn broadcast(&self, packet: &Packet, opts: &BroadcastOptions) -> Result<(), AdapterError>;
async fn broadcast(&self, packet: &Packet, opts: &BroadcastOptions)
-> Result<(), AdapterError>;
async fn add(&self, sid: &str, room: &str, ns: &str) -> Result<(), AdapterError>;
async fn del(&self, sid: &str, room: &str, ns: &str) -> Result<(), AdapterError>;
async fn del_all(&self, sid: &str, ns: &str) -> Result<(), AdapterError>;
async fn fetch_sockets(&self, opts: &BroadcastOptions) -> Result<Vec<SocketInfo>, AdapterError>;
async fn fetch_sockets(&self, opts: &BroadcastOptions)
-> Result<Vec<SocketInfo>, AdapterError>;
async fn socket_rooms(&self, sid: &str) -> Result<HashSet<String>, AdapterError>;
fn server_id(&self) -> &str;
async fn close(&self) -> Result<(), AdapterError>;
@@ -84,7 +92,12 @@ pub trait Adapter: Send + Sync + 'static {
/// Register a socket SID → engine SID mapping in the adapter.
/// Must be called when a socket first connects, before any room operations.
/// The `ns` parameter is the namespace path this socket belongs to.
async fn register(&self, _socket_sid: &str, _engine_sid: &str, _ns: &str) -> Result<(), AdapterError> {
async fn register(
&self,
_socket_sid: &str,
_engine_sid: &str,
_ns: &str,
) -> Result<(), AdapterError> {
Ok(())
}
@@ -95,5 +108,5 @@ pub trait Adapter: Send + Sync + 'static {
}
pub use local::LocalAdapter;
pub use nats::NatsAdapter;
pub use redis::RedisAdapter;
pub use nats::NatsAdapter;
socket/adapter/nats.rs
+90 -32
View File
@@ -5,7 +5,9 @@ use async_trait::async_trait;
use dashmap::DashMap;
use tokio::sync::mpsc;
use crate::socket::adapter::{Adapter, AdapterError, BroadcastOptions, BusMessage, SocketInfo};
use crate::socket::adapter::{
Adapter, AdapterError, BroadcastOptions, BusMessage, LocalBroadcastFn, SocketInfo,
};
use crate::socket::message_bus::MessageBus;
use crate::socket::packet::Packet;
use crate::socket::parser;
@@ -15,11 +17,16 @@ use crate::socket::socket::Socket;
/// Only performs local dispatch — no remote state writes needed.
async fn handle_bus_message(
msg: BusMessage,
on_local_broadcast: &Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
on_local_broadcast: &LocalBroadcastFn,
server_id: &str,
) {
match msg {
BusMessage::Broadcast { namespace: _, packet, opts, server_id: sender_id } => {
BusMessage::Broadcast {
namespace: _,
packet,
opts,
server_id: sender_id,
} => {
if sender_id == server_id {
return;
}
@@ -29,13 +36,18 @@ async fn handle_bus_message(
}
// NATS adapter manages room state locally; cross-server join/leave/disconnect
// are informational only and don't require duplicate state writes.
BusMessage::SocketJoin { server_id: sender_id, .. }
| BusMessage::SocketLeave { server_id: sender_id, .. }
| BusMessage::SocketDisconnect { server_id: sender_id, .. } => {
if sender_id == server_id {
return;
}
BusMessage::SocketJoin {
server_id: sender_id,
..
}
| BusMessage::SocketLeave {
server_id: sender_id,
..
}
| BusMessage::SocketDisconnect {
server_id: sender_id,
..
} => if sender_id == server_id {},
}
}
@@ -51,7 +63,7 @@ pub struct NatsAdapter {
sockets: DashMap<String, Arc<Socket>>,
server_id: String,
namespace: String,
on_local_broadcast: Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
on_local_broadcast: LocalBroadcastFn,
}
impl NatsAdapter {
@@ -59,7 +71,7 @@ impl NatsAdapter {
message_bus: Arc<dyn MessageBus>,
server_id: String,
namespace: String,
on_local_broadcast: Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
on_local_broadcast: LocalBroadcastFn,
) -> Self {
Self {
message_bus,
@@ -133,7 +145,11 @@ impl NatsAdapter {
#[async_trait]
impl Adapter for NatsAdapter {
async fn broadcast(&self, packet: &Packet, opts: &BroadcastOptions) -> Result<(), AdapterError> {
async fn broadcast(
&self,
packet: &Packet,
opts: &BroadcastOptions,
) -> Result<(), AdapterError> {
if opts.flags.local_only {
(self.on_local_broadcast)(packet, opts);
return Ok(());
@@ -146,8 +162,8 @@ impl Adapter for NatsAdapter {
server_id: self.server_id.clone(),
};
let payload = serde_json::to_vec(&msg)
.map_err(|e| AdapterError::Serialization(e.to_string()))?;
let payload =
serde_json::to_vec(&msg).map_err(|e| AdapterError::Serialization(e.to_string()))?;
self.message_bus
.publish(&format!("socket.io:{}:broadcast", self.namespace), &payload)
@@ -158,20 +174,30 @@ impl Adapter for NatsAdapter {
Ok(())
}
async fn register(&self, socket_sid: &str, engine_sid: &str, _ns: &str) -> Result<(), AdapterError> {
self.socket_sids.insert(socket_sid.to_string(), engine_sid.to_string());
async fn register(
&self,
socket_sid: &str,
engine_sid: &str,
_ns: &str,
) -> Result<(), AdapterError> {
self.socket_sids
.insert(socket_sid.to_string(), engine_sid.to_string());
Ok(())
}
async fn add(&self, sid: &str, room: &str, _ns: &str) -> Result<(), AdapterError> {
self.socket_rooms
.entry(sid.to_string())
.and_modify(|set| { set.insert(room.to_string()); })
.and_modify(|set| {
set.insert(room.to_string());
})
.or_insert_with(|| HashSet::from([room.to_string()]));
self.rooms
.entry(room.to_string())
.and_modify(|set| { set.insert(sid.to_string()); })
.and_modify(|set| {
set.insert(sid.to_string());
})
.or_insert_with(|| HashSet::from([sid.to_string()]));
let msg = BusMessage::SocketJoin {
@@ -181,8 +207,8 @@ impl Adapter for NatsAdapter {
server_id: self.server_id.clone(),
};
let payload = serde_json::to_vec(&msg)
.map_err(|e| AdapterError::Serialization(e.to_string()))?;
let payload =
serde_json::to_vec(&msg).map_err(|e| AdapterError::Serialization(e.to_string()))?;
self.message_bus
.publish(&format!("socket.io:{}:join", self.namespace), &payload)
@@ -196,14 +222,24 @@ impl Adapter for NatsAdapter {
if let Some(mut entry) = self.socket_rooms.get_mut(sid) {
entry.value_mut().remove(room);
}
if self.socket_rooms.get(sid).map(|e| e.value().is_empty()).unwrap_or(true) {
if self
.socket_rooms
.get(sid)
.map(|e| e.value().is_empty())
.unwrap_or(true)
{
self.socket_rooms.remove(sid);
}
if let Some(mut entry) = self.rooms.get_mut(room) {
entry.value_mut().remove(sid);
}
if self.rooms.get(room).map(|e| e.value().is_empty()).unwrap_or(true) {
if self
.rooms
.get(room)
.map(|e| e.value().is_empty())
.unwrap_or(true)
{
self.rooms.remove(room);
}
@@ -214,8 +250,8 @@ impl Adapter for NatsAdapter {
server_id: self.server_id.clone(),
};
let payload = serde_json::to_vec(&msg)
.map_err(|e| AdapterError::Serialization(e.to_string()))?;
let payload =
serde_json::to_vec(&msg).map_err(|e| AdapterError::Serialization(e.to_string()))?;
self.message_bus
.publish(&format!("socket.io:{}:leave", self.namespace), &payload)
@@ -231,7 +267,12 @@ impl Adapter for NatsAdapter {
if let Some(mut entry) = self.rooms.get_mut(room) {
entry.value_mut().remove(sid);
}
if self.rooms.get(room).map(|e| e.value().is_empty()).unwrap_or(true) {
if self
.rooms
.get(room)
.map(|e| e.value().is_empty())
.unwrap_or(true)
{
self.rooms.remove(room);
}
}
@@ -246,18 +287,24 @@ impl Adapter for NatsAdapter {
server_id: self.server_id.clone(),
};
let payload = serde_json::to_vec(&msg)
.map_err(|e| AdapterError::Serialization(e.to_string()))?;
let payload =
serde_json::to_vec(&msg).map_err(|e| AdapterError::Serialization(e.to_string()))?;
self.message_bus
.publish(&format!("socket.io:{}:disconnect", self.namespace), &payload)
.publish(
&format!("socket.io:{}:disconnect", self.namespace),
&payload,
)
.await
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
Ok(())
}
async fn fetch_sockets(&self, opts: &BroadcastOptions) -> Result<Vec<SocketInfo>, AdapterError> {
async fn fetch_sockets(
&self,
opts: &BroadcastOptions,
) -> Result<Vec<SocketInfo>, AdapterError> {
let mut result = Vec::new();
let target_sids: HashSet<String> = if opts.rooms.is_empty() {
@@ -276,7 +323,11 @@ impl Adapter for NatsAdapter {
if opts.except.contains(&sid) {
continue;
}
let rooms = self.socket_rooms.get(&sid).map(|e| e.value().clone()).unwrap_or_default();
let rooms = self
.socket_rooms
.get(&sid)
.map(|e| e.value().clone())
.unwrap_or_default();
result.push(SocketInfo {
sid: sid.clone(),
namespace: self.namespace.clone(),
@@ -288,7 +339,11 @@ impl Adapter for NatsAdapter {
}
async fn socket_rooms(&self, sid: &str) -> Result<HashSet<String>, AdapterError> {
Ok(self.socket_rooms.get(sid).map(|e| e.value().clone()).unwrap_or_default())
Ok(self
.socket_rooms
.get(sid)
.map(|e| e.value().clone())
.unwrap_or_default())
}
fn server_id(&self) -> &str {
@@ -296,7 +351,10 @@ impl Adapter for NatsAdapter {
}
async fn close(&self) -> Result<(), AdapterError> {
self.message_bus.close().await.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
self.message_bus
.close()
.await
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
Ok(())
}
}
socket/adapter/redis.rs
+107 -34
View File
@@ -7,7 +7,9 @@ use fred::clients::Client;
use fred::interfaces::{KeysInterface, SetsInterface};
use tokio::sync::mpsc;
use crate::socket::adapter::{Adapter, AdapterError, BroadcastOptions, BusMessage, SocketInfo};
use crate::socket::adapter::{
Adapter, AdapterError, BroadcastOptions, BusMessage, LocalBroadcastFn, SocketInfo,
};
use crate::socket::message_bus::MessageBus;
use crate::socket::packet::Packet;
use crate::socket::parser;
@@ -28,11 +30,16 @@ fn socket_rooms_key(ns: &str, sid: &str) -> String {
/// Only performs local state updates — the remote server already wrote to Redis.
async fn handle_bus_message(
msg: BusMessage,
on_local_broadcast: &Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
on_local_broadcast: &LocalBroadcastFn,
server_id: &str,
) {
match msg {
BusMessage::Broadcast { namespace: _, packet, opts, server_id: sender_id } => {
BusMessage::Broadcast {
namespace: _,
packet,
opts,
server_id: sender_id,
} => {
if sender_id == server_id {
return;
}
@@ -40,13 +47,20 @@ async fn handle_bus_message(
on_local_broadcast(&decoded_packet, &opts);
}
}
BusMessage::SocketJoin { server_id: sender_id, .. }
| BusMessage::SocketLeave { server_id: sender_id, .. }
| BusMessage::SocketDisconnect { server_id: sender_id, .. } => {
BusMessage::SocketJoin {
server_id: sender_id,
..
}
| BusMessage::SocketLeave {
server_id: sender_id,
..
}
| BusMessage::SocketDisconnect {
server_id: sender_id,
..
} => {
// Skip messages from this server; remote server already updated Redis
if sender_id == server_id {
return;
}
if sender_id == server_id {}
// No duplicate Redis writes — the sender already persisted the state change
}
}
@@ -58,10 +72,12 @@ pub struct RedisAdapter {
room_subscribers: DashMap<String, mpsc::Receiver<Vec<u8>>>,
socket_rooms: DashMap<String, HashSet<String>>,
rooms: DashMap<String, HashSet<String>>,
/// socket_sid → engine_sid mapping for local inspection.
socket_sids: DashMap<String, String>,
sockets: DashMap<String, Arc<Socket>>,
server_id: String,
namespace: String,
on_local_broadcast: Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
on_local_broadcast: LocalBroadcastFn,
}
impl RedisAdapter {
@@ -70,7 +86,7 @@ impl RedisAdapter {
redis_client: Client,
server_id: String,
namespace: String,
on_local_broadcast: Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
on_local_broadcast: LocalBroadcastFn,
) -> Self {
Self {
message_bus,
@@ -81,6 +97,7 @@ impl RedisAdapter {
room_subscribers: DashMap::new(),
socket_rooms: DashMap::new(),
rooms: DashMap::new(),
socket_sids: DashMap::new(),
sockets: DashMap::new(),
}
}
@@ -144,7 +161,11 @@ impl RedisAdapter {
#[async_trait]
impl Adapter for RedisAdapter {
async fn broadcast(&self, packet: &Packet, opts: &BroadcastOptions) -> Result<(), AdapterError> {
async fn broadcast(
&self,
packet: &Packet,
opts: &BroadcastOptions,
) -> Result<(), AdapterError> {
if opts.flags.local_only {
(self.on_local_broadcast)(packet, opts);
return Ok(());
@@ -157,11 +178,11 @@ impl Adapter for RedisAdapter {
server_id: self.server_id.clone(),
};
let payload = serde_json::to_vec(&msg)
.map_err(|e| AdapterError::Serialization(e.to_string()))?;
let payload =
serde_json::to_vec(&msg).map_err(|e| AdapterError::Serialization(e.to_string()))?;
self.message_bus
.publish(&format!("socket.io:{}:broadcast", packet.namespace), &payload)
.publish(&format!("socket.io:{}:broadcast", self.namespace), &payload)
.await
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
@@ -185,12 +206,16 @@ impl Adapter for RedisAdapter {
self.socket_rooms
.entry(sid.to_string())
.and_modify(|set| { set.insert(room.to_string()); })
.and_modify(|set| {
set.insert(room.to_string());
})
.or_insert_with(|| HashSet::from([room.to_string()]));
self.rooms
.entry(room.to_string())
.and_modify(|set| { set.insert(sid.to_string()); })
.and_modify(|set| {
set.insert(sid.to_string());
})
.or_insert_with(|| HashSet::from([sid.to_string()]));
let msg = BusMessage::SocketJoin {
@@ -200,11 +225,11 @@ impl Adapter for RedisAdapter {
server_id: self.server_id.clone(),
};
let payload = serde_json::to_vec(&msg)
.map_err(|e| AdapterError::Serialization(e.to_string()))?;
let payload =
serde_json::to_vec(&msg).map_err(|e| AdapterError::Serialization(e.to_string()))?;
self.message_bus
.publish(&format!("socket.io:{}:join", ns), &payload)
.publish(&format!("socket.io:{}:join", self.namespace), &payload)
.await
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
@@ -228,14 +253,24 @@ impl Adapter for RedisAdapter {
if let Some(mut entry) = self.socket_rooms.get_mut(sid) {
entry.value_mut().remove(room);
}
if self.socket_rooms.get(sid).map(|e| e.value().is_empty()).unwrap_or(true) {
if self
.socket_rooms
.get(sid)
.map(|e| e.value().is_empty())
.unwrap_or(true)
{
self.socket_rooms.remove(sid);
}
if let Some(mut entry) = self.rooms.get_mut(room) {
entry.value_mut().remove(sid);
}
if self.rooms.get(room).map(|e| e.value().is_empty()).unwrap_or(true) {
if self
.rooms
.get(room)
.map(|e| e.value().is_empty())
.unwrap_or(true)
{
self.rooms.remove(room);
}
@@ -246,11 +281,11 @@ impl Adapter for RedisAdapter {
server_id: self.server_id.clone(),
};
let payload = serde_json::to_vec(&msg)
.map_err(|e| AdapterError::Serialization(e.to_string()))?;
let payload =
serde_json::to_vec(&msg).map_err(|e| AdapterError::Serialization(e.to_string()))?;
self.message_bus
.publish(&format!("socket.io:{}:leave", ns), &payload)
.publish(&format!("socket.io:{}:leave", self.namespace), &payload)
.await
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
@@ -263,7 +298,12 @@ impl Adapter for RedisAdapter {
if let Some(mut entry) = self.rooms.get_mut(room) {
entry.value_mut().remove(sid);
}
if self.rooms.get(room).map(|e| e.value().is_empty()).unwrap_or(true) {
if self
.rooms
.get(room)
.map(|e| e.value().is_empty())
.unwrap_or(true)
{
self.rooms.remove(room);
}
@@ -280,6 +320,7 @@ impl Adapter for RedisAdapter {
.await
.map_err(|e| AdapterError::Redis(e.to_string()))?;
self.socket_sids.remove(sid);
self.sockets.remove(sid);
let msg = BusMessage::SocketDisconnect {
@@ -288,22 +329,43 @@ impl Adapter for RedisAdapter {
server_id: self.server_id.clone(),
};
let payload = serde_json::to_vec(&msg)
.map_err(|e| AdapterError::Serialization(e.to_string()))?;
let payload =
serde_json::to_vec(&msg).map_err(|e| AdapterError::Serialization(e.to_string()))?;
self.message_bus
.publish(&format!("socket.io:{}:disconnect", ns), &payload)
.publish(
&format!("socket.io:{}:disconnect", self.namespace),
&payload,
)
.await
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
Ok(())
}
async fn fetch_sockets(&self, opts: &BroadcastOptions) -> Result<Vec<SocketInfo>, AdapterError> {
async fn register(
&self,
socket_sid: &str,
engine_sid: &str,
_ns: &str,
) -> Result<(), AdapterError> {
self.socket_sids
.insert(socket_sid.to_string(), engine_sid.to_string());
Ok(())
}
async fn unregister(&self, socket_sid: &str, ns: &str) -> Result<(), AdapterError> {
self.del_all(socket_sid, ns).await
}
async fn fetch_sockets(
&self,
opts: &BroadcastOptions,
) -> Result<Vec<SocketInfo>, AdapterError> {
let mut result = Vec::new();
let target_sids: HashSet<String> = if opts.rooms.is_empty() {
self.sockets.iter().map(|e| e.key().clone()).collect()
self.socket_sids.iter().map(|e| e.key().clone()).collect()
} else {
let mut sids = HashSet::new();
for room in &opts.rooms {
@@ -318,7 +380,11 @@ impl Adapter for RedisAdapter {
if opts.except.contains(&sid) {
continue;
}
let rooms = self.socket_rooms.get(&sid).map(|e| e.value().clone()).unwrap_or_default();
let rooms = self
.socket_rooms
.get(&sid)
.map(|e| e.value().clone())
.unwrap_or_default();
result.push(SocketInfo {
sid: sid.clone(),
namespace: self.namespace.clone(),
@@ -330,7 +396,11 @@ impl Adapter for RedisAdapter {
}
async fn socket_rooms(&self, sid: &str) -> Result<HashSet<String>, AdapterError> {
Ok(self.socket_rooms.get(sid).map(|e| e.value().clone()).unwrap_or_default())
Ok(self
.socket_rooms
.get(sid)
.map(|e| e.value().clone())
.unwrap_or_default())
}
fn server_id(&self) -> &str {
@@ -338,7 +408,10 @@ impl Adapter for RedisAdapter {
}
async fn close(&self) -> Result<(), AdapterError> {
self.message_bus.close().await.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
self.message_bus
.close()
.await
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
Ok(())
}
}
socket/message_bus/mod.rs
+2 -2
View File
@@ -1,5 +1,5 @@
pub mod redis;
pub mod nats;
pub mod redis;
use async_trait::async_trait;
use thiserror::Error;
@@ -27,5 +27,5 @@ pub trait MessageBus: Send + Sync + 'static {
async fn close(&self) -> Result<(), MessageBusError>;
}
pub use nats::NatsMessageBus;
pub use redis::RedisMessageBus;
pub use nats::NatsMessageBus;
socket/message_bus/nats.rs
+3 -2
View File
@@ -34,7 +34,8 @@ impl MessageBus for NatsMessageBus {
async fn subscribe(&self, channel: &str) -> Result<mpsc::Receiver<Vec<u8>>, MessageBusError> {
let (tx, rx) = mpsc::channel::<Vec<u8>>(256);
let mut subscriber = self.client
let mut subscriber = self
.client
.subscribe(channel.to_string())
.await
.map_err(|e| MessageBusError::Nats(e.to_string()))?;
@@ -85,4 +86,4 @@ impl MessageBus for NatsMessageBus {
self.shutdowns.clear();
Ok(())
}
}
}
socket/message_bus/redis.rs
+5 -6
View File
@@ -13,8 +13,8 @@ pub struct RedisMessageBus {
impl RedisMessageBus {
pub async fn new(redis_url: &str) -> Result<Self, MessageBusError> {
let config = Config::from_url(redis_url)
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
let config =
Config::from_url(redis_url).map_err(|e| MessageBusError::Redis(e.to_string()))?;
let client = Client::new(config.clone(), None, None, None);
let subscriber = SubscriberClient::new(config, None, None, None);
@@ -64,9 +64,8 @@ impl MessageBus for RedisMessageBus {
tokio::spawn(async move {
while let Ok(message) = message_rx.recv().await {
if &message.channel == &channel_owned {
let data: Vec<u8> = FromValue::from_value(message.value)
.unwrap_or_default();
if message.channel == channel_owned {
let data: Vec<u8> = FromValue::from_value(message.value).unwrap_or_default();
if tx.send(data).await.is_err() {
break;
}
@@ -96,4 +95,4 @@ impl MessageBus for RedisMessageBus {
.map_err(|e| MessageBusError::Redis(e.to_string()))?;
Ok(())
}
}
}
socket/mod.rs
+11 -5
View File
@@ -5,12 +5,18 @@ pub mod packet;
pub mod parser;
pub mod server;
pub mod session_store;
#[allow(clippy::module_inception)]
pub mod socket;
pub use adapter::{Adapter, AdapterError, BroadcastOptions, BroadcastFlags, BusMessage, LocalAdapter, RedisAdapter, NatsAdapter, SocketInfo};
pub use message_bus::{MessageBus, MessageBusError, RedisMessageBus, NatsMessageBus};
pub use namespace::{is_valid_namespace, Namespace, NamespaceManager};
pub use adapter::{
Adapter, AdapterError, BroadcastFlags, BroadcastOptions, BusMessage, LocalAdapter, NatsAdapter,
RedisAdapter, SocketInfo,
};
pub use message_bus::{MessageBus, MessageBusError, NatsMessageBus, RedisMessageBus};
pub use namespace::{Namespace, NamespaceManager, is_valid_namespace};
pub use packet::{Packet, PacketType};
pub use server::{SocketServer, SocketServerBuilder};
pub use session_store::{InMemorySessionStore, RedisSessionStore, SessionError, SessionInfo, SessionStoreTrait};
pub use socket::Socket;
pub use session_store::{
InMemorySessionStore, RedisSessionStore, SessionError, SessionInfo, SessionStoreTrait,
};
pub use socket::Socket;
socket/namespace.rs
+148 -21
View File
@@ -4,12 +4,13 @@ use std::sync::Arc;
use dashmap::DashMap;
use tokio::sync::RwLock;
use crate::socket::adapter::{Adapter, BroadcastOptions, BroadcastFlags};
use crate::socket::adapter::{Adapter, BroadcastFlags, BroadcastOptions};
use crate::socket::packet::Packet;
use crate::socket::socket::Socket;
pub type EventHandler = Arc<dyn Fn(&Socket, &serde_json::Value) + Send + Sync>;
type ConnectHandler = Arc<dyn Fn(&Socket, Option<&serde_json::Value>) -> Result<(), String> + Send + Sync>;
pub type EventHandler = Arc<dyn Fn(Arc<Socket>, &serde_json::Value) + Send + Sync>;
type ConnectHandler =
Arc<dyn Fn(&Socket, Option<&serde_json::Value>) -> Result<(), String> + Send + Sync>;
pub struct Namespace {
pub path: String,
@@ -19,6 +20,8 @@ pub struct Namespace {
engine_to_socket: DashMap<String, String>,
handlers: RwLock<HashMap<String, Vec<EventHandler>>>,
connect_handler: RwLock<Option<ConnectHandler>>,
rooms: DashMap<String, HashSet<String>>,
socket_rooms: DashMap<String, HashSet<String>>,
pub(crate) adapter: RwLock<Option<Arc<dyn Adapter>>>,
}
@@ -30,6 +33,8 @@ impl Namespace {
engine_to_socket: DashMap::new(),
handlers: RwLock::new(HashMap::new()),
connect_handler: RwLock::new(None),
rooms: DashMap::new(),
socket_rooms: DashMap::new(),
adapter: RwLock::new(None),
}
}
@@ -40,11 +45,15 @@ impl Namespace {
}
/// Add a socket to this namespace. Returns Err if the connect handler rejects.
pub async fn add_socket(&self, socket: Arc<Socket>) -> Result<(), String> {
pub async fn add_socket(
&self,
socket: Arc<Socket>,
auth_data: Option<&serde_json::Value>,
) -> Result<(), String> {
// Run connect handler before adding to storage
let handler = self.connect_handler.read().await;
if let Some(ref h) = *handler {
h(&socket, None)?;
h(&socket, auth_data)?;
}
drop(handler);
@@ -53,10 +62,10 @@ impl Namespace {
// Register with adapter (socket_sid → engine_sid mapping)
let adapter = self.adapter.read().await;
if let Some(ref adapter) = *adapter {
if let Err(e) = adapter.register(&socket_sid, &engine_sid, &self.path).await {
tracing::warn!("Adapter register error for socket {}: {}", socket_sid, e);
}
if let Some(ref adapter) = *adapter
&& let Err(e) = adapter.register(&socket_sid, &engine_sid, &self.path).await
{
tracing::warn!("Adapter register error for socket {}: {}", socket_sid, e);
}
// Store socket by socket_sid, plus reverse index
@@ -69,12 +78,13 @@ impl Namespace {
pub async fn remove_socket_by_sid(&self, socket_sid: &str) {
if let Some((_, socket)) = self.sockets.remove(socket_sid) {
self.engine_to_socket.remove(&socket.engine_sid);
self.remove_socket_from_local_rooms(socket_sid);
let adapter = self.adapter.read().await;
if let Some(ref adapter) = *adapter {
if let Err(e) = adapter.del_all(socket_sid, &self.path).await {
tracing::warn!("Adapter del_all error for socket {}: {}", socket_sid, e);
}
if let Some(ref adapter) = *adapter
&& let Err(e) = adapter.del_all(socket_sid, &self.path).await
{
tracing::warn!("Adapter del_all error for socket {}: {}", socket_sid, e);
}
}
}
@@ -130,7 +140,12 @@ impl Namespace {
}
}
pub async fn emit_to_room(&self, room: &str, event: impl Into<String>, data: serde_json::Value) {
pub async fn emit_to_room(
&self,
room: &str,
event: impl Into<String>,
data: serde_json::Value,
) {
let event_name = event.into();
let packet = Packet::event(&self.path, serde_json::json!([event_name, data]), None);
@@ -145,20 +160,64 @@ impl Namespace {
tracing::warn!("Adapter broadcast to room error: {}", e);
}
} else {
self.emit_local(&packet);
self.emit_local_to_room(&packet, room, &HashSet::new());
}
}
pub fn emit_local(&self, packet: &Packet) {
for entry in self.sockets.iter() {
let socket = entry.value();
if socket.send_packet(packet).is_err() {
tracing::warn!("Failed to send event to socket {}", socket.sid);
self.send_local_packet(entry.value(), packet);
}
}
pub fn emit_local_filtered(&self, packet: &Packet, opts: &BroadcastOptions) {
if opts.rooms.is_empty() {
for entry in self.sockets.iter() {
if !opts.except.contains(entry.key()) {
self.send_local_packet(entry.value(), packet);
}
}
return;
}
let mut target_sids = HashSet::new();
for room in &opts.rooms {
if let Some(room_sids) = self.rooms.get(room) {
target_sids.extend(room_sids.value().iter().cloned());
}
}
for sid in target_sids {
if opts.except.contains(&sid) {
continue;
}
if let Some(socket) = self.get_socket(&sid) {
self.send_local_packet(&socket, packet);
}
}
}
pub async fn emit_to(&self, socket_sid: &str, event: impl Into<String>, data: serde_json::Value) {
fn emit_local_to_room(&self, packet: &Packet, room: &str, except: &HashSet<String>) {
let opts = BroadcastOptions {
rooms: HashSet::from([room.to_string()]),
except: except.clone(),
flags: BroadcastFlags::default(),
};
self.emit_local_filtered(packet, &opts);
}
fn send_local_packet(&self, socket: &Socket, packet: &Packet) {
if socket.send_packet(packet).is_err() {
tracing::warn!("Failed to send event to socket {}", socket.sid);
}
}
pub async fn emit_to(
&self,
socket_sid: &str,
event: impl Into<String>,
data: serde_json::Value,
) {
if let Some(socket) = self.get_socket(socket_sid) {
let event_name = event.into();
let packet = Packet::event(&self.path, serde_json::json!([event_name, data]), None);
@@ -168,11 +227,79 @@ impl Namespace {
}
}
pub async fn handle_event(&self, socket: &Socket, event: &str, data: &serde_json::Value) {
pub async fn handle_event(&self, socket: Arc<Socket>, event: &str, data: &serde_json::Value) {
let handlers = self.handlers.read().await;
if let Some(event_handlers) = handlers.get(event) {
for handler in event_handlers {
handler(socket, data);
handler(Arc::clone(&socket), data);
}
}
}
pub async fn join_room(&self, socket_sid: &str, room: &str) -> crate::ImksResult<()> {
if !self.sockets.contains_key(socket_sid) {
return Err(crate::ImksError::SocketNotFound(socket_sid.to_string()));
}
self.rooms
.entry(room.to_string())
.or_default()
.value_mut()
.insert(socket_sid.to_string());
self.socket_rooms
.entry(socket_sid.to_string())
.or_default()
.value_mut()
.insert(room.to_string());
let adapter = self.adapter.read().await;
if let Some(ref adapter) = *adapter
&& let Err(e) = adapter.add(socket_sid, room, &self.path).await
{
self.remove_local_room(socket_sid, room);
return Err(e.into());
}
Ok(())
}
pub async fn leave_room(&self, socket_sid: &str, room: &str) -> crate::ImksResult<()> {
let adapter = self.adapter.read().await;
if let Some(ref adapter) = *adapter {
adapter.del(socket_sid, room, &self.path).await?;
}
self.remove_local_room(socket_sid, room);
Ok(())
}
fn remove_local_room(&self, socket_sid: &str, room: &str) {
if let Some(mut sids) = self.rooms.get_mut(room) {
sids.value_mut().remove(socket_sid);
if sids.value().is_empty() {
drop(sids);
self.rooms.remove(room);
}
}
if let Some(mut rooms) = self.socket_rooms.get_mut(socket_sid) {
rooms.value_mut().remove(room);
if rooms.value().is_empty() {
drop(rooms);
self.socket_rooms.remove(socket_sid);
}
}
}
fn remove_socket_from_local_rooms(&self, socket_sid: &str) {
if let Some((_, rooms)) = self.socket_rooms.remove(socket_sid) {
for room in rooms {
if let Some(mut sids) = self.rooms.get_mut(&room) {
sids.value_mut().remove(socket_sid);
if sids.value().is_empty() {
drop(sids);
self.rooms.remove(&room);
}
}
}
}
}
socket/parser.rs
+31 -29
View File
@@ -24,19 +24,18 @@ pub fn encode(packet: &Packet) -> String {
if let Some(ref data) = packet.data {
if packet.has_binary() {
let data_with_placeholders = replace_binary_with_placeholders(data, packet.attachment_count());
let encoded_data = serde_json::to_string(&data_with_placeholders)
.unwrap_or_else(|e| {
tracing::error!("Failed to serialize socket packet data: {}", e);
"null".to_string()
});
let data_with_placeholders =
replace_binary_with_placeholders(data, packet.attachment_count());
let encoded_data = serde_json::to_string(&data_with_placeholders).unwrap_or_else(|e| {
tracing::error!("Failed to serialize socket packet data: {}", e);
"null".to_string()
});
result.push_str(&encoded_data);
} else {
let encoded_data = serde_json::to_string(data)
.unwrap_or_else(|e| {
tracing::error!("Failed to serialize socket packet data: {}", e);
"null".to_string()
});
let encoded_data = serde_json::to_string(data).unwrap_or_else(|e| {
tracing::error!("Failed to serialize socket packet data: {}", e);
"null".to_string()
});
result.push_str(&encoded_data);
}
}
@@ -67,7 +66,8 @@ pub fn decode(input: &str) -> Result<Packet, PacketError> {
let type_char = chars.next().ok_or(PacketError::Empty)?;
let packet_type = PacketType::try_from(type_char)?;
let attachment_count = if matches!(packet_type, PacketType::BinaryEvent | PacketType::BinaryAck) {
let attachment_count = if matches!(packet_type, PacketType::BinaryEvent | PacketType::BinaryAck)
{
let mut count_str = String::new();
while let Some(&c) = chars.peek() {
if c == '-' {
@@ -126,7 +126,11 @@ pub fn decode(input: &str) -> Result<Packet, PacketError> {
id,
attachments: Vec::new(),
// Store attachment_count for binary packets; actual attachments come via decode_with_attachments
expected_attachments: if attachment_count > 0 { Some(attachment_count) } else { None },
expected_attachments: if attachment_count > 0 {
Some(attachment_count)
} else {
None
},
})
}
@@ -144,10 +148,10 @@ pub fn decode_with_attachments(
packet.attachments = attachments;
packet.expected_attachments = None;
if packet.has_binary() {
if let Some(ref data) = packet.data {
packet.data = Some(replace_placeholders_with_binary(data, &packet.attachments));
}
if packet.has_binary()
&& let Some(ref data) = packet.data
{
packet.data = Some(replace_placeholders_with_binary(data, &packet.attachments));
}
Ok(packet)
@@ -204,12 +208,12 @@ fn replace_binary_with_placeholders(value: &Value, total_attachments: usize) ->
}
}
fn replace_binary_with_placeholders_inner(value: &Value, placeholder_idx: &mut usize) -> Value {
fn replace_binary_with_placeholders_inner(value: &Value, _placeholder_idx: &mut usize) -> Value {
match value {
Value::Array(arr) => {
let new_arr: Vec<Value> = arr
.iter()
.map(|v| replace_binary_with_placeholders_inner(v, placeholder_idx))
.map(|v| replace_binary_with_placeholders_inner(v, _placeholder_idx))
.collect();
Value::Array(new_arr)
}
@@ -218,7 +222,7 @@ fn replace_binary_with_placeholders_inner(value: &Value, placeholder_idx: &mut u
for (k, v) in map {
new_map.insert(
k.clone(),
replace_binary_with_placeholders_inner(v, placeholder_idx),
replace_binary_with_placeholders_inner(v, _placeholder_idx),
);
}
Value::Object(new_map)
@@ -236,15 +240,13 @@ fn replace_placeholders_with_binary(value: &Value, attachments: &[Vec<u8>]) -> V
// Check if this is a placeholder object: { "_placeholder": true, "num": N }
if let (Some(Value::Bool(true)), Some(Value::Number(num))) =
(map.get("_placeholder"), map.get("num"))
&& let Some(idx) = num.as_u64()
&& let Some(attachment) = attachments.get(idx as usize)
{
if let Some(idx) = num.as_u64() {
if let Some(attachment) = attachments.get(idx as usize) {
return Value::String(base64::Engine::encode(
&base64::engine::general_purpose::STANDARD,
attachment,
));
}
}
return Value::String(base64::Engine::encode(
&base64::engine::general_purpose::STANDARD,
attachment,
));
}
let mut new_map = serde_json::Map::new();
@@ -389,4 +391,4 @@ mod tests {
assert_eq!(packet.expected_attachments, Some(1));
assert_eq!(packet.namespace, "/");
}
}
}
socket/server.rs
+55 -41
View File
@@ -103,8 +103,14 @@ impl SocketServerBuilder {
let adapter = adapter_clone.clone();
tokio::spawn(async move {
handle_engine_message(
sid, engine_packet, &namespaces, &socket_txs, &engine_store, &adapter,
).await;
sid,
engine_packet,
&namespaces,
&socket_txs,
&engine_store,
&adapter,
)
.await;
});
},
));
@@ -136,10 +142,18 @@ async fn handle_engine_message(
adapter: &Arc<dyn Adapter>,
) {
if let EnginePacketData::Text(ref text) = engine_packet.data {
if let Ok(socket_packet) = parser::decode(text) {
match socket_packet.packet_type {
match parser::decode(text) {
Ok(socket_packet) => match socket_packet.packet_type {
PacketType::Connect => {
handle_connect(&engine_sid, &socket_packet, namespaces, socket_txs, engine_store, adapter).await;
handle_connect(
&engine_sid,
&socket_packet,
namespaces,
socket_txs,
engine_store,
adapter,
)
.await;
}
PacketType::Disconnect => {
handle_disconnect(&engine_sid, &socket_packet, namespaces, socket_txs);
@@ -151,6 +165,9 @@ async fn handle_engine_message(
handle_ack(&engine_sid, &socket_packet);
}
_ => {}
},
Err(e) => {
tracing::warn!(engine_sid = %engine_sid, error = %e, "Invalid Socket.IO packet");
}
}
}
@@ -166,22 +183,21 @@ async fn handle_connect(
) {
// Validate namespace path to prevent DoS via arbitrary namespace creation
if !crate::socket::namespace::is_valid_namespace(&packet.namespace) {
tracing::warn!("Rejected connect with invalid namespace: {}", packet.namespace);
tracing::warn!(
"Rejected connect with invalid namespace: {}",
packet.namespace
);
return;
}
let namespace = namespaces.get_or_create_namespace(&packet.namespace);
// Ensure newly created namespaces get the shared adapter
// Ensure newly created namespaces get the shared adapter before registration.
{
let ns_adapter = namespace.adapter.read().await;
if ns_adapter.is_none() {
drop(ns_adapter);
let adapter_ref = adapter.clone();
let ns_clone = namespace.clone();
tokio::spawn(async move {
ns_clone.set_adapter(adapter_ref).await;
});
namespace.set_adapter(adapter.clone()).await;
}
}
@@ -198,7 +214,10 @@ async fn handle_connect(
// Run connect handler and add to namespace.
// If the handler rejects, clean up and do NOT send a Connect response.
if let Err(msg) = namespace.add_socket(socket.clone()).await {
if let Err(msg) = namespace
.add_socket(socket.clone(), packet.data.as_ref())
.await
{
tracing::warn!("Socket {} connection rejected: {}", socket_sid, msg);
socket_txs.remove(&socket_sid);
return;
@@ -227,7 +246,9 @@ async fn handle_connect(
}
// Forwarding task ended — ensure socket is cleaned up from namespace
socket_txs_clone.remove(&socket_sid_clone);
namespace_clone.remove_socket_by_sid(&socket_sid_clone).await;
namespace_clone
.remove_socket_by_sid(&socket_sid_clone)
.await;
});
// Send Connect response (only after handler passed)
@@ -260,34 +281,27 @@ fn handle_disconnect(
}
}
fn handle_event(
engine_sid: &str,
packet: &Packet,
namespaces: &Arc<NamespaceManager>,
) {
if let Some(namespace) = namespaces.get_namespace(&packet.namespace) {
if let Some(socket) = namespace.get_socket_by_engine_sid(engine_sid) {
if let Some(ref data) = packet.data {
if let Some(arr) = data.as_array() {
if let Some(event) = arr.first().and_then(|v| v.as_str()) {
let event_data = if arr.len() > 1 {
serde_json::Value::Array(arr[1..].to_vec())
} else {
serde_json::Value::Null
};
fn handle_event(engine_sid: &str, packet: &Packet, namespaces: &Arc<NamespaceManager>) {
if let Some(namespace) = namespaces.get_namespace(&packet.namespace)
&& let Some(socket) = namespace.get_socket_by_engine_sid(engine_sid)
&& let Some(ref data) = packet.data
&& let Some(arr) = data.as_array()
&& let Some(event) = arr.first().and_then(|v| v.as_str())
{
let event_data = if arr.len() > 1 {
serde_json::Value::Array(arr[1..].to_vec())
} else {
serde_json::Value::Null
};
let namespace_clone = namespace.clone();
let event = event.to_string();
let socket_clone = socket.clone();
tokio::spawn(async move {
namespace_clone
.handle_event(&socket_clone, &event, &event_data)
.await;
});
}
}
}
}
let namespace_clone = namespace.clone();
let event = event.to_string();
let socket_clone = socket.clone();
tokio::spawn(async move {
namespace_clone
.handle_event(socket_clone, &event, &event_data)
.await;
});
}
}
socket/session_store/memory.rs
+7 -2
View File
@@ -33,7 +33,12 @@ fn now_millis() -> u64 {
#[async_trait]
impl SessionStoreTrait for InMemorySessionStore {
async fn create(&self, sid: &str, transport: &str, server_id: &str) -> Result<(), SessionError> {
async fn create(
&self,
sid: &str,
transport: &str,
server_id: &str,
) -> Result<(), SessionError> {
let info = SessionInfo {
sid: sid.to_string(),
transport: transport.to_string(),
@@ -85,4 +90,4 @@ impl SessionStoreTrait for InMemorySessionStore {
async fn exists(&self, sid: &str) -> Result<bool, SessionError> {
Ok(self.sessions.contains_key(sid))
}
}
}
socket/session_store/mod.rs
+3 -2
View File
@@ -28,7 +28,8 @@ pub struct SessionInfo {
#[async_trait]
pub trait SessionStoreTrait: Send + Sync + 'static {
async fn create(&self, sid: &str, transport: &str, server_id: &str) -> Result<(), SessionError>;
async fn create(&self, sid: &str, transport: &str, server_id: &str)
-> Result<(), SessionError>;
async fn get(&self, sid: &str) -> Result<Option<SessionInfo>, SessionError>;
async fn set_state(&self, sid: &str, state: &str) -> Result<(), SessionError>;
async fn set_transport(&self, sid: &str, transport: &str) -> Result<(), SessionError>;
@@ -38,4 +39,4 @@ pub trait SessionStoreTrait: Send + Sync + 'static {
}
pub use memory::InMemorySessionStore;
pub use redis::RedisSessionStore;
pub use redis::RedisSessionStore;
socket/session_store/redis.rs
+19 -6
View File
@@ -36,7 +36,12 @@ impl RedisSessionStore {
#[async_trait]
impl SessionStoreTrait for RedisSessionStore {
async fn create(&self, sid: &str, transport: &str, server_id: &str) -> Result<(), SessionError> {
async fn create(
&self,
sid: &str,
transport: &str,
server_id: &str,
) -> Result<(), SessionError> {
let key = self.key(sid);
let now = now_millis();
@@ -67,7 +72,8 @@ impl SessionStoreTrait for RedisSessionStore {
// Use hgetall directly — if the key doesn't exist Redis returns an empty map.
// This avoids the TOCTOU race between EXISTS and HGETALL.
let values: std::collections::HashMap<String, String> = self.client
let values: std::collections::HashMap<String, String> = self
.client
.hgetall::<std::collections::HashMap<String, String>, _>(&key)
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
@@ -81,8 +87,14 @@ impl SessionStoreTrait for RedisSessionStore {
transport: values.get("transport").cloned().unwrap_or_default(),
state: values.get("state").cloned().unwrap_or_default(),
server_id: values.get("server_id").cloned().unwrap_or_default(),
created_at: values.get("created_at").and_then(|v| v.parse::<u64>().ok()).unwrap_or(0),
last_ping: values.get("last_ping").and_then(|v| v.parse::<u64>().ok()).unwrap_or(0),
created_at: values
.get("created_at")
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(0),
last_ping: values
.get("last_ping")
.and_then(|v| v.parse::<u64>().ok())
.unwrap_or(0),
};
Ok(Some(info))
@@ -154,11 +166,12 @@ impl SessionStoreTrait for RedisSessionStore {
async fn exists(&self, sid: &str) -> Result<bool, SessionError> {
let key = self.key(sid);
let exists: bool = self.client
let exists: bool = self
.client
.exists::<bool, _>(&key)
.await
.map_err(|e| SessionError::Redis(e.to_string()))?;
Ok(exists)
}
}
}
socket/socket.rs
+27 -2
View File
@@ -1,6 +1,8 @@
use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::mpsc;
use uuid::Uuid;
use crate::socket::packet::Packet;
@@ -8,10 +10,13 @@ pub struct Socket {
pub sid: String,
pub namespace: String,
pub engine_sid: String,
/// Authenticated user ID, set once during `on_connect`.
user_id: OnceLock<Uuid>,
ack_id: AtomicU64,
tx: mpsc::Sender<Packet>,
}
#[allow(clippy::result_large_err)]
impl Socket {
pub fn new(
sid: String,
@@ -24,10 +29,22 @@ impl Socket {
namespace,
engine_sid,
ack_id: AtomicU64::new(0),
user_id: OnceLock::new(),
tx,
}
}
/// Set the authenticated user ID after JWT verification.
/// Safe to call once; subsequent calls are ignored.
pub fn set_user_id(&self, id: Uuid) {
let _ = self.user_id.set(id);
}
/// Get the authenticated user ID, if set.
pub fn user_id(&self) -> Option<Uuid> {
self.user_id.get().copied()
}
pub fn next_ack_id(&self) -> u64 {
self.ack_id.fetch_add(1, Ordering::SeqCst)
}
@@ -36,7 +53,11 @@ impl Socket {
self.tx.try_send(packet.clone())
}
pub fn emit(&self, event: impl Into<String>, data: serde_json::Value) -> Result<(), mpsc::error::TrySendError<Packet>> {
pub fn emit(
&self,
event: impl Into<String>,
data: serde_json::Value,
) -> Result<(), mpsc::error::TrySendError<Packet>> {
let packet = Packet::event(
&self.namespace,
serde_json::json!([event.into(), data]),
@@ -65,7 +86,11 @@ impl Socket {
self.send_packet(&packet)
}
pub fn send_ack(&self, id: u64, data: serde_json::Value) -> Result<(), mpsc::error::TrySendError<Packet>> {
pub fn send_ack(
&self,
id: u64,
data: serde_json::Value,
) -> Result<(), mpsc::error::TrySendError<Packet>> {
let packet = Packet::ack(&self.namespace, data, id);
self.send_packet(&packet)
}