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
93
94
95
96
97
98
99
100
101
102
103

//! Abstraction over WebSocket implementations.
//!
//! Use the [`fastwebsockets`] implementation of these traits as an example for implementing them
//! for other WebSocket implementations.
//!
//! [`fastwebsockets`]: https://github.com/MercuryWorkshop/epoxy-tls/blob/multiplexed/wisp/src/fastwebsockets.rs
use bytes::Bytes;
use futures::lock::Mutex;
use std::sync::Arc;

/// Opcode of the WebSocket frame.
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum OpCode {
    /// Text frame.
    Text,
    /// Binary frame.
    Binary,
    /// Close frame.
    Close,
    /// Ping frame.
    Ping,
    /// Pong frame.
    Pong,
}

/// WebSocket frame.
#[derive(Debug, Clone)]
pub struct Frame {
    /// Whether the frame is finished or not.
    pub finished: bool,
    /// Opcode of the WebSocket frame.
    pub opcode: OpCode,
    /// Payload of the WebSocket frame.
    pub payload: Bytes,
}

impl Frame {
    /// Create a new text frame.
    pub fn text(payload: Bytes) -> Self {
        Self {
            finished: true,
            opcode: OpCode::Text,
            payload,
        }
    }

    /// Create a new binary frame.
    pub fn binary(payload: Bytes) -> Self {
        Self {
            finished: true,
            opcode: OpCode::Binary,
            payload,
        }
    }

    /// Create a new close frame.
    pub fn close(payload: Bytes) -> Self {
        Self {
            finished: true,
            opcode: OpCode::Close,
            payload,
        }
    }
}

/// Generic WebSocket read trait.
pub trait WebSocketRead {
    /// Read a frame from the socket.
    fn wisp_read_frame(
        &mut self,
        tx: &crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite>,
    ) -> impl std::future::Future<Output = Result<Frame, crate::WispError>>;
}

/// Generic WebSocket write trait.
pub trait WebSocketWrite {
    /// Write a frame to the socket.
    fn wisp_write_frame(
        &mut self,
        frame: Frame,
    ) -> impl std::future::Future<Output = Result<(), crate::WispError>>;
}

/// Locked WebSocket that can be shared between threads.
pub struct LockedWebSocketWrite<S>(Arc<Mutex<S>>);

impl<S: WebSocketWrite> LockedWebSocketWrite<S> {
    /// Create a new locked websocket.
    pub fn new(ws: S) -> Self {
        Self(Arc::new(Mutex::new(ws)))
    }

    /// Write a frame to the websocket.
    pub async fn write_frame(&self, frame: Frame) -> Result<(), crate::WispError> {
        self.0.lock().await.wisp_write_frame(frame).await
    }
}

impl<S: WebSocketWrite> Clone for LockedWebSocketWrite<S> {
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }
}