paritytech · niklasad1 · Oct 5, 2021 · Jun 30, 2021 · Jun 30, 2021 · Jun 30, 2021
@@ -31,7 +31,7 @@ use crate::types::{
 };
 use crate::HttpClientBuilder;
 use jsonrpsee_test_utils::helpers::*;
-use jsonrpsee_test_utils::types::Id;
+use jsonrpsee_test_utils::mocks::Id;
 use jsonrpsee_test_utils::TimeoutFutureExt;
 
 #[tokio::test]

@@ -32,7 +32,7 @@ use crate::types::error::{CallError, Error};
 use crate::{server::StopHandle, HttpServerBuilder, RpcModule};
 
 use jsonrpsee_test_utils::helpers::*;
-use jsonrpsee_test_utils::types::{Id, StatusCode, TestContext};
+use jsonrpsee_test_utils::mocks::{Id, StatusCode, TestContext};
 use jsonrpsee_test_utils::TimeoutFutureExt;
 use serde_json::Value as JsonValue;
 use tokio::task::JoinHandle;

@@ -15,6 +15,6 @@ hyper = { version = "0.14.10", features = ["full"] }
 log = "0.4"
 serde = { version = "1", default-features = false, features = ["derive"] }
 serde_json = "1"
-soketto = "0.6"
+soketto = { version = "0.7", features = ["http"] }
 tokio = { version = "1", features = ["net", "rt-multi-thread", "macros", "time"] }
-tokio-util = { version = "0.6", features = ["compat"] }
+tokio-util = { version = "0.6", features = ["compat"] }
@@ -24,7 +24,7 @@
 // IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
-use crate::types::{Body, HttpResponse, Id, Uri};
+use crate::mocks::{Body, HttpResponse, Id, Uri};
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Request, Response, Server};
 use serde_json::Value;

@@ -32,7 +32,7 @@ use std::{future::Future, time::Duration};
 use tokio::time::{timeout, Timeout};
 
 pub mod helpers;
-pub mod types;
+pub mod mocks;
 
 /// Helper extension trait which allows to limit execution time for the futures.
 /// It is helpful in tests to ensure that no future will ever get stuck forever.

@@ -34,10 +34,8 @@ use futures_util::{
 	stream::{self, StreamExt},
 };
 use serde::{Deserialize, Serialize};
-use soketto::handshake::{self, server::Response, Error as SokettoError, Server};
-use std::io;
-use std::net::SocketAddr;
-use std::time::Duration;
+use soketto::handshake::{self, http::is_upgrade_request, server::Response, Error as SokettoError, Server};
+use std::{io, net::SocketAddr, path::Path, time::Duration};
 use tokio::net::TcpStream;
 use tokio_util::compat::{Compat, TokioAsyncReadCompatExt};
 
@@ -314,3 +312,61 @@ async fn connection_task(socket: tokio::net::TcpStream, mode: ServerMode, mut ex
 		}
 	}
 }
+
+// Run a WebSocket server running on localhost that redirects requests for testing.
+// Requests to any url except for `/myblock/two` will redirect one or two times (HTTP 301) and eventually end up in `/myblock/two`.
+pub fn ws_server_with_redirect(other_server: String) -> String {
+	let addr = ([127, 0, 0, 1], 0).into();
+
+	let service = hyper::service::make_service_fn(move |_| {
+		let other_server = other_server.clone();
+		async move {
+			Ok::<_, hyper::Error>(hyper::service::service_fn(move |req| {
+				let other_server = other_server.clone();
+				async move { handler(req, other_server).await }
+			}))
+		}
+	});
+	let server = hyper::Server::bind(&addr).serve(service);
+	let addr = server.local_addr();
+
+	tokio::spawn(async move { server.await });
+	format!("ws://{}", addr)
+}
+
+/// Handle incoming HTTP Requests.
+async fn handler(
+	req: hyper::Request<Body>,
+	other_server: String,
+) -> Result<hyper::Response<Body>, soketto::BoxedError> {
+	if is_upgrade_request(&req) {
+		let path = req.uri().path();
+		log::error!("{:?}", path);
+
+		match path {
+			"/myblock/two" => {
+				let response = hyper::Response::builder()
+					.status(301)
+					.header("Location", other_server)
+					.body(Body::empty())
+					.unwrap();
+				Ok(response)
+			}
+			"/myblock/one" => {
+				let response =
+					hyper::Response::builder().status(301).header("Location", "two").body(Body::empty()).unwrap();
+				Ok(response)
+			}
+			_ => {
+				let response = hyper::Response::builder()
+					.status(301)
+					.header("Location", "/myblock/one")
+					.body(Body::empty())
+					.unwrap();
+				Ok(response)
+			}
+		}
+	} else {
+		panic!("expect upgrade to WS");
+	}
+}
diff --git a/types/Cargo.toml b/types/Cargo.toml
@@ -19,5 +19,5 @@ log = { version = "0.4", default-features = false }
 serde = { version = "1", default-features = false, features = ["derive"] }
 serde_json = { version = "1", default-features = false, features = ["alloc", "raw_value", "std"] }
 thiserror = "1.0"
-soketto = "0.6"
+soketto = "0.7"
 hyper = "0.14.10"
@@ -21,7 +21,7 @@ jsonrpsee-types = { path = "../types", version = "0.3.0" }
 log = "0.4"
 serde = "1"
 serde_json = "1"
-soketto = "0.6"
+soketto = "0.7"
 pin-project = "1"
 thiserror = "1"
 url = "2"
@@ -30,3 +30,4 @@ rustls-native-certs = "0.5.0"
 
 [dev-dependencies]
 jsonrpsee-test-utils = { path = "../test-utils" }
+env_logger = "0.9"
@@ -49,7 +49,9 @@ use futures::{
 use tokio::sync::Mutex;
 
 use serde::de::DeserializeOwned;
-use std::{borrow::Cow, time::Duration};
+use std::time::Duration;
+
+pub use soketto::handshake::client::Header;
 
 /// Wrapper over a [`oneshot::Receiver`](futures::channel::oneshot::Receiver) that reads
 /// the underlying channel once and then stores the result in String.
@@ -106,9 +108,10 @@ pub struct WsClientBuilder<'a> {
 	max_request_body_size: u32,
 	request_timeout: Duration,
 	connection_timeout: Duration,
-	origin_header: Option<Cow<'a, str>>,
+	origin_header: Option<Header<'a>>,
 	max_concurrent_requests: usize,
 	max_notifs_per_subscription: usize,
+	max_redirections: usize,
 }
 
 impl<'a> Default for WsClientBuilder<'a> {
@@ -121,6 +124,7 @@ impl<'a> Default for WsClientBuilder<'a> {
 			origin_header: None,
 			max_concurrent_requests: 256,
 			max_notifs_per_subscription: 1024,
+			max_redirections: 10,
 		}
 	}
 }
@@ -151,8 +155,8 @@ impl<'a> WsClientBuilder<'a> {
 	}
 
 	/// Set origin header to pass during the handshake.
-	pub fn origin_header(mut self, origin: &'a str) -> Self {
-		self.origin_header = Some(Cow::Borrowed(origin));
+	pub fn origin_header(mut self, origin: Header<'a>) -> Self {
+		self.origin_header = Some(origin);
 		self
 	}
 
@@ -176,6 +180,12 @@ impl<'a> WsClientBuilder<'a> {
 		self
 	}
 
+	/// Set the max number of redirections to perform until a connection is regarded as failed.
+	pub fn max_redirections(mut self, redirect: usize) -> Self {
+		self.max_redirections = redirect;
+		self
+	}
+
 	/// Build the client with specified URL to connect to.
 	/// If the port number is missing from the URL, the default port number is used.
 	///
@@ -201,6 +211,7 @@ impl<'a> WsClientBuilder<'a> {
 			timeout: self.connection_timeout,
 			origin_header: self.origin_header,
 			max_request_body_size: self.max_request_body_size,
+			max_redirections: self.max_redirections,
 		};
 
 		let (sender, receiver) = builder.build().await.map_err(|e| Error::Transport(e.into()))?;

@@ -32,7 +32,7 @@ use crate::types::{
 };
 use crate::WsClientBuilder;
 use jsonrpsee_test_utils::helpers::*;
-use jsonrpsee_test_utils::types::{Id, WebSocketTestServer};
+use jsonrpsee_test_utils::mocks::{Id, WebSocketTestServer};
 use jsonrpsee_test_utils::TimeoutFutureExt;
 use serde_json::Value as JsonValue;
 
@@ -263,3 +263,35 @@ fn assert_error_response(err: Error, exp: ErrorObject) {
 		e => panic!("Expected error: \"{}\", got: {:?}", err, e),
 	};
 }
+
+//#[cfg_attr(target_os = "windows", ignore)]
+#[tokio::test]
+async fn redirections() {
+	env_logger::try_init();
+	let expected = "abc 123";
+	let server = WebSocketTestServer::with_hardcoded_response(
+		"127.0.0.1:0".parse().unwrap(),
+		ok_response(expected.into(), Id::Num(0)),
+	)
+	.with_default_timeout()
+	.await
+	.unwrap();
+
+	let server_url = format!("ws://{}", server.local_addr());
+	let redirect_url = jsonrpsee_test_utils::mocks::ws_server_with_redirect(server_url);
+
+	// The client will first connect to a server that only performs re-directions and finally
+	// redirect to another server to complete the handshake.
+	let client = WsClientBuilder::default().build(&redirect_url).with_default_timeout().await;
+	// It's an ok client
+	let client = match client {
+		Ok(Ok(client)) => client,
+		Ok(Err(e)) => panic!("WsClient builder failed with: {:?}", e),
+		Err(e) => panic!("WsClient builder timed out with: {:?}", e),
+	};
+	// It's connected
+	assert!(client.is_connected());
+	// It works
+	let response = client.request::<String>("anything", ParamsSer::NoParams).with_default_timeout().await.unwrap();
+	assert_eq!(response.unwrap(), String::from(expected));
+}