From 521f1a0767339a7e455cdc92580aa477900a50d7 Mon Sep 17 00:00:00 2001
From: Daniel Abramov <dabramov@snapview.de>
Date: Mon, 16 Nov 2020 19:25:08 +0100
Subject: [PATCH] clean up the redirect logic a bit

---
 src/client.rs       | 63 +++++++++++++++++++++------------------------
 src/protocol/mod.rs |  4 ---
 2 files changed, 30 insertions(+), 37 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 49ed656..b28cf3c 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -4,7 +4,7 @@ use std::io::{Read, Write};
 use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
 use std::result::Result as StdResult;
 
-use http::Uri;
+use http::{Uri, request::Parts};
 use log::*;
 
 use url::Url;
@@ -89,10 +89,12 @@ use crate::stream::{Mode, NoDelay};
 pub fn connect_with_config<Req: IntoClientRequest>(
     request: Req,
     config: Option<WebSocketConfig>,
+    max_redirects: u8,
 ) -> Result<(WebSocket<AutoStream>, Response)> {
-    let mut request: Request = request.into_client_request()?;
 
-    fn inner(request: Request, config: Option<WebSocketConfig>) -> Result<(WebSocket<AutoStream>, Response)> {
+    fn try_client_handshake(request: Request, config: Option<WebSocketConfig>)
+        -> Result<(WebSocket<AutoStream>, Response)>
+    {
         let uri = request.uri();
         let mode = uri_mode(uri)?;
         let host = request
@@ -106,44 +108,39 @@ pub fn connect_with_config<Req: IntoClientRequest>(
         let addrs = (host, port).to_socket_addrs()?;
         let mut stream = connect_to_some(addrs.as_slice(), &request.uri(), mode)?;
         NoDelay::set_nodelay(&mut stream, true)?;
-        client_with_config(request, stream, config).map_err(|e| match e{
+        client_with_config(request, stream, config).map_err(|e| match e {
             HandshakeError::Failure(f) => f,
             HandshakeError::Interrupted(_) => panic!("Bug: blocking handshake not blocked"),
         })
     }
 
-    let max_redirects = config.as_ref().and_then(|c| c.max_redirects).unwrap_or(0);
-    let mut redirects = 0;
-
-    loop {
-        // Copy all the fields from the initial reqeust **except** the URI. This will be used in the event of a redirection code
-        // Have to manually clone Method because there is one field that contains a Box,
-        // but in the case of normal request methods it is Copy
-        let request2 = Request::builder()
-            .method(request.method().clone())
-            .version(request.version());
-
-        match inner(request, config) {
-            Ok(r) => return Ok(r),
-            Err(e) => match e {
-                Error::Http(res) =>  {
-                    if res.status().is_redirection() {
-                        let uri = res.headers().get("Location").ok_or(Error::NoLocation)?;
-                        debug!("Redirecting to {:?}", uri);
-                        request = request2.uri(uri.to_str()?.parse::<Uri>()?).body(()).unwrap();
-                        redirects += 1;
-                        if redirects > max_redirects {
-                            return Err(Error::Http(res));
-                        }
-                    } else {
-                        return Err(Error::Http(res));
-                    }
-                }
-                _ => return Err(e),
+    fn create_request(parts: &Parts, uri: &Uri) -> Request {
+        let mut builder = Request::builder()
+            .uri(uri.clone())
+            .method(parts.method.clone())
+            .version(parts.version.clone());
+        *builder.headers_mut().expect("Failed to create `Request`") = parts.headers.clone();
+        builder.body(()).expect("Failed to create `Request`")
+    }
+
+    let (parts, _) = request.into_client_request()?.into_parts();
+    let mut uri = parts.uri.clone();
+
+    for attempt in 0..(max_redirects + 1) {
+        let request = create_request(&parts, &uri);
+
+        match try_client_handshake(request, config) {
+            Err(Error::Http(res)) if res.status().is_redirection() && attempt < max_redirects => {
+                let location = res.headers().get("Location").ok_or(Error::NoLocation)?;
+                uri = location.to_str()?.parse::<Uri>()?;
+                debug!("Redirecting to {:?}", uri);
+                continue;
             }
+            other => return other,
         }
     }
 
+    unreachable!("Bug in a redirect handling logic")
 }
 
 /// Connect to the given WebSocket in blocking mode.
@@ -159,7 +156,7 @@ pub fn connect_with_config<Req: IntoClientRequest>(
 /// use `client` instead. There is no need to enable the "tls" feature if you don't call
 /// `connect` since it's the only function that uses native_tls.
 pub fn connect<Req: IntoClientRequest>(request: Req) -> Result<(WebSocket<AutoStream>, Response)> {
-    connect_with_config(request, None)
+    connect_with_config(request, None, 0)
 }
 
 fn connect_to_some(addrs: &[SocketAddr], uri: &Uri, mode: Mode) -> Result<AutoStream> {
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 505dddd..8137393 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -43,9 +43,6 @@ pub struct WebSocketConfig {
     /// be reasonably big for all normal use-cases but small enough to prevent memory eating
     /// by a malicious user.
     pub max_frame_size: Option<usize>,
-    /// The max number of redirects the client should follow before aborting the connection.
-    /// The default value is 3. `None` here means that the client will not attempt to follow redirects.
-    pub max_redirects: Option<u8>,
 }
 
 impl Default for WebSocketConfig {
@@ -54,7 +51,6 @@ impl Default for WebSocketConfig {
             max_send_queue: None,
             max_message_size: Some(64 << 20),
             max_frame_size: Some(16 << 20),
-            max_redirects: Some(3)
         }
     }
 }