From 07f89984c485a4a6a4d8586f0004638fec71bd3f Mon Sep 17 00:00:00 2001
From: Johann Tuffe <tafia973@gmail.com>
Date: Mon, 18 Dec 2017 17:53:17 +0800
Subject: [PATCH] store intercept to set_proxy on requests

---
 Cargo.toml               |  3 ++-
 src/async_impl/client.rs | 10 ++++++++++
 src/proxy.rs             |  2 +-
 tests/support/server.rs  |  3 ++-
 4 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 74ec12d..6e3bf95 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,7 +26,8 @@ tokio-io = "0.1"
 tokio-tls = "0.1"
 url = "1.2"
 uuid = { version = "0.5", features = ["v4"] }
-hyper-proxy = "0.2.0"
+hyper-proxy = { version = "0.2.0", path = "../hyper-proxy" }
+#hyper-proxy = "0.2.0"
 
 [dev-dependencies]
 env_logger = "0.4"
diff --git a/src/async_impl/client.rs b/src/async_impl/client.rs
index 4f2e880..b99b169 100644
--- a/src/async_impl/client.rs
+++ b/src/async_impl/client.rs
@@ -7,6 +7,7 @@ use futures::{Async, Future, Poll};
 use hyper::client::{Connect, FutureResponse, HttpConnector};
 use hyper::header::{Headers, Location, Referer, UserAgent, Accept, Encoding,
                     AcceptEncoding, Range, qitem};
+use hyper_proxy::Intercept;
 use hyper_tls::HttpsConnector;
 use native_tls::{TlsConnector, TlsConnectorBuilder};
 use tokio_core::reactor::Handle;
@@ -128,6 +129,7 @@ impl ClientBuilder {
         if !config.hostname_verification {
             https_connector.danger_disable_hostname_verification(true);
         }
+        let intercept = config.proxy.inner.intercept().clone();
         let mut connector = config.proxy.inner.with_connector(https_connector);
         let tls = try_!(config.tls.build());
         connector.set_tls(Some(tls));
@@ -140,6 +142,7 @@ impl ClientBuilder {
             inner: Arc::new(ClientRef {
                 gzip: config.gzip,
                 hyper: hyper_client,
+                intercept: intercept,
                 headers: config.headers,
                 redirect_policy: config.redirect_policy,
                 referer: config.referer,
@@ -412,6 +415,9 @@ impl Client {
             reusable
         });
 
+        if self.inner.intercept.matches(&uri) && uri.scheme() == Some("http") {
+            req.set_proxy(true);
+        }
         let in_flight = self.inner.hyper.request(req);
 
         Pending {
@@ -451,6 +457,7 @@ impl fmt::Debug for ClientBuilder {
 struct ClientRef {
     gzip: bool,
     headers: Headers,
+    intercept: Intercept,
     hyper: HyperClient,
     redirect_policy: RedirectPolicy,
     referer: bool,
@@ -549,6 +556,9 @@ impl Future for PendingRequest {
                             if let Some(Some(ref body)) = self.body {
                                 req.set_body(body.clone());
                             }
+                            if self.client.intercept.matches(&uri) && uri.scheme() == Some("http") {
+                                req.set_proxy(true);
+                            }
                             self.in_flight = self.client.hyper.request(req);
                             continue;
                         },
diff --git a/src/proxy.rs b/src/proxy.rs
index 78ec807..d1f84e3 100644
--- a/src/proxy.rs
+++ b/src/proxy.rs
@@ -113,7 +113,7 @@ impl Proxy {
     /// # fn main() {}
     pub fn custom<F, U: IntoUrl>(fun: F, url: U) -> ::Result<Proxy>
     where F: Fn(&Uri) -> bool + 'static + Send + Sync {
-        Proxy::new(Intercept::Custom(Box::new(fun)), url)
+        Proxy::new(Intercept::Custom(fun.into()), url)
     }
 
     /*
diff --git a/tests/support/server.rs b/tests/support/server.rs
index 145bbf0..e882e29 100644
--- a/tests/support/server.rs
+++ b/tests/support/server.rs
@@ -52,7 +52,8 @@ pub fn spawn(txns: Vec<Txn>) -> Server {
             let mut n = 0;
             while n < expected.len() {
                 match socket.read(&mut buf[n..]) {
-                    Ok(0) | Err(_) => break,
+                    Ok(0) => break,
+                    Err(e) => panic!(e),
                     Ok(nread) => n += nread,
                 }
             }