use crate::chunker;
use std::io::{self, Write};
use std::mem;
pub struct BodyWriter<D, E>(Inner<D, E>)
where
D: From<Vec<u8>> + Send + 'static,
E: Send + 'static;
enum Inner<D, E>
where
D: From<Vec<u8>> + Send + 'static,
E: Send + 'static,
{
Raw(chunker::BodyWriter<D, E>),
Gzipped(flate2::write::GzEncoder<chunker::BodyWriter<D, E>>),
Dead,
}
impl<D, E> BodyWriter<D, E>
where
D: From<Vec<u8>> + Send + 'static,
E: Send + 'static,
{
pub(crate) fn raw(raw: chunker::BodyWriter<D, E>) -> Self {
BodyWriter(Inner::Raw(raw))
}
pub(crate) fn gzipped(raw: chunker::BodyWriter<D, E>, level: flate2::Compression) -> Self {
BodyWriter(Inner::Gzipped(flate2::GzBuilder::new().write(raw, level)))
}
pub fn abort(&mut self, error: E) {
match mem::replace(&mut self.0, Inner::Dead) {
Inner::Dead => (),
Inner::Raw(ref mut w) => w.abort(error),
Inner::Gzipped(ref mut g) => g.get_mut().abort(error),
};
}
}
impl<D, E> Write for BodyWriter<D, E>
where
D: From<Vec<u8>> + Send + 'static,
E: Send + 'static,
{
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let r = match self.0 {
Inner::Dead => return Err(io::Error::new(io::ErrorKind::BrokenPipe, "body is dead")),
Inner::Raw(ref mut w) => w.write(buf),
Inner::Gzipped(ref mut w) => w.write(buf),
};
if r.is_err() {
self.0 = Inner::Dead;
}
r
}
fn flush(&mut self) -> io::Result<()> {
let r = match self.0 {
Inner::Dead => return Err(io::Error::new(io::ErrorKind::BrokenPipe, "body is dead")),
Inner::Raw(ref mut w) => w.flush(),
Inner::Gzipped(ref mut w) => w.flush(),
};
if r.is_err() {
self.0 = Inner::Dead;
}
r
}
}