Skip to content

Commit

Permalink
Add (configurable) automatic gzip decompression.
Browse files Browse the repository at this point in the history
  • Loading branch information
Jason Schein committed Feb 27, 2017
1 parent 1d3c134 commit ab5e477
Show file tree
Hide file tree
Showing 4 changed files with 159 additions and 19 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ serde = "0.9"
serde_json = "0.9"
serde_urlencoded = "0.4"
url = "1.2"
libflate = "0.1.3"

[dev-dependencies]
env_logger = "0.3"
133 changes: 114 additions & 19 deletions src/client.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
use std::fmt;
use std::io::{self, Read};
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, Ordering};

use hyper::client::IntoUrl;
use hyper::header::{Headers, ContentType, Location, Referer, UserAgent, Accept};
use hyper::header::{Headers, ContentType, Location, Referer, UserAgent, Accept, ContentEncoding, Encoding, ContentLength};
use hyper::method::Method;
use hyper::status::StatusCode;
use hyper::version::HttpVersion;
Expand Down Expand Up @@ -38,10 +39,16 @@ impl Client {
inner: Arc::new(ClientRef {
hyper: client,
redirect_policy: Mutex::new(RedirectPolicy::default()),
auto_ungzip: AtomicBool::new(true),
}),
})
}

/// Enable auto gzip decompression by checking the ContentEncoding response header.
pub fn gzip(&mut self, enable: bool) {
self.inner.auto_ungzip.store(enable, Ordering::Relaxed);
}

/// Set a `RedirectPolicy` for this client.
pub fn redirect(&mut self, policy: RedirectPolicy) {
*self.inner.redirect_policy.lock().unwrap() = policy;
Expand Down Expand Up @@ -94,13 +101,15 @@ impl fmt::Debug for Client {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("Client")
.field("redirect_policy", &self.inner.redirect_policy)
.field("auto_ungzip", &self.inner.auto_ungzip)
.finish()
}
}

struct ClientRef {
hyper: ::hyper::Client,
redirect_policy: Mutex<RedirectPolicy>,
auto_ungzip: AtomicBool,
}

fn new_hyper_client() -> ::Result<::hyper::Client> {
Expand Down Expand Up @@ -268,7 +277,7 @@ impl RequestBuilder {
loc
} else {
return Ok(Response {
inner: res
inner: Decoder::from_hyper_response(res, client.auto_ungzip.load(Ordering::Relaxed))
});
}
};
Expand All @@ -282,14 +291,14 @@ impl RequestBuilder {
} else {
debug!("redirect_policy disallowed redirection to '{}'", loc);
return Ok(Response {
inner: res
inner: Decoder::from_hyper_response(res, client.auto_ungzip.load(Ordering::Relaxed))
})
}
},
Err(e) => {
debug!("Location header had invalid URI: {:?}", e);
return Ok(Response {
inner: res
inner: Decoder::from_hyper_response(res, client.auto_ungzip.load(Ordering::Relaxed))
})
}
};
Expand All @@ -299,7 +308,7 @@ impl RequestBuilder {
//TODO: removeSensitiveHeaders(&mut headers, &url);
} else {
return Ok(Response {
inner: res
inner: Decoder::from_hyper_response(res, client.auto_ungzip.load(Ordering::Relaxed))
});
}
}
Expand All @@ -318,26 +327,56 @@ impl fmt::Debug for RequestBuilder {

/// A Response to a submitted `Request`.
pub struct Response {
inner: ::hyper::client::Response,
inner: Decoder,
}

impl fmt::Debug for Response {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
return match &self.inner {
&Decoder::PlainText(ref hyper_response) => {
f.debug_struct("Response")
.field("status", &hyper_response.status)
.field("headers", &hyper_response.headers)
.field("version", &hyper_response.version)
.finish()
},
&Decoder::Gzip{ref status, ref version, ref headers, ..} => {
f.debug_struct("Response")
.field("status", &status)
.field("headers", &headers)
.field("version", &version)
.finish()
}
}
}
}

impl Response {
/// Get the `StatusCode`.
#[inline]
pub fn status(&self) -> &StatusCode {
&self.inner.status
match &self.inner {
&Decoder::PlainText(ref hyper_response) => &hyper_response.status,
&Decoder::Gzip{ref status, ..} => status
}
}

/// Get the `Headers`.
#[inline]
pub fn headers(&self) -> &Headers {
&self.inner.headers
match &self.inner {
&Decoder::PlainText(ref hyper_response) => &hyper_response.headers,
&Decoder::Gzip{ref headers, ..} => headers
}
}

/// Get the `HttpVersion`.
#[inline]
pub fn version(&self) -> &HttpVersion {
&self.inner.version
match &self.inner {
&Decoder::PlainText(ref hyper_response) => &hyper_response.version,
&Decoder::Gzip{ref version, ..} => version
}
}

/// Try and deserialize the response body as JSON.
Expand All @@ -347,6 +386,72 @@ impl Response {
}
}

enum Decoder {
/// A `PlainText` decoder just returns the response content as is.
PlainText(::hyper::client::Response),
/// A `Gzip` decoder will uncompress the gziped response content before returning it.
Gzip {
decoder: ::libflate::gzip::Decoder<::hyper::client::Response>,
headers: ::hyper::header::Headers,
version: ::hyper::version::HttpVersion,
status: ::hyper::status::StatusCode,
}
}

impl Decoder {
/// Constructs a Decoder from a hyper request.
///
/// A decoder is just a wrapper around the hyper request that knows
/// how to decode the content body of the request.
///
/// Uses the correct variant by inspecting the Content-Encoding header.
fn from_hyper_response(res: ::hyper::client::Response, check_gzip: bool) -> Self {
if !check_gzip {
return Decoder::PlainText(res);
}

let mut is_gzip = false;
match res.headers.get::<ContentEncoding>() {
Some(encoding_types) => {
if encoding_types.contains(&Encoding::Gzip) {
is_gzip = true;
}
if let Some(content_length) = res.headers.get::<ContentLength>() {
if content_length.0 == 0 {
warn!("GZipped response with content-length of 0");
is_gzip = false;
}
}
}
_ => {}
}

if is_gzip {
return Decoder::Gzip {
status: res.status.clone(),
version: res.version.clone(),
headers: res.headers.clone(),
decoder: ::libflate::gzip::Decoder::new(res).unwrap(),
};
} else {
return Decoder::PlainText(res);
}
}
}

impl Read for Decoder {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
&mut Decoder::PlainText(ref mut hyper_response) => {
hyper_response.read(buf)
},
&mut Decoder::Gzip{ref mut decoder, ..} => {
decoder.read(buf)
}
}
}
}

/// Read the body of the Response.
impl Read for Response {
#[inline]
Expand All @@ -355,16 +460,6 @@ impl Read for Response {
}
}

impl fmt::Debug for Response {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("Response")
.field("status", self.status())
.field("headers", self.headers())
.field("version", self.version())
.finish()
}
}

#[cfg(test)]
mod tests {
use super::*;
Expand Down
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,7 @@
extern crate hyper;

#[macro_use] extern crate log;
extern crate libflate;
extern crate hyper_native_tls;
extern crate serde;
extern crate serde_json;
Expand Down
43 changes: 43 additions & 0 deletions tests/client.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
extern crate reqwest;
extern crate libflate;

#[macro_use] mod server;

use std::io::Read;
use std::io::prelude::*;

#[test]
fn test_get() {
Expand Down Expand Up @@ -248,3 +250,44 @@ fn test_accept_header_is_not_changed_if_set() {

assert_eq!(res.status(), &reqwest::StatusCode::Ok);
}

#[test]
fn test_gzip_response() {
let mut encoder = ::libflate::gzip::Encoder::new(Vec::new()).unwrap();
match encoder.write(b"test request") {
Ok(n) => assert!(n > 0, "Failed to write to encoder."),
_ => panic!("Failed to gzip encode string.")
};

let gzipped_content = encoder.finish().into_result().unwrap();

let mut response = format!("\
HTTP/1.1 200 OK\r\n\
Server: test-accept\r\n\
Content-Encoding: gzip\r\n\
Content-Length: {}\r\n\
\r\n", &gzipped_content.len())
.into_bytes();
response.extend(&gzipped_content);

let server = server! {
request: b"\
GET /gzip HTTP/1.1\r\n\
Host: $HOST\r\n\
User-Agent: $USERAGENT\r\n\
Accept: */*\r\n\
\r\n\
",
response: response
};
let mut res = reqwest::get(&format!("http://{}/gzip", server.addr()))
.unwrap();

let mut body = ::std::string::String::new();
match res.read_to_string(&mut body) {
Ok(n) => assert!(n > 0, "Failed to write to buffer."),
_ => panic!("Failed to write to buffer.")
};

assert_eq!(body, "test request");
}

0 comments on commit ab5e477

Please sign in to comment.