rainier_websocket/

socket.rs

//! The handle a handler holds — [`Socket`].

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;

use serde::Serialize;
use tokio::sync::mpsc;

use rainier_support::{Error, Result};

use crate::message::Message;

/// A unique id for one connection, for the life of the process.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct SocketId(u64);

impl SocketId {
    /// The next id. Monotonic, so ordering by it is arrival order.
    pub fn next() -> Self {
        static NEXT: AtomicU64 = AtomicU64::new(1);
        Self(NEXT.fetch_add(1, Ordering::Relaxed))
    }

    /// The id as a number, for logs and for keying a map.
    pub fn get(self) -> u64 {
        self.0
    }
}

impl std::fmt::Display for SocketId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// One connected client.
///
/// Cheap to clone and safe to keep: sending goes through a channel, so a handle
/// stored in a room registry does not hold the connection's task or its
/// buffers. Sending to a socket that has gone away is an error, not a panic.
///
/// ```ignore
/// socket.send("hello").await?;
/// socket.send_json(&Update { count: 3 }).await?;
/// socket.close().await;
/// ```
#[derive(Debug, Clone)]
pub struct Socket {
    id: SocketId,
    path: Arc<str>,
    params: Arc<Vec<(String, String)>>,
    outbound: mpsc::UnboundedSender<Outbound>,
    closed: Arc<AtomicBool>,
}

/// What the connection task pulls off the queue.
#[derive(Debug)]
pub enum Outbound {
    /// Send this frame.
    Send(Message),
    /// Close the connection, with an optional reason.
    Close(Option<String>),
}

impl Socket {
    /// Build a handle over `outbound`.
    ///
    /// The transport owns the receiving half; a handler is only ever given
    /// this side.
    pub fn new(
        id: SocketId,
        path: impl Into<Arc<str>>,
        params: Vec<(String, String)>,
        outbound: mpsc::UnboundedSender<Outbound>,
    ) -> Self {
        Self {
            id,
            path: path.into(),
            params: Arc::new(params),
            outbound,
            closed: Arc::new(AtomicBool::new(false)),
        }
    }

    /// This connection's id.
    pub fn id(&self) -> SocketId {
        self.id
    }

    /// The path it connected to.
    pub fn path(&self) -> &str {
        &self.path
    }

    /// A route parameter captured from the path — `rooms/{room}` gives
    /// `room`.
    pub fn param(&self, name: &str) -> Option<&str> {
        self.params.iter().find(|(key, _)| key == name).map(|(_, value)| value.as_str())
    }

    /// A route parameter, parsed.
    ///
    /// A **bad request** when it will not parse: the value came from the URL
    /// the client asked for.
    pub fn parse_param<T>(&self, name: &str) -> Result<T>
    where
        T: std::str::FromStr,
        T::Err: std::fmt::Display,
    {
        let raw = self
            .param(name)
            .ok_or_else(|| Error::internal(format!("the route captured no `{name}`")))?;

        raw.parse::<T>()
            .map_err(|e| Error::bad_request(format!("`{name}` in the path is invalid: {e}")))
    }

    /// Every captured parameter.
    pub fn params(&self) -> &[(String, String)] {
        &self.params
    }

    /// Send a frame.
    ///
    /// Queued rather than written: a handler that awaited the socket would
    /// block on a slow client, and one slow client must not hold up the task
    /// reading everyone else's messages.
    pub fn send(&self, message: impl Into<Message>) -> Result<()> {
        self.push(Outbound::Send(message.into()))
    }

    /// Send a value as a JSON text frame.
    pub fn send_json(&self, value: &impl Serialize) -> Result<()> {
        self.send(Message::json(value)?)
    }

    /// Close the connection.
    ///
    /// Idempotent: closing twice is not an error, because two halves of a
    /// handler deciding to close at once is normal.
    pub fn close(&self) -> Result<()> {
        if self.closed.swap(true, Ordering::SeqCst) {
            return Ok(());
        }
        self.push(Outbound::Close(None))
    }

    /// Close with a reason the peer will see.
    pub fn close_with(&self, reason: impl Into<String>) -> Result<()> {
        if self.closed.swap(true, Ordering::SeqCst) {
            return Ok(());
        }
        self.push(Outbound::Close(Some(reason.into())))
    }

    /// Whether this handle has been asked to close, or the connection has
    /// gone.
    pub fn is_closed(&self) -> bool {
        self.closed.load(Ordering::SeqCst) || self.outbound.is_closed()
    }

    fn push(&self, outbound: Outbound) -> Result<()> {
        self.outbound.send(outbound).map_err(|_| {
            // The peer disconnected. Routine — a registry holding stale
            // handles finds out this way — so it is an error to handle, not a
            // panic and not a silent success.
            Error::internal(format!("socket {} has gone away", self.id))
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn socket() -> (Socket, mpsc::UnboundedReceiver<Outbound>) {
        let (tx, rx) = mpsc::unbounded_channel();
        (Socket::new(SocketId::next(), "/ws/rooms/7", vec![("room".into(), "7".into())], tx), rx)
    }

    #[test]
    fn a_send_is_queued_not_written() {
        let (socket, mut rx) = socket();

        socket.send("hello").unwrap();

        match rx.try_recv().unwrap() {
            Outbound::Send(Message::Text(text)) => assert_eq!(text, "hello"),
            other => panic!("{other:?}"),
        }
    }

    #[test]
    fn route_parameters_come_through() {
        let (socket, _rx) = socket();

        assert_eq!(socket.param("room"), Some("7"));
        assert_eq!(socket.parse_param::<u64>("room").unwrap(), 7);
        assert_eq!(socket.param("nope"), None);
    }

    #[test]
    fn a_parameter_that_will_not_parse_is_a_bad_request() {
        let (tx, _rx) = mpsc::unbounded_channel();
        let socket =
            Socket::new(SocketId::next(), "/ws/rooms/x", vec![("room".into(), "x".into())], tx);

        assert_eq!(socket.parse_param::<u64>("room").unwrap_err().status(), 400);
    }

    #[test]
    fn closing_twice_is_not_an_error() {
        let (socket, mut rx) = socket();

        socket.close().unwrap();
        socket.close().unwrap();

        assert!(matches!(rx.try_recv(), Ok(Outbound::Close(None))));
        assert!(rx.try_recv().is_err(), "only one close should be queued");
        assert!(socket.is_closed());
    }

    #[test]
    fn sending_to_a_socket_that_went_away_is_an_error_not_a_panic() {
        let (socket, rx) = socket();
        drop(rx);

        let err = socket.send("hello").unwrap_err();
        assert!(err.message().contains("gone away"), "{}", err.message());
        assert!(socket.is_closed());
    }

    #[test]
    fn ids_are_unique_and_ordered() {
        let first = SocketId::next();
        let second = SocketId::next();

        assert!(second > first);
    }
}