Skip to content
Clone
main
206 lines 6.4 KB Rust
use serde_json::Value;
use tokio::sync::oneshot;
use super::io::Input;
use super::state::Session;
use crate::packet::{self, Kind, Packet as SocketPacket};
use crate::{EmitPacket, SocketEvent, SocketPayload};
const RESERVED_EVENTS: [&str; 6] = [
"connect",
"connect_error",
"disconnect",
"disconnecting",
"newListener",
"removeListener",
];
impl Session {
pub(super) fn dispatch(&mut self, packet: SocketPacket, binaries: Vec<Vec<u8>>) -> bool {
let mut values = match packet.data {
Value::Array(values) => values.into_iter(),
_ => return false,
};
let Some(Value::String(event)) = values.next() else {
return false;
};
if RESERVED_EVENTS.contains(&event.as_str()) {
return false;
}
let session_id = self
.namespaces
.get(&packet.namespace)
.cloned()
.unwrap_or_default();
let ack = packet.id.map(|id| {
let (sender, receiver) = oneshot::channel::<SocketPayload>();
let input = self.input.clone();
let namespace = packet.namespace.clone();
tokio::spawn(async move {
if let Ok(Ok(payload)) =
tokio::time::timeout(std::time::Duration::from_secs(30), receiver).await
{
let _ = tokio::time::timeout(
std::time::Duration::from_secs(5),
input.send(Input::Ack {
namespace,
id,
payload,
}),
)
.await;
}
});
sender
});
let event = SocketEvent {
user_id: self.user_id,
device_id: self.device_id.clone(),
access_token: self.access_token.clone(),
namespace: packet.namespace,
session_id,
event,
data: values.collect(),
binaries,
ack,
};
self.hub.events().try_send(event).is_ok()
}
pub(super) fn handle_binary(&mut self, data: &[u8]) -> bool {
let mut overflow = false;
{
let Some(pending) = self.pending_binary.as_mut() else {
return true;
};
match pending.bytes.checked_add(data.len()) {
Some(total) if total <= self.max_payload => {
pending.bytes = total;
pending.data.push(data.to_vec());
}
_ => overflow = true,
}
}
if overflow {
self.pending_binary = None;
return false;
}
let complete = self
.pending_binary
.as_ref()
.is_some_and(|pending| pending.data.len() == pending.packet.attachments);
if !complete {
return true;
}
let pending = self.pending_binary.take().expect("pending binary packet");
match pending.packet.kind {
Kind::BinaryEvent => return self.dispatch(pending.packet, pending.data),
Kind::BinaryAck => {
let payload = payload(pending.packet.data, pending.data);
self.resolve_ack(&pending.packet.namespace, pending.packet.id, payload);
}
_ => {}
}
true
}
pub(super) fn resolve_ack(&mut self, namespace: &str, id: Option<u64>, payload: SocketPayload) {
let Some(id) = id else {
return;
};
if self
.pending_acks
.get(&id)
.is_some_and(|(pending_namespace, deadline, _)| {
pending_namespace == namespace && *deadline > tokio::time::Instant::now()
})
{
if let Some((_, _, sender)) = self.pending_acks.remove(&id) {
let _ = sender.send(payload);
}
}
}
pub(super) async fn emit(
&mut self,
packet: EmitPacket,
ack: Option<(oneshot::Sender<SocketPayload>, tokio::time::Instant)>,
) -> bool {
if !self.namespaces.contains_key(&packet.namespace) {
return true;
}
if ack.as_ref().is_some_and(|(sender, deadline)| {
sender.is_closed() || *deadline <= tokio::time::Instant::now()
}) {
return true;
}
let id = if ack.is_some() {
let id = self.next_ack;
self.next_ack = self.next_ack.wrapping_add(1);
Some(id)
} else {
None
};
if let (Some(id), Some((ack, deadline))) = (id, ack) {
self.pending_acks
.insert(id, (packet.namespace.clone(), deadline, ack));
}
let mut data = Vec::with_capacity(packet.args.len() + 1);
data.push(Value::String(packet.event));
data.extend(packet.args);
let text = if packet.binaries.is_empty() {
packet::encode_event(&packet.namespace, id, &data)
} else {
packet::encode_binary_event(&packet.namespace, id, &data, packet.binaries.len())
};
let mut packets = Vec::with_capacity(packet.binaries.len() + 1);
packets.push(crate::engine::text(text));
packets.extend(packet.binaries.into_iter().map(crate::engine::binary));
let sent = self.output.send(packets).await;
if !sent {
if let Some(id) = id {
self.pending_acks.remove(&id);
}
}
sent
}
pub(super) async fn send_ack(
&mut self,
namespace: String,
id: u64,
payload: SocketPayload,
) -> bool {
if !self.namespaces.contains_key(&namespace) {
return true;
}
let SocketPayload { data, binaries } = payload;
let text = if binaries.is_empty() {
packet::encode_ack(&namespace, id, &data)
} else {
packet::encode_binary_ack(&namespace, id, &data, binaries.len())
};
let mut packets = Vec::with_capacity(binaries.len() + 1);
packets.push(crate::engine::text(text));
packets.extend(binaries.into_iter().map(crate::engine::binary));
self.output.send(packets).await
}
}
pub(super) fn payload(data: Value, binaries: Vec<Vec<u8>>) -> SocketPayload {
let data = match data {
Value::Array(values) => values,
_ => Vec::new(),
};
SocketPayload { data, binaries }
}