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

//! Methods to connect to an WebSocket as a client.

use std::net::{TcpStream, SocketAddr, ToSocketAddrs};
use std::result::Result as StdResult;
use std::io::{Read, Write};

use url::Url;

#[cfg(feature="tls")]
mod encryption {
    use std::net::TcpStream;
    use native_tls::{TlsConnector, HandshakeError as TlsHandshakeError};
    pub use native_tls::TlsStream;

    pub use stream::Stream as StreamSwitcher;
    /// TCP stream switcher (plain/TLS).
    pub type AutoStream = StreamSwitcher<TcpStream, TlsStream<TcpStream>>;

    use stream::Mode;
    use error::Result;

    pub fn wrap_stream(stream: TcpStream, domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(StreamSwitcher::Plain(stream)),
            Mode::Tls => {
                let connector = TlsConnector::builder()?.build()?;
                connector.connect(domain, stream)
                    .map_err(|e| match e {
                        TlsHandshakeError::Failure(f) => f.into(),
                        TlsHandshakeError::Interrupted(_) => panic!("Bug: TLS handshake not blocked"),
                    })
                    .map(StreamSwitcher::Tls)
            }
        }
    }
}

#[cfg(not(feature="tls"))]
mod encryption {
    use std::net::TcpStream;

    use stream::Mode;
    use error::{Error, Result};

    /// TLS support is nod compiled in, this is just standard `TcpStream`.
    pub type AutoStream = TcpStream;

    pub fn wrap_stream(stream: TcpStream, _domain: &str, mode: Mode) -> Result<AutoStream> {
        match mode {
            Mode::Plain => Ok(stream),
            Mode::Tls => Err(Error::Url("TLS support not compiled in.".into())),
        }
    }
}

pub use self::encryption::AutoStream;
use self::encryption::wrap_stream;

use protocol::WebSocket;
use handshake::HandshakeError;
use handshake::client::{ClientHandshake, Request};
use stream::{NoDelay, Mode};
use error::{Error, Result};


/// Connect to the given WebSocket in blocking mode.
///
/// The URL may be either ws:// or wss://.
/// To support wss:// URLs, feature "tls" must be turned on.
///
/// This function "just works" for those who wants a simple blocking solution
/// similar to `std::net::TcpStream`. If you want a non-blocking or other
/// custom stream, call `client` instead.
///
/// This function uses `native_tls` to do TLS. If you want to use other TLS libraries,
/// use `client` instead. There is no need to enable the "tls" feature if you don't call
/// `connect` since it's the only function that uses native_tls.
pub fn connect<'t, Req: Into<Request<'t>>>(request: Req) -> Result<WebSocket<AutoStream>> {
    let request: Request = request.into();
    let mode = url_mode(&request.url)?;
    let addrs = request.url.to_socket_addrs()?;
    let mut stream = connect_to_some(addrs, &request.url, mode)?;
    NoDelay::set_nodelay(&mut stream, true)?;
    client(request, stream)
        .map_err(|e| match e {
            HandshakeError::Failure(f) => f,
            HandshakeError::Interrupted(_) => panic!("Bug: blocking handshake not blocked"),
        })
}

fn connect_to_some<A>(addrs: A, url: &Url, mode: Mode) -> Result<AutoStream>
    where A: Iterator<Item=SocketAddr>
{
    let domain = url.host_str().ok_or_else(|| Error::Url("No host name in the URL".into()))?;
    for addr in addrs {
        debug!("Trying to contact {} at {}...", url, addr);
        if let Ok(raw_stream) = TcpStream::connect(addr) {
            if let Ok(stream) = wrap_stream(raw_stream, domain, mode) {
                return Ok(stream)
            }
        }
    }
    Err(Error::Url(format!("Unable to connect to {}", url).into()))
}

/// Get the mode of the given URL.
///
/// This function may be used to ease the creation of custom TLS streams
/// in non-blocking algorithmss or for use with TLS libraries other than `native_tls`.
pub fn url_mode(url: &Url) -> Result<Mode> {
    match url.scheme() {
        "ws" => Ok(Mode::Plain),
        "wss" => Ok(Mode::Tls),
        _ => Err(Error::Url("URL scheme not supported".into()))
    }
}

/// Do the client handshake over the given stream.
///
/// Use this function if you need a nonblocking handshake support or if you
/// want to use a custom stream like `mio::tcp::TcpStream` or `openssl::ssl::SslStream`.
/// Any stream supporting `Read + Write` will do.
pub fn client<'t, Stream, Req>(request: Req, stream: Stream)
    -> StdResult<WebSocket<Stream>, HandshakeError<Stream, ClientHandshake>>
where Stream: Read + Write,
      Req: Into<Request<'t>>,
{
    ClientHandshake::start(stream, request.into()).handshake()
}