Skip to content

Commit

Permalink
Merge pull request #771 from hyperium/allow-proxy
Browse files Browse the repository at this point in the history
feat(client): add Proxy support
  • Loading branch information
seanmonstar committed Apr 26, 2016
2 parents 4828437 + 25010fc commit 3a3e086
Show file tree
Hide file tree
Showing 4 changed files with 115 additions and 16 deletions.
64 changes: 56 additions & 8 deletions src/client/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,9 +55,9 @@
//! clone2.post("http://example.domain/post").body("foo=bar").send().unwrap();
//! });
//! ```
use std::borrow::Cow;
use std::default::Default;
use std::io::{self, copy, Read};
use std::iter::Extend;
use std::fmt;

use std::time::Duration;
Expand All @@ -66,7 +66,7 @@ use url::Url;
use url::ParseError as UrlError;

use header::{Headers, Header, HeaderFormat};
use header::{ContentLength, Location};
use header::{ContentLength, Host, Location};
use method::Method;
use net::{NetworkConnector, NetworkStream};
use Error;
Expand All @@ -90,6 +90,7 @@ pub struct Client {
redirect_policy: RedirectPolicy,
read_timeout: Option<Duration>,
write_timeout: Option<Duration>,
proxy: Option<(Cow<'static, str>, Cow<'static, str>, u16)>
}

impl fmt::Debug for Client {
Expand All @@ -98,6 +99,7 @@ impl fmt::Debug for Client {
.field("redirect_policy", &self.redirect_policy)
.field("read_timeout", &self.read_timeout)
.field("write_timeout", &self.write_timeout)
.field("proxy", &self.proxy)
.finish()
}
}
Expand Down Expand Up @@ -127,6 +129,7 @@ impl Client {
redirect_policy: Default::default(),
read_timeout: None,
write_timeout: None,
proxy: None,
}
}

Expand All @@ -145,6 +148,12 @@ impl Client {
self.write_timeout = dur;
}

/// Set a proxy for requests of this Client.
pub fn set_proxy<S, H>(&mut self, scheme: S, host: H, port: u16)
where S: Into<Cow<'static, str>>, H: Into<Cow<'static, str>> {
self.proxy = Some((scheme.into(), host.into(), port));
}

/// Build a Get request.
pub fn get<U: IntoUrl>(&self, url: U) -> RequestBuilder {
self.request(Method::Get, url)
Expand Down Expand Up @@ -247,7 +256,7 @@ impl<'a> RequestBuilder<'a> {
pub fn send(self) -> ::Result<Response> {
let RequestBuilder { client, method, url, headers, body } = self;
let mut url = try!(url);
trace!("send {:?} {:?}", method, url);
trace!("send method={:?}, url={:?}, client={:?}", method, url, client);

let can_have_body = match method {
Method::Get | Method::Head => false,
Expand All @@ -261,12 +270,25 @@ impl<'a> RequestBuilder<'a> {
};

loop {
let message = {
let (host, port) = try!(get_host_and_port(&url));
try!(client.protocol.new_message(&host, port, url.scheme()))
let mut req = {
let (scheme, host, port) = match client.proxy {
Some(ref proxy) => (proxy.0.as_ref(), proxy.1.as_ref(), proxy.2),
None => {
let hp = try!(get_host_and_port(&url));
(url.scheme(), hp.0, hp.1)
}
};
let mut headers = match headers {
Some(ref headers) => headers.clone(),
None => Headers::new(),
};
headers.set(Host {
hostname: host.to_owned(),
port: Some(port),
});
let message = try!(client.protocol.new_message(&host, port, scheme));
Request::with_headers_and_message(method.clone(), url.clone(), headers, message)
};
let mut req = try!(Request::with_message(method.clone(), url.clone(), message));
headers.as_ref().map(|headers| req.headers_mut().extend(headers.iter()));

try!(req.set_write_timeout(client.write_timeout));
try!(req.set_read_timeout(client.read_timeout));
Expand Down Expand Up @@ -456,6 +478,8 @@ fn get_host_and_port(url: &Url) -> ::Result<(&str, u16)> {
mod tests {
use std::io::Read;
use header::Server;
use http::h1::Http11Message;
use mock::{MockStream};
use super::{Client, RedirectPolicy};
use super::pool::Pool;
use url::Url;
Expand All @@ -477,6 +501,30 @@ mod tests {
"
});


#[test]
fn test_proxy() {
use super::pool::PooledStream;
mock_connector!(ProxyConnector {
b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n"
});
let mut client = Client::with_connector(Pool::with_connector(Default::default(), ProxyConnector));
client.set_proxy("http", "example.proxy", 8008);
let mut dump = vec![];
client.get("http://127.0.0.1/foo/bar").send().unwrap().read_to_end(&mut dump).unwrap();

{
let box_message = client.protocol.new_message("example.proxy", 8008, "http").unwrap();
let message = box_message.downcast::<Http11Message>().unwrap();
let stream = message.into_inner().downcast::<PooledStream<MockStream>>().unwrap().into_inner();
let s = ::std::str::from_utf8(&stream.write).unwrap();
let request_line = "GET http://127.0.0.1/foo/bar HTTP/1.1\r\n";
assert_eq!(&s[..request_line.len()], request_line);
assert!(s.contains("Host: example.proxy:8008\r\n"));
}

}

#[test]
fn test_redirect_followall() {
let mut client = Client::with_connector(MockRedirectPolicy);
Expand Down
14 changes: 10 additions & 4 deletions src/client/pool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -102,22 +102,21 @@ impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector fo
type Stream = PooledStream<S>;
fn connect(&self, host: &str, port: u16, scheme: &str) -> ::Result<PooledStream<S>> {
let key = key(host, port, scheme);
let mut locked = self.inner.lock().unwrap();
let mut should_remove = false;
let inner = match locked.conns.get_mut(&key) {
let inner = match self.inner.lock().unwrap().conns.get_mut(&key) {
Some(ref mut vec) => {
trace!("Pool had connection, using");
should_remove = vec.len() == 1;
vec.pop().unwrap()
}
_ => PooledStreamInner {
None => PooledStreamInner {
key: key.clone(),
stream: try!(self.connector.connect(host, port, scheme)),
previous_response_expected_no_content: false,
}
};
if should_remove {
locked.conns.remove(&key);
self.inner.lock().unwrap().conns.remove(&key);
}
Ok(PooledStream {
inner: Some(inner),
Expand All @@ -134,6 +133,13 @@ pub struct PooledStream<S> {
pool: Arc<Mutex<PoolImpl<S>>>,
}

impl<S: NetworkStream> PooledStream<S> {
/// Take the wrapped stream out of the pool completely.
pub fn into_inner(mut self) -> S {
self.inner.take().expect("PooledStream lost its inner stream").stream
}
}

#[derive(Debug)]
struct PooledStreamInner<S> {
key: Key,
Expand Down
38 changes: 36 additions & 2 deletions src/client/request.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,14 +70,20 @@ impl Request<Fresh> {
});
}

Ok(Request {
Ok(Request::with_headers_and_message(method, url, headers, message))
}

#[doc(hidden)]
pub fn with_headers_and_message(method: Method, url: Url, headers: Headers, message: Box<HttpMessage>)
-> Request<Fresh> {
Request {
method: method,
headers: headers,
url: url,
version: version::HttpVersion::Http11,
message: message,
_marker: PhantomData,
})
}
}

/// Create a new client request.
Expand Down Expand Up @@ -129,6 +135,8 @@ impl Request<Fresh> {
pub fn headers_mut(&mut self) -> &mut Headers { &mut self.headers }
}



impl Request<Streaming> {
/// Completes writing the request, and returns a response to read from.
///
Expand Down Expand Up @@ -246,6 +254,32 @@ mod tests {
assert!(!s.contains("Content-Length:"));
}

#[test]
fn test_host_header() {
let url = Url::parse("http://example.dom").unwrap();
let req = Request::with_connector(
Get, url, &mut MockConnector
).unwrap();
let bytes = run_request(req);
let s = from_utf8(&bytes[..]).unwrap();
assert!(s.contains("Host: example.dom"));
}

#[test]
fn test_proxy() {
let url = Url::parse("http://example.dom").unwrap();
let proxy_url = Url::parse("http://pro.xy").unwrap();
let mut req = Request::with_connector(
Get, proxy_url, &mut MockConnector
).unwrap();
req.url = url;
let bytes = run_request(req);
let s = from_utf8(&bytes[..]).unwrap();
let request_line = "GET http://example.dom/ HTTP/1.1";
assert_eq!(&s[..request_line.len()], request_line);
assert!(s.contains("Host: pro.xy"));
}

#[test]
fn test_post_chunked_with_encoding() {
let url = Url::parse("http://example.dom").unwrap();
Expand Down
15 changes: 13 additions & 2 deletions src/http/h1.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use url::Position as UrlPosition;

use buffer::BufReader;
use Error;
use header::{Headers, ContentLength, TransferEncoding};
use header::{Headers, Host, ContentLength, TransferEncoding};
use header::Encoding::Chunked;
use method::{Method};
use net::{NetworkConnector, NetworkStream};
Expand Down Expand Up @@ -144,7 +144,18 @@ impl HttpMessage for Http11Message {
let mut stream = BufWriter::new(stream);

{
let uri = &head.url[UrlPosition::BeforePath..UrlPosition::AfterQuery];
let uri = match head.headers.get::<Host>() {
Some(host)
if Some(&*host.hostname) == head.url.host_str()
&& host.port == head.url.port_or_known_default() => {
&head.url[UrlPosition::BeforePath..UrlPosition::AfterQuery]
},
_ => {
trace!("url and host header dont match, using absolute uri form");
head.url.as_ref()
}

};

let version = version::HttpVersion::Http11;
debug!("request line: {:?} {:?} {:?}", head.method, uri, version);
Expand Down

0 comments on commit 3a3e086

Please sign in to comment.