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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
/**
 * Meows is a simpmle library for making it easy to implement websocket message
 * handlers, built on top of async-tungstenite and the async ecosystem
 */

#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate smol;

use async_tungstenite::WebSocketStream;
use futures::future::BoxFuture;
use futures::prelude::*;
use log::*;
use serde::de::DeserializeOwned;
use smol::{Async, Task};
use std::collections::HashMap;
use std::net::{TcpListener, TcpStream};
use std::sync::{Arc, RwLock};

pub use serde_json::Value;
/** Re-exporting for convenience */
pub use tungstenite::Message;

/**
 * The Envelope handles the serialization/deserialization of the outer part of a
 * websocket message.
 *
 * All websocket messages are expected to have the basic format of:
 *  ```json
 *  {
 *      "type" : "foo",
 *      "value": {}
 *  }
 *  ```
 *  The contents of `value` can be completely arbitrary and are expected to be
 *  deserializable into whatever the `type` value string is , e.g. `Foo` in this
 *  example.
 */
#[derive(Debug, Deserialize, Serialize)]
pub struct Envelope {
    #[serde(rename = "type")]
    pub ttype: String,
    pub value: Value,
}

/**
 * Support converting an Envelope directly into a String type
 *
 * ```
 * use meows::Envelope;
 *
 * let e = Envelope { ttype: String::from("rust"), value: serde_json::Value::Null };
 * let into: String = e.into();
 * assert_eq!(r#"{"type":"rust","value":null}"#, into);
 * ```
 */
impl Into<String> for Envelope {
    fn into(self) -> String {
        serde_json::to_string(&self).expect("Curiosly failed to serialize an envelope")
    }
}

/**
 * THe Request struct brings message specific information into the handlers
 */
pub struct Request<ServerState, ClientState> {
    pub env: Envelope,
    pub state: Arc<ServerState>,
    pub client_state: Arc<RwLock<ClientState>>,
}

impl<ServerState, ClientState> Request<ServerState, ClientState> {
    pub fn from_value<ValueType: DeserializeOwned>(&mut self) -> Option<ValueType> {
        serde_json::from_value(self.env.value.take()).map_or(None, |v| Some(v))
    }
}

/**
 * Endpoint comes from tide, and I'm still not sure how this magic works
 */
pub trait Endpoint<ServerState, ClientState>: Send + Sync + 'static {
    /// Invoke the endpoint within the given context
    fn call<'a>(&'a self, req: Request<ServerState, ClientState>) -> BoxFuture<'a, Option<Message>>;
}

impl<ServerState, ClientState, F: Send + Sync + 'static, Fut> Endpoint<ServerState, ClientState> for F
where
    F: Fn(Request<ServerState, ClientState>) -> Fut,
    Fut: Future<Output = Option<Message>> + Send + 'static,
{
    fn call<'a>(&'a self, req: Request<ServerState, ClientState>) -> BoxFuture<'a, Option<Message>> {
        let fut = (self)(req);
        Box::pin(fut)
    }
}

pub trait DefaultEndpoint<ServerState, ClientState>: Send + Sync + 'static {
    /// Invoke the endpoint within the given context
    fn call<'a>(&'a self, msg: String, state: Arc<ServerState>) -> BoxFuture<'a, Option<Message>>;
}

impl<ServerState, ClientState, F: Send + Sync + 'static, Fut> DefaultEndpoint<ServerState, ClientState> for F
where
    F: Fn(String, Arc<ServerState>) -> Fut,
    Fut: Future<Output = Option<Message>> + Send + 'static,
{
    fn call<'a>(&'a self, msg: String, state: Arc<ServerState>) -> BoxFuture<'a, Option<Message>> {
        let fut = (self)(msg, state);
        Box::pin(fut)
    }
}

type Callback<ServerState, ClientState> = Arc<Box<dyn Endpoint<ServerState, ClientState>>>;
type DefaultCallback<ServerState, ClientState> = Arc<Box<dyn DefaultEndpoint<ServerState, ClientState>>>;

/**
 * The Server is the primary means of listening for messages
 */
pub struct Server<ServerState, ClientState> {
    state: Arc<ServerState>,
    handlers: Arc<RwLock<HashMap<String, Callback<ServerState, ClientState>>>>,
    default: DefaultCallback<ServerState, ClientState>,
}

impl<ServerState: 'static + Send + Sync, ClientState: 'static + Default + Send + Sync> Server<ServerState, ClientState> {

    /**
     * with_state will construct the Server with the given state object
     */
    pub fn with_state(state: ServerState) -> Self {
        Server {
            state: Arc::new(state),
            handlers: Arc::new(RwLock::new(HashMap::default())),
            default: Arc::new(Box::new(Server::<ServerState, ClientState>::default_handler)),
        }
    }

    /**
     * Add a handler for a specific message type
     *
     * ```
     * use meows::*;
     * #[macro_use]
     * extern crate serde_derive;
     *
     * #[derive(Debug, Deserialize, Serialize)]
     * struct Ping {
     *     msg: String,
     * }
     *
     * async fn handle_ping(mut req: Request<(), ()>) -> Option<Message> {
     *   if let Some(ping) = req.from_value::<Ping>() {
     *       println!("Ping received: {:?}", ping);
     *   }
     *   Some(Message::text("pong"))
     * }
     *
     * # fn main() {
     * let mut server = Server::new();
     * server.on("ping", handle_ping);
     * # }
     * ```
     */
    pub fn on(&mut self, message_type: &str, invoke: impl Endpoint<ServerState, ClientState>) {
        if let Ok(mut h) = self.handlers.write() {
            h.insert(message_type.to_owned(), Arc::new(Box::new(invoke)));
        }
    }

    /**
     * Set the default message handler, which will be invoked any time that a
     * message is received that cannot be deserialized as a Meows Envelope
     *
     * ```
     * use meows::*;
     * use std::sync::Arc;
     *
     * async fn my_default(message: String, _state: Arc<()>) -> Option<Message> {
     *   None
     * }
     *
     * let mut server = Server::new();
     * server.default(my_default);
     * ```
     */
    pub fn default(&mut self, invoke: impl DefaultEndpoint<ServerState, ClientState>) {
        self.default = Arc::new(Box::new(invoke));
    }

    /**
     * Default handler which is used if the user doesn't specify a handler
     * that should be used for messages Meows doesn't understand
     */
    async fn default_handler(_msg: String, _state: Arc<ServerState>) -> Option<Message> {
        None
    }

    /**
     * The serve() function will listen for inbound webhook connections
     *
     *
     * ```no_run
     * use meows::*;
     * use smol;
     *
     * fn main() -> Result<(), std::io::Error> {
     *   let mut server = Server::new();
     *   smol::run(async move {
     *     server.serve("127.0.0.1:8105".to_string()).await
     *   })
     * }
     * ```
     */
    pub async fn serve(&self, listen_on: String) -> Result<(), std::io::Error> {
        debug!("Starting to listen on: {}", &listen_on);
        let listener = Async::<TcpListener>::bind(listen_on)?;

        loop {
            let (stream, _) = listener.accept().await?;

            match async_tungstenite::accept_async(stream).await {
                Ok(ws) => {
                    let state = self.state.clone();
                    let handlers = self.handlers.clone();
                    let default = self.default.clone();
                    Task::spawn(async move {
                        Server::<ServerState, ClientState>::handle_connection(state, default, handlers, ws)
                            .await;
                    })
                    .detach();
                }
                Err(e) => {
                    error!("Failed to process WebSocket handshake: {}", e);
                }
            }
        }
    }

    /**
     * Handle connection is invoked in its own task for each new WebSocket,
     * from which it will read messages and invoke the appropriate handlers
     */
    async fn handle_connection(
        state: Arc<ServerState>,
        default: DefaultCallback<ServerState, ClientState>,
        handlers: Arc<RwLock<HashMap<String, Callback<ServerState, ClientState>>>>,
        mut stream: WebSocketStream<Async<TcpStream>>,
    ) -> Result<(), std::io::Error> {

        let client_state = Arc::new(RwLock::new(ClientState::default()));

        while let Some(raw) = stream.next().await {
            let client_state = client_state.clone();

            trace!("WebSocket message received: {:?}", raw);
            match raw {
                Ok(message) => {
                    let message = message.to_string();

                    if let Ok(envelope) = serde_json::from_str::<Envelope>(&message) {
                        debug!("Envelope deserialized: {:?}", envelope);

                        let handler = match handlers.read() {
                            Ok(h) => {
                                if let Some(handler) = h.get(&envelope.ttype) {
                                    Some(handler.clone())
                                } else {
                                    debug!("No handler found for message type: {}", envelope.ttype);
                                    None
                                }
                            }
                            _ => None,
                        };

                        if let Some(handler) = handler {
                            let req = Request {
                                env: envelope,
                                state: state.clone(),
                                client_state: client_state.clone(),
                            };

                            if let Some(response) = handler.call(req).await {
                                stream.send(response).await;
                            }
                        }
                    } else {
                        if let Some(response) = default.call(message, state.clone()).await {
                            stream.send(response).await;
                        }
                    }
                }
                Err(e) => {
                    error!("Error receiving message: {}", e);
                }
            }
        }
        Ok(())
    }
}

impl Server<(), ()> {
    pub fn new() -> Self {
        Server {
            state: Arc::new(()),
            handlers: Arc::new(RwLock::new(HashMap::default())),
            default: Arc::new(Box::new(Server::<(), ()>::default_handler)),
        }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
        assert_eq!(2 + 2, 4);
    }
}