From 05527b3e512a61015608099a2c401609238602fe Mon Sep 17 00:00:00 2001
From: Thiago Arrais <thiago.arrais@gmail.com>
Date: Fri, 4 Dec 2020 13:57:39 -0300
Subject: [PATCH] respect max_redirects configuration

---
 awc/src/builder.rs |  1 +
 awc/src/connect.rs | 25 ++++++++++++++++---------
 awc/src/lib.rs     |  2 ++
 awc/src/sender.rs  |  2 +-
 4 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/awc/src/builder.rs b/awc/src/builder.rs
index 40fe89f65..c00e83bb2 100644
--- a/awc/src/builder.rs
+++ b/awc/src/builder.rs
@@ -191,6 +191,7 @@ impl ClientBuilder {
         let config = ClientConfig {
             headers: self.headers,
             timeout: self.timeout,
+            max_redirects: self.max_redirects,
             connector,
         };
         Client(Rc::new(config))
diff --git a/awc/src/connect.rs b/awc/src/connect.rs
index 810864891..4b5f130e8 100644
--- a/awc/src/connect.rs
+++ b/awc/src/connect.rs
@@ -5,6 +5,7 @@ use std::pin::Pin;
 use std::rc::Rc;
 use std::task::{Context, Poll};
 use std::{fmt, io, mem, net};
+use url::Url;
 
 use actix_codec::{AsyncRead, AsyncWrite, Framed};
 use actix_http::body::Body;
@@ -12,7 +13,7 @@ use actix_http::client::{
     Connect as ClientConnect, ConnectError, Connection, SendRequestError,
 };
 use actix_http::h1::ClientCodec;
-use actix_http::http::{HeaderMap, Uri};
+use actix_http::http::HeaderMap;
 use actix_http::Extensions;
 use actix_http::{RequestHead, RequestHeadType, ResponseHead};
 use actix_service::Service;
@@ -27,6 +28,7 @@ pub(crate) trait Connect {
         head: RequestHead,
         body: Body,
         addr: Option<net::SocketAddr>,
+        max_redirect: usize,
     ) -> Pin<Box<dyn Future<Output = Result<ClientResponse, SendRequestError>>>>;
 
     fn send_request_extra(
@@ -85,12 +87,15 @@ where
         head: RequestHead,
         body: Body,
         addr: Option<net::SocketAddr>,
+        max_redirects: usize,
     ) -> Pin<Box<dyn Future<Output = Result<ClientResponse, SendRequestError>>>> {
         fn deal_with_redirects<S>(
             backend: Rc<RefCell<S>>,
             head: RequestHead,
             body: Body,
             addr: Option<net::SocketAddr>,
+            max_redirects: usize,
+            redirect_count: usize,
         ) -> Pin<Box<dyn Future<Output = Result<ClientResponse, SendRequestError>>>>
         where
             S: Service<Request = ClientConnect, Error = ConnectError> + 'static,
@@ -139,23 +144,25 @@ where
                             if resphead.status.is_redirection();
                             if let Some(location_value) = resphead.headers.get(actix_http::http::header::LOCATION);
                             if let Ok(location_str) = location_value.to_str();
-                            if let Ok(location_uri) = location_str.parse::<Uri>();
                             then {
+                                if redirect_count >= max_redirects {
+                                     // TODO: need a better error
+                                    return Err(SendRequestError::Timeout);
+                                }
                                 if resphead.status == actix_http::http::StatusCode::SEE_OTHER {
                                     reqhead.method = actix_http::http::Method::GET;
                                     reqbody = Body::None;
                                 }
-                                let mut parts = location_uri.clone().into_parts();
-                                if location_uri.authority().is_none() {
-                                    parts.scheme = Some(uri.scheme().unwrap().clone());
-                                    parts.authority = Some(uri.authority().unwrap().clone());
-                                }
-                                reqhead.uri = Uri::from_parts(parts).unwrap();
+                                let base_url = Url::parse(&uri.to_string()).unwrap();
+                                let url_parser = Url::options().base_url(Some(&base_url));
+                                reqhead.uri = url_parser.parse(location_str).unwrap().as_str().parse().unwrap();
                                 return deal_with_redirects(
                                     backend.clone(),
                                     reqhead,
                                     reqbody,
                                     addr,
+                                    max_redirects,
+                                    redirect_count + 1,
                                 )
                                 .await;
                             }
@@ -167,7 +174,7 @@ where
             })
         }
 
-        deal_with_redirects(self.0.clone(), head, body, addr)
+        deal_with_redirects(self.0.clone(), head, body, addr, max_redirects, 0)
     }
 
     fn send_request_extra(
diff --git a/awc/src/lib.rs b/awc/src/lib.rs
index fb391fdab..7459b1fd9 100644
--- a/awc/src/lib.rs
+++ b/awc/src/lib.rs
@@ -148,6 +148,7 @@ pub(crate) struct ClientConfig {
     pub(crate) connector: RefCell<Box<dyn Connect>>,
     pub(crate) headers: HeaderMap,
     pub(crate) timeout: Option<Duration>,
+    pub(crate) max_redirects: usize,
 }
 
 impl Default for Client {
@@ -158,6 +159,7 @@ impl Default for Client {
             ))))),
             headers: HeaderMap::new(),
             timeout: Some(Duration::from_secs(5)),
+            max_redirects: 10,
         }))
     }
 }
diff --git a/awc/src/sender.rs b/awc/src/sender.rs
index 0bcdf4307..9a254211d 100644
--- a/awc/src/sender.rs
+++ b/awc/src/sender.rs
@@ -186,7 +186,7 @@ impl RequestSender {
 
         let fut = match self {
             RequestSender::Owned(head) => {
-                connector.send_request(head, body.into(), addr)
+                connector.send_request(head, body.into(), addr, config.max_redirects)
             }
             RequestSender::Rc(head, extra_headers) => {
                 connector.send_request_extra(head, extra_headers, body.into(), addr)