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
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
use log::Level::Error as ErrorLevel;
#[cfg(feature = "nativetls")]
use native_tls::{TlsConnector, TlsStream as SslStream};
#[cfg(feature = "ssl")]
use openssl::ssl::{SslConnector, SslMethod, SslStream};
use url;

use frame::Frame;
use handshake::{Handshake, Request, Response};
use message::Message;
use protocol::CloseCode;
use result::{Error, Kind, Result};
use util::{Timeout, Token};

#[cfg(any(feature = "ssl", feature = "nativetls"))]
use util::TcpStream;

/// The core trait of this library.
/// Implementing this trait provides the business logic of the WebSocket application.
pub trait Handler {
    // general

    /// Called when a request to shutdown all connections has been received.
    #[inline]
    fn on_shutdown(&mut self) {
        debug!("Handler received WebSocket shutdown request.");
    }

    // WebSocket events

    /// Called when the WebSocket handshake is successful and the connection is open for sending
    /// and receiving messages.
    fn on_open(&mut self, shake: Handshake) -> Result<()> {
        if let Some(addr) = shake.remote_addr()? {
            debug!("Connection with {} now open", addr);
        }
        Ok(())
    }

    /// Called on incoming messages.
    fn on_message(&mut self, msg: Message) -> Result<()> {
        debug!("Received message {:?}", msg);
        Ok(())
    }

    /// Called any time this endpoint receives a close control frame.
    /// This may be because the other endpoint is initiating a closing handshake,
    /// or it may be the other endpoint confirming the handshake initiated by this endpoint.
    fn on_close(&mut self, code: CloseCode, reason: &str) {
        debug!("Connection closing due to ({:?}) {}", code, reason);
    }

    /// Called when an error occurs on the WebSocket.
    fn on_error(&mut self, err: Error) {
        // Ignore connection reset errors by default, but allow library clients to see them by
        // overriding this method if they want
        if let Kind::Io(ref err) = err.kind {
            if let Some(104) = err.raw_os_error() {
                return;
            }
        }

        error!("{:?}", err);
        if !log_enabled!(ErrorLevel) {
            println!(
                "Encountered an error: {}\nEnable a logger to see more information.",
                err
            );
        }
    }

    // handshake events

    /// A method for handling the low-level workings of the request portion of the WebSocket
    /// handshake.
    ///
    /// Implementors should select a WebSocket protocol and extensions where they are supported.
    ///
    /// Implementors can inspect the Request and must return a Response or an error
    /// indicating that the handshake failed. The default implementation provides conformance with
    /// the WebSocket protocol, and implementors should use the `Response::from_request` method and
    /// then modify the resulting response as necessary in order to maintain conformance.
    ///
    /// This method will not be called when the handler represents a client endpoint. Use
    /// `build_request` to provide an initial handshake request.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let mut res = try!(Response::from_request(req));
    /// if try!(req.extensions()).iter().find(|&&ext| ext.contains("myextension-name")).is_some() {
    ///     res.add_extension("myextension-name")
    /// }
    /// Ok(res)
    /// ```
    #[inline]
    fn on_request(&mut self, req: &Request) -> Result<Response> {
        debug!("Handler received request:\n{}", req);
        Response::from_request(req)
    }

    /// A method for handling the low-level workings of the response portion of the WebSocket
    /// handshake.
    ///
    /// Implementors can inspect the Response and choose to fail the connection by
    /// returning an error. This method will not be called when the handler represents a server
    /// endpoint. The response should indicate which WebSocket protocol and extensions the server
    /// has agreed to if any.
    #[inline]
    fn on_response(&mut self, res: &Response) -> Result<()> {
        debug!("Handler received response:\n{}", res);
        Ok(())
    }

    // timeout events

    /// Called when a timeout is triggered.
    ///
    /// This method will be called when the eventloop encounters a timeout on the specified
    /// token. To schedule a timeout with your specific token use the `Sender::timeout` method.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// const GRATI: Token = Token(1);
    ///
    /// ... Handler
    ///
    /// fn on_open(&mut self, _: Handshake) -> Result<()> {
    ///     // schedule a timeout to send a gratuitous pong every 5 seconds
    ///     self.ws.timeout(5_000, GRATI)
    /// }
    ///
    /// fn on_timeout(&mut self, event: Token) -> Result<()> {
    ///     if event == GRATI {
    ///         // send gratuitous pong
    ///         try!(self.ws.pong(vec![]))
    ///         // reschedule the timeout
    ///         self.ws.timeout(5_000, GRATI)
    ///     } else {
    ///         Err(Error::new(ErrorKind::Internal, "Invalid timeout token encountered!"))
    ///     }
    /// }
    /// ```
    #[inline]
    fn on_timeout(&mut self, event: Token) -> Result<()> {
        debug!("Handler received timeout token: {:?}", event);
        Ok(())
    }

    /// Called when a timeout has been scheduled on the eventloop.
    ///
    /// This method is the hook for obtaining a Timeout object that may be used to cancel a
    /// timeout. This is a noop by default.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// const PING: Token = Token(1);
    /// const EXPIRE: Token = Token(2);
    ///
    /// ... Handler
    ///
    /// fn on_open(&mut self, _: Handshake) -> Result<()> {
    ///     // schedule a timeout to send a ping every 5 seconds
    ///     try!(self.ws.timeout(5_000, PING));
    ///     // schedule a timeout to close the connection if there is no activity for 30 seconds
    ///     self.ws.timeout(30_000, EXPIRE)
    /// }
    ///
    /// fn on_timeout(&mut self, event: Token) -> Result<()> {
    ///     match event {
    ///         PING => {
    ///             self.ws.ping(vec![]);
    ///             self.ws.timeout(5_000, PING)
    ///         }
    ///         EXPIRE => self.ws.close(CloseCode::Away),
    ///         _ => Err(Error::new(ErrorKind::Internal, "Invalid timeout token encountered!")),
    ///     }
    /// }
    ///
    /// fn on_new_timeout(&mut self, event: Token, timeout: Timeout) -> Result<()> {
    ///     if event == EXPIRE {
    ///         if let Some(t) = self.timeout.take() {
    ///             try!(self.ws.cancel(t))
    ///         }
    ///         self.timeout = Some(timeout)
    ///     }
    ///     Ok(())
    /// }
    ///
    /// fn on_frame(&mut self, frame: Frame) -> Result<Option<Frame>> {
    ///     // some activity has occurred, let's reset the expiration
    ///     try!(self.ws.timeout(30_000, EXPIRE));
    ///     Ok(Some(frame))
    /// }
    /// ```
    #[inline]
    fn on_new_timeout(&mut self, _: Token, _: Timeout) -> Result<()> {
        // default implementation discards the timeout handle
        Ok(())
    }

    // frame events

    /// A method for handling incoming frames.
    ///
    /// This method provides very low-level access to the details of the WebSocket protocol. It may
    /// be necessary to implement this method in order to provide a particular extension, but
    /// incorrect implementation may cause the other endpoint to fail the connection.
    ///
    /// Returning `Ok(None)` will cause the connection to forget about a particular frame. This is
    /// useful if you want ot filter out a frame or if you don't want any of the default handler
    /// methods to run.
    ///
    /// By default this method simply ensures that no reserved bits are set.
    #[inline]
    fn on_frame(&mut self, frame: Frame) -> Result<Option<Frame>> {
        debug!("Handler received: {}", frame);
        // default implementation doesn't allow for reserved bits to be set
        if frame.has_rsv1() || frame.has_rsv2() || frame.has_rsv3() {
            Err(Error::new(
                Kind::Protocol,
                "Encountered frame with reserved bits set.",
            ))
        } else {
            Ok(Some(frame))
        }
    }

    /// A method for handling outgoing frames.
    ///
    /// This method provides very low-level access to the details of the WebSocket protocol. It may
    /// be necessary to implement this method in order to provide a particular extension, but
    /// incorrect implementation may cause the other endpoint to fail the connection.
    ///
    /// Returning `Ok(None)` will cause the connection to forget about a particular frame, meaning
    /// that it will not be sent. You can use this approach to merge multiple frames into a single
    /// frame before sending the message.
    ///
    /// For messages, this method will be called with a single complete, final frame before any
    /// fragmentation is performed. Automatic fragmentation will be performed on the returned
    /// frame, if any, based on the `fragment_size` setting.
    ///
    /// By default this method simply ensures that no reserved bits are set.
    #[inline]
    fn on_send_frame(&mut self, frame: Frame) -> Result<Option<Frame>> {
        trace!("Handler will send: {}", frame);
        // default implementation doesn't allow for reserved bits to be set
        if frame.has_rsv1() || frame.has_rsv2() || frame.has_rsv3() {
            Err(Error::new(
                Kind::Protocol,
                "Encountered frame with reserved bits set.",
            ))
        } else {
            Ok(Some(frame))
        }
    }

    // constructors

    /// A method for creating the initial handshake request for WebSocket clients.
    ///
    /// The default implementation provides conformance with the WebSocket protocol, but this
    /// method may be overridden. In order to facilitate conformance,
    /// implementors should use the `Request::from_url` method and then modify the resulting
    /// request as necessary.
    ///
    /// Implementors should indicate any available WebSocket extensions here.
    ///
    /// # Examples
    /// ```ignore
    /// let mut req = try!(Request::from_url(url));
    /// req.add_extension("permessage-deflate; client_max_window_bits");
    /// Ok(req)
    /// ```
    #[inline]
    fn build_request(&mut self, url: &url::Url) -> Result<Request> {
        trace!("Handler is building request to {}.", url);
        Request::from_url(url)
    }

    /// A method for wrapping a client TcpStream with Ssl Authentication machinery
    ///
    /// Override this method to customize how the connection is encrypted. By default
    /// this will use the Server Name Indication extension in conformance with RFC6455.
    #[inline]
    #[cfg(feature = "ssl")]
    fn upgrade_ssl_client(
        &mut self,
        stream: TcpStream,
        url: &url::Url,
    ) -> Result<SslStream<TcpStream>> {
        let domain = url.domain().ok_or(Error::new(
            Kind::Protocol,
            format!("Unable to parse domain from {}. Needed for SSL.", url),
        ))?;
        let connector = SslConnector::builder(SslMethod::tls())
            .map_err(|e| {
                Error::new(
                    Kind::Internal,
                    format!("Failed to upgrade client to SSL: {}", e),
                )
            })?
            .build();
        connector.connect(domain, stream).map_err(Error::from)
    }

    #[inline]
    #[cfg(feature = "nativetls")]
    fn upgrade_ssl_client(
        &mut self,
        stream: TcpStream,
        url: &url::Url,
    ) -> Result<SslStream<TcpStream>> {
        let domain = url.domain().ok_or(Error::new(
            Kind::Protocol,
            format!("Unable to parse domain from {}. Needed for SSL.", url),
        ))?;

        let connector = TlsConnector::new().map_err(|e| {
            Error::new(
                Kind::Internal,
                format!("Failed to upgrade client to SSL: {}", e),
            )
        })?;

        connector.connect(domain, stream).map_err(Error::from)
    }
    /// A method for wrapping a server TcpStream with Ssl Authentication machinery
    ///
    /// Override this method to customize how the connection is encrypted. By default
    /// this method is not implemented.
    #[inline]
    #[cfg(any(feature = "ssl", feature = "nativetls"))]
    fn upgrade_ssl_server(&mut self, _: TcpStream) -> Result<SslStream<TcpStream>> {
        unimplemented!()
    }
}

impl<F> Handler for F
where
    F: Fn(Message) -> Result<()>,
{
    fn on_message(&mut self, msg: Message) -> Result<()> {
        self(msg)
    }
}

mod test {
    #![allow(unused_imports, unused_variables, dead_code)]
    use super::*;
    use frame;
    use handshake::{Handshake, Request, Response};
    use message;
    use mio;
    use protocol::CloseCode;
    use result::Result;
    use url;

    #[derive(Debug, Eq, PartialEq)]
    struct M;
    impl Handler for M {
        fn on_message(&mut self, _: message::Message) -> Result<()> {
            println!("test");
            Ok(())
        }

        fn on_frame(&mut self, f: frame::Frame) -> Result<Option<frame::Frame>> {
            Ok(None)
        }
    }

    #[test]
    fn handler() {
        struct H;

        impl Handler for H {
            fn on_open(&mut self, shake: Handshake) -> Result<()> {
                assert!(shake.request.key().is_ok());
                assert!(shake.response.key().is_ok());
                Ok(())
            }

            fn on_message(&mut self, msg: message::Message) -> Result<()> {
                Ok(assert_eq!(
                    msg,
                    message::Message::Text(String::from("testme"))
                ))
            }

            fn on_close(&mut self, code: CloseCode, _: &str) {
                assert_eq!(code, CloseCode::Normal)
            }
        }

        let mut h = H;
        let url = url::Url::parse("wss://127.0.0.1:3012").unwrap();
        let req = Request::from_url(&url).unwrap();
        let res = Response::from_request(&req).unwrap();
        h.on_open(Handshake {
            request: req,
            response: res,
            peer_addr: None,
            local_addr: None,
        }).unwrap();
        h.on_message(message::Message::Text("testme".to_owned()))
            .unwrap();
        h.on_close(CloseCode::Normal, "");
    }

    #[test]
    fn closure_handler() {
        let mut close = |msg| {
            assert_eq!(msg, message::Message::Binary(vec![1, 2, 3]));
            Ok(())
        };

        close
            .on_message(message::Message::Binary(vec![1, 2, 3]))
            .unwrap();
    }
}