2022-05-01 18:36:09 +00:00
|
|
|
use hyper::client::connect::dns::GaiResolver;
|
|
|
|
|
use hyper::client::HttpConnector;
|
2022-06-18 22:46:50 +00:00
|
|
|
use hyper::header::{CONNECTION, HOST, UPGRADE};
|
2022-04-07 12:31:27 +00:00
|
|
|
use hyper::server::conn::AddrStream;
|
2022-04-12 14:41:51 +00:00
|
|
|
use hyper::service::{make_service_fn, service_fn};
|
2022-06-18 22:46:50 +00:00
|
|
|
use hyper::{Body, Client, HeaderMap, Request, Response, Server, StatusCode, Uri};
|
2022-05-01 18:36:09 +00:00
|
|
|
use hyper_reverse_proxy::ReverseProxy;
|
2022-04-07 12:31:27 +00:00
|
|
|
use std::convert::Infallible;
|
|
|
|
|
use std::net::{IpAddr, SocketAddr};
|
2022-04-12 14:41:51 +00:00
|
|
|
use test_context::test_context;
|
|
|
|
|
use test_context::AsyncTestContext;
|
|
|
|
|
use tokio::sync::oneshot::Sender;
|
|
|
|
|
use tokio::task::JoinHandle;
|
2022-04-07 12:31:27 +00:00
|
|
|
use tokiotest_httpserver::handler::HandlerBuilder;
|
2022-04-12 14:41:51 +00:00
|
|
|
use tokiotest_httpserver::{take_port, HttpTestContext};
|
2022-04-07 12:31:27 +00:00
|
|
|
|
2022-05-01 18:36:09 +00:00
|
|
|
lazy_static::lazy_static! {
|
|
|
|
|
static ref PROXY_CLIENT: ReverseProxy<HttpConnector<GaiResolver>> = {
|
|
|
|
|
ReverseProxy::new(
|
|
|
|
|
hyper::Client::new(),
|
|
|
|
|
)
|
|
|
|
|
};
|
|
|
|
|
}
|
|
|
|
|
|
2022-04-07 12:31:27 +00:00
|
|
|
struct ProxyTestContext {
|
|
|
|
|
sender: Sender<()>,
|
|
|
|
|
proxy_handler: JoinHandle<Result<(), hyper::Error>>,
|
|
|
|
|
http_back: HttpTestContext,
|
2022-04-12 14:41:51 +00:00
|
|
|
port: u16,
|
2022-04-07 12:31:27 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test_context(ProxyTestContext)]
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn test_get_error_500(ctx: &mut ProxyTestContext) {
|
2022-04-12 14:41:51 +00:00
|
|
|
let client = Client::new();
|
|
|
|
|
let resp = client
|
|
|
|
|
.request(
|
|
|
|
|
Request::builder()
|
|
|
|
|
.header("keep-alive", "treu")
|
|
|
|
|
.method("GET")
|
|
|
|
|
.uri(ctx.uri("/500"))
|
|
|
|
|
.body(Body::from(""))
|
|
|
|
|
.unwrap(),
|
|
|
|
|
)
|
|
|
|
|
.await
|
|
|
|
|
.unwrap();
|
2022-04-07 12:31:27 +00:00
|
|
|
assert_eq!(500, resp.status());
|
|
|
|
|
}
|
|
|
|
|
|
2022-05-10 07:02:21 +00:00
|
|
|
#[test_context(ProxyTestContext)]
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn test_upgrade_mismatch(ctx: &mut ProxyTestContext) {
|
|
|
|
|
ctx.http_back.add(
|
2022-05-10 07:05:33 +00:00
|
|
|
HandlerBuilder::new("/ws")
|
2022-05-10 07:02:21 +00:00
|
|
|
.status_code(StatusCode::SWITCHING_PROTOCOLS)
|
|
|
|
|
.build(),
|
|
|
|
|
);
|
|
|
|
|
let resp = Client::new()
|
|
|
|
|
.request(
|
|
|
|
|
Request::builder()
|
|
|
|
|
.header(CONNECTION, "Upgrade")
|
|
|
|
|
.header(UPGRADE, "websocket")
|
|
|
|
|
.method("GET")
|
2022-05-10 07:05:33 +00:00
|
|
|
.uri(ctx.uri("/ws"))
|
2022-05-10 07:02:21 +00:00
|
|
|
.body(Body::from(""))
|
|
|
|
|
.unwrap(),
|
|
|
|
|
)
|
|
|
|
|
.await
|
|
|
|
|
.unwrap();
|
|
|
|
|
assert_eq!(resp.status(), 502);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test_context(ProxyTestContext)]
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn test_upgrade_unrequested(ctx: &mut ProxyTestContext) {
|
|
|
|
|
ctx.http_back.add(
|
2022-05-10 07:05:33 +00:00
|
|
|
HandlerBuilder::new("/wrong_switch")
|
2022-05-10 07:02:21 +00:00
|
|
|
.status_code(StatusCode::SWITCHING_PROTOCOLS)
|
|
|
|
|
.build(),
|
|
|
|
|
);
|
2022-05-10 07:05:33 +00:00
|
|
|
let resp = Client::new().get(ctx.uri("/wrong_switch")).await.unwrap();
|
2022-05-10 07:02:21 +00:00
|
|
|
assert_eq!(resp.status(), 502);
|
|
|
|
|
}
|
|
|
|
|
|
2022-04-07 12:31:27 +00:00
|
|
|
#[test_context(ProxyTestContext)]
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn test_get(ctx: &mut ProxyTestContext) {
|
2022-06-18 22:46:50 +00:00
|
|
|
let mut headers = HeaderMap::new();
|
|
|
|
|
headers.insert(
|
|
|
|
|
HOST,
|
|
|
|
|
format!("127.0.0.1:{}", ctx.http_back.port).parse().unwrap(),
|
|
|
|
|
);
|
2022-06-18 23:20:59 +00:00
|
|
|
|
|
|
|
|
ctx.http_back.add(
|
|
|
|
|
HandlerBuilder::new("/foo")
|
|
|
|
|
.status_code(StatusCode::OK)
|
|
|
|
|
.headers(headers)
|
|
|
|
|
.build(),
|
|
|
|
|
);
|
|
|
|
|
let resp = Client::new().get(ctx.uri("/foo")).await.unwrap();
|
|
|
|
|
assert_eq!(200, resp.status());
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[test_context(ProxyTestContext)]
|
|
|
|
|
#[tokio::test]
|
|
|
|
|
async fn test_headers(ctx: &mut ProxyTestContext) {
|
|
|
|
|
let mut headers = HeaderMap::new();
|
|
|
|
|
headers.insert("x-forwarded-for", "127.0.0.1".parse().unwrap());
|
|
|
|
|
headers.insert(
|
|
|
|
|
"x-forwarded-host",
|
|
|
|
|
format!("localhost:{}", ctx.port).parse().unwrap(),
|
|
|
|
|
);
|
|
|
|
|
|
2022-04-12 14:41:51 +00:00
|
|
|
ctx.http_back.add(
|
|
|
|
|
HandlerBuilder::new("/foo")
|
|
|
|
|
.status_code(StatusCode::OK)
|
2022-06-18 22:46:50 +00:00
|
|
|
.headers(headers)
|
2022-04-12 14:41:51 +00:00
|
|
|
.build(),
|
|
|
|
|
);
|
2022-04-07 12:31:27 +00:00
|
|
|
let resp = Client::new().get(ctx.uri("/foo")).await.unwrap();
|
|
|
|
|
assert_eq!(200, resp.status());
|
|
|
|
|
}
|
|
|
|
|
|
2022-04-12 14:41:51 +00:00
|
|
|
async fn handle(
|
|
|
|
|
client_ip: IpAddr,
|
|
|
|
|
req: Request<Body>,
|
|
|
|
|
backend_port: u16,
|
|
|
|
|
) -> Result<Response<Body>, Infallible> {
|
2022-05-01 18:36:09 +00:00
|
|
|
match PROXY_CLIENT
|
|
|
|
|
.call(
|
|
|
|
|
client_ip,
|
|
|
|
|
format!("http://127.0.0.1:{}", backend_port).as_str(),
|
|
|
|
|
req,
|
|
|
|
|
)
|
|
|
|
|
.await
|
2022-04-12 14:41:51 +00:00
|
|
|
{
|
|
|
|
|
Ok(response) => Ok(response),
|
|
|
|
|
Err(_) => Ok(Response::builder().status(502).body(Body::empty()).unwrap()),
|
2022-04-07 12:31:27 +00:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
#[async_trait::async_trait]
|
2022-05-01 18:36:09 +00:00
|
|
|
impl<'a> AsyncTestContext for ProxyTestContext {
|
2022-04-07 12:31:27 +00:00
|
|
|
async fn setup() -> ProxyTestContext {
|
|
|
|
|
let http_back: HttpTestContext = AsyncTestContext::setup().await;
|
|
|
|
|
let (sender, receiver) = tokio::sync::oneshot::channel::<()>();
|
|
|
|
|
let bp_to_move = http_back.port;
|
2022-05-01 18:36:09 +00:00
|
|
|
|
2022-04-07 12:31:27 +00:00
|
|
|
let make_svc = make_service_fn(move |conn: &AddrStream| {
|
|
|
|
|
let remote_addr = conn.remote_addr().ip();
|
|
|
|
|
let back_port = bp_to_move;
|
|
|
|
|
async move {
|
|
|
|
|
Ok::<_, Infallible>(service_fn(move |req| handle(remote_addr, req, back_port)))
|
|
|
|
|
}
|
|
|
|
|
});
|
|
|
|
|
let port = take_port();
|
|
|
|
|
let addr = SocketAddr::new("127.0.0.1".parse().unwrap(), port);
|
2022-04-12 14:41:51 +00:00
|
|
|
let server = Server::bind(&addr)
|
|
|
|
|
.serve(make_svc)
|
|
|
|
|
.with_graceful_shutdown(async {
|
|
|
|
|
receiver.await.ok();
|
|
|
|
|
});
|
2022-04-07 12:31:27 +00:00
|
|
|
let proxy_handler = tokio::spawn(server);
|
|
|
|
|
ProxyTestContext {
|
|
|
|
|
sender,
|
|
|
|
|
proxy_handler,
|
|
|
|
|
http_back,
|
2022-04-12 14:41:51 +00:00
|
|
|
port,
|
2022-04-07 12:31:27 +00:00
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
async fn teardown(self) {
|
|
|
|
|
let _ = AsyncTestContext::teardown(self.http_back);
|
|
|
|
|
let _ = self.sender.send(()).unwrap();
|
|
|
|
|
let _ = tokio::join!(self.proxy_handler);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
impl ProxyTestContext {
|
|
|
|
|
pub fn uri(&self, path: &str) -> Uri {
|
2022-04-12 14:41:51 +00:00
|
|
|
format!("http://{}:{}{}", "localhost", self.port, path)
|
|
|
|
|
.parse::<Uri>()
|
|
|
|
|
.unwrap()
|
2022-04-07 12:31:27 +00:00
|
|
|
}
|
2022-04-12 14:41:51 +00:00
|
|
|
}
|
