gitdataai/lib/socketio/socket.rs

use std::fmt;
use std::{collections::HashSet, sync::Arc};

use serde::Serialize;
use serde_json::Value;

use crate::{
    adapter::BroadcastOptions, error::Result, packet::Packet, server::SocketIo,
    session::Session,
};

#[derive(Clone)]
pub struct Socket {
    pub io: SocketIo,
    pub session: Arc<Session>,
    pub namespace: String,
    pub sid: String,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DisconnectReason {
    Client,
    TransportClosed,
    PingTimeout,
    Server,
}

#[derive(Clone)]
pub struct AckSender {
    session: Arc<Session>,
    namespace: String,
    id: u64,
}

impl fmt::Debug for AckSender {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AckSender")
            .field("namespace", &self.namespace)
            .field("id", &self.id)
            .finish_non_exhaustive()
    }
}

impl AckSender {
    pub fn new(session: Arc<Session>, namespace: String, id: u64) -> Self {
        Self {
            session,
            namespace,
            id,
        }
    }

    pub async fn send<T: Serialize>(&self, data: T) -> Result<()> {
        let args = match serde_json::to_value(data)? {
            Value::Array(values) => values,
            value => vec![value],
        };
        self.session
            .enqueue_socket_packet(Packet::ack(&self.namespace, self.id, args))
            .await;
        Ok(())
    }
}

impl Socket {
    pub fn id(&self) -> &str {
        &self.sid
    }

    pub fn namespace(&self) -> &str {
        &self.namespace
    }

    pub fn session_user(&self) -> Option<uuid::Uuid> {
        self.session
            .user
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .clone()
    }

    pub fn set_user(&self, user: uuid::Uuid) {
        *self.session.user.lock().unwrap_or_else(|e| e.into_inner()) =
            Some(user);
    }

    pub async fn rooms(&self) -> HashSet<String> {
        self.session
            .namespaces
            .lock()
            .await
            .get(&self.namespace)
            .map(|state| state.rooms.clone())
            .unwrap_or_default()
    }

    pub async fn auth(&self) -> Option<Value> {
        self.session
            .namespaces
            .lock()
            .await
            .get(&self.namespace)
            .and_then(|state| state.auth.clone())
    }

    pub async fn emit<T: Serialize>(&self, event: &str, data: T) -> Result<()> {
        self.io
            .emit_to_sid(&self.namespace, &self.session.engine_id, event, data)
            .await
    }

    pub async fn emit_with_ack<T: Serialize>(
        &self,
        event: &str,
        data: T,
    ) -> Result<Vec<Value>> {
        self.io
            .emit_to_sid_with_ack(
                &self.namespace,
                &self.session.engine_id,
                event,
                data,
            )
            .await
    }

    pub async fn emit_binary(
        &self,
        event: &str,
        args: Vec<Value>,
        binary: Vec<Vec<u8>>,
    ) -> Result<()> {
        self.io
            .emit_binary_to_sid(
                &self.namespace,
                &self.session.engine_id,
                event,
                args,
                binary,
            )
            .await
    }

    pub async fn broadcast<T: Serialize>(
        &self,
        event: &str,
        data: T,
    ) -> Result<()> {
        let opts = BroadcastOptions {
            namespace: self.namespace.clone(),
            skip_sid: Some(self.session.engine_id.clone()),
            ..BroadcastOptions::default()
        };
        self.io.broadcast_with_opts(opts, event, data).await
    }

    pub async fn join(&self, room: impl Into<String>) -> Result<()> {
        self.io
            .join(&self.namespace, &self.session.engine_id, room.into())
            .await
    }

    pub async fn leave(&self, room: &str) -> Result<()> {
        self.io
            .leave(&self.namespace, &self.session.engine_id, room)
            .await
    }

    pub async fn disconnect(&self) -> Result<()> {
        self.session
            .enqueue_socket_packet(Packet::new(
                crate::packet::PacketType::Disconnect,
                &self.namespace,
                None,
                None,
            ))
            .await;
        self.io
            .disconnect_socket(
                &self.namespace,
                &self.session,
                DisconnectReason::Server,
            )
            .await
    }
}