twisty 0.3.1

//! twisty tokio-proto implementation
use config::Config;
use futures::Future;
use futures::future;
use native_tls::TlsConnector;
use std::error::Error;
use std::io::{self, Write};
use std::net::SocketAddr;
use term;
use tokio_core::net::TcpStream;
use tokio_core::reactor::Handle;
use tokio_proto::TcpClient;
use tokio_proto::pipeline::ClientService;
use tokio_service::Service;
use tokio_tls::proto::Client as TlsClientProto;
use twist::client::{BaseFrame, OpCode, WebSocketFrame, WebSocketProtocol};
use twist_lz4::ClientLz4;

/// Generate an `io::ErrorKind::Other` error with the given description.
pub fn other(desc: &str) -> io::Error {
    io::Error::new(io::ErrorKind::Other, desc)
}

/// tokio-service for unsecure twisty connections.
pub struct Client {
    /// inner
    inner: ClientService<TcpStream, WebSocketProtocol>,
}

/// Generate the twisty prompt.
fn prompt() -> Result<(), io::Error> {
    let mut t = term::stdout().ok_or_else(|| other("invalid term"))?;
    t.fg(term::color::BRIGHT_GREEN)?;
    write!(t, "twisty > ")?;
    t.reset()?;
    t.flush()?;
    Ok(())
}

impl Client {
    /// Establish a connection to a server at the provided `addr`.
    pub fn connect(config: &Config,
                   socket_addr: SocketAddr,
                   handle: &Handle)
                   -> Box<Future<Item = Client, Error = io::Error>> {
        let mut lz4: ClientLz4 = Default::default();
        lz4.stdout(config.dual()).stderr(config.dual());
        lz4.set_enabled(true);
        let mut ws_proto: WebSocketProtocol = Default::default();
        ws_proto.stdout(config.dual());
        ws_proto.stderr(config.dual());
        ws_proto.client(true);
        ws_proto.per_message(lz4);
        let ret = TcpClient::new(ws_proto)
            .connect(&socket_addr, handle)
            .map(|client_service| Client { inner: client_service });
        Box::new(ret)
    }

    /// Read a line from stdin generate a call future as a text base frame.
    pub fn read_line(&self) -> Box<Future<Item = WebSocketFrame, Error = io::Error>> {
        let mut buf = String::new();
        if prompt().is_err() {
            out!("twisty > ");
        }

        match io::stdin().read_line(&mut buf) {
            Ok(_) => {
                if buf.starts_with("exit") {
                    return Box::new(future::err(other("exit")));
                }
                let mut frame: WebSocketFrame = Default::default();
                let mut base: BaseFrame = Default::default();
                base.set_fin(true);
                base.set_masked(true);
                base.set_mask(0);
                base.set_opcode(OpCode::Text);
                base.set_payload_length(buf.len() as u64);

                // Compress if longer than 20 bytes.
                if buf.len() > 20 {
                    base.set_rsv2(true);
                }

                base.set_application_data(buf.as_bytes().to_vec());
                frame.set_base(base);
                self.call(frame)
            }
            Err(e) => Box::new(future::err(e)),
        }
    }
}

impl Service for Client {
    type Request = WebSocketFrame;
    type Response = WebSocketFrame;
    type Error = io::Error;
    // For simplicity, box the future.
    type Future = Box<Future<Item = WebSocketFrame, Error = io::Error>>;

    fn call(&self, req: WebSocketFrame) -> Self::Future {
        Box::new(self.inner.call(req).then(|result| match result {
                                               Ok(resp) => Box::new(future::done(Ok(resp))),
                                               Err(e) => Box::new(future::err(e)),
                                           }))
    }
}

/// tokio-service for secure twisty connections.
pub struct Tls {
    /// Inner
    inner: ClientService<TcpStream, TlsClientProto<WebSocketProtocol>>,
}

impl Tls {
    /// Establish a connection to a server at the provided `addr`.
    pub fn connect(config: &Config,
                   socket_addr: SocketAddr,
                   handle: &Handle)
                   -> Box<Future<Item = Tls, Error = io::Error>> {
        let builder = match TlsConnector::builder() {
            Ok(builder) => builder,
            Err(e) => return Box::new(future::err(other(e.description()))),
        };
        let connector = match builder.build() {
            Ok(connector) => connector,
            Err(e) => return Box::new(future::err(other(e.description()))),
        };
        let mut lz4: ClientLz4 = Default::default();
        lz4.stdout(config.dual()).stderr(config.dual());
        lz4.set_enabled(true);
        let mut ws_proto: WebSocketProtocol = Default::default();
        ws_proto.stdout(config.dual());
        ws_proto.stderr(config.dual());
        ws_proto.client(true);
        ws_proto.per_message(lz4);
        let tls_proto = TlsClientProto::new(ws_proto, connector, config.host());
        let ret = TcpClient::new(tls_proto)
            .connect(&socket_addr, handle)
            .map(|client_service| Tls { inner: client_service });
        Box::new(ret)
    }

    /// Read a line from stdin generate a call future as a text base frame.
    pub fn read_line(&self) -> Box<Future<Item = WebSocketFrame, Error = io::Error>> {
        let mut buf = String::new();
        if prompt().is_err() {
            out!("twisty > ");
        }

        match io::stdin().read_line(&mut buf) {
            Ok(_) => {
                if buf.starts_with("exit") {
                    return Box::new(future::err(other("exit")));
                }
                let mut frame: WebSocketFrame = Default::default();
                let mut base: BaseFrame = Default::default();
                base.set_fin(true);
                base.set_masked(true);
                base.set_mask(0);
                base.set_opcode(OpCode::Text);
                base.set_payload_length(buf.len() as u64);

                // Compress if longer than 20 bytes.
                if buf.len() > 20 {
                    base.set_rsv2(true);
                }

                base.set_application_data(buf.as_bytes().to_vec());
                frame.set_base(base);
                self.call(frame)
            }
            Err(e) => Box::new(future::err(e)),
        }
    }
}

impl Service for Tls {
    type Request = WebSocketFrame;
    type Response = WebSocketFrame;
    type Error = io::Error;
    // For simplicity, box the future.
    type Future = Box<Future<Item = WebSocketFrame, Error = io::Error>>;

    fn call(&self, req: WebSocketFrame) -> Self::Future {
        Box::new(self.inner.call(req).then(|result| match result {
                                               Ok(resp) => Box::new(future::done(Ok(resp))),
                                               Err(e) => Box::new(future::err(e)),
                                           }))
    }
}