feat(auth): add authentication protocol definitions and build configuration
- Add TokenClaims message for JWT payload structure with user id, issuer, timestamps, and scopes - Implement IssueTokenRequest/Response for creating access and refresh tokens with TTL support - Create RefreshTokenRequest/Response for token rotation functionality - Define RevokeTokenRequest/Response with support for single token or user-wide revocation - Add VerifyTokenRequest/Response for validating JWT tokens with detailed claims information - Implement signing key distribution system with GetSigningKeysRequest/Response - Create TokenService gRPC service with IssueToken, RefreshToken, RevokeToken, VerifyToken, and GetSigningKeys methods - Add build.rs configuration to compile proto files using tonic_prost_build - Include channel, channel_settings, member, and permission protocol definitions for IM services - Generate Rust code bindings through pb/core.rs and pb/im.rs modules
This commit is contained in:
zhenyi
@@ -0,0 +1,199 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use dashmap::DashMap;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::socket::adapter::{Adapter, AdapterError, BroadcastOptions, SocketInfo};
|
||||
use crate::socket::packet::Packet;
|
||||
|
||||
pub struct LocalAdapter {
|
||||
server_id: String,
|
||||
rooms: Arc<DashMap<String, HashSet<String>>>,
|
||||
socket_rooms: Arc<DashMap<String, HashSet<String>>>,
|
||||
/// socket_sid → engine_sid
|
||||
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>,
|
||||
}
|
||||
|
||||
impl LocalAdapter {
|
||||
pub fn new(
|
||||
send_fn: impl Fn(&str, &Packet) -> Result<(), String> + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
server_id: Uuid::new_v4().to_string(),
|
||||
rooms: Arc::new(DashMap::new()),
|
||||
socket_rooms: Arc::new(DashMap::new()),
|
||||
socket_sids: Arc::new(DashMap::new()),
|
||||
socket_namespace: Arc::new(DashMap::new()),
|
||||
send_fn: Arc::new(send_fn),
|
||||
}
|
||||
}
|
||||
|
||||
fn room_key(ns: &str, room: &str) -> String {
|
||||
format!("{}:{}", ns, room)
|
||||
}
|
||||
|
||||
/// Collect socket SIDs matching the broadcast options, scoped to the given namespace.
|
||||
fn collect_matching_sids(&self, opts: &BroadcastOptions, namespace: &str) -> Vec<String> {
|
||||
if opts.rooms.is_empty() {
|
||||
// Broadcast to all sockets in this namespace only
|
||||
self.socket_sids
|
||||
.iter()
|
||||
.filter(|e| {
|
||||
self.socket_namespace
|
||||
.get(e.key())
|
||||
.map(|ns| ns.value() == namespace)
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.map(|e| e.key().clone())
|
||||
.collect()
|
||||
} else {
|
||||
let mut sids = HashSet::new();
|
||||
for room in &opts.rooms {
|
||||
let key = Self::room_key(namespace, room);
|
||||
if let Some(entry) = self.rooms.get(&key) {
|
||||
for sid in entry.value() {
|
||||
sids.insert(sid.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
sids.into_iter().collect()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Adapter for LocalAdapter {
|
||||
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 {
|
||||
if opts.except.contains(sid) {
|
||||
continue;
|
||||
}
|
||||
// socket_sids maps socket SID -> engine SID
|
||||
if let Some(entry) = self.socket_sids.get(sid) {
|
||||
let engine_sid = entry.value();
|
||||
let result = (self.send_fn)(engine_sid, packet);
|
||||
if let Err(e) = result {
|
||||
tracing::warn!("Failed to broadcast to {}: {}", sid, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
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());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn unregister(&self, socket_sid: &str, ns: &str) -> Result<(), AdapterError> {
|
||||
self.del_all(socket_sid, ns).await
|
||||
}
|
||||
|
||||
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());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn del(&self, sid: &str, room: &str, ns: &str) -> Result<(), AdapterError> {
|
||||
let key = Self::room_key(ns, room);
|
||||
if let Some(mut room_sids) = self.rooms.get_mut(&key) {
|
||||
room_sids.value_mut().remove(sid);
|
||||
if room_sids.value_mut().is_empty() {
|
||||
drop(room_sids);
|
||||
self.rooms.remove(&key);
|
||||
}
|
||||
}
|
||||
if let Some(mut rooms) = self.socket_rooms.get_mut(sid) {
|
||||
rooms.value_mut().remove(room);
|
||||
if rooms.value_mut().is_empty() {
|
||||
drop(rooms);
|
||||
self.socket_rooms.remove(sid);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn del_all(&self, sid: &str, ns: &str) -> Result<(), AdapterError> {
|
||||
if let Some((_, rooms)) = self.socket_rooms.remove(sid) {
|
||||
for room in &rooms {
|
||||
let key = Self::room_key(ns, room);
|
||||
if let Some(mut room_sids) = self.rooms.get_mut(&key) {
|
||||
room_sids.value_mut().remove(sid);
|
||||
if room_sids.value_mut().is_empty() {
|
||||
drop(room_sids);
|
||||
self.rooms.remove(&key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
self.socket_sids.remove(sid);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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() {
|
||||
self.socket_sids.iter().map(|e| e.key().clone()).collect()
|
||||
} else {
|
||||
let mut sids_set = HashSet::new();
|
||||
for room in &opts.rooms {
|
||||
for entry in self.rooms.iter() {
|
||||
if entry.key().ends_with(&format!(":{}", room)) {
|
||||
for sid in entry.value() {
|
||||
sids_set.insert(sid.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
sids_set.into_iter().collect()
|
||||
};
|
||||
let mut result = Vec::new();
|
||||
for sid in &sids {
|
||||
if opts.except.contains(sid) {
|
||||
continue;
|
||||
}
|
||||
if self.socket_sids.contains_key(sid) {
|
||||
let namespace = self.socket_namespace
|
||||
.get(sid)
|
||||
.map(|r| r.value().clone())
|
||||
.unwrap_or_default();
|
||||
let rooms = self.socket_rooms
|
||||
.get(sid)
|
||||
.map(|r| r.value().clone())
|
||||
.unwrap_or_default();
|
||||
result.push(SocketInfo {
|
||||
sid: sid.clone(),
|
||||
namespace,
|
||||
rooms,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn socket_rooms(&self, sid: &str) -> Result<HashSet<String>, AdapterError> {
|
||||
Ok(self.socket_rooms
|
||||
.get(sid)
|
||||
.map(|r| r.value().clone())
|
||||
.unwrap_or_default())
|
||||
}
|
||||
|
||||
fn server_id(&self) -> &str {
|
||||
&self.server_id
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), AdapterError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
pub mod local;
|
||||
pub mod redis;
|
||||
pub mod nats;
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::socket::packet::Packet;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum AdapterError {
|
||||
#[error("Redis error: {0}")]
|
||||
Redis(String),
|
||||
#[error("NATS error: {0}")]
|
||||
Nats(String),
|
||||
#[error("Message bus error: {0}")]
|
||||
MessageBus(String),
|
||||
#[error("Serialization error: {0}")]
|
||||
Serialization(String),
|
||||
#[error("Room error: {0}")]
|
||||
Room(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize, PartialEq)]
|
||||
pub struct BroadcastOptions {
|
||||
pub rooms: HashSet<String>,
|
||||
pub except: HashSet<String>,
|
||||
pub flags: BroadcastFlags,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize, PartialEq)]
|
||||
pub struct BroadcastFlags {
|
||||
pub local_only: bool,
|
||||
pub broadcast: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct SocketInfo {
|
||||
pub sid: String,
|
||||
pub namespace: String,
|
||||
pub rooms: HashSet<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq)]
|
||||
pub enum BusMessage {
|
||||
Broadcast {
|
||||
namespace: String,
|
||||
packet: String,
|
||||
opts: BroadcastOptions,
|
||||
server_id: String,
|
||||
},
|
||||
SocketJoin {
|
||||
namespace: String,
|
||||
sid: String,
|
||||
room: String,
|
||||
server_id: String,
|
||||
},
|
||||
SocketLeave {
|
||||
namespace: String,
|
||||
sid: String,
|
||||
room: String,
|
||||
server_id: String,
|
||||
},
|
||||
SocketDisconnect {
|
||||
namespace: String,
|
||||
sid: String,
|
||||
server_id: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait Adapter: Send + Sync + 'static {
|
||||
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 socket_rooms(&self, sid: &str) -> Result<HashSet<String>, AdapterError>;
|
||||
fn server_id(&self) -> &str;
|
||||
async fn close(&self) -> Result<(), AdapterError>;
|
||||
|
||||
/// 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> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Unregister a socket from the adapter, removing all local mappings.
|
||||
async fn unregister(&self, _socket_sid: &str, _ns: &str) -> Result<(), AdapterError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub use local::LocalAdapter;
|
||||
pub use redis::RedisAdapter;
|
||||
pub use nats::NatsAdapter;
|
||||
@@ -0,0 +1,302 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use dashmap::DashMap;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::socket::adapter::{Adapter, AdapterError, BroadcastOptions, BusMessage, SocketInfo};
|
||||
use crate::socket::message_bus::MessageBus;
|
||||
use crate::socket::packet::Packet;
|
||||
use crate::socket::parser;
|
||||
use crate::socket::socket::Socket;
|
||||
|
||||
/// Handle incoming bus messages from other servers.
|
||||
/// 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>,
|
||||
server_id: &str,
|
||||
) {
|
||||
match msg {
|
||||
BusMessage::Broadcast { namespace: _, packet, opts, server_id: sender_id } => {
|
||||
if sender_id == server_id {
|
||||
return;
|
||||
}
|
||||
if let Ok(decoded_packet) = parser::decode(&packet) {
|
||||
on_local_broadcast(&decoded_packet, &opts);
|
||||
}
|
||||
}
|
||||
// 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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// NATS-based adapter that manages room state locally and uses NATS
|
||||
/// for cross-server broadcast only. Does NOT depend on Redis.
|
||||
pub struct NatsAdapter {
|
||||
message_bus: Arc<dyn MessageBus>,
|
||||
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 delivery
|
||||
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>,
|
||||
}
|
||||
|
||||
impl NatsAdapter {
|
||||
pub fn new(
|
||||
message_bus: Arc<dyn MessageBus>,
|
||||
server_id: String,
|
||||
namespace: String,
|
||||
on_local_broadcast: Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
|
||||
) -> Self {
|
||||
Self {
|
||||
message_bus,
|
||||
server_id,
|
||||
namespace,
|
||||
on_local_broadcast,
|
||||
room_subscribers: DashMap::new(),
|
||||
socket_rooms: DashMap::new(),
|
||||
rooms: DashMap::new(),
|
||||
socket_sids: DashMap::new(),
|
||||
sockets: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn init(&self) -> Result<(), AdapterError> {
|
||||
let channels = ["broadcast", "join", "leave", "disconnect"];
|
||||
let prefix = format!("socket.io:{}:", self.namespace);
|
||||
|
||||
for channel_type in channels {
|
||||
let subject = format!("{}{}", prefix, channel_type);
|
||||
match self.message_bus.subscribe(&subject).await {
|
||||
Ok(rx) => {
|
||||
self.room_subscribers.insert(channel_type.to_string(), rx);
|
||||
}
|
||||
Err(e) => return Err(AdapterError::MessageBus(e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
self.spawn_listener();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn spawn_listener(&self) {
|
||||
let server_id = self.server_id.clone();
|
||||
let on_local_broadcast = self.on_local_broadcast.clone();
|
||||
|
||||
let mut broadcast_rx = self.room_subscribers.remove("broadcast").map(|(_, rx)| rx);
|
||||
let mut join_rx = self.room_subscribers.remove("join").map(|(_, rx)| rx);
|
||||
let mut leave_rx = self.room_subscribers.remove("leave").map(|(_, rx)| rx);
|
||||
let mut disconnect_rx = self.room_subscribers.remove("disconnect").map(|(_, rx)| rx);
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
Some(data) = async { broadcast_rx.as_mut()?.recv().await } => {
|
||||
if let Ok(msg) = serde_json::from_slice::<BusMessage>(&data) {
|
||||
handle_bus_message(msg, &on_local_broadcast, &server_id).await;
|
||||
}
|
||||
}
|
||||
Some(data) = async { join_rx.as_mut()?.recv().await } => {
|
||||
if let Ok(msg) = serde_json::from_slice::<BusMessage>(&data) {
|
||||
handle_bus_message(msg, &on_local_broadcast, &server_id).await;
|
||||
}
|
||||
}
|
||||
Some(data) = async { leave_rx.as_mut()?.recv().await } => {
|
||||
if let Ok(msg) = serde_json::from_slice::<BusMessage>(&data) {
|
||||
handle_bus_message(msg, &on_local_broadcast, &server_id).await;
|
||||
}
|
||||
}
|
||||
Some(data) = async { disconnect_rx.as_mut()?.recv().await } => {
|
||||
if let Ok(msg) = serde_json::from_slice::<BusMessage>(&data) {
|
||||
handle_bus_message(msg, &on_local_broadcast, &server_id).await;
|
||||
}
|
||||
}
|
||||
else => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Adapter for NatsAdapter {
|
||||
async fn broadcast(&self, packet: &Packet, opts: &BroadcastOptions) -> Result<(), AdapterError> {
|
||||
if opts.flags.local_only {
|
||||
(self.on_local_broadcast)(packet, opts);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let msg = BusMessage::Broadcast {
|
||||
namespace: self.namespace.clone(),
|
||||
packet: parser::encode(packet),
|
||||
opts: opts.clone(),
|
||||
server_id: self.server_id.clone(),
|
||||
};
|
||||
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
|
||||
(self.on_local_broadcast)(packet, opts);
|
||||
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());
|
||||
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()); })
|
||||
.or_insert_with(|| HashSet::from([room.to_string()]));
|
||||
|
||||
self.rooms
|
||||
.entry(room.to_string())
|
||||
.and_modify(|set| { set.insert(sid.to_string()); })
|
||||
.or_insert_with(|| HashSet::from([sid.to_string()]));
|
||||
|
||||
let msg = BusMessage::SocketJoin {
|
||||
namespace: self.namespace.clone(),
|
||||
sid: sid.to_string(),
|
||||
room: room.to_string(),
|
||||
server_id: self.server_id.clone(),
|
||||
};
|
||||
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn del(&self, sid: &str, room: &str, _ns: &str) -> Result<(), AdapterError> {
|
||||
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) {
|
||||
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) {
|
||||
self.rooms.remove(room);
|
||||
}
|
||||
|
||||
let msg = BusMessage::SocketLeave {
|
||||
namespace: self.namespace.clone(),
|
||||
sid: sid.to_string(),
|
||||
room: room.to_string(),
|
||||
server_id: self.server_id.clone(),
|
||||
};
|
||||
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn del_all(&self, sid: &str, _ns: &str) -> Result<(), AdapterError> {
|
||||
if let Some((_, rooms)) = self.socket_rooms.remove(sid) {
|
||||
for room in &rooms {
|
||||
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) {
|
||||
self.rooms.remove(room);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.socket_sids.remove(sid);
|
||||
self.sockets.remove(sid);
|
||||
|
||||
let msg = BusMessage::SocketDisconnect {
|
||||
namespace: self.namespace.clone(),
|
||||
sid: sid.to_string(),
|
||||
server_id: self.server_id.clone(),
|
||||
};
|
||||
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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.socket_sids.iter().map(|e| e.key().clone()).collect()
|
||||
} else {
|
||||
let mut sids = HashSet::new();
|
||||
for room in &opts.rooms {
|
||||
if let Some(entry) = self.rooms.get(room) {
|
||||
sids.extend(entry.value().iter().cloned());
|
||||
}
|
||||
}
|
||||
sids
|
||||
};
|
||||
|
||||
for sid in target_sids {
|
||||
if opts.except.contains(&sid) {
|
||||
continue;
|
||||
}
|
||||
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(),
|
||||
rooms,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
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())
|
||||
}
|
||||
|
||||
fn server_id(&self) -> &str {
|
||||
&self.server_id
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), AdapterError> {
|
||||
self.message_bus.close().await.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,344 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use dashmap::DashMap;
|
||||
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::message_bus::MessageBus;
|
||||
use crate::socket::packet::Packet;
|
||||
use crate::socket::parser;
|
||||
use crate::socket::socket::Socket;
|
||||
|
||||
const KEY_PREFIX_ROOMS: &str = "socket.io:rooms";
|
||||
const KEY_PREFIX_SOCKET_ROOMS: &str = "socket.io:socket_rooms";
|
||||
|
||||
fn room_key(ns: &str, room: &str) -> String {
|
||||
format!("{}:{}:{}", KEY_PREFIX_ROOMS, ns, room)
|
||||
}
|
||||
|
||||
fn socket_rooms_key(ns: &str, sid: &str) -> String {
|
||||
format!("{}:{}:{}", KEY_PREFIX_SOCKET_ROOMS, ns, sid)
|
||||
}
|
||||
|
||||
/// Handle incoming bus messages from other servers.
|
||||
/// 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>,
|
||||
server_id: &str,
|
||||
) {
|
||||
match msg {
|
||||
BusMessage::Broadcast { namespace: _, packet, opts, server_id: sender_id } => {
|
||||
if sender_id == server_id {
|
||||
return;
|
||||
}
|
||||
if let Ok(decoded_packet) = parser::decode(&packet) {
|
||||
on_local_broadcast(&decoded_packet, &opts);
|
||||
}
|
||||
}
|
||||
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;
|
||||
}
|
||||
// No duplicate Redis writes — the sender already persisted the state change
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RedisAdapter {
|
||||
message_bus: Arc<dyn MessageBus>,
|
||||
redis_client: Client,
|
||||
room_subscribers: DashMap<String, mpsc::Receiver<Vec<u8>>>,
|
||||
socket_rooms: DashMap<String, HashSet<String>>,
|
||||
rooms: DashMap<String, HashSet<String>>,
|
||||
sockets: DashMap<String, Arc<Socket>>,
|
||||
server_id: String,
|
||||
namespace: String,
|
||||
on_local_broadcast: Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
|
||||
}
|
||||
|
||||
impl RedisAdapter {
|
||||
pub fn new(
|
||||
message_bus: Arc<dyn MessageBus>,
|
||||
redis_client: Client,
|
||||
server_id: String,
|
||||
namespace: String,
|
||||
on_local_broadcast: Arc<dyn Fn(&Packet, &BroadcastOptions) + Send + Sync + 'static>,
|
||||
) -> Self {
|
||||
Self {
|
||||
message_bus,
|
||||
redis_client,
|
||||
server_id,
|
||||
namespace,
|
||||
on_local_broadcast,
|
||||
room_subscribers: DashMap::new(),
|
||||
socket_rooms: DashMap::new(),
|
||||
rooms: DashMap::new(),
|
||||
sockets: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn init(&self) -> Result<(), AdapterError> {
|
||||
let channels = ["broadcast", "join", "leave", "disconnect"];
|
||||
let prefix = format!("socket.io:{}:", self.namespace);
|
||||
|
||||
for channel_type in channels {
|
||||
let channel = format!("{}{}", prefix, channel_type);
|
||||
match self.message_bus.subscribe(&channel).await {
|
||||
Ok(rx) => {
|
||||
self.room_subscribers.insert(channel_type.to_string(), rx);
|
||||
}
|
||||
Err(e) => return Err(AdapterError::MessageBus(e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
self.spawn_listener();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn spawn_listener(&self) {
|
||||
let server_id = self.server_id.clone();
|
||||
let on_local_broadcast = self.on_local_broadcast.clone();
|
||||
|
||||
let mut broadcast_rx = self.room_subscribers.remove("broadcast").map(|(_, rx)| rx);
|
||||
let mut join_rx = self.room_subscribers.remove("join").map(|(_, rx)| rx);
|
||||
let mut leave_rx = self.room_subscribers.remove("leave").map(|(_, rx)| rx);
|
||||
let mut disconnect_rx = self.room_subscribers.remove("disconnect").map(|(_, rx)| rx);
|
||||
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
Some(data) = async { broadcast_rx.as_mut()?.recv().await } => {
|
||||
if let Ok(msg) = serde_json::from_slice::<BusMessage>(&data) {
|
||||
handle_bus_message(msg, &on_local_broadcast, &server_id).await;
|
||||
}
|
||||
}
|
||||
Some(data) = async { join_rx.as_mut()?.recv().await } => {
|
||||
if let Ok(msg) = serde_json::from_slice::<BusMessage>(&data) {
|
||||
handle_bus_message(msg, &on_local_broadcast, &server_id).await;
|
||||
}
|
||||
}
|
||||
Some(data) = async { leave_rx.as_mut()?.recv().await } => {
|
||||
if let Ok(msg) = serde_json::from_slice::<BusMessage>(&data) {
|
||||
handle_bus_message(msg, &on_local_broadcast, &server_id).await;
|
||||
}
|
||||
}
|
||||
Some(data) = async { disconnect_rx.as_mut()?.recv().await } => {
|
||||
if let Ok(msg) = serde_json::from_slice::<BusMessage>(&data) {
|
||||
handle_bus_message(msg, &on_local_broadcast, &server_id).await;
|
||||
}
|
||||
}
|
||||
else => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Adapter for RedisAdapter {
|
||||
async fn broadcast(&self, packet: &Packet, opts: &BroadcastOptions) -> Result<(), AdapterError> {
|
||||
if opts.flags.local_only {
|
||||
(self.on_local_broadcast)(packet, opts);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let msg = BusMessage::Broadcast {
|
||||
namespace: packet.namespace.clone(),
|
||||
packet: parser::encode(packet),
|
||||
opts: opts.clone(),
|
||||
server_id: self.server_id.clone(),
|
||||
};
|
||||
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
|
||||
(self.on_local_broadcast)(packet, opts);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn add(&self, sid: &str, room: &str, ns: &str) -> Result<(), AdapterError> {
|
||||
let rk = room_key(ns, room);
|
||||
let srk = socket_rooms_key(ns, sid);
|
||||
|
||||
self.redis_client
|
||||
.sadd::<(), _, _>(&rk, sid)
|
||||
.await
|
||||
.map_err(|e| AdapterError::Redis(e.to_string()))?;
|
||||
|
||||
self.redis_client
|
||||
.sadd::<(), _, _>(&srk, room)
|
||||
.await
|
||||
.map_err(|e| AdapterError::Redis(e.to_string()))?;
|
||||
|
||||
self.socket_rooms
|
||||
.entry(sid.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()); })
|
||||
.or_insert_with(|| HashSet::from([sid.to_string()]));
|
||||
|
||||
let msg = BusMessage::SocketJoin {
|
||||
namespace: ns.to_string(),
|
||||
sid: sid.to_string(),
|
||||
room: room.to_string(),
|
||||
server_id: self.server_id.clone(),
|
||||
};
|
||||
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn del(&self, sid: &str, room: &str, ns: &str) -> Result<(), AdapterError> {
|
||||
let rk = room_key(ns, room);
|
||||
let srk = socket_rooms_key(ns, sid);
|
||||
|
||||
self.redis_client
|
||||
.srem::<(), _, _>(&rk, sid)
|
||||
.await
|
||||
.map_err(|e| AdapterError::Redis(e.to_string()))?;
|
||||
|
||||
self.redis_client
|
||||
.srem::<(), _, _>(&srk, room)
|
||||
.await
|
||||
.map_err(|e| AdapterError::Redis(e.to_string()))?;
|
||||
|
||||
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) {
|
||||
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) {
|
||||
self.rooms.remove(room);
|
||||
}
|
||||
|
||||
let msg = BusMessage::SocketLeave {
|
||||
namespace: ns.to_string(),
|
||||
sid: sid.to_string(),
|
||||
room: room.to_string(),
|
||||
server_id: self.server_id.clone(),
|
||||
};
|
||||
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn del_all(&self, sid: &str, ns: &str) -> Result<(), AdapterError> {
|
||||
if let Some((_, rooms)) = self.socket_rooms.remove(sid) {
|
||||
for room in &rooms {
|
||||
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) {
|
||||
self.rooms.remove(room);
|
||||
}
|
||||
|
||||
let rk = room_key(ns, room);
|
||||
if let Err(e) = self.redis_client.srem::<(), _, _>(&rk, sid).await {
|
||||
tracing::warn!("Redis SREM room error: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let srk = socket_rooms_key(ns, sid);
|
||||
self.redis_client
|
||||
.del::<(), _>(&srk)
|
||||
.await
|
||||
.map_err(|e| AdapterError::Redis(e.to_string()))?;
|
||||
|
||||
self.sockets.remove(sid);
|
||||
|
||||
let msg = BusMessage::SocketDisconnect {
|
||||
namespace: ns.to_string(),
|
||||
sid: sid.to_string(),
|
||||
server_id: self.server_id.clone(),
|
||||
};
|
||||
|
||||
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)
|
||||
.await
|
||||
.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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()
|
||||
} else {
|
||||
let mut sids = HashSet::new();
|
||||
for room in &opts.rooms {
|
||||
if let Some(entry) = self.rooms.get(room) {
|
||||
sids.extend(entry.value().iter().cloned());
|
||||
}
|
||||
}
|
||||
sids
|
||||
};
|
||||
|
||||
for sid in target_sids {
|
||||
if opts.except.contains(&sid) {
|
||||
continue;
|
||||
}
|
||||
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(),
|
||||
rooms,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
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())
|
||||
}
|
||||
|
||||
fn server_id(&self) -> &str {
|
||||
&self.server_id
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), AdapterError> {
|
||||
self.message_bus.close().await.map_err(|e| AdapterError::MessageBus(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user