From e7ec0f9fd7f5a4a95ce03c3207b6521225c65e9e Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Mon, 19 Mar 2018 09:30:58 -0700
Subject: [PATCH] ws tests and proper write payload ref

---
 src/body.rs            |  2 +-
 src/server/h1writer.rs |  7 ++++---
 tests/test_ws.rs       | 40 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/src/body.rs b/src/body.rs
index cdf8c0391..57df35287 100644
--- a/src/body.rs
+++ b/src/body.rs
@@ -238,7 +238,7 @@ impl AsRef<[u8]> for Binary {
     #[inline]
     fn as_ref(&self) -> &[u8] {
         match *self {
-            Binary::Bytes(ref bytes) => &bytes[..],
+            Binary::Bytes(ref bytes) => bytes.as_ref(),
             Binary::Slice(slice) => slice,
             Binary::SharedString(ref s) => s.as_bytes(),
             Binary::ArcSharedString(ref s) => s.as_bytes(),
diff --git a/src/server/h1writer.rs b/src/server/h1writer.rs
index 46c3cf1a2..68ce71af4 100644
--- a/src/server/h1writer.rs
+++ b/src/server/h1writer.rs
@@ -227,9 +227,10 @@ impl<T: AsyncWrite, H: 'static> Writer for H1Writer<T, H> {
                 // shortcut for upgraded connection
                 if self.flags.contains(Flags::UPGRADE) {
                     if self.buffer.is_empty() {
-                        let n = self.write_data(payload.as_ref())?;
-                        if payload.len() < n {
-                            self.buffer.extend_from_slice(&payload.as_ref()[n..]);
+                        let pl: &[u8] = payload.as_ref();
+                        let n = self.write_data(pl)?;
+                        if pl.len() < n {
+                            self.buffer.extend_from_slice(&pl[n..]);
                             return Ok(WriterState::Done);
                         }
                     } else {
diff --git a/tests/test_ws.rs b/tests/test_ws.rs
index edda3f64b..c1319b264 100644
--- a/tests/test_ws.rs
+++ b/tests/test_ws.rs
@@ -3,9 +3,11 @@ extern crate actix_web;
 extern crate futures;
 extern crate http;
 extern crate bytes;
+extern crate rand;
 
 use bytes::Bytes;
 use futures::Stream;
+use rand::Rng;
 
 use actix_web::*;
 use actix::prelude::*;
@@ -51,3 +53,41 @@ fn test_simple() {
     let (item, _) = srv.execute(reader.into_future()).unwrap();
     assert_eq!(item, Some(ws::Message::Close(ws::CloseCode::Normal)));
 }
+
+#[test]
+fn test_large_text() {
+    let data = rand::thread_rng()
+        .gen_ascii_chars()
+        .take(65_536)
+        .collect::<String>();
+
+    let mut srv = test::TestServer::new(
+        |app| app.handler(|req| ws::start(req, Ws)));
+    let (mut reader, mut writer) = srv.ws().unwrap();
+
+    for _ in 0..100 {
+        writer.text(data.clone());
+        let (item, r) = srv.execute(reader.into_future()).unwrap();
+        reader = r;
+        assert_eq!(item, Some(ws::Message::Text(data.clone())));
+    }
+}
+
+#[test]
+fn test_large_bin() {
+    let data = rand::thread_rng()
+        .gen_ascii_chars()
+        .take(65_536)
+        .collect::<String>();
+
+    let mut srv = test::TestServer::new(
+        |app| app.handler(|req| ws::start(req, Ws)));
+    let (mut reader, mut writer) = srv.ws().unwrap();
+
+    for _ in 0..100 {
+        writer.binary(data.clone());
+        let (item, r) = srv.execute(reader.into_future()).unwrap();
+        reader = r;
+        assert_eq!(item, Some(ws::Message::Binary(Binary::from(data.clone()))));
+    }
+}