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,301 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::engine::packet::Packet as EnginePacket;
|
||||
use crate::engine::packet::PacketData as EnginePacketData;
|
||||
use crate::engine::server::{EngineConfig, EngineServer};
|
||||
use crate::engine::session::SessionStore;
|
||||
use crate::socket::adapter::{Adapter, LocalAdapter};
|
||||
use crate::socket::namespace::NamespaceManager;
|
||||
use crate::socket::packet::{Packet, PacketType};
|
||||
use crate::socket::parser;
|
||||
use crate::socket::socket::Socket;
|
||||
|
||||
pub struct SocketServer {
|
||||
pub engine: Arc<EngineServer>,
|
||||
pub namespaces: Arc<NamespaceManager>,
|
||||
pub adapter: Arc<dyn Adapter>,
|
||||
socket_txs: Arc<DashMap<String, mpsc::Sender<Packet>>>,
|
||||
}
|
||||
|
||||
impl SocketServer {
|
||||
pub fn new(config: EngineConfig) -> Self {
|
||||
SocketServerBuilder::new(config).build()
|
||||
}
|
||||
|
||||
pub fn builder(config: EngineConfig) -> SocketServerBuilder {
|
||||
SocketServerBuilder::new(config)
|
||||
}
|
||||
|
||||
pub fn of(&self, path: impl Into<String>) -> Arc<crate::socket::namespace::Namespace> {
|
||||
self.namespaces.get_or_create_namespace(&path.into())
|
||||
}
|
||||
|
||||
pub async fn run_http(self: Arc<Self>, addr: &str) -> std::io::Result<()> {
|
||||
self.engine.clone().run_http(addr).await
|
||||
}
|
||||
|
||||
pub fn register_socket(&self, sid: String, tx: mpsc::Sender<Packet>) {
|
||||
self.socket_txs.insert(sid, tx);
|
||||
}
|
||||
|
||||
pub fn unregister_socket(&self, sid: &str) {
|
||||
self.socket_txs.remove(sid);
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SocketServerBuilder {
|
||||
config: EngineConfig,
|
||||
adapter: Option<Arc<dyn Adapter>>,
|
||||
}
|
||||
|
||||
impl SocketServerBuilder {
|
||||
pub fn new(config: EngineConfig) -> Self {
|
||||
Self {
|
||||
config,
|
||||
adapter: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn adapter(mut self, adapter: Arc<dyn Adapter>) -> Self {
|
||||
self.adapter = Some(adapter);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn build(self) -> SocketServer {
|
||||
let namespaces = Arc::new(NamespaceManager::new());
|
||||
let socket_txs: Arc<DashMap<String, mpsc::Sender<Packet>>> = Arc::new(DashMap::new());
|
||||
let engine_store = SessionStore::new();
|
||||
|
||||
let namespaces_clone = namespaces.clone();
|
||||
let socket_txs_clone = socket_txs.clone();
|
||||
let engine_store_clone = engine_store.clone();
|
||||
|
||||
let adapter: Arc<dyn Adapter> = self.adapter.unwrap_or_else(|| {
|
||||
let ns_clone = namespaces.clone();
|
||||
let send_fn = move |engine_sid: &str, packet: &Packet| {
|
||||
if let Some(ns) = ns_clone.get_namespace(&packet.namespace) {
|
||||
if let Some(socket) = ns.get_socket_by_engine_sid(engine_sid) {
|
||||
socket.send_packet(packet).map_err(|e| e.to_string())
|
||||
} else {
|
||||
Err(format!(
|
||||
"Socket with engine_sid {} not found in namespace {}",
|
||||
engine_sid, packet.namespace
|
||||
))
|
||||
}
|
||||
} else {
|
||||
Err(format!("Namespace {} not found", packet.namespace))
|
||||
}
|
||||
};
|
||||
Arc::new(LocalAdapter::new(send_fn))
|
||||
});
|
||||
|
||||
let adapter_clone = adapter.clone();
|
||||
let engine = Arc::new(EngineServer::with_store(
|
||||
self.config,
|
||||
engine_store,
|
||||
move |sid, engine_packet| {
|
||||
let namespaces = namespaces_clone.clone();
|
||||
let socket_txs = socket_txs_clone.clone();
|
||||
let engine_store = engine_store_clone.clone();
|
||||
let adapter = adapter_clone.clone();
|
||||
tokio::spawn(async move {
|
||||
handle_engine_message(
|
||||
sid, engine_packet, &namespaces, &socket_txs, &engine_store, &adapter,
|
||||
).await;
|
||||
});
|
||||
},
|
||||
));
|
||||
|
||||
let server = SocketServer {
|
||||
engine,
|
||||
namespaces,
|
||||
adapter,
|
||||
socket_txs,
|
||||
};
|
||||
|
||||
for ns in server.namespaces.all_namespaces() {
|
||||
let adapter_ref = server.adapter.clone();
|
||||
tokio::spawn(async move {
|
||||
ns.set_adapter(adapter_ref).await;
|
||||
});
|
||||
}
|
||||
|
||||
server
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_engine_message(
|
||||
engine_sid: String,
|
||||
engine_packet: EnginePacket,
|
||||
namespaces: &Arc<NamespaceManager>,
|
||||
socket_txs: &Arc<DashMap<String, mpsc::Sender<Packet>>>,
|
||||
engine_store: &SessionStore,
|
||||
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 {
|
||||
PacketType::Connect => {
|
||||
handle_connect(&engine_sid, &socket_packet, namespaces, socket_txs, engine_store, adapter).await;
|
||||
}
|
||||
PacketType::Disconnect => {
|
||||
handle_disconnect(&engine_sid, &socket_packet, namespaces, socket_txs);
|
||||
}
|
||||
PacketType::Event => {
|
||||
handle_event(&engine_sid, &socket_packet, namespaces);
|
||||
}
|
||||
PacketType::Ack => {
|
||||
handle_ack(&engine_sid, &socket_packet);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_connect(
|
||||
engine_sid: &str,
|
||||
packet: &Packet,
|
||||
namespaces: &Arc<NamespaceManager>,
|
||||
socket_txs: &Arc<DashMap<String, mpsc::Sender<Packet>>>,
|
||||
engine_store: &SessionStore,
|
||||
adapter: &Arc<dyn Adapter>,
|
||||
) {
|
||||
// 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);
|
||||
return;
|
||||
}
|
||||
|
||||
let namespace = namespaces.get_or_create_namespace(&packet.namespace);
|
||||
|
||||
// Ensure newly created namespaces get the shared adapter
|
||||
{
|
||||
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;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let socket_sid = crate::engine::session::generate_sid();
|
||||
let (tx, mut rx) = mpsc::channel::<Packet>(256);
|
||||
socket_txs.insert(socket_sid.clone(), tx.clone());
|
||||
|
||||
let socket = Arc::new(Socket::new(
|
||||
socket_sid.clone(),
|
||||
packet.namespace.clone(),
|
||||
engine_sid.to_string(),
|
||||
tx,
|
||||
));
|
||||
|
||||
// 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 {
|
||||
tracing::warn!("Socket {} connection rejected: {}", socket_sid, msg);
|
||||
socket_txs.remove(&socket_sid);
|
||||
return;
|
||||
}
|
||||
|
||||
// Connect handler passed — spawn forwarding task
|
||||
let engine_store_clone = engine_store.clone();
|
||||
let engine_sid_clone = engine_sid.to_string();
|
||||
let socket_sid_clone = socket_sid.clone();
|
||||
let socket_txs_clone = socket_txs.clone();
|
||||
let namespace_clone = namespace.clone();
|
||||
tokio::spawn(async move {
|
||||
while let Some(socket_packet) = rx.recv().await {
|
||||
let encoded = parser::encode(&socket_packet);
|
||||
let engine_packet = EnginePacket::message_text(encoded);
|
||||
|
||||
if let Some(session) = engine_store_clone.get(&engine_sid_clone) {
|
||||
let mut s = session.write().await;
|
||||
if s.state == crate::engine::session::SessionState::Closed {
|
||||
break;
|
||||
}
|
||||
s.push_packet(engine_packet);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
// 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;
|
||||
});
|
||||
|
||||
// Send Connect response (only after handler passed)
|
||||
let response = Packet::connect(
|
||||
&socket.namespace,
|
||||
Some(serde_json::json!({ "sid": &socket.sid })),
|
||||
);
|
||||
|
||||
if socket.send_packet(&response).is_err() {
|
||||
tracing::warn!("Failed to send connect response to socket {}", socket.sid);
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_disconnect(
|
||||
engine_sid: &str,
|
||||
packet: &Packet,
|
||||
namespaces: &Arc<NamespaceManager>,
|
||||
socket_txs: &Arc<DashMap<String, mpsc::Sender<Packet>>>,
|
||||
) {
|
||||
if let Some(namespace) = namespaces.get_namespace(&packet.namespace) {
|
||||
// Look up socket by engine_sid, then remove by socket_sid
|
||||
if let Some(socket) = namespace.get_socket_by_engine_sid(engine_sid) {
|
||||
socket_txs.remove(&socket.sid);
|
||||
let socket_sid = socket.sid.clone();
|
||||
let ns_clone = namespace.clone();
|
||||
tokio::spawn(async move {
|
||||
ns_clone.remove_socket_by_sid(&socket_sid).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
};
|
||||
|
||||
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;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_ack(engine_sid: &str, packet: &Packet) {
|
||||
tracing::debug!(
|
||||
"Received ACK from {} for namespace {} with id {:?}",
|
||||
engine_sid,
|
||||
packet.namespace,
|
||||
packet.id
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user