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
//! Lightweight, event-driven WebSockets for Rust.
#![allow(deprecated)]
#![deny(missing_copy_implementations, trivial_casts, trivial_numeric_casts, unstable_features,
        unused_import_braces)]

extern crate byteorder;
extern crate bytes;
extern crate httparse;
extern crate mio;
extern crate mio_extras;
#[cfg(feature = "ssl")]
extern crate openssl;
#[cfg(feature = "nativetls")]
extern crate native_tls;
extern crate rand;
extern crate sha1;
extern crate slab;
extern crate url;
#[macro_use]
extern crate log;

mod communication;
mod connection;
mod factory;
mod frame;
mod handler;
mod handshake;
mod io;
mod message;
mod protocol;
mod result;
mod stream;

#[cfg(feature = "permessage-deflate")]
pub mod deflate;

pub mod util;

pub use factory::Factory;
pub use handler::Handler;

pub use communication::Sender;
pub use frame::Frame;
pub use handshake::{Handshake, Request, Response};
pub use message::Message;
pub use protocol::{CloseCode, OpCode};
pub use result::Kind as ErrorKind;
pub use result::{Error, Result};

use std::borrow::Borrow;
use std::default::Default;
use std::fmt;
use std::net::{SocketAddr, ToSocketAddrs};

use mio::Poll;

/// A utility function for setting up a WebSocket server.
///
/// # Safety
///
/// This function blocks until the event loop finishes running. Avoid calling this method within
/// another WebSocket handler.
///
/// # Examples
///
/// ```no_run
/// use ws::listen;
///
/// listen("127.0.0.1:3012", |out| {
///     move |msg| {
///        out.send(msg)
///    }
/// }).unwrap()
/// ```
///
pub fn listen<A, F, H>(addr: A, factory: F) -> Result<()>
where
    A: ToSocketAddrs + fmt::Debug,
    F: FnMut(Sender) -> H,
    H: Handler,
{
    let ws = WebSocket::new(factory)?;
    ws.listen(addr)?;
    Ok(())
}

/// A utility function for setting up a WebSocket client.
///
/// # Safety
///
/// This function blocks until the event loop finishes running. Avoid calling this method within
/// another WebSocket handler. If you need to establish a connection from inside of a handler,
/// use the `connect` method on the Sender.
///
/// # Examples
///
/// ```no_run
/// use ws::{connect, CloseCode};
///
/// connect("ws://127.0.0.1:3012", |out| {
///     out.send("Hello WebSocket").unwrap();
///
///     move |msg| {
///         println!("Got message: {}", msg);
///         out.close(CloseCode::Normal)
///     }
/// }).unwrap()
/// ```
///
pub fn connect<U, F, H>(url: U, factory: F) -> Result<()>
where
    U: Borrow<str>,
    F: FnMut(Sender) -> H,
    H: Handler,
{
    let mut ws = WebSocket::new(factory)?;
    let parsed = url::Url::parse(url.borrow()).map_err(|err| {
        Error::new(
            ErrorKind::Internal,
            format!("Unable to parse {} as url due to {:?}", url.borrow(), err),
        )
    })?;
    ws.connect(parsed)?;
    ws.run()?;
    Ok(())
}

/// WebSocket settings
#[derive(Debug, Clone, Copy)]
pub struct Settings {
    /// The maximum number of connections that this WebSocket will support.
    /// The default setting is low and should be increased when expecting more
    /// connections because this is a hard limit and no new connections beyond
    /// this limit can be made until an old connection is dropped.
    /// Default: 100
    pub max_connections: usize,
    /// The number of events anticipated per connection. The event loop queue size will
    /// be `queue_size` * `max_connections`. In order to avoid an overflow error,
    /// `queue_size` * `max_connections` must be less than or equal to `usize::max_value()`.
    /// The queue is shared between connections, which means that a connection may schedule
    /// more events than `queue_size` provided that another connection is using less than
    /// `queue_size`. However, if the queue is maxed out a Queue error will occur.
    /// Default: 5
    pub queue_size: usize,
    /// Whether to panic when unable to establish a new TCP connection.
    /// Default: false
    pub panic_on_new_connection: bool,
    /// Whether to panic when a shutdown of the WebSocket is requested.
    /// Default: false
    pub panic_on_shutdown: bool,
    /// The maximum number of fragments the connection can handle without reallocating.
    /// Default: 10
    pub fragments_capacity: usize,
    /// Whether to reallocate when `fragments_capacity` is reached. If this is false,
    /// a Capacity error will be triggered instead.
    /// Default: true
    pub fragments_grow: bool,
    /// The maximum length of outgoing frames. Messages longer than this will be fragmented.
    /// Default: 65,535
    pub fragment_size: usize,
    /// The size of the incoming buffer. A larger buffer uses more memory but will allow for fewer
    /// reallocations.
    /// Default: 2048
    pub in_buffer_capacity: usize,
    /// Whether to reallocate the incoming buffer when `in_buffer_capacity` is reached. If this is
    /// false, a Capacity error will be triggered instead.
    /// Default: true
    pub in_buffer_grow: bool,
    /// The size of the outgoing buffer. A larger buffer uses more memory but will allow for fewer
    /// reallocations.
    /// Default: 2048
    pub out_buffer_capacity: usize,
    /// Whether to reallocate the incoming buffer when `out_buffer_capacity` is reached. If this is
    /// false, a Capacity error will be triggered instead.
    /// Default: true
    pub out_buffer_grow: bool,
    /// Whether to panic when an Internal error is encountered. Internal errors should generally
    /// not occur, so this setting defaults to true as a debug measure, whereas production
    /// applications should consider setting it to false.
    /// Default: true
    pub panic_on_internal: bool,
    /// Whether to panic when a Capacity error is encountered.
    /// Default: false
    pub panic_on_capacity: bool,
    /// Whether to panic when a Protocol error is encountered.
    /// Default: false
    pub panic_on_protocol: bool,
    /// Whether to panic when an Encoding error is encountered.
    /// Default: false
    pub panic_on_encoding: bool,
    /// Whether to panic when a Queue error is encountered.
    /// Default: false
    pub panic_on_queue: bool,
    /// Whether to panic when an Io error is encountered.
    /// Default: false
    pub panic_on_io: bool,
    /// Whether to panic when a Timer error is encountered.
    /// Default: false
    pub panic_on_timeout: bool,
    /// Whether to shutdown the eventloop when an interrupt is received.
    /// Default: true
    pub shutdown_on_interrupt: bool,
    /// The WebSocket protocol requires frames sent from client endpoints to be masked as a
    /// security and sanity precaution. Enforcing this requirement, which may be removed at some
    /// point may cause incompatibilities. If you need the extra security, set this to true.
    /// Default: false
    pub masking_strict: bool,
    /// The WebSocket protocol requires clients to verify the key returned by a server to ensure
    /// that the server and all intermediaries can perform the protocol. Verifying the key will
    /// consume processing time and other resources with the benefit that we can fail the
    /// connection early. The default in WS-RS is to accept any key from the server and instead
    /// fail late if a protocol error occurs. Change this setting to enable key verification.
    /// Default: false
    pub key_strict: bool,
    /// The WebSocket protocol requires clients to perform an opening handshake using the HTTP
    /// GET method for the request. However, since only WebSockets are supported on the connection,
    /// verifying the method of handshake requests is not always necessary. To enforce the
    /// requirement that handshakes begin with a GET method, set this to true.
    /// Default: false
    pub method_strict: bool,
    /// Indicate whether server connections should use ssl encryption when accepting connections.
    /// Setting this to true means that clients should use the `wss` scheme to connect to this
    /// server. Note that using this flag will in general necessitate overriding the
    /// `Handler::upgrade_ssl_server` method in order to provide the details of the ssl context. It may be
    /// simpler for most users to use a reverse proxy such as nginx to provide server side
    /// encryption.
    ///
    /// Default: false
    pub encrypt_server: bool,
    /// Disables Nagle's algorithm.
    /// Usually tcp socket tries to accumulate packets to send them all together (every 200ms).
    /// When enabled socket will try to send packet as fast as possible.
    ///
    /// Default: false
    pub tcp_nodelay: bool,
    /// Size of mio's timer each tick in milliseconds.
    ///
    /// Default: 100
    pub timer_tick_millis: u64
}

impl Default for Settings {
    fn default() -> Settings {
        Settings {
            max_connections: 100,
            queue_size: 5,
            panic_on_new_connection: false,
            panic_on_shutdown: false,
            fragments_capacity: 10,
            fragments_grow: true,
            fragment_size: u16::max_value() as usize,
            in_buffer_capacity: 2048,
            in_buffer_grow: true,
            out_buffer_capacity: 2048,
            out_buffer_grow: true,
            panic_on_internal: true,
            panic_on_capacity: false,
            panic_on_protocol: false,
            panic_on_encoding: false,
            panic_on_queue: false,
            panic_on_io: false,
            panic_on_timeout: false,
            shutdown_on_interrupt: true,
            masking_strict: false,
            key_strict: false,
            method_strict: false,
            encrypt_server: false,
            tcp_nodelay: false,
            timer_tick_millis: 100
        }
    }
}

/// The WebSocket struct. A WebSocket can support multiple incoming and outgoing connections.
pub struct WebSocket<F>
where
    F: Factory,
{
    poll: Poll,
    handler: io::Handler<F>,
}

impl<F> WebSocket<F>
where
    F: Factory,
{
    /// Create a new WebSocket using the given Factory to create handlers.
    pub fn new(factory: F) -> Result<WebSocket<F>> {
        Builder::new().build(factory)
    }

    /// Consume the WebSocket and bind to the specified address.
    /// If the `addr_spec` yields multiple addresses this will return after the
    /// first successful bind. `local_addr` can be called to determine which
    /// address it ended up binding to.
    /// After the server is successfully bound you should start it using `run`.
    pub fn bind<A>(mut self, addr_spec: A) -> Result<WebSocket<F>>
    where
        A: ToSocketAddrs,
    {
        let mut last_error = Error::new(ErrorKind::Internal, "No address given");

        for addr in addr_spec.to_socket_addrs()? {
            if let Err(e) = self.handler.listen(&mut self.poll, &addr) {
                error!("Unable to listen on {}", addr);
                last_error = e;
            } else {
                let actual_addr = self.handler.local_addr().unwrap_or(addr);
                info!("Listening for new connections on {}.", actual_addr);
                return Ok(self);
            }
        }

        Err(last_error)
    }

    /// Consume the WebSocket and listen for new connections on the specified address.
    ///
    /// # Safety
    ///
    /// This method will block until the event loop finishes running.
    pub fn listen<A>(self, addr_spec: A) -> Result<WebSocket<F>>
    where
        A: ToSocketAddrs,
    {
        self.bind(addr_spec).and_then(|server| server.run())
    }

    /// Queue an outgoing connection on this WebSocket. This method may be called multiple times,
    /// but the actual connections will not be established until `run` is called.
    pub fn connect(&mut self, url: url::Url) -> Result<&mut WebSocket<F>> {
        let sender = self.handler.sender();
        info!("Queuing connection to {}", url);
        sender.connect(url)?;
        Ok(self)
    }

    /// Run the WebSocket. This will run the encapsulated event loop blocking the calling thread until
    /// the WebSocket is shutdown.
    pub fn run(mut self) -> Result<WebSocket<F>> {
        self.handler.run(&mut self.poll)?;
        Ok(self)
    }

    /// Get a Sender that can be used to send messages on all connections.
    /// Calling `send` on this Sender is equivalent to calling `broadcast`.
    /// Calling `shutdown` on this Sender will shutdown the WebSocket even if no connections have
    /// been established.
    #[inline]
    pub fn broadcaster(&self) -> Sender {
        self.handler.sender()
    }

    /// Get the local socket address this socket is bound to. Will return an error
    /// if the backend returns an error. Will return a `NotFound` error if
    /// this WebSocket is not a listening socket.
    pub fn local_addr(&self) -> ::std::io::Result<SocketAddr> {
        self.handler.local_addr()
    }
}

/// Utility for constructing a WebSocket from various settings.
#[derive(Debug, Default, Clone, Copy)]
pub struct Builder {
    settings: Settings,
}

// TODO: add convenience methods for each setting
impl Builder {
    /// Create a new Builder with default settings.
    pub fn new() -> Builder {
        Builder::default()
    }

    /// Build a WebSocket using this builder and a factory.
    /// It is possible to use the same builder to create multiple WebSockets.
    pub fn build<F>(&self, factory: F) -> Result<WebSocket<F>>
    where
        F: Factory,
    {
        Ok(WebSocket {
            poll: Poll::new()?,
            handler: io::Handler::new(factory, self.settings),
        })
    }

    /// Set the WebSocket settings to use.
    pub fn with_settings(&mut self, settings: Settings) -> &mut Builder {
        self.settings = settings;
        self
    }
}