basic websocket client

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "actix-web"
-version = "0.3.3"
+version = "0.3.4"
 authors = ["Nikolay Kim <fafhrd91@gmail.com>"]
 description = "Actix web framework"
 readme = "README.md"
@@ -34,28 +34,29 @@ tls = ["native-tls", "tokio-tls"]
 alpn = ["openssl", "openssl/v102", "openssl/v110", "tokio-openssl"]
 
 [dependencies]
-log = "0.4"
-failure = "0.1"
-failure_derive = "0.1"
+base64 = "0.9"
+bitflags = "1.0"
+brotli2 = "^0.3.2"
+failure = "0.1.1"
+flate2 = "1.0"
 h2 = "0.1"
 http = "^0.1.2"
 httparse = "1.2"
 http-range = "0.1"
-time = "0.1"
+libc = "0.2"
+log = "0.4"
 mime = "0.3"
 mime_guess = "1.8"
+num_cpus = "1.0"
+percent-encoding = "1.0"
+rand = "0.4"
 regex = "0.2"
-sha1 = "0.4"
-url = "1.6"
-libc = "0.2"
 serde = "1.0"
 serde_json = "1.0"
-brotli2 = "^0.3.2"
-percent-encoding = "1.0"
+sha1 = "0.4"
 smallvec = "0.6"
-bitflags = "1.0"
-num_cpus = "1.0"
-flate2 = "1.0"
+time = "0.1"
+url = "1.6"
 cookie = { version="0.10", features=["percent-encode", "secure"] }
 
 # io
@@ -65,6 +66,7 @@ bytes = "0.4"
 futures = "0.1"
 tokio-io = "0.1"
 tokio-core = "0.1"
+trust-dns-resolver = "0.7"
 
 # native-tls
 native-tls = { version="0.1", optional = true }

examples/websocket/Cargo.toml

@@ -7,8 +7,14 @@ authors = ["Nikolay Kim <fafhrd91@gmail.com>"]
 name = "server"
 path = "src/main.rs"
 
+[[bin]]
+name = "client"
+path = "src/client.rs"
+
 [dependencies]
 env_logger = "*"
 futures = "0.1"
-actix = "^0.4.2"
-actix-web = { git = "https://github.com/actix/actix-web.git" }
+tokio-core = "0.1"
+#actix = "^0.4.2"
+actix = { git = "https://github.com/actix/actix.git" }
+actix-web = { path="../../" }

examples/websocket/src/client.rs

@@ -0,0 +1,108 @@
+//! Simple websocket client.
+
+#![allow(unused_variables)]
+extern crate actix;
+extern crate actix_web;
+extern crate env_logger;
+extern crate futures;
+extern crate tokio_core;
+
+use std::{io, thread};
+use std::time::Duration;
+
+use actix::*;
+use futures::Future;
+use tokio_core::net::TcpStream;
+use actix_web::ws::{client, Message, WsClientError};
+
+
+fn main() {
+    ::std::env::set_var("RUST_LOG", "actix_web=info");
+    let _ = env_logger::init();
+    let sys = actix::System::new("ws-example");
+
+    Arbiter::handle().spawn(
+        client::WsClient::new("http://127.0.0.1:8080/ws/")
+            .connect().unwrap()
+            .map_err(|e| {
+                println!("Error: {}", e);
+                ()
+            })
+            .map(|(reader, writer)| {
+                let addr: SyncAddress<_> = ChatClient::create(|ctx| {
+                    ctx.add_stream(reader);
+                    ChatClient(writer)
+                });
+
+                // start console loop
+                thread::spawn(move|| {
+                    loop {
+                        let mut cmd = String::new();
+                        if io::stdin().read_line(&mut cmd).is_err() {
+                            println!("error");
+                            return
+                        }
+                        addr.send(ClientCommand(cmd));
+                    }
+                });
+
+                ()
+            })
+    );
+
+    let _ = sys.run();
+}
+
+
+struct ChatClient(client::WsWriter<TcpStream>);
+
+#[derive(Message)]
+struct ClientCommand(String);
+
+impl Actor for ChatClient {
+    type Context = Context<Self>;
+
+    fn started(&mut self, ctx: &mut Context<Self>) {
+        // start heartbeats otherwise server will disconnect after 10 seconds
+        self.hb(ctx)
+    }
+
+    fn stopping(&mut self, _: &mut Context<Self>) -> bool {
+        println!("Disconnected");
+
+        // Stop application on disconnect
+        Arbiter::system().send(actix::msgs::SystemExit(0));
+
+        true
+    }
+}
+
+impl ChatClient {
+    fn hb(&self, ctx: &mut Context<Self>) {
+        ctx.run_later(Duration::new(1, 0), |act, ctx| {
+            act.0.ping("");
+            act.hb(ctx);
+        });
+    }
+}
+
+/// Handle stdin commands
+impl Handler<ClientCommand> for ChatClient {
+    type Result = ();
+
+    fn handle(&mut self, msg: ClientCommand, ctx: &mut Context<Self>) {
+        self.0.text(msg.0.as_str())
+    }
+}
+
+/// Handle server websocket messages
+impl Handler<Result<Message, WsClientError>> for ChatClient {
+    type Result = ();
+
+    fn handle(&mut self, msg: Result<Message, WsClientError>, ctx: &mut Context<Self>) {
+        match msg {
+            Ok(Message::Text(txt)) => println!("Server: {:?}", txt),
+            _ => ()
+        }
+    }
+}

examples/websocket/src/main.rs

@@ -44,7 +44,7 @@ impl Handler<ws::Message> for MyWebSocket {
 }
 
 fn main() {
-    ::std::env::set_var("RUST_LOG", "actix_web=trace");
+    ::std::env::set_var("RUST_LOG", "actix_web=info");
     let _ = env_logger::init();
     let sys = actix::System::new("ws-example");
 

src/client/mod.rs

@@ -0,0 +1,5 @@
+mod parser;
+mod response;
+
+pub use self::response::ClientResponse;
+pub use self::parser::{HttpResponseParser, HttpResponseParserError};

src/client/parser.rs

@@ -0,0 +1,239 @@
+use std::mem;
+use httparse;
+use http::{Version, HttpTryFrom, HeaderMap, StatusCode};
+use http::header::{self, HeaderName, HeaderValue};
+use bytes::BytesMut;
+use futures::{Poll, Async};
+
+use error::{ParseError, PayloadError};
+use payload::{Payload, PayloadWriter, DEFAULT_BUFFER_SIZE};
+
+use server::{utils, IoStream};
+use server::h1::{Decoder, chunked};
+use server::encoding::PayloadType;
+
+use super::ClientResponse;
+
+const MAX_BUFFER_SIZE: usize = 131_072;
+const MAX_HEADERS: usize = 96;
+
+
+pub struct HttpResponseParser {
+    payload: Option<PayloadInfo>,
+}
+
+enum Decoding {
+    Paused,
+    Ready,
+    NotReady,
+}
+
+struct PayloadInfo {
+    tx: PayloadType,
+    decoder: Decoder,
+}
+
+#[derive(Debug)]
+pub enum HttpResponseParserError {
+    Disconnect,
+    Payload,
+    Error(ParseError),
+}
+
+impl HttpResponseParser {
+    pub fn new() -> HttpResponseParser {
+        HttpResponseParser {
+            payload: None,
+        }
+    }
+
+    fn decode(&mut self, buf: &mut BytesMut) -> Result<Decoding, HttpResponseParserError> {
+        if let Some(ref mut payload) = self.payload {
+            if payload.tx.capacity() > DEFAULT_BUFFER_SIZE {
+                return Ok(Decoding::Paused)
+            }
+            loop {
+                match payload.decoder.decode(buf) {
+                    Ok(Async::Ready(Some(bytes))) => {
+                        payload.tx.feed_data(bytes)
+                    },
+                    Ok(Async::Ready(None)) => {
+                        payload.tx.feed_eof();
+                        return Ok(Decoding::Ready)
+                    },
+                    Ok(Async::NotReady) => return Ok(Decoding::NotReady),
+                    Err(err) => {
+                        payload.tx.set_error(err.into());
+                        return Err(HttpResponseParserError::Payload)
+                    }
+                }
+            }
+        } else {
+            return Ok(Decoding::Ready)
+        }
+    }
+
+    pub fn parse<T>(&mut self, io: &mut T, buf: &mut BytesMut)
+                    -> Poll<ClientResponse, HttpResponseParserError>
+        where T: IoStream
+    {
+        // read payload
+        if self.payload.is_some() {
+            match utils::read_from_io(io, buf) {
+                Ok(Async::Ready(0)) => {
+                    if let Some(ref mut payload) = self.payload {
+                        payload.tx.set_error(PayloadError::Incomplete);
+                    }
+                    // http channel should not deal with payload errors
+                    return Err(HttpResponseParserError::Payload)
+                },
+                Err(err) => {
+                    if let Some(ref mut payload) = self.payload {
+                        payload.tx.set_error(err.into());
+                    }
+                    // http channel should not deal with payload errors
+                    return Err(HttpResponseParserError::Payload)
+                }
+                _ => (),
+            }
+            match self.decode(buf)? {
+                Decoding::Ready => self.payload = None,
+                Decoding::Paused | Decoding::NotReady => return Ok(Async::NotReady),
+            }
+        }
+
+        // if buf is empty parse_message will always return NotReady, let's avoid that
+        let read = if buf.is_empty() {
+            match utils::read_from_io(io, buf) {
+                Ok(Async::Ready(0)) => {
+                    // debug!("Ignored premature client disconnection");
+                    return Err(HttpResponseParserError::Disconnect);
+                },
+                Ok(Async::Ready(_)) => (),
+                Ok(Async::NotReady) =>
+                    return Ok(Async::NotReady),
+                Err(err) =>
+                    return Err(HttpResponseParserError::Error(err.into()))
+            }
+            false
+        } else {
+            true
+        };
+
+        loop {
+            match HttpResponseParser::parse_message(buf).map_err(HttpResponseParserError::Error)? {
+                Async::Ready((msg, decoder)) => {
+                    // process payload
+                    if let Some(payload) = decoder {
+                        self.payload = Some(payload);
+                        match self.decode(buf)? {
+                            Decoding::Paused | Decoding::NotReady => (),
+                            Decoding::Ready => self.payload = None,
+                        }
+                    }
+                    return Ok(Async::Ready(msg));
+                },
+                Async::NotReady => {
+                    if buf.capacity() >= MAX_BUFFER_SIZE {
+                        error!("MAX_BUFFER_SIZE unprocessed data reached, closing");
+                        return Err(HttpResponseParserError::Error(ParseError::TooLarge));
+                    }
+                    if read {
+                        match utils::read_from_io(io, buf) {
+                            Ok(Async::Ready(0)) => {
+                                debug!("Ignored premature client disconnection");
+                                return Err(HttpResponseParserError::Disconnect);
+                            },
+                            Ok(Async::Ready(_)) => (),
+                            Ok(Async::NotReady) =>
+                                return Ok(Async::NotReady),
+                            Err(err) =>
+                                return Err(HttpResponseParserError::Error(err.into()))
+                        }
+                    } else {
+                        return Ok(Async::NotReady)
+                    }
+                },
+            }
+        }
+    }
+
+    fn parse_message(buf: &mut BytesMut) -> Poll<(ClientResponse, Option<PayloadInfo>), ParseError>
+    {
+        // Parse http message
+        let bytes_ptr = buf.as_ref().as_ptr() as usize;
+        let mut headers: [httparse::Header; MAX_HEADERS] =
+            unsafe{mem::uninitialized()};
+
+        let (len, version, status, headers_len) = {
+            let b = unsafe{ let b: &[u8] = buf; mem::transmute(b) };
+            let mut resp = httparse::Response::new(&mut headers);
+            match resp.parse(b)? {
+                httparse::Status::Complete(len) => {
+                    let version = if resp.version.unwrap() == 1 {
+                        Version::HTTP_11
+                    } else {
+                        Version::HTTP_10
+                    };
+                    let status = StatusCode::from_u16(resp.code.unwrap())
+                        .map_err(|_| ParseError::Status)?;
+
+                    (len, version, status, resp.headers.len())
+                }
+                httparse::Status::Partial => return Ok(Async::NotReady),
+            }
+        };
+
+
+        let slice = buf.split_to(len).freeze();
+
+        // convert headers
+        let mut hdrs = HeaderMap::new();
+        for header in headers[..headers_len].iter() {
+            if let Ok(name) = HeaderName::try_from(header.name) {
+                let v_start = header.value.as_ptr() as usize - bytes_ptr;
+                let v_end = v_start + header.value.len();
+                let value = unsafe {
+                    HeaderValue::from_shared_unchecked(slice.slice(v_start, v_end)) };
+                hdrs.append(name, value);
+            } else {
+                return Err(ParseError::Header)
+            }
+        }
+
+        let decoder = if let Some(len) = hdrs.get(header::CONTENT_LENGTH) {
+            // Content-Length
+            if let Ok(s) = len.to_str() {
+                if let Ok(len) = s.parse::<u64>() {
+                    Some(Decoder::length(len))
+                } else {
+                    debug!("illegal Content-Length: {:?}", len);
+                    return Err(ParseError::Header)
+                }
+            } else {
+                debug!("illegal Content-Length: {:?}", len);
+                return Err(ParseError::Header)
+            }
+        } else if chunked(&hdrs)? {
+            // Chunked encoding
+            Some(Decoder::chunked())
+        } else if hdrs.contains_key(header::UPGRADE) {
+            Some(Decoder::eof())
+        } else {
+            None
+        };
+
+        if let Some(decoder) = decoder {
+            let (psender, payload) = Payload::new(false);
+            let info = PayloadInfo {
+                tx: PayloadType::new(&hdrs, psender),
+                decoder: decoder,
+            };
+            Ok(Async::Ready(
+                (ClientResponse::new(status, version, hdrs, Some(payload)), Some(info))))
+        } else {
+            Ok(Async::Ready(
+                (ClientResponse::new(status, version, hdrs, None), None)))
+        }
+    }
+}

src/client/response.rs

@@ -0,0 +1,76 @@
+#![allow(dead_code)]
+use std::fmt;
+use http::{HeaderMap, StatusCode, Version};
+use http::header::HeaderValue;
+
+use payload::Payload;
+
+
+pub struct ClientResponse {
+    /// The response's status
+    status: StatusCode,
+
+    /// The response's version
+    version: Version,
+
+    /// The response's headers
+    headers: HeaderMap<HeaderValue>,
+
+    payload: Option<Payload>,
+}
+
+impl ClientResponse {
+    pub fn new(status: StatusCode, version: Version,
+               headers: HeaderMap<HeaderValue>, payload: Option<Payload>) -> Self {
+        ClientResponse {
+            status: status, version: version, headers: headers, payload: payload
+        }
+    }
+
+ /// Get the HTTP version of this response.
+    #[inline]
+    pub fn version(&self) -> Version {
+        self.version
+    }
+
+    /// Get the headers from the response.
+    #[inline]
+    pub fn headers(&self) -> &HeaderMap {
+        &self.headers
+    }
+
+    /// Get a mutable reference to the headers.
+    #[inline]
+    pub fn headers_mut(&mut self) -> &mut HeaderMap {
+        &mut self.headers
+    }
+
+    /// Get the status from the server.
+    #[inline]
+    pub fn status(&self) -> StatusCode {
+        self.status
+    }
+
+    /// Set the `StatusCode` for this response.
+    #[inline]
+    pub fn status_mut(&mut self) -> &mut StatusCode {
+        &mut self.status
+    }
+}
+
+impl fmt::Debug for ClientResponse {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let res = write!(
+            f, "\nClientResponse {:?} {}\n", self.version, self.status);
+        let _ = write!(f, "  headers:\n");
+        for key in self.headers.keys() {
+            let vals: Vec<_> = self.headers.get_all(key).iter().collect();
+            if vals.len() > 1 {
+                let _ = write!(f, "    {:?}: {:?}\n", key, vals);
+            } else {
+                let _ = write!(f, "    {:?}: {:?}\n", key, vals[0]);
+            }
+        }
+        res
+    }
+}

src/error.rs

@@ -16,7 +16,7 @@ use http::{header, StatusCode, Error as HttpError};
 use http::uri::InvalidUriBytes;
 use http_range::HttpRangeParseError;
 use serde_json::error::Error as JsonError;
-use url::ParseError as UrlParseError;
+pub use url::ParseError as UrlParseError;
 
 // re-exports
 pub use cookie::{ParseError as CookieParseError};
@@ -106,6 +106,9 @@ default impl<T: StdError + Sync + Send + 'static> ResponseError for T {
 /// `InternalServerError` for `JsonError`
 impl ResponseError for JsonError {}
 
+/// `InternalServerError` for `UrlParseError`
+impl ResponseError for UrlParseError {}
+
 /// Return `InternalServerError` for `HttpError`,
 /// Response generation can return `HttpError`, so it is internal error
 impl ResponseError for HttpError {}

src/httpresponse.rs

@@ -1,17 +1,15 @@
-//! Pieces pertaining to the HTTP response.
+//! Http response
 use std::{mem, str, fmt};
 use std::io::Write;
 use std::cell::RefCell;
-use std::convert::Into;
 use std::collections::VecDeque;
 
-use cookie::CookieJar;
+use cookie::{Cookie, CookieJar};
 use bytes::{Bytes, BytesMut, BufMut};
 use http::{StatusCode, Version, HeaderMap, HttpTryFrom, Error as HttpError};
 use http::header::{self, HeaderName, HeaderValue};
 use serde_json;
 use serde::Serialize;
-use cookie::Cookie;
 
 use body::Body;
 use error::Error;
@@ -261,7 +259,6 @@ impl HttpResponseBuilder {
     /// }
     /// fn main() {}
     /// ```
-    #[inline]
     pub fn header<K, V>(&mut self, key: K, value: V) -> &mut Self
         where HeaderName: HttpTryFrom<K>,
               HeaderValue: HttpTryFrom<V>

src/lib.rs

@@ -47,6 +47,7 @@
 #[macro_use]
 extern crate log;
 extern crate time;
+extern crate base64;
 extern crate bytes;
 extern crate sha1;
 extern crate regex;
@@ -66,6 +67,7 @@ extern crate httparse;
 extern crate http_range;
 extern crate mime;
 extern crate mime_guess;
+extern crate rand;
 extern crate url;
 extern crate libc;
 extern crate serde;
@@ -76,6 +78,7 @@ extern crate percent_encoding;
 extern crate smallvec;
 extern crate num_cpus;
 extern crate h2 as http2;
+extern crate trust_dns_resolver;
 #[macro_use] extern crate actix;
 
 #[cfg(test)]
@@ -106,6 +109,8 @@ mod resource;
 mod handler;
 mod pipeline;
 
+mod client;
+
 pub mod fs;
 pub mod ws;
 pub mod error;

src/server/h1.rs

@@ -550,7 +550,7 @@ impl Reader {
 }
 
 /// Check if request has chunked transfer encoding
-fn chunked(headers: &HeaderMap) -> Result<bool, ParseError> {
+pub fn chunked(headers: &HeaderMap) -> Result<bool, ParseError> {
     if let Some(encodings) = headers.get(header::TRANSFER_ENCODING) {
         if let Ok(s) = encodings.to_str() {
             Ok(s.to_lowercase().contains("chunked"))
@@ -567,7 +567,7 @@ fn chunked(headers: &HeaderMap) -> Result<bool, ParseError> {
 /// If a message body does not include a Transfer-Encoding, it *should*
 /// include a Content-Length header.
 #[derive(Debug, Clone, PartialEq)]
-struct Decoder {
+pub struct Decoder {
     kind: Kind,
 }
 

src/server/mod.rs

@@ -9,14 +9,14 @@ use tokio_core::net::TcpStream;
 mod srv;
 mod worker;
 mod channel;
-mod encoding;
-mod h1;
+pub(crate) mod encoding;
+pub(crate) mod h1;
 mod h2;
 mod h1writer;
 mod h2writer;
 mod settings;
-mod shared;
-mod utils;
+pub(crate) mod shared;
+pub(crate) mod utils;
 
 pub use self::srv::HttpServer;
 pub use self::settings::ServerSettings;

src/ws/client.rs

@@ -0,0 +1,561 @@
+//! Http client request
+use std::{fmt, io, str};
+use std::rc::Rc;
+use std::time::Duration;
+use std::cell::UnsafeCell;
+
+use base64;
+use rand;
+use cookie::{Cookie, CookieJar};
+use bytes::BytesMut;
+use http::{Method, Version, HeaderMap, HttpTryFrom, StatusCode, Error as HttpError};
+use http::header::{self, HeaderName, HeaderValue};
+use url::Url;
+use sha1::Sha1;
+use futures::{Async, Future, Poll, Stream};
+// use futures::unsync::oneshot;
+use tokio_core::net::TcpStream;
+
+use body::{Body, Binary};
+use error::UrlParseError;
+use headers::ContentEncoding;
+use server::shared::SharedBytes;
+
+use server::{utils, IoStream};
+use client::{HttpResponseParser, HttpResponseParserError};
+
+use super::Message;
+use super::proto::{CloseCode, OpCode};
+use super::frame::Frame;
+use super::writer::Writer;
+use super::connect::{TcpConnector, TcpConnectorError};
+
+/// Websockt client error
+#[derive(Fail, Debug)]
+pub enum WsClientError {
+    #[fail(display="Invalid url")]
+    InvalidUrl,
+    #[fail(display="Invalid response status")]
+    InvalidResponseStatus,
+    #[fail(display="Invalid upgrade header")]
+    InvalidUpgradeHeader,
+    #[fail(display="Invalid connection header")]
+    InvalidConnectionHeader,
+    #[fail(display="Invalid challenge response")]
+    InvalidChallengeResponse,
+    #[fail(display="Http parsing error")]
+    Http(HttpError),
+    #[fail(display="Url parsing error")]
+    Url(UrlParseError),
+    #[fail(display="Response parsing error")]
+    ResponseParseError(HttpResponseParserError),
+    #[fail(display="{}", _0)]
+    Connection(TcpConnectorError),
+    #[fail(display="{}", _0)]
+    Io(io::Error),
+    #[fail(display="Disconnected")]
+    Disconnected,
+}
+
+impl From<HttpError> for WsClientError {
+    fn from(err: HttpError) -> WsClientError {
+        WsClientError::Http(err)
+    }
+}
+
+impl From<UrlParseError> for WsClientError {
+    fn from(err: UrlParseError) -> WsClientError {
+        WsClientError::Url(err)
+    }
+}
+
+impl From<TcpConnectorError> for WsClientError {
+    fn from(err: TcpConnectorError) -> WsClientError {
+        WsClientError::Connection(err)
+    }
+}
+
+impl From<io::Error> for WsClientError {
+    fn from(err: io::Error) -> WsClientError {
+        WsClientError::Io(err)
+    }
+}
+
+impl From<HttpResponseParserError> for WsClientError {
+    fn from(err: HttpResponseParserError) -> WsClientError {
+        WsClientError::ResponseParseError(err)
+    }
+}
+
+type WsFuture<T> = Future<Item=(WsReader<T>, WsWriter<T>), Error=WsClientError>;
+
+/// Websockt client
+pub struct WsClient {
+    request: Option<ClientRequest>,
+    err: Option<WsClientError>,
+    http_err: Option<HttpError>,
+    cookies: Option<CookieJar>,
+    origin: Option<HeaderValue>,
+    protocols: Option<String>,
+}
+
+impl WsClient {
+
+    pub fn new<S: AsRef<str>>(url: S) -> WsClient {
+        let mut cl = WsClient {
+            request: None,
+            err: None,
+            http_err: None,
+            cookies: None,
+            origin: None,
+            protocols: None };
+
+        match Url::parse(url.as_ref()) {
+            Ok(url) => {
+                if url.scheme() != "http" && url.scheme() != "https" &&
+                    url.scheme() != "ws" && url.scheme() != "wss" || !url.has_host() {
+                        cl.err = Some(WsClientError::InvalidUrl);
+                } else {
+                    cl.request = Some(ClientRequest::new(Method::GET, url));
+                }
+            },
+            Err(err) => cl.err = Some(err.into()),
+        }
+        cl
+    }
+
+    pub fn protocols<U, V>(&mut self, protos: U) -> &mut Self
+        where U: IntoIterator<Item=V> + 'static,
+              V: AsRef<str>
+    {
+        let mut protos = protos.into_iter()
+            .fold(String::new(), |acc, s| {acc + s.as_ref() + ","});
+        protos.pop();
+        self.protocols = Some(protos);
+        self
+    }
+
+    pub fn cookie<'c>(&mut self, cookie: Cookie<'c>) -> &mut Self {
+        if self.cookies.is_none() {
+            let mut jar = CookieJar::new();
+            jar.add(cookie.into_owned());
+            self.cookies = Some(jar)
+        } else {
+            self.cookies.as_mut().unwrap().add(cookie.into_owned());
+        }
+        self
+    }
+
+    /// Set request Origin
+    pub fn origin<V>(&mut self, origin: V) -> &mut Self
+        where HeaderValue: HttpTryFrom<V>
+    {
+        match HeaderValue::try_from(origin) {
+            Ok(value) => self.origin = Some(value),
+            Err(e) => self.http_err = Some(e.into()),
+        }
+        self
+    }
+
+    pub fn header<K, V>(&mut self, key: K, value: V) -> &mut Self
+        where HeaderName: HttpTryFrom<K>,
+              HeaderValue: HttpTryFrom<V>
+    {
+        if let Some(parts) = parts(&mut self.request, &self.err, &self.http_err) {
+            match HeaderName::try_from(key) {
+                Ok(key) => {
+                    match HeaderValue::try_from(value) {
+                        Ok(value) => { parts.headers.append(key, value); }
+                        Err(e) => self.http_err = Some(e.into()),
+                    }
+                },
+                Err(e) => self.http_err = Some(e.into()),
+            };
+        }
+        self
+    }
+
+    pub fn connect(&mut self) -> Result<Box<WsFuture<TcpStream>>, WsClientError> {
+        if let Some(e) = self.err.take() {
+            return Err(e)
+        }
+        if let Some(e) = self.http_err.take() {
+            return Err(e.into())
+        }
+        let mut request = self.request.take().expect("cannot reuse request builder");
+
+        // headers
+        if let Some(ref jar) = self.cookies {
+            for cookie in jar.delta() {
+                request.headers.append(
+                    header::SET_COOKIE,
+                    HeaderValue::from_str(&cookie.to_string()).map_err(HttpError::from)?);
+            }
+        }
+
+        // origin
+        if let Some(origin) = self.origin.take() {
+            request.headers.insert(header::ORIGIN, origin);
+        }
+
+        request.headers.insert(header::UPGRADE, HeaderValue::from_static("websocket"));
+        request.headers.insert(header::CONNECTION, HeaderValue::from_static("upgrade"));
+        request.headers.insert(
+            HeaderName::try_from("SEC-WEBSOCKET-VERSION").unwrap(),
+            HeaderValue::from_static("13"));
+
+        if let Some(protocols) = self.protocols.take() {
+            request.headers.insert(
+                HeaderName::try_from("SEC-WEBSOCKET-PROTOCOL").unwrap(),
+                HeaderValue::try_from(protocols.as_str()).unwrap());
+        }
+
+        let connect = TcpConnector::new(
+            request.url.host_str().unwrap(),
+            request.url.port().unwrap_or(80), Duration::from_secs(5));
+
+        Ok(Box::new(
+            connect
+                .from_err()
+                .and_then(move |stream| WsHandshake::new(stream, request))))
+    }
+}
+
+#[inline]
+fn parts<'a>(parts: &'a mut Option<ClientRequest>,
+             err: &Option<WsClientError>,
+             http_err: &Option<HttpError>) -> Option<&'a mut ClientRequest>
+{
+    if err.is_some() || http_err.is_some() {
+        return None
+    }
+    parts.as_mut()
+}
+
+pub(crate) struct ClientRequest {
+    pub url: Url,
+    pub method: Method,
+    pub version: Version,
+    pub headers: HeaderMap,
+    pub body: Body,
+    pub chunked: Option<bool>,
+    pub encoding: ContentEncoding,
+}
+
+impl ClientRequest {
+
+    #[inline]
+    fn new(method: Method, url: Url) -> ClientRequest {
+        ClientRequest {
+            url: url,
+            method: method,
+            version: Version::HTTP_11,
+            headers: HeaderMap::with_capacity(16),
+            body: Body::Empty,
+            chunked: None,
+            encoding: ContentEncoding::Auto,
+        }
+    }
+}
+
+impl fmt::Debug for ClientRequest {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let res = write!(f, "\nClientRequest {:?} {}:{}\n",
+                         self.version, self.method, self.url);
+        let _ = write!(f, "  headers:\n");
+        for key in self.headers.keys() {
+            let vals: Vec<_> = self.headers.get_all(key).iter().collect();
+            if vals.len() > 1 {
+                let _ = write!(f, "    {:?}: {:?}\n", key, vals);
+            } else {
+                let _ = write!(f, "    {:?}: {:?}\n", key, vals[0]);
+            }
+        }
+        res
+    }
+}
+
+struct WsInner<T> {
+    stream: T,
+    writer: Writer,
+    parser: HttpResponseParser,
+    parser_buf: BytesMut,
+    closed: bool,
+    error_sent: bool,
+}
+
+struct WsHandshake<T> {
+    inner: Option<WsInner<T>>,
+    request: ClientRequest,
+    sent: bool,
+    key: String,
+}
+
+impl<T: IoStream> WsHandshake<T> {
+    fn new(stream: T, mut request: ClientRequest) -> WsHandshake<T> {
+        // Generate a random key for the `Sec-WebSocket-Key` header.
+        // a base64-encoded (see Section 4 of [RFC4648]) value that,
+        // when decoded, is 16 bytes in length (RFC 6455)
+        let sec_key: [u8; 16] = rand::random();
+        let key = base64::encode(&sec_key);
+
+        request.headers.insert(
+            HeaderName::try_from("SEC-WEBSOCKET-KEY").unwrap(),
+            HeaderValue::try_from(key.as_str()).unwrap());
+
+        let inner = WsInner {
+            stream: stream,
+            writer: Writer::new(SharedBytes::default()),
+            parser: HttpResponseParser::new(),
+            parser_buf: BytesMut::new(),
+            closed: false,
+            error_sent: false,
+        };
+
+        WsHandshake {
+            key: key,
+            inner: Some(inner),
+            request: request,
+            sent: false,
+        }
+    }
+}
+
+impl<T: IoStream> Future for WsHandshake<T> {
+    type Item = (WsReader<T>, WsWriter<T>);
+    type Error = WsClientError;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        let mut inner = self.inner.take().unwrap();
+
+        if !self.sent {
+            self.sent = true;
+            inner.writer.start(&mut self.request);
+        }
+        if let Err(err) = inner.writer.poll_completed(&mut inner.stream, false) {
+            return Err(err.into())
+        }
+
+        match inner.parser.parse(&mut inner.stream, &mut inner.parser_buf) {
+            Ok(Async::Ready(resp)) => {
+                // verify response
+                if resp.status() != StatusCode::SWITCHING_PROTOCOLS {
+                    return Err(WsClientError::InvalidResponseStatus)
+                }
+                // Check for "UPGRADE" to websocket header
+                let has_hdr = if let Some(hdr) = resp.headers().get(header::UPGRADE) {
+                    if let Ok(s) = hdr.to_str() {
+                        s.to_lowercase().contains("websocket")
+                    } else {
+                        false
+                    }
+                } else {
+                    false
+                };
+                if !has_hdr {
+                    return Err(WsClientError::InvalidUpgradeHeader)
+                }
+                // Check for "CONNECTION" header
+                let has_hdr = if let Some(conn) = resp.headers().get(header::CONNECTION) {
+                    if let Ok(s) = conn.to_str() {
+                        s.to_lowercase().contains("upgrade")
+                    } else { false }
+                } else { false };
+                if !has_hdr {
+                    return Err(WsClientError::InvalidConnectionHeader)
+                }
+
+                let match_key = if let Some(key) = resp.headers().get(
+                    HeaderName::try_from("SEC-WEBSOCKET-ACCEPT").unwrap())
+                {
+                    // ... field is constructed by concatenating /key/ ...
+                    // ... with the string "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" (RFC 6455)
+                    const WS_GUID: &'static [u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
+                    let mut sha1 = Sha1::new();
+                    sha1.update(self.key.as_ref());
+                    sha1.update(WS_GUID);
+                    key.as_bytes() == base64::encode(&sha1.digest().bytes()).as_bytes()
+                } else {
+                    false
+                };
+                if !match_key {
+                    return Err(WsClientError::InvalidChallengeResponse)
+                }
+
+                let inner = Rc::new(UnsafeCell::new(Inner{inner: inner}));
+                Ok(Async::Ready(
+                    (WsReader{inner: Rc::clone(&inner)},
+                     WsWriter{inner: inner})))
+            },
+            Ok(Async::NotReady) => {
+                self.inner = Some(inner);
+                Ok(Async::NotReady)
+            },
+            Err(err) => Err(err.into())
+        }
+    }
+}
+
+
+struct Inner<T> {
+    inner: WsInner<T>,
+}
+
+pub struct WsReader<T> {
+    inner: Rc<UnsafeCell<Inner<T>>>
+}
+
+impl<T> WsReader<T> {
+    #[inline]
+    fn as_mut(&mut self) -> &mut Inner<T> {
+        unsafe{ &mut *self.inner.get() }
+    }
+}
+
+impl<T: IoStream> Stream for WsReader<T> {
+    type Item = Message;
+    type Error = WsClientError;
+
+    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        let inner = self.as_mut();
+        let mut done = false;
+
+        match utils::read_from_io(&mut inner.inner.stream, &mut inner.inner.parser_buf) {
+            Ok(Async::Ready(0)) => {
+                done = true;
+                inner.inner.closed = true;
+            },
+            Ok(Async::Ready(_)) | Ok(Async::NotReady) => (),
+            Err(err) =>
+                return Err(err.into())
+        }
+
+        // write
+        let _ = inner.inner.writer.poll_completed(&mut inner.inner.stream, false);
+
+        // read
+        match Frame::parse(&mut inner.inner.parser_buf) {
+            Ok(Some(frame)) => {
+                // trace!("WsFrame {}", frame);
+                let (_finished, opcode, payload) = frame.unpack();
+
+                match opcode {
+                    OpCode::Continue => unimplemented!(),
+                    OpCode::Bad =>
+                        Ok(Async::Ready(Some(Message::Error))),
+                    OpCode::Close => {
+                        inner.inner.closed = true;
+                        inner.inner.error_sent = true;
+                        Ok(Async::Ready(Some(Message::Closed)))
+                    },
+                    OpCode::Ping =>
+                        Ok(Async::Ready(Some(
+                            Message::Ping(
+                                String::from_utf8_lossy(payload.as_ref()).into())))),
+                    OpCode::Pong =>
+                        Ok(Async::Ready(Some(
+                            Message::Pong(
+                                String::from_utf8_lossy(payload.as_ref()).into())))),
+                    OpCode::Binary =>
+                        Ok(Async::Ready(Some(Message::Binary(payload)))),
+                    OpCode::Text => {
+                        let tmp = Vec::from(payload.as_ref());
+                        match String::from_utf8(tmp) {
+                            Ok(s) =>
+                                Ok(Async::Ready(Some(Message::Text(s)))),
+                            Err(_) =>
+                                Ok(Async::Ready(Some(Message::Error))),
+                        }
+                    }
+                }
+            }
+            Ok(None) => {
+                if done {
+                    Ok(Async::Ready(None))
+                } else if inner.inner.closed {
+                    if !inner.inner.error_sent {
+                        inner.inner.error_sent = true;
+                        Ok(Async::Ready(Some(Message::Closed)))
+                    } else {
+                        Ok(Async::Ready(None))
+                    }
+                } else {
+                    Ok(Async::NotReady)
+                }
+            },
+            Err(err) => {
+                inner.inner.closed = true;
+                inner.inner.error_sent = true;
+                Err(err.into())
+            }
+        }
+    }
+}
+
+pub struct WsWriter<T> {
+    inner: Rc<UnsafeCell<Inner<T>>>
+}
+
+impl<T: IoStream> WsWriter<T> {
+    #[inline]
+    fn as_mut(&mut self) -> &mut Inner<T> {
+        unsafe{ &mut *self.inner.get() }
+    }
+}
+
+impl<T: IoStream> WsWriter<T> {
+
+    /// Write payload
+    #[inline]
+    fn write<B: Into<Binary>>(&mut self, data: B) {
+        if !self.as_mut().inner.closed {
+            let _ = self.as_mut().inner.writer.write(&data.into());
+        } else {
+            warn!("Trying to write to disconnected response");
+        }
+    }
+
+    /// Send text frame
+    pub fn text(&mut self, text: &str) {
+        let mut frame = Frame::message(Vec::from(text), OpCode::Text, true);
+        let mut buf = Vec::new();
+        frame.format(&mut buf).unwrap();
+
+        self.write(buf);
+    }
+
+    /// Send binary frame
+    pub fn binary<B: Into<Binary>>(&mut self, data: B) {
+        let mut frame = Frame::message(data, OpCode::Binary, true);
+        let mut buf = Vec::new();
+        frame.format(&mut buf).unwrap();
+
+        self.write(buf);
+    }
+
+    /// Send ping frame
+    pub fn ping(&mut self, message: &str) {
+        let mut frame = Frame::message(Vec::from(message), OpCode::Ping, true);
+        let mut buf = Vec::new();
+        frame.format(&mut buf).unwrap();
+
+        self.write(buf);
+    }
+
+    /// Send pong frame
+    pub fn pong(&mut self, message: &str) {
+        let mut frame = Frame::message(Vec::from(message), OpCode::Pong, true);
+        let mut buf = Vec::new();
+        frame.format(&mut buf).unwrap();
+
+        self.write(buf);
+    }
+
+    /// Send close frame
+    pub fn close(&mut self, code: CloseCode, reason: &str) {
+        let mut frame = Frame::close(code, reason);
+        let mut buf = Vec::new();
+        frame.format(&mut buf).unwrap();
+        self.write(buf);
+    }
+}

src/ws/connect.rs

@@ -0,0 +1,141 @@
+use std::io;
+use std::net::SocketAddr;
+use std::collections::VecDeque;
+use std::time::Duration;
+
+use actix::Arbiter;
+use trust_dns_resolver::ResolverFuture;
+use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};
+use trust_dns_resolver::lookup_ip::LookupIpFuture;
+use futures::{Async, Future, Poll};
+use tokio_core::reactor::Timeout;
+use tokio_core::net::{TcpStream, TcpStreamNew};
+
+
+#[derive(Fail, Debug)]
+pub enum TcpConnectorError {
+    /// Failed to resolve the hostname
+    #[fail(display = "Failed resolving hostname: {}", _0)]
+    Dns(String),
+
+    /// Address is invalid
+    #[fail(display = "Invalid input: {}", _0)]
+    InvalidInput(&'static str),
+
+    /// Connecting took too long
+    #[fail(display = "Timeout out while establishing connection")]
+    Timeout,
+
+    /// Connection io error
+    #[fail(display = "{}", _0)]
+    IoError(io::Error),
+}
+
+pub struct TcpConnector {
+    lookup: Option<LookupIpFuture>,
+    port: u16,
+    ips: VecDeque<SocketAddr>,
+    error: Option<TcpConnectorError>,
+    timeout: Timeout,
+    stream: Option<TcpStreamNew>,
+}
+
+impl TcpConnector {
+
+    pub fn new<S: AsRef<str>>(addr: S, port: u16, timeout: Duration) -> TcpConnector {
+        println!("TES: {:?} {:?}", addr.as_ref(), port);
+
+        // try to parse as a regular SocketAddr first
+        if let Ok(addr) = addr.as_ref().parse() {
+            let mut ips = VecDeque::new();
+            ips.push_back(addr);
+
+            TcpConnector {
+                lookup: None,
+                port: port,
+                ips: ips,
+                error: None,
+                stream: None,
+                timeout: Timeout::new(timeout, Arbiter::handle()).unwrap() }
+        } else {
+            // we need to do dns resolution
+            let resolve = match ResolverFuture::from_system_conf(Arbiter::handle()) {
+                Ok(resolve) => resolve,
+                Err(err) => {
+                    warn!("Can not create system dns resolver: {}", err);
+                    ResolverFuture::new(
+                        ResolverConfig::default(),
+                        ResolverOpts::default(),
+                        Arbiter::handle())
+                }
+            };
+
+            TcpConnector {
+                lookup: Some(resolve.lookup_ip(addr.as_ref())),
+                port: port,
+                ips: VecDeque::new(),
+                error: None,
+                stream: None,
+                timeout: Timeout::new(timeout, Arbiter::handle()).unwrap() }
+        }
+    }
+}
+
+impl Future for TcpConnector {
+    type Item = TcpStream;
+    type Error = TcpConnectorError;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        if let Some(err) = self.error.take() {
+            Err(err)
+        } else {
+            // timeout
+            if let Ok(Async::Ready(_)) = self.timeout.poll() {
+                return Err(TcpConnectorError::Timeout)
+            }
+
+            // lookip ips
+            if let Some(mut lookup) = self.lookup.take() {
+                match lookup.poll() {
+                    Ok(Async::NotReady) => {
+                        self.lookup = Some(lookup);
+                        return Ok(Async::NotReady)
+                    },
+                    Ok(Async::Ready(ips)) => {
+                        let port = self.port;
+                        let ips = ips.iter().map(|ip| SocketAddr::new(ip, port));
+                        self.ips.extend(ips);
+                        if self.ips.is_empty() {
+                            return Err(TcpConnectorError::Dns(
+                                "Expect at least one A dns record".to_owned()))
+                        }
+                    },
+                    Err(err) => return Err(TcpConnectorError::Dns(format!("{}", err))),
+                }
+            }
+
+            // connect
+            loop {
+                if let Some(mut new) = self.stream.take() {
+                    match new.poll() {
+                        Ok(Async::Ready(sock)) =>
+                            return Ok(Async::Ready(sock)),
+                        Ok(Async::NotReady) => {
+                            self.stream = Some(new);
+                            return Ok(Async::NotReady)
+                        },
+                        Err(err) => {
+                            if self.ips.is_empty() {
+                                return Err(TcpConnectorError::IoError(err))
+                            }
+                        }
+                    }
+                }
+
+                // try to connect
+                let addr = self.ips.pop_front().unwrap();
+                self.stream = Some(TcpStream::connect(&addr, Arbiter::handle()));
+            }
+        }
+    }
+}

src/ws/mod.rs

@@ -60,11 +60,18 @@ mod proto;
 mod context;
 mod mask;
 
+mod connect;
+mod writer;
+
+pub mod client;
+
 use ws::frame::Frame;
 use ws::proto::{hash_key, OpCode};
 pub use ws::proto::CloseCode;
 pub use ws::context::WebsocketContext;
 
+pub use self::client::{WsClient, WsClientError, WsReader, WsWriter};
+
 const SEC_WEBSOCKET_ACCEPT: &str = "SEC-WEBSOCKET-ACCEPT";
 const SEC_WEBSOCKET_KEY: &str = "SEC-WEBSOCKET-KEY";
 const SEC_WEBSOCKET_VERSION: &str = "SEC-WEBSOCKET-VERSION";

src/ws/proto.rs

@@ -1,7 +1,7 @@
 use std::fmt;
 use std::convert::{Into, From};
 use sha1;
-
+use base64;
 
 use self::OpCode::*;
 /// Operation codes as part of rfc6455.
@@ -188,10 +188,7 @@ impl From<u16> for CloseCode {
     }
 }
 
-
 static WS_GUID: &'static str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
-static BASE64: &'static [u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
-
 
 // TODO: hash is always same size, we dont need String
 pub(crate) fn hash_key(key: &[u8]) -> String {
@@ -200,48 +197,7 @@ pub(crate) fn hash_key(key: &[u8]) -> String {
     hasher.update(key);
     hasher.update(WS_GUID.as_bytes());
 
-    encode_base64(&hasher.digest().bytes())
-}
-
-
-// This code is based on rustc_serialize base64 STANDARD
-fn encode_base64(data: &[u8]) -> String {
-    let len = data.len();
-    let mod_len = len % 3;
-
-    let mut encoded = vec![b'='; (len + 2) / 3 * 4];
-    {
-        let mut in_iter = data[..len - mod_len].iter().map(|&c| u32::from(c));
-        let mut out_iter = encoded.iter_mut();
-
-        let enc = |val| BASE64[val as usize];
-        let mut write = |val| *out_iter.next().unwrap() = val;
-
-        while let (Some(one), Some(two), Some(three)) = (in_iter.next(), in_iter.next(), in_iter.next()) {
-            let g24 = one << 16 | two << 8 | three;
-            write(enc((g24 >> 18) & 63));
-            write(enc((g24 >> 12) & 63));
-            write(enc((g24 >> 6 ) & 63));
-            write(enc(g24 & 63));
-        }
-
-        match mod_len {
-            1 => {
-                let pad = u32::from(data[len-1]) << 16;
-                write(enc((pad >> 18) & 63));
-                write(enc((pad >> 12) & 63));
-            }
-            2 => {
-                let pad = u32::from(data[len-2]) << 16 | u32::from(data[len-1]) << 8;
-                write(enc((pad >> 18) & 63));
-                write(enc((pad >> 12) & 63));
-                write(enc((pad >> 6) & 63));
-            }
-            _ => (),
-        }
-    }
-
-    String::from_utf8(encoded).unwrap()
+    base64::encode(&hasher.digest().bytes())
 }
 
 

src/ws/writer.rs

@@ -0,0 +1,152 @@
+#![allow(dead_code)]
+use std::io;
+use bytes::BufMut;
+use futures::{Async, Poll};
+use tokio_io::AsyncWrite;
+// use http::header::{HeaderValue, CONNECTION, DATE};
+
+use body::Binary;
+use server::{WriterState, MAX_WRITE_BUFFER_SIZE};
+use server::shared::SharedBytes;
+
+use super::client::ClientRequest;
+
+
+const AVERAGE_HEADER_SIZE: usize = 30; // totally scientific
+
+bitflags! {
+    struct Flags: u8 {
+        const STARTED = 0b0000_0001;
+        const UPGRADE = 0b0000_0010;
+        const KEEPALIVE = 0b0000_0100;
+        const DISCONNECTED = 0b0000_1000;
+    }
+}
+
+pub(crate) struct Writer {
+    flags: Flags,
+    written: u64,
+    headers_size: u32,
+    buffer: SharedBytes,
+}
+
+impl Writer {
+
+    pub fn new(buf: SharedBytes) -> Writer {
+        Writer {
+            flags: Flags::empty(),
+            written: 0,
+            headers_size: 0,
+            buffer: buf,
+        }
+    }
+
+    pub fn disconnected(&mut self) {
+        self.buffer.take();
+    }
+
+    pub fn keepalive(&self) -> bool {
+        self.flags.contains(Flags::KEEPALIVE) && !self.flags.contains(Flags::UPGRADE)
+    }
+
+    fn write_to_stream<T: AsyncWrite>(&mut self, stream: &mut T) -> io::Result<WriterState> {
+        while !self.buffer.is_empty() {
+            match stream.write(self.buffer.as_ref()) {
+                Ok(0) => {
+                    self.disconnected();
+                    return Ok(WriterState::Done);
+                },
+                Ok(n) => {
+                    let _ = self.buffer.split_to(n);
+                },
+                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+                    if self.buffer.len() > MAX_WRITE_BUFFER_SIZE {
+                        return Ok(WriterState::Pause)
+                    } else {
+                        return Ok(WriterState::Done)
+                    }
+                }
+                Err(err) => return Err(err),
+            }
+        }
+        Ok(WriterState::Done)
+    }
+}
+
+impl Writer {
+
+    pub fn start(&mut self, msg: &mut ClientRequest) {
+        // prepare task
+        self.flags.insert(Flags::STARTED);
+
+        // render message
+        {
+            let buffer = self.buffer.get_mut();
+            buffer.reserve(256 + msg.headers.len() * AVERAGE_HEADER_SIZE);
+
+            // status line
+            // helpers::write_status_line(version, msg.status().as_u16(), &mut buffer);
+            // buffer.extend_from_slice(msg.reason().as_bytes());
+            buffer.extend_from_slice(b"GET /ws/ HTTP/1.1");
+            buffer.extend_from_slice(b"\r\n");
+
+            // write headers
+            for (key, value) in &msg.headers {
+                let v = value.as_ref();
+                let k = key.as_str().as_bytes();
+                buffer.reserve(k.len() + v.len() + 4);
+                buffer.put_slice(k);
+                buffer.put_slice(b": ");
+                buffer.put_slice(v);
+                buffer.put_slice(b"\r\n");
+            }
+
+            // using helpers::date is quite a lot faster
+            //if !msg.headers.contains_key(DATE) {
+            //    helpers::date(&mut buffer);
+            //} else {
+                // msg eof
+                buffer.extend_from_slice(b"\r\n");
+            //}
+            self.headers_size = buffer.len() as u32;
+        }
+    }
+
+    pub fn write(&mut self, payload: &Binary) -> io::Result<WriterState> {
+        self.written += payload.len() as u64;
+        if !self.flags.contains(Flags::DISCONNECTED) {
+            self.buffer.extend_from_slice(payload.as_ref())
+        }
+
+        if self.buffer.len() > MAX_WRITE_BUFFER_SIZE {
+            Ok(WriterState::Pause)
+        } else {
+            Ok(WriterState::Done)
+        }
+    }
+
+    pub fn write_eof(&mut self) -> io::Result<WriterState> {
+        if self.buffer.len() > MAX_WRITE_BUFFER_SIZE {
+            Ok(WriterState::Pause)
+        } else {
+            Ok(WriterState::Done)
+        }
+    }
+
+    #[inline]
+    pub fn poll_completed<T: AsyncWrite>(&mut self, stream: &mut T, shutdown: bool)
+                                         -> Poll<(), io::Error>
+    {
+        match self.write_to_stream(stream) {
+            Ok(WriterState::Done) => {
+                if shutdown {
+                    stream.shutdown()
+                } else {
+                    Ok(Async::Ready(()))
+                }
+            },
+            Ok(WriterState::Pause) => Ok(Async::NotReady),
+            Err(err) => Err(err)
+        }
+    }
+}