1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

use std::future::Future;
use std::pin::Pin;

use crate::Error;

/// Abstrction around a websocket connection.
///
/// Built in implementations are provided for `ws_stream_wasm` & `async_tungstenite`.
///
/// If users wish to add support for a new client they should implement this trait.
pub trait Connection {
    /// Receive the next message on this connection.
    fn receive(&mut self) -> impl Future<Output = Option<Message>> + Send;

    /// Send a message with on connection
    fn send(&mut self, message: Message) -> impl Future<Output = Result<(), Error>> + Send;
}

/// A websocket message
///
/// Websocket client libraries usually provide their own version of this struct.
/// The [Connection] trait for a given client should handle translation to & from this enum.
pub enum Message {
    /// A message containing the given text payload
    Text(String),
    /// A message that closes the connection with the given code & reason
    Close {
        /// The status code for this close message
        code: Option<u16>,
        /// Some text explaining the reason the connection is being closed
        reason: Option<String>,
    },
    /// A ping
    Ping,
    /// A reply to a ping
    Pong,
}

impl Message {
    pub(crate) fn deserialize<T>(self) -> Result<T, Error>
    where
        T: serde::de::DeserializeOwned,
    {
        let Message::Text(text) = self else {
            panic!("Don't call deserialize on non-text messages");
        };

        serde_json::from_str(&text).map_err(|error| Error::Decode(error.to_string()))
    }

    pub(crate) fn init(payload: Option<serde_json::Value>) -> Self {
        Self::Text(
            serde_json::to_string(&crate::protocol::ConnectionInit::new(payload))
                .expect("payload is already serialized so this shouldn't fail"),
        )
    }

    pub(crate) fn graphql_pong() -> Self {
        Self::Text(serde_json::to_string(&crate::protocol::Message::Pong::<()>).unwrap())
    }

    pub(crate) fn complete(id: usize) -> Self {
        Self::Text(
            serde_json::to_string(&crate::protocol::Message::Complete::<()> { id: id.to_string() })
                .unwrap(),
        )
    }
}

/// An object safe wrapper around the Connection trait, allowing us
/// to use it dynamically
pub(crate) trait ObjectSafeConnection: Send {
    fn receive(&mut self) -> Pin<Box<dyn Future<Output = Option<Message>> + Send + '_>>;

    fn send(
        &mut self,
        message: Message,
    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>>;
}

impl<T: Connection + Sized + Send> ObjectSafeConnection for T {
    fn receive(&mut self) -> Pin<Box<dyn Future<Output = Option<Message>> + Send + '_>> {
        Box::pin(Connection::receive(self))
    }

    fn send(
        &mut self,
        message: Message,
    ) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
        Box::pin(Connection::send(self, message))
    }
}