Skip to content

Commit ae244b1

Browse files
LarkoooLarkooo
authored
fix(proxy): use h1 for websocket upgrade (#184)
* fix(proxy): use h1 for websocket upgrade * fmt * support upgrades over tls
1 parent 7c7f144 commit ae244b1

4 files changed

Lines changed: 35 additions & 9 deletions

File tree

‎crates/server/src/handlers/graphql.rs‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,19 @@ impl Handler for GraphQLHandler {
2828
async fn handle(&self, req: Request<Body>, client_addr: IpAddr) -> Response<Body> {
2929
if let Some(addr) = self.graphql_addr {
3030
let graphql_addr = format!("http://{}", addr);
31-
match crate::proxy::GRAPHQL_PROXY_CLIENT
32-
.call(client_addr, &graphql_addr, req)
33-
.await
34-
{
31+
32+
// Use WebSocket-compatible client for WebSocket upgrade requests
33+
let result = if crate::proxy::is_websocket_upgrade(&req) {
34+
crate::proxy::WEBSOCKET_PROXY_CLIENT
35+
.call(client_addr, &graphql_addr, req)
36+
.await
37+
} else {
38+
crate::proxy::PROXY_CLIENT
39+
.call(client_addr, &graphql_addr, req)
40+
.await
41+
};
42+
43+
match result {
3544
Ok(response) => response,
3645
Err(_error) => {
3746
error!(target: LOG_TARGET, "GraphQL proxy error: {:?}", _error);

‎crates/server/src/handlers/grpc.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ impl Handler for GrpcHandler {
3232
async fn handle(&self, req: Request<Body>, client_addr: IpAddr) -> Response<Body> {
3333
if let Some(grpc_addr) = self.grpc_addr {
3434
let grpc_addr = format!("http://{}", grpc_addr);
35-
match crate::proxy::GRPC_PROXY_CLIENT
35+
match crate::proxy::PROXY_CLIENT
3636
.call(client_addr, &grpc_addr, req)
3737
.await
3838
{

‎crates/server/src/handlers/static_files.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ impl Handler for StaticHandler {
2727
async fn handle(&self, req: Request<Body>, client_addr: IpAddr) -> Response<Body> {
2828
if let Some(artifacts_addr) = self.artifacts_addr {
2929
let artifacts_addr = format!("http://{}", artifacts_addr);
30-
match crate::proxy::GRAPHQL_PROXY_CLIENT
30+
match crate::proxy::PROXY_CLIENT
3131
.call(client_addr, &artifacts_addr, req)
3232
.await
3333
{

‎crates/server/src/proxy.rs‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,15 +56,15 @@ const DEFAULT_EXPOSED_HEADERS: [&str; 4] = [
5656
const DEFAULT_MAX_AGE: Duration = Duration::from_secs(24 * 60 * 60);
5757

5858
lazy_static::lazy_static! {
59-
pub(crate) static ref GRAPHQL_PROXY_CLIENT: ReverseProxy<HttpConnector<GaiResolver>> = {
59+
pub(crate) static ref WEBSOCKET_PROXY_CLIENT: ReverseProxy<HttpConnector<GaiResolver>> = {
6060
ReverseProxy::new(
6161
Client::builder()
62-
.http2_only(true)
6362
.build_http(),
6463
)
6564
};
6665

67-
pub(crate) static ref GRPC_PROXY_CLIENT: ReverseProxy<HttpConnector<GaiResolver>> = {
66+
67+
pub(crate) static ref PROXY_CLIENT: ReverseProxy<HttpConnector<GaiResolver>> = {
6868
ReverseProxy::new(
6969
Client::builder()
7070
.http2_only(true)
@@ -73,6 +73,21 @@ lazy_static::lazy_static! {
7373
};
7474
}
7575

76+
// Helper function to check if a request is a WebSocket upgrade request
77+
pub fn is_websocket_upgrade(req: &Request<Body>) -> bool {
78+
req.headers()
79+
.get("upgrade")
80+
.and_then(|v| v.to_str().ok())
81+
.map(|v| v.to_lowercase() == "websocket")
82+
.unwrap_or(false)
83+
&& req
84+
.headers()
85+
.get("connection")
86+
.and_then(|v| v.to_str().ok())
87+
.map(|v| v.to_lowercase().contains("upgrade"))
88+
.unwrap_or(false)
89+
}
90+
7691
#[derive(Debug)]
7792
pub struct Proxy {
7893
addr: SocketAddr,
@@ -238,8 +253,10 @@ impl Proxy {
238253
}
239254
});
240255

256+
241257
if let Err(e) = hyper::server::conn::Http::new()
242258
.serve_connection(tls_stream, service)
259+
.with_upgrades() // Enable connection upgrades for WebSocket over TLS
243260
.await
244261
{
245262
error!(target: LOG_TARGET, error = ?e, "Serving connection.");

0 commit comments

Comments
 (0)