diff --git a/src/body.rs b/src/body.rs index 7eaa54468..beecf8bfa 100644 --- a/src/body.rs +++ b/src/body.rs @@ -1,4 +1,4 @@ -use std::fmt; +use std::{fmt, mem}; use std::rc::Rc; use std::sync::Arc; use bytes::{Bytes, BytesMut}; @@ -122,6 +122,22 @@ impl Binary { pub fn from_slice(s: &[u8]) -> Binary { Binary::Bytes(Bytes::from(s)) } + + /// Convert Binary to a Bytes instance + pub fn take(&mut self) -> Bytes { + mem::replace(self, Binary::Slice(b"")).into() + } +} + +impl Clone for Binary { + fn clone(&self) -> Binary { + match *self { + Binary::Bytes(ref bytes) => Binary::Bytes(bytes.clone()), + Binary::Slice(slice) => Binary::Bytes(Bytes::from(slice)), + Binary::SharedString(ref s) => Binary::Bytes(Bytes::from(s.as_str())), + Binary::ArcSharedString(ref s) => Binary::Bytes(Bytes::from(s.as_str())), + } + } } impl Into for Binary { diff --git a/src/server/encoding.rs b/src/server/encoding.rs index c0cd6fd32..271c9f55c 100644 --- a/src/server/encoding.rs +++ b/src/server/encoding.rs @@ -396,7 +396,7 @@ impl PayloadEncoder { ContentEncoding::Auto => unreachable!() }; // TODO return error! - let _ = enc.write(bytes.as_ref()); + let _ = enc.write(bytes.clone()); let _ = enc.write_eof(); *bytes = Binary::from(tmp.get_mut().take()); @@ -520,7 +520,7 @@ impl PayloadEncoder { #[cfg_attr(feature = "cargo-clippy", allow(inline_always))] #[inline(always)] - pub fn write(&mut self, payload: &[u8]) -> Result<(), io::Error> { + pub fn write(&mut self, payload: Binary) -> Result<(), io::Error> { self.0.write(payload) } @@ -629,10 +629,10 @@ impl ContentEncoder { #[cfg_attr(feature = "cargo-clippy", allow(inline_always))] #[inline(always)] - pub fn write(&mut self, data: &[u8]) -> Result<(), io::Error> { + pub fn write(&mut self, data: Binary) -> Result<(), io::Error> { match *self { ContentEncoder::Br(ref mut encoder) => { - match encoder.write(data) { + match encoder.write(data.as_ref()) { Ok(_) => encoder.flush(), Err(err) => { @@ -642,7 +642,7 @@ impl ContentEncoder { } }, ContentEncoder::Gzip(ref mut encoder) => { - match encoder.write(data) { + match encoder.write(data.as_ref()) { Ok(_) => encoder.flush(), Err(err) => { @@ -652,7 +652,7 @@ impl ContentEncoder { } } ContentEncoder::Deflate(ref mut encoder) => { - match encoder.write(data) { + match encoder.write(data.as_ref()) { Ok(_) => encoder.flush(), Err(err) => { @@ -727,10 +727,10 @@ impl TransferEncoding { /// Encode message. Return `EOF` state of encoder #[inline] - pub fn encode(&mut self, msg: &[u8]) -> io::Result { + pub fn encode(&mut self, msg: Binary) -> io::Result { match self.kind { TransferEncodingKind::Eof => { - self.buffer.get_mut().extend_from_slice(msg); + self.buffer.get_mut().extend_from_slice(msg.as_ref()); Ok(msg.is_empty()) }, TransferEncodingKind::Chunked(ref mut eof) => { @@ -744,7 +744,7 @@ impl TransferEncoding { } else { write!(self.buffer.get_mut(), "{:X}\r\n", msg.len()) .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?; - self.buffer.get_mut().extend_from_slice(msg); + self.buffer.get_mut().extend_from_slice(msg.as_ref()); self.buffer.get_mut().extend_from_slice(b"\r\n"); } Ok(*eof) @@ -754,7 +754,7 @@ impl TransferEncoding { return Ok(*remaining == 0) } let max = cmp::min(*remaining, msg.len() as u64); - self.buffer.get_mut().extend_from_slice(msg[..max as usize].as_ref()); + self.buffer.get_mut().extend_from_slice(msg.as_ref()[..max as usize].as_ref()); *remaining -= max as u64; Ok(*remaining == 0) @@ -781,7 +781,7 @@ impl io::Write for TransferEncoding { #[inline] fn write(&mut self, buf: &[u8]) -> io::Result { - self.encode(buf)?; + self.encode(Binary::from_slice(buf))?; Ok(buf.len()) } diff --git a/src/server/h1writer.rs b/src/server/h1writer.rs index 0befa9ad3..c448313bc 100644 --- a/src/server/h1writer.rs +++ b/src/server/h1writer.rs @@ -174,7 +174,7 @@ impl Writer for H1Writer { if let Body::Binary(bytes) = body { self.written = bytes.len() as u64; - self.encoder.write(bytes.as_ref())?; + self.encoder.write(bytes)?; } else { msg.replace_body(body); } @@ -186,7 +186,7 @@ impl Writer for H1Writer { if !self.flags.contains(Flags::DISCONNECTED) { if self.flags.contains(Flags::STARTED) { // TODO: add warning, write after EOF - self.encoder.write(payload.as_ref())?; + self.encoder.write(payload)?; return Ok(WriterState::Done) } else { // might be response to EXCEPT diff --git a/src/server/h2writer.rs b/src/server/h2writer.rs index 1cf1b0b61..9b522b38d 100644 --- a/src/server/h2writer.rs +++ b/src/server/h2writer.rs @@ -168,7 +168,7 @@ impl Writer for H2Writer { if let Body::Binary(bytes) = body { self.flags.insert(Flags::EOF); self.written = bytes.len() as u64; - self.encoder.write(bytes.as_ref())?; + self.encoder.write(bytes)?; if let Some(ref mut stream) = self.stream { stream.reserve_capacity(cmp::min(self.encoder.len(), CHUNK_SIZE)); } @@ -185,7 +185,7 @@ impl Writer for H2Writer { if !self.flags.contains(Flags::DISCONNECTED) { if self.flags.contains(Flags::STARTED) { // TODO: add warning, write after EOF - self.encoder.write(payload.as_ref())?; + self.encoder.write(payload)?; } else { // might be response for EXCEPT self.encoder.get_mut().extend_from_slice(payload.as_ref())