ws client redirections

ws-client/src/transport.rs

@@ -24,10 +24,11 @@
 // IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
+use crate::stream::EitherStream;
 use futures::io::{BufReader, BufWriter};
-use futures::prelude::*;
 use soketto::connection;
-use soketto::handshake::client::{Client as WsRawClient, ServerResponse};
+use soketto::handshake::client::{Client as WsHandshakeClient, ServerResponse};
+use std::path::{Path, PathBuf};
 use std::{borrow::Cow, io, net::SocketAddr, sync::Arc, time::Duration};
 use thiserror::Error;
 use tokio::net::TcpStream;
@@ -37,7 +38,10 @@ use tokio_rustls::{
 	TlsConnector,
 };
 
-type TlsOrPlain = crate::stream::EitherStream<TcpStream, TlsStream<TcpStream>>;
+type TlsOrPlain = EitherStream<TcpStream, TlsStream<TcpStream>>;
+
+// TODO(niklasad1): make this configurable.
+const MAX_REDIRECTIONS_ALLOWED: usize = 5;
 
 /// Sending end of WebSocket transport.
 #[derive(Debug)]
@@ -112,7 +116,7 @@ pub enum WsHandshakeError {
 	#[error("Invalid DNS name: {}", 0)]
 	InvalidDnsName(#[source] InvalidDNSNameError),
 
-	/// RawServer rejected our handshake.
+	/// Server rejected the handshake.
 	#[error("Connection rejected with status code: {}", status_code)]
 	Rejected {
 		/// HTTP status code that the server returned.
@@ -179,74 +183,118 @@ impl<'a> WsTransportClientBuilder<'a> {
 			Mode::Plain => None,
 		};
 
-		let mut err = None;
-		for sockaddr in &self.target.sockaddrs {
-			match self.try_connect(*sockaddr, &connector).await {
-				Ok(res) => return Ok(res),
-				Err(e) => {
-					log::debug!("Failed to connect to sockaddr: {:?} with err: {:?}", sockaddr, e);
-					err = Some(Err(e));
-				}
-			}
-		}
-		// NOTE(niklasad1): this is most likely unreachable because [`Url::socket_addrs`] doesn't
-		// return an empty `Vec` if no socket address was found for the host name.
-		err.unwrap_or(Err(WsHandshakeError::NoAddressFound(self.target.host)))
+		self.try_connect(connector).await
 	}
 
 	async fn try_connect(
-		&self,
-		sockaddr: SocketAddr,
-		tls_connector: &Option<TlsConnector>,
+		self,
+		mut tls_connector: Option<TlsConnector>,
 	) -> Result<(Sender, Receiver), WsHandshakeError> {
-		// Try establish the TCP connection.
-		let tcp_stream = {
-			let socket = TcpStream::connect(sockaddr);
-			let timeout = tokio::time::sleep(self.timeout);
-			futures::pin_mut!(socket, timeout);
-			match future::select(socket, timeout).await {
-				future::Either::Left((socket, _)) => {
-					let socket = socket?;
-					if let Err(err) = socket.set_nodelay(true) {
-						log::warn!("set nodelay failed: {:?}", err);
-					}
-					match tls_connector {
-						None => TlsOrPlain::Plain(socket),
-						Some(connector) => {
-							let dns_name = DNSNameRef::try_from_ascii_str(&self.target.host)?;
-							let tls_stream = connector.connect(dns_name, socket).await?;
-							TlsOrPlain::Tls(tls_stream)
-						}
-					}
-				}
-				future::Either::Right((_, _)) => return Err(WsHandshakeError::Timeout(self.timeout)),
+		let mut sockaddrs = self.target.sockaddrs;
+		let mut path_and_query = PathBuf::from(self.target.path_and_query);
+		let mut host = self.target.host;
+		let mut host_header = self.target.host_header;
+
+		let mut err = None;
+
+		for _ in 0..MAX_REDIRECTIONS_ALLOWED {
+			let sockaddr = match sockaddrs.pop() {
+				Some(addr) => addr,
+				None => return err.unwrap_or(Err(WsHandshakeError::NoAddressFound(host))),
+			};
+
+			let tcp_stream = connect(sockaddr, self.timeout, &host, &tls_connector).await?;
+			let mut client = WsHandshakeClient::new(
+				BufReader::new(BufWriter::new(tcp_stream)),
+				&host_header,
+				path_and_query.to_str().expect("valid UTF-8 checked by Url::parse; qed"),
+			);
+			if let Some(origin) = self.origin_header.as_ref() {
+				client.set_origin(origin);
 			}
-		};
+			// Perform the initial handshake.
+			match client.handshake().await? {
+				ServerResponse::Accepted { .. } => {
+					let mut builder = client.into_builder();
+					builder.set_max_message_size(self.max_request_body_size as usize);
+					let (sender, receiver) = builder.finish();
+					return Ok((Sender { inner: sender }, Receiver { inner: receiver }));
+				}
 
-		log::debug!("Connecting to target: {:?}", self.target);
-		let mut client = WsRawClient::new(
-			BufReader::new(BufWriter::new(tcp_stream)),
-			&self.target.host_header,
-			&self.target.path_and_query,
-		);
-		if let Some(origin) = self.origin_header.as_ref() {
-			client.set_origin(origin);
+				ServerResponse::Rejected { status_code } => {
+					err = Some(Err(WsHandshakeError::Rejected { status_code }));
+				}
+				ServerResponse::Redirect { status_code, location } => {
+					log::trace!("recv redirection: status_code: {}, location: {}", status_code, location);
+					match url::Url::parse(&location) {
+						// redirection with absolute path => need to lookup.
+						Ok(url) => {
+							let target = Target::parse(url)?;
+							sockaddrs = target.sockaddrs;
+							path_and_query = PathBuf::from(target.path_and_query);
+							host = target.host;
+							host_header = target.host_header;
+							tls_connector = match target.mode {
+								Mode::Tls => {
+									let mut client_config = rustls::ClientConfig::default();
+									if let CertificateStore::Native = self.certificate_store {
+										client_config.root_store = rustls_native_certs::load_native_certs()
+											.map_err(|(_, e)| WsHandshakeError::CertificateStore(e))?;
+									}
+									Some(Arc::new(client_config).into())
+								}
+								Mode::Plain => None,
+							};
+						}
+						// redirection is relative, either `/baz` or `bar`.
+						Err(
+							url::ParseError::RelativeUrlWithoutBase | url::ParseError::RelativeUrlWithCannotBeABaseBase,
+						) => {
+							// replace the entire path if `location` is `/`.
+							if location.starts_with('/') {
+								path_and_query = location.into();
+							} else {
+								// join paths such that the leaf is replaced with `location`.
+								let strip_last_child =
+									Path::new(&path_and_query).ancestors().nth(1).unwrap_or_else(|| Path::new("/"));
+								path_and_query = strip_last_child.join(location);
+							}
+						}
+						Err(e) => {
+							err = Some(Err(WsHandshakeError::Url(e.to_string().into())));
+						}
+					};
+				}
+			};
 		}
+		err.unwrap_or(Err(WsHandshakeError::NoAddressFound(host)))
+	}
+}
 
-		// Perform the initial handshake.
-		match client.handshake().await? {
-			ServerResponse::Accepted { .. } => {}
-			ServerResponse::Rejected { status_code } | ServerResponse::Redirect { status_code, .. } => {
-				// TODO: HTTP redirects also lead here #339.
-				return Err(WsHandshakeError::Rejected { status_code });
+async fn connect(
+	sockaddr: SocketAddr,
+	timeout_dur: Duration,
+	host: &str,
+	tls_connector: &Option<TlsConnector>,
+) -> Result<EitherStream<TcpStream, TlsStream<TcpStream>>, WsHandshakeError> {
+	let socket = TcpStream::connect(sockaddr);
+	let timeout = tokio::time::sleep(timeout_dur);
+	tokio::select! {
+		socket = socket => {
+			let socket = socket?;
+			if let Err(err) = socket.set_nodelay(true) {
+				log::warn!("set nodelay failed: {:?}", err);
+			}
+			match tls_connector {
+				None => Ok(TlsOrPlain::Plain(socket)),
+				Some(connector) => {
+					let dns_name = DNSNameRef::try_from_ascii_str(host)?;
+					let tls_stream = connector.connect(dns_name, socket).await?;
+					Ok(TlsOrPlain::Tls(tls_stream))
+				}
 			}
 		}
-
-		// If the handshake succeeded, return.
-		let mut builder = client.into_builder();
-		builder.set_max_message_size(self.max_request_body_size as usize);
-		let (sender, receiver) = builder.finish();
-		Ok((Sender { inner: sender }, Receiver { inner: receiver }))
+		_ = timeout => Err(WsHandshakeError::Timeout(timeout_dur))
 	}
 }