micropelago/hyper-reverse-proxy
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

tests: ws

Browse Source
This commit is contained in:
chesedo 2022-05-04 10:25:11 +02:00 committed by Felipe Noronha
parent 96a398de85
commit 20dbf00931
2 changed files with 127 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
Cargo.toml
View File

@ -28,6 +28,7 @@ hyper = { version = "0.14.18", features = ["server"] }
tokio = { version = "1.17.0", features = ["full"] } tokio = { version = "1.17.0", features = ["full"] }
futures = "0.3.21" futures = "0.3.21"
async-trait = "0.1.53" async-trait = "0.1.53"
async-tungstenite = { version = "0.17", features = ["tokio-runtime"] }
tokio-test = "0.4.2" tokio-test = "0.4.2"
test-context = "0.1.3" test-context = "0.1.3"
tokiotest-httpserver = "0.2.1" tokiotest-httpserver = "0.2.1"
@ -38,3 +39,5 @@ hyper-trust-dns = { version = "0.4.2", features = [
"rustls-webpki" "rustls-webpki"
] } ] }
rand = "0.8.5" rand = "0.8.5"
tungstenite = "0.17"
url = "2.2"

124
tests/test_websocket.rs Normal file
View File

@ -0,0 +1,124 @@
use std::{
convert::Infallible,
net::{IpAddr, SocketAddr},
process::exit,
time::Duration,
};
use async_tungstenite::tokio::{accept_async, connect_async};
use futures::{SinkExt, StreamExt};
use hyper::{
client::{connect::dns::GaiResolver, HttpConnector},
server::conn::AddrStream,
service::{make_service_fn, service_fn},
Body, Request, Response, Server,
};
use hyper_reverse_proxy::ReverseProxy;
use test_context::{test_context, AsyncTestContext};
use tokio::{net::TcpListener, sync::oneshot::Sender, task::JoinHandle};
use tokiotest_httpserver::take_port;
use tungstenite::Message;
use url::Url;
lazy_static::lazy_static! {
static ref PROXY_CLIENT: ReverseProxy<HttpConnector<GaiResolver>> = {
ReverseProxy::new(
hyper::Client::new(),
)
};
}
struct ProxyTestContext {
sender: Sender<()>,
proxy_handler: JoinHandle<Result<(), hyper::Error>>,
ws_handler: JoinHandle<()>,
port: u16,
}
#[test_context(ProxyTestContext)]
#[tokio::test]
async fn test_websocket(ctx: &mut ProxyTestContext) {
println!("making client connection");
let (mut client, _) =
connect_async(Url::parse(&format!("ws://127.0.0.1:{}", ctx.port)).unwrap())
.await
.unwrap();
println!("made client connection");
client.send(Message::Ping("ping".into())).await.unwrap();
let msg = client.next().await.unwrap().unwrap();
assert!(matches!(msg, Message::Pong(inner) if inner == "pong".as_bytes()));
}
async fn handle(
client_ip: IpAddr,
req: Request<Body>,
backend_port: u16,
) -> Result<Response<Body>, Infallible> {
match PROXY_CLIENT
.call(
client_ip,
format!("http://127.0.0.1:{}", backend_port).as_str(),
req,
)
.await
{
Ok(response) => Ok(response),
Err(err) => panic!("did not expect error: {:?}", err),
}
}
#[async_trait::async_trait]
impl<'a> AsyncTestContext for ProxyTestContext {
async fn setup() -> ProxyTestContext {
tokio::spawn(async {
tokio::time::sleep(Duration::from_secs(5)).await;
println!("Unit test executed too long, perhaps its stuck...");
exit(1);
});
let (sender, receiver) = tokio::sync::oneshot::channel::<()>();
let ws_port = take_port();
let ws_handler = tokio::spawn(async move {
let ws_server = TcpListener::bind(("127.0.0.1", ws_port)).await.unwrap();
while let Ok((stream, addr)) = ws_server.accept().await {
println!("incoming connection: {addr}");
let mut websocket = accept_async(stream).await.unwrap();
let msg = websocket.next().await.unwrap().unwrap();
assert!(matches!(msg, Message::Ping(inner) if inner == "ping".as_bytes()));
println!("past ping");
websocket.send(Message::Pong("pong".into())).await.unwrap();
println!("past pong");
}
});
let make_svc = make_service_fn(move |conn: &AddrStream| {
let remote_addr = conn.remote_addr().ip();
async move { Ok::<_, Infallible>(service_fn(move |req| handle(remote_addr, req, ws_port))) }
});
let port = take_port();
let addr = SocketAddr::new("127.0.0.1".parse().unwrap(), port);
let server = Server::bind(&addr)
.serve(make_svc)
.with_graceful_shutdown(async {
receiver.await.ok();
});
let proxy_handler = tokio::spawn(server);
ProxyTestContext {
sender,
proxy_handler,
ws_handler,
port,
}
}
async fn teardown(self) {
let _ = self.sender.send(()).unwrap();
let _ = tokio::join!(self.proxy_handler, self.ws_handler);
}
}