273bd1b
use std::collections::HashMap;
use std::time::Instant;
use tokio::sync::{mpsc, oneshot};
use uuid::Uuid;
use super::events::payload;
use super::io::{Input, Output};
use crate::backend::{ConnectedSession, ConnectionPhase, DeliveryContext};
use crate::engine::{self, Data, Packet as EnginePacket};
use crate::packet::{self, Kind, Packet as SocketPacket};
use crate::{SocketHub, SocketPayload};
pub(super) struct Session {
pub(super) hub: SocketHub,
pub(super) sid: String,
pub(super) user_id: i64,
pub(super) access_token: String,
pub(super) device_id: String,
pub(super) deliveries: HashMap<String, (oneshot::Sender<()>, tokio::task::JoinHandle<()>)>,
pub(super) input: mpsc::Sender<Input>,
pub(super) output: Output,
pub(super) namespaces: HashMap<String, String>,
pub(super) pending_binary: Option<PendingBinary>,
pub(super) pending_acks:
HashMap<u64, (String, tokio::time::Instant, oneshot::Sender<SocketPayload>)>,
pub(super) next_ack: u64,
pub(super) last_pong: Instant,
pub(super) last_actor_heartbeat: Instant,
pub(super) max_payload: usize,
pub(super) auth_deadline: Option<tokio::time::Instant>,
}
pub(super) struct PendingBinary {
pub(super) packet: SocketPacket,
pub(super) data: Vec<Vec<u8>>,
pub(super) bytes: usize,
}
impl Session {
pub(super) async fn handle(&mut self, input: Input) -> bool {
match input {
Input::Packets(packets) => {
for packet in packets {
if !self.handle_packet(packet).await {
return false;
}
}
true
}
Input::Poll(responder) => self.output.poll(responder),
Input::Upgrade { tx, result } => {
if result.is_closed() {
return true;
}
if matches!(&self.output, Output::WebSocket { .. }) {
let _ = result.send(false);
return true;
}
let upgraded = self.output.upgrade(tx).await;
let _ = result.send(upgraded);
upgraded
}
Input::Emit {
packet,
ack,
accepted,
} => {
let connected = self.namespaces.contains_key(&packet.namespace);
let sent = self.emit(packet, ack).await;
if let Some(accepted) = accepted {
let _ = accepted.send(connected && sent);
}
sent
}
Input::Ack {
namespace,
id,
payload,
} => self.send_ack(namespace, id, payload).await,
Input::Close => false,
}
}
async fn handle_packet(&mut self, packet: EnginePacket) -> bool {
match packet {
EnginePacket::Ping(payload) => {
self.output.send(vec![EnginePacket::Pong(payload)]).await
}
EnginePacket::Pong(_) => {
self.last_pong = Instant::now();
if !self.namespaces.is_empty()
&& self.last_actor_heartbeat.elapsed() >= self.hub.options().heartbeat_interval
{
self.last_actor_heartbeat = Instant::now();
self.heartbeat_actor();
}
true
}
EnginePacket::Message(Data::Text(text)) => self.handle_socket(&text).await,
EnginePacket::Message(Data::Binary(data)) => self.handle_binary(&data),
EnginePacket::Close => false,
EnginePacket::Open(_) | EnginePacket::Upgrade(_) | EnginePacket::Noop => true,
}
}
async fn handle_socket(&mut self, text: &str) -> bool {
if self.pending_binary.is_some() {
return false;
}
let Some(packet) = packet::decode(text) else {
return false;
};
match packet.kind {
Kind::Connect => self.connect(packet).await,
Kind::Disconnect => {
self.disconnect(&packet.namespace).await;
true
}
Kind::Event => {
if !self.namespaces.contains_key(&packet.namespace) {
return self
.connect_error(&packet.namespace, "namespace is not connected")
.await;
}
self.dispatch(packet, Vec::new())
}
Kind::Ack => {
self.resolve_ack(
&packet.namespace,
packet.id,
payload(packet.data, Vec::new()),
);
true
}
Kind::BinaryEvent => {
if !self.namespaces.contains_key(&packet.namespace) {
return self
.connect_error(&packet.namespace, "namespace is not connected")
.await;
}
self.pending_binary = Some(PendingBinary {
packet,
data: Vec::new(),
bytes: 0,
});
true
}
Kind::BinaryAck => {
self.pending_binary = Some(PendingBinary {
packet,
data: Vec::new(),
bytes: 0,
});
true
}
Kind::ConnectError => true,
}
}
async fn connect(&mut self, packet: SocketPacket) -> bool {
if self.user_id == 0 {
let token = packet
.data
.get("token")
.or_else(|| packet.data.get("access_token"))
.and_then(serde_json::Value::as_str);
let Some(token) = token else {
return self
.connect_error(&packet.namespace, "missing access token")
.await;
};
match self.hub.authenticate_token(token).await {
Ok(user_id) => {
self.user_id = user_id;
self.access_token = token.to_string();
self.auth_deadline = None;
}
Err(_) => {
return self
.connect_error(&packet.namespace, "authentication failed")
.await;
}
}
}
let namespace = packet.namespace;
let device = packet
.data
.get("device_id")
.and_then(serde_json::Value::as_str);
if self.device_id.is_empty() {
self.device_id = match device {
Some(device) => match Uuid::parse_str(device) {
Ok(device) => device.to_string(),
Err(_) => {
return self
.connect_error(&namespace, "device_id must be a UUID")
.await;
}
},
None => Uuid::new_v4().to_string(),
};
} else if device.is_some_and(|device| {
Uuid::parse_str(device)
.map(|id| id.to_string())
.ok()
.as_deref()
!= Some(self.device_id.as_str())
}) {
return self
.connect_error(&namespace, "device_id cannot change within a session")
.await;
}
if self.namespaces.len() >= 32 || namespace.len() > 256 {
return self
.connect_error(&namespace, "namespace limit exceeded")
.await;
}
if self.namespaces.contains_key(&namespace) {
return self
.connect_error(&namespace, "namespace already connected")
.await;
}
let session_id = Uuid::new_v4().to_string();
let handle = self.hub.session_handle(self.input.clone());
let connected = ConnectedSession {
user_id: self.user_id,
device_id: self.device_id.clone(),
namespace: namespace.clone(),
session_id: session_id.clone(),
session_sid: self.sid.clone(),
access_token: self.access_token.clone(),
first_namespace: self.namespaces.is_empty(),
hub: self.hub.clone(),
handle: handle.clone(),
};
if let Err(error) = self.hub.backend().on_connect(&connected).await {
return self
.connect_error(&namespace, &format!("connection setup failed: {error}"))
.await;
}
self.hub
.register(
self.user_id,
&namespace,
session_id.clone(),
self.input.clone(),
)
.await;
self.namespaces
.insert(namespace.clone(), session_id.clone());
self.hub.backend().heartbeat(self.user_id).await;
self.hub
.report_connection(ConnectionPhase::Open, self.user_id, &self.sid)
.await;
let packet = packet::encode_connect(&namespace, &session_id, &self.device_id);
self.last_actor_heartbeat = Instant::now();
if !self.output.send(vec![engine::text(packet)]).await {
return false;
}
let (stop, stopped) = oneshot::channel();
let backend = self.hub.backend().clone();
let context = DeliveryContext {
user_id: self.user_id,
device_id: self.device_id.clone(),
namespace: namespace.clone(),
session_id: session_id.clone(),
session_sid: self.sid.clone(),
access_token: self.access_token.clone(),
hub: self.hub.clone(),
handle,
stop: stopped,
};
let task = tokio::spawn(async move {
backend.run_delivery(context).await;
});
self.deliveries.insert(namespace, (stop, task));
true
}
async fn connect_error(&mut self, namespace: &str, message: &str) -> bool {
let packet = packet::encode_connect_error(namespace, message);
self.output.send(vec![engine::text(packet)]).await
}
async fn disconnect(&mut self, namespace: &str) {
if let Some((stop, _)) = self.deliveries.remove(namespace) {
let _ = stop.send(());
}
self.pending_acks
.retain(|_, (pending_namespace, _, _)| pending_namespace != namespace);
if let Some(session_id) = self.namespaces.remove(namespace) {
self.hub
.unregister(self.user_id, namespace, &session_id)
.await;
}
}
fn heartbeat_actor(&self) {
let backend = self.hub.backend().clone();
let hub = self.hub.clone();
let user_id = self.user_id;
let session = self.sid.clone();
tokio::spawn(async move {
let _ = tokio::time::timeout(
std::time::Duration::from_secs(5),
backend.heartbeat(user_id),
)
.await;
hub.report_connection(ConnectionPhase::Heartbeat, user_id, &session)
.await;
});
}
pub(super) async fn cleanup(&mut self) {
let mut tasks = Vec::new();
for (_, (stop, task)) in self.deliveries.drain() {
let _ = stop.send(());
tasks.push(task);
}
let _ = tokio::time::timeout(
std::time::Duration::from_secs(5),
futures_util::future::join_all(tasks),
)
.await;
for (namespace, session_id) in self.namespaces.drain() {
self.hub
.unregister(self.user_id, &namespace, &session_id)
.await;
}
self.pending_acks.clear();
self.output.close().await;
self.hub
.report_connection(ConnectionPhase::Close, self.user_id, &self.sid)
.await;
self.hub.remove_session(&self.sid).await;
}
}