mirror of
https://github.com/fafhrd91/actix-web
synced 2025-06-24 22:37:35 +02:00
add content-encoding decompression
This commit is contained in:
Nikolay Kim
191
actix-http/src/encoding/decoder.rs
Normal file
191
actix-http/src/encoding/decoder.rs
Normal file
@ -0,0 +1,191 @@
|
||||
use std::io::{self, Write};
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::{Async, Poll, Stream};
|
||||
|
||||
#[cfg(feature = "brotli")]
|
||||
use brotli2::write::BrotliDecoder;
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
use flate2::write::{GzDecoder, ZlibDecoder};
|
||||
|
||||
use super::Writer;
|
||||
use crate::error::PayloadError;
|
||||
use crate::http::header::{ContentEncoding, HeaderMap, CONTENT_ENCODING};
|
||||
|
||||
pub struct Decoder<T> {
|
||||
stream: T,
|
||||
decoder: Option<ContentDecoder>,
|
||||
}
|
||||
|
||||
impl<T> Decoder<T>
|
||||
where
|
||||
T: Stream<Item = Bytes, Error = PayloadError>,
|
||||
{
|
||||
pub fn new(stream: T, encoding: ContentEncoding) -> Self {
|
||||
let decoder = match encoding {
|
||||
#[cfg(feature = "brotli")]
|
||||
ContentEncoding::Br => Some(ContentDecoder::Br(Box::new(
|
||||
BrotliDecoder::new(Writer::new()),
|
||||
))),
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoding::Deflate => Some(ContentDecoder::Deflate(Box::new(
|
||||
ZlibDecoder::new(Writer::new()),
|
||||
))),
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoding::Gzip => Some(ContentDecoder::Gzip(Box::new(
|
||||
GzDecoder::new(Writer::new()),
|
||||
))),
|
||||
_ => None,
|
||||
};
|
||||
Decoder { stream, decoder }
|
||||
}
|
||||
|
||||
pub fn from_headers(headers: &HeaderMap, stream: T) -> Self {
|
||||
// check content-encoding
|
||||
let encoding = if let Some(enc) = headers.get(CONTENT_ENCODING) {
|
||||
if let Ok(enc) = enc.to_str() {
|
||||
ContentEncoding::from(enc)
|
||||
} else {
|
||||
ContentEncoding::Identity
|
||||
}
|
||||
} else {
|
||||
ContentEncoding::Identity
|
||||
};
|
||||
|
||||
Self::new(stream, encoding)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Stream for Decoder<T>
|
||||
where
|
||||
T: Stream<Item = Bytes, Error = PayloadError>,
|
||||
{
|
||||
type Item = Bytes;
|
||||
type Error = PayloadError;
|
||||
|
||||
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
|
||||
loop {
|
||||
match self.stream.poll()? {
|
||||
Async::Ready(Some(chunk)) => {
|
||||
if let Some(ref mut decoder) = self.decoder {
|
||||
match decoder.feed_data(chunk) {
|
||||
Ok(Some(chunk)) => return Ok(Async::Ready(Some(chunk))),
|
||||
Ok(None) => continue,
|
||||
Err(e) => return Err(e.into()),
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Async::Ready(None) => {
|
||||
return if let Some(mut decoder) = self.decoder.take() {
|
||||
match decoder.feed_eof() {
|
||||
Ok(chunk) => Ok(Async::Ready(chunk)),
|
||||
Err(e) => Err(e.into()),
|
||||
}
|
||||
} else {
|
||||
Ok(Async::Ready(None))
|
||||
};
|
||||
}
|
||||
Async::NotReady => break,
|
||||
}
|
||||
}
|
||||
Ok(Async::NotReady)
|
||||
}
|
||||
}
|
||||
|
||||
enum ContentDecoder {
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
Deflate(Box<ZlibDecoder<Writer>>),
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
Gzip(Box<GzDecoder<Writer>>),
|
||||
#[cfg(feature = "brotli")]
|
||||
Br(Box<BrotliDecoder<Writer>>),
|
||||
}
|
||||
|
||||
impl ContentDecoder {
|
||||
fn feed_eof(&mut self) -> io::Result<Option<Bytes>> {
|
||||
match self {
|
||||
#[cfg(feature = "brotli")]
|
||||
ContentDecoder::Br(ref mut decoder) => match decoder.finish() {
|
||||
Ok(mut writer) => {
|
||||
let b = writer.take();
|
||||
if !b.is_empty() {
|
||||
Ok(Some(b))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
},
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentDecoder::Gzip(ref mut decoder) => match decoder.try_finish() {
|
||||
Ok(_) => {
|
||||
let b = decoder.get_mut().take();
|
||||
if !b.is_empty() {
|
||||
Ok(Some(b))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
},
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentDecoder::Deflate(ref mut decoder) => match decoder.try_finish() {
|
||||
Ok(_) => {
|
||||
let b = decoder.get_mut().take();
|
||||
if !b.is_empty() {
|
||||
Ok(Some(b))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn feed_data(&mut self, data: Bytes) -> io::Result<Option<Bytes>> {
|
||||
match self {
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentDecoder::Br(ref mut decoder) => match decoder.write_all(&data) {
|
||||
Ok(_) => {
|
||||
decoder.flush()?;
|
||||
let b = decoder.get_mut().take();
|
||||
if !b.is_empty() {
|
||||
Ok(Some(b))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
},
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentDecoder::Gzip(ref mut decoder) => match decoder.write_all(&data) {
|
||||
Ok(_) => {
|
||||
decoder.flush()?;
|
||||
let b = decoder.get_mut().take();
|
||||
if !b.is_empty() {
|
||||
Ok(Some(b))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
},
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentDecoder::Deflate(ref mut decoder) => match decoder.write_all(&data) {
|
||||
Ok(_) => {
|
||||
decoder.flush()?;
|
||||
let b = decoder.get_mut().take();
|
||||
if !b.is_empty() {
|
||||
Ok(Some(b))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
234
actix-http/src/encoding/encoder.rs
Normal file
234
actix-http/src/encoding/encoder.rs
Normal file
@ -0,0 +1,234 @@
|
||||
//! Stream encoder
|
||||
use std::io::{self, Write};
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures::{Async, Poll};
|
||||
|
||||
#[cfg(feature = "brotli")]
|
||||
use brotli2::write::BrotliEncoder;
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
use flate2::write::{GzEncoder, ZlibEncoder};
|
||||
|
||||
use crate::body::{Body, BodyLength, MessageBody, ResponseBody};
|
||||
use crate::http::header::{ContentEncoding, CONTENT_ENCODING};
|
||||
use crate::http::{HeaderValue, HttpTryFrom, StatusCode};
|
||||
use crate::{Error, Head, ResponseHead};
|
||||
|
||||
use super::Writer;
|
||||
|
||||
pub struct Encoder<B> {
|
||||
body: EncoderBody<B>,
|
||||
encoder: Option<ContentEncoder>,
|
||||
}
|
||||
|
||||
impl<B: MessageBody> Encoder<B> {
|
||||
pub fn response(
|
||||
encoding: ContentEncoding,
|
||||
head: &mut ResponseHead,
|
||||
body: ResponseBody<B>,
|
||||
) -> ResponseBody<Encoder<B>> {
|
||||
let has_ce = head.headers().contains_key(CONTENT_ENCODING);
|
||||
match body {
|
||||
ResponseBody::Other(b) => match b {
|
||||
Body::None => ResponseBody::Other(Body::None),
|
||||
Body::Empty => ResponseBody::Other(Body::Empty),
|
||||
Body::Bytes(buf) => {
|
||||
if !(has_ce
|
||||
|| encoding == ContentEncoding::Identity
|
||||
|| encoding == ContentEncoding::Auto)
|
||||
{
|
||||
let mut enc = ContentEncoder::encoder(encoding).unwrap();
|
||||
|
||||
// TODO return error!
|
||||
let _ = enc.write(buf.as_ref());
|
||||
let body = enc.finish().unwrap();
|
||||
update_head(encoding, head);
|
||||
ResponseBody::Other(Body::Bytes(body))
|
||||
} else {
|
||||
ResponseBody::Other(Body::Bytes(buf))
|
||||
}
|
||||
}
|
||||
Body::Message(stream) => {
|
||||
if has_ce || head.status == StatusCode::SWITCHING_PROTOCOLS {
|
||||
ResponseBody::Body(Encoder {
|
||||
body: EncoderBody::Other(stream),
|
||||
encoder: None,
|
||||
})
|
||||
} else {
|
||||
update_head(encoding, head);
|
||||
head.no_chunking = false;
|
||||
ResponseBody::Body(Encoder {
|
||||
body: EncoderBody::Other(stream),
|
||||
encoder: ContentEncoder::encoder(encoding),
|
||||
})
|
||||
}
|
||||
}
|
||||
},
|
||||
ResponseBody::Body(stream) => {
|
||||
if has_ce || head.status == StatusCode::SWITCHING_PROTOCOLS {
|
||||
ResponseBody::Body(Encoder {
|
||||
body: EncoderBody::Body(stream),
|
||||
encoder: None,
|
||||
})
|
||||
} else {
|
||||
update_head(encoding, head);
|
||||
head.no_chunking = false;
|
||||
ResponseBody::Body(Encoder {
|
||||
body: EncoderBody::Body(stream),
|
||||
encoder: ContentEncoder::encoder(encoding),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum EncoderBody<B> {
|
||||
Body(B),
|
||||
Other(Box<dyn MessageBody>),
|
||||
}
|
||||
|
||||
impl<B: MessageBody> MessageBody for Encoder<B> {
|
||||
fn length(&self) -> BodyLength {
|
||||
if self.encoder.is_none() {
|
||||
match self.body {
|
||||
EncoderBody::Body(ref b) => b.length(),
|
||||
EncoderBody::Other(ref b) => b.length(),
|
||||
}
|
||||
} else {
|
||||
BodyLength::Stream
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_next(&mut self) -> Poll<Option<Bytes>, Error> {
|
||||
loop {
|
||||
let result = match self.body {
|
||||
EncoderBody::Body(ref mut b) => b.poll_next()?,
|
||||
EncoderBody::Other(ref mut b) => b.poll_next()?,
|
||||
};
|
||||
match result {
|
||||
Async::NotReady => return Ok(Async::NotReady),
|
||||
Async::Ready(Some(chunk)) => {
|
||||
if let Some(ref mut encoder) = self.encoder {
|
||||
if encoder.write(&chunk)? {
|
||||
return Ok(Async::Ready(Some(encoder.take())));
|
||||
}
|
||||
} else {
|
||||
return Ok(Async::Ready(Some(chunk)));
|
||||
}
|
||||
}
|
||||
Async::Ready(None) => {
|
||||
if let Some(encoder) = self.encoder.take() {
|
||||
let chunk = encoder.finish()?;
|
||||
if chunk.is_empty() {
|
||||
return Ok(Async::Ready(None));
|
||||
} else {
|
||||
return Ok(Async::Ready(Some(chunk)));
|
||||
}
|
||||
} else {
|
||||
return Ok(Async::Ready(None));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn update_head(encoding: ContentEncoding, head: &mut ResponseHead) {
|
||||
head.headers_mut().insert(
|
||||
CONTENT_ENCODING,
|
||||
HeaderValue::try_from(Bytes::from_static(encoding.as_str().as_bytes())).unwrap(),
|
||||
);
|
||||
}
|
||||
|
||||
enum ContentEncoder {
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
Deflate(ZlibEncoder<Writer>),
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
Gzip(GzEncoder<Writer>),
|
||||
#[cfg(feature = "brotli")]
|
||||
Br(BrotliEncoder<Writer>),
|
||||
}
|
||||
|
||||
impl ContentEncoder {
|
||||
fn encoder(encoding: ContentEncoding) -> Option<Self> {
|
||||
match encoding {
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoding::Deflate => Some(ContentEncoder::Deflate(ZlibEncoder::new(
|
||||
Writer::new(),
|
||||
flate2::Compression::fast(),
|
||||
))),
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoding::Gzip => Some(ContentEncoder::Gzip(GzEncoder::new(
|
||||
Writer::new(),
|
||||
flate2::Compression::fast(),
|
||||
))),
|
||||
#[cfg(feature = "brotli")]
|
||||
ContentEncoding::Br => {
|
||||
Some(ContentEncoder::Br(BrotliEncoder::new(Writer::new(), 3)))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn take(&mut self) -> Bytes {
|
||||
match *self {
|
||||
#[cfg(feature = "brotli")]
|
||||
ContentEncoder::Br(ref mut encoder) => encoder.get_mut().take(),
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoder::Deflate(ref mut encoder) => encoder.get_mut().take(),
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoder::Gzip(ref mut encoder) => encoder.get_mut().take(),
|
||||
}
|
||||
}
|
||||
|
||||
fn finish(self) -> Result<Bytes, io::Error> {
|
||||
match self {
|
||||
#[cfg(feature = "brotli")]
|
||||
ContentEncoder::Br(encoder) => match encoder.finish() {
|
||||
Ok(writer) => Ok(writer.buf.freeze()),
|
||||
Err(err) => Err(err),
|
||||
},
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoder::Gzip(encoder) => match encoder.finish() {
|
||||
Ok(writer) => Ok(writer.buf.freeze()),
|
||||
Err(err) => Err(err),
|
||||
},
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoder::Deflate(encoder) => match encoder.finish() {
|
||||
Ok(writer) => Ok(writer.buf.freeze()),
|
||||
Err(err) => Err(err),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn write(&mut self, data: &[u8]) -> Result<bool, io::Error> {
|
||||
match *self {
|
||||
#[cfg(feature = "brotli")]
|
||||
ContentEncoder::Br(ref mut encoder) => match encoder.write_all(data) {
|
||||
Ok(_) => Ok(!encoder.get_ref().buf.is_empty()),
|
||||
Err(err) => {
|
||||
trace!("Error decoding br encoding: {}", err);
|
||||
Err(err)
|
||||
}
|
||||
},
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoder::Gzip(ref mut encoder) => match encoder.write_all(data) {
|
||||
Ok(_) => Ok(!encoder.get_ref().buf.is_empty()),
|
||||
Err(err) => {
|
||||
trace!("Error decoding gzip encoding: {}", err);
|
||||
Err(err)
|
||||
}
|
||||
},
|
||||
#[cfg(any(feature = "flate2-c", feature = "flate2-rust"))]
|
||||
ContentEncoder::Deflate(ref mut encoder) => match encoder.write_all(data) {
|
||||
Ok(_) => Ok(!encoder.get_ref().buf.is_empty()),
|
||||
Err(err) => {
|
||||
trace!("Error decoding deflate encoding: {}", err);
|
||||
Err(err)
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
35
actix-http/src/encoding/mod.rs
Normal file
35
actix-http/src/encoding/mod.rs
Normal file
@ -0,0 +1,35 @@
|
||||
//! Content-Encoding support
|
||||
use std::io;
|
||||
|
||||
use bytes::{Bytes, BytesMut};
|
||||
|
||||
mod decoder;
|
||||
mod encoder;
|
||||
|
||||
pub use self::decoder::Decoder;
|
||||
pub use self::encoder::Encoder;
|
||||
|
||||
pub(self) struct Writer {
|
||||
buf: BytesMut,
|
||||
}
|
||||
|
||||
impl Writer {
|
||||
fn new() -> Writer {
|
||||
Writer {
|
||||
buf: BytesMut::with_capacity(8192),
|
||||
}
|
||||
}
|
||||
fn take(&mut self) -> Bytes {
|
||||
self.buf.take().freeze()
|
||||
}
|
||||
}
|
||||
|
||||
impl io::Write for Writer {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.buf.extend_from_slice(buf);
|
||||
Ok(buf.len())
|
||||
}
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@ -12,6 +12,7 @@ pub mod body;
|
||||
mod builder;
|
||||
pub mod client;
|
||||
mod config;
|
||||
pub mod encoding;
|
||||
mod extensions;
|
||||
mod header;
|
||||
mod helpers;
|
||||
|
||||
Reference in New Issue
Block a user