refacto: simplify server request handling

src/tunnel/server/handler_http2.rs

@@ -1,25 +1,20 @@
 use crate::restrictions::types::RestrictionsRules;
-use crate::tunnel::server::utils::{
+use crate::tunnel::server::utils::{bad_request, inject_cookie};
-    extract_path_prefix, extract_tunnel_info, extract_x_forwarded_for, inject_cookie, validate_tunnel,
-};
 use crate::tunnel::server::WsServer;
+use crate::tunnel::transport;
 use crate::tunnel::transport::http2::{Http2TunnelRead, Http2TunnelWrite};
-use crate::tunnel::{transport, RemoteAddr};
-use crate::LocalProtocol;
 use bytes::Bytes;
 use futures_util::StreamExt;
 use http_body_util::combinators::BoxBody;
 use http_body_util::{BodyStream, Either, StreamBody};
 use hyper::body::{Frame, Incoming};
 use hyper::header::CONTENT_TYPE;
-use hyper::{http, Request, Response, StatusCode};
+use hyper::{Request, Response, StatusCode};
 use std::net::SocketAddr;
-use std::pin::Pin;
 use std::sync::Arc;
-use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::sync::{mpsc, oneshot};
 use tokio_stream::wrappers::ReceiverStream;
-use tracing::{info, warn, Instrument, Span};
+use tracing::{Instrument, Span};
 
 pub(super) async fn http_server_upgrade(
     server: WsServer,
@@ -28,11 +23,13 @@ pub(super) async fn http_server_upgrade(
     client_addr: SocketAddr,
     mut req: Request<Incoming>,
 ) -> Response<Either<String, BoxBody<Bytes, anyhow::Error>>> {
-    let (remote_addr, local_rx, local_tx, need_cookie) =
+    let (remote_addr, local_rx, local_tx, need_cookie) = match server
-        match exec_tunnel_request(server, restrictions, restrict_path_prefix, client_addr, &req).await {
+        .handle_tunnel_request(restrictions, restrict_path_prefix, client_addr, &req)
-            Ok(ret) => ret,
+        .await
-            Err(err) => return err,
+    {
-        };
+        Ok(ret) => ret,
+        Err(err) => return err,
+    };
 
     let req_content_type = req.headers_mut().remove(CONTENT_TYPE);
     let ws_rx = BodyStream::new(req.into_body());
@@ -60,10 +57,8 @@ pub(super) async fn http_server_upgrade(
         .instrument(Span::current()),
     );
 
-    if need_cookie {
+    if need_cookie && inject_cookie(&mut response, &remote_addr).is_err() {
-        if let Err(response) = inject_cookie(&mut response, &remote_addr, Either::Left) {
+        return bad_request();
-            return response;
-        }
     }
 
     if let Some(content_type) = req_content_type {
@@ -72,93 +67,3 @@ pub(super) async fn http_server_upgrade(
 
     response
 }
-
-pub(super) async fn exec_tunnel_request(
-    server: WsServer,
-    restrictions: Arc<RestrictionsRules>,
-    restrict_path_prefix: Option<String>,
-    mut client_addr: SocketAddr,
-    req: &Request<Incoming>,
-) -> Result<
-    (
-        RemoteAddr,
-        Pin<Box<dyn AsyncRead + Send>>,
-        Pin<Box<dyn AsyncWrite + Send>>,
-        bool,
-    ),
-    Response<Either<String, BoxBody<Bytes, anyhow::Error>>>,
-> {
-    match extract_x_forwarded_for(req) {
-        Ok(Some((x_forward_for, x_forward_for_str))) => {
-            info!("Request X-Forwarded-For: {:?}", x_forward_for);
-            Span::current().record("forwarded_for", x_forward_for_str);
-            client_addr.set_ip(x_forward_for);
-        }
-        Ok(_) => {}
-        Err(err) => return Err(err.map(Either::Left)),
-    };
-
-    let path_prefix = match extract_path_prefix(req) {
-        Ok(p) => p,
-        Err(err) => return Err(err.map(Either::Left)),
-    };
-
-    if let Some(restrict_path) = restrict_path_prefix {
-        if path_prefix != restrict_path {
-            warn!(
-                "Client requested upgrade path '{}' does not match upgrade path restriction '{}' (mTLS, etc.)",
-                path_prefix, restrict_path
-            );
-            return Err(http::Response::builder()
-                .status(StatusCode::BAD_REQUEST)
-                .body(Either::Left("Invalid upgrade request".to_string()))
-                .unwrap());
-        }
-    }
-
-    let jwt = match extract_tunnel_info(req) {
-        Ok(jwt) => jwt,
-        Err(err) => return Err(err.map(Either::Left)),
-    };
-
-    Span::current().record("id", &jwt.claims.id);
-    Span::current().record("remote", format!("{}:{}", jwt.claims.r, jwt.claims.rp));
-    let remote = match RemoteAddr::try_from(jwt.claims) {
-        Ok(remote) => remote,
-        Err(err) => {
-            warn!("Rejecting connection with bad tunnel info: {} {}", err, req.uri());
-            return Err(http::Response::builder()
-                .status(StatusCode::BAD_REQUEST)
-                .body(Either::Left("Invalid upgrade request".to_string()))
-                .unwrap());
-        }
-    };
-
-    let restriction = match validate_tunnel(&remote, path_prefix, &restrictions) {
-        Ok(matched_restriction) => {
-            info!("Tunnel accepted due to matched restriction: {}", matched_restriction.name);
-            matched_restriction
-        }
-        Err(err) => return Err(err.map(Either::Left)),
-    };
-
-    let req_protocol = remote.protocol.clone();
-    let inject_cookie = matches!(
-        req_protocol,
-        LocalProtocol::ReverseSocks5 { .. } | LocalProtocol::ReverseHttpProxy { .. }
-    );
-    let tunnel = match server.run_tunnel(restriction, remote, client_addr).await {
-        Ok(ret) => ret,
-        Err(err) => {
-            warn!("Rejecting connection with bad upgrade request: {} {}", err, req.uri());
-            return Err(http::Response::builder()
-                .status(StatusCode::BAD_REQUEST)
-                .body(Either::Left("Invalid upgrade request".to_string()))
-                .unwrap());
-        }
-    };
-
-    let (remote_addr, local_rx, local_tx) = tunnel;
-    info!("connected to {:?} {}:{}", req_protocol, remote_addr.host, remote_addr.port);
-    Ok((remote_addr, local_rx, local_tx, inject_cookie))
-}

src/tunnel/server/handler_websocket.rs

@@ -1,6 +1,5 @@
 use crate::restrictions::types::RestrictionsRules;
-use crate::tunnel::server::handler_http2::exec_tunnel_request;
+use crate::tunnel::server::utils::{bad_request, inject_cookie};
-use crate::tunnel::server::utils::inject_cookie;
 use crate::tunnel::server::WsServer;
 use crate::tunnel::transport;
 use crate::tunnel::transport::websocket::{WebsocketTunnelRead, WebsocketTunnelWrite};
@@ -9,7 +8,7 @@ use http_body_util::combinators::BoxBody;
 use http_body_util::Either;
 use hyper::body::Incoming;
 use hyper::header::{HeaderValue, SEC_WEBSOCKET_PROTOCOL};
-use hyper::{http, Request, Response, StatusCode};
+use hyper::{Request, Response};
 use std::net::SocketAddr;
 use std::sync::Arc;
 use tokio::sync::oneshot;
@@ -24,27 +23,23 @@ pub(super) async fn ws_server_upgrade(
 ) -> Response<Either<String, BoxBody<Bytes, anyhow::Error>>> {
     if !fastwebsockets::upgrade::is_upgrade_request(&req) {
         warn!("Rejecting connection with bad upgrade request: {}", req.uri());
-        return http::Response::builder()
+        return bad_request();
-            .status(StatusCode::BAD_REQUEST)
-            .body(Either::Left("Invalid upgrade request".to_string()))
-            .unwrap();
     }
 
     let mask_frame = server.config.websocket_mask_frame;
-    let (remote_addr, local_rx, local_tx, need_cookie) =
+    let (remote_addr, local_rx, local_tx, need_cookie) = match server
-        match exec_tunnel_request(server, restrictions, restrict_path_prefix, client_addr, &req).await {
+        .handle_tunnel_request(restrictions, restrict_path_prefix, client_addr, &req)
-            Ok(ret) => ret,
+        .await
-            Err(err) => return err,
+    {
-        };
+        Ok(ret) => ret,
+        Err(err) => return err,
+    };
 
     let (response, fut) = match fastwebsockets::upgrade::upgrade(&mut req) {
         Ok(ret) => ret,
         Err(err) => {
             warn!("Rejecting connection with bad upgrade request: {} {}", err, req.uri());
-            return http::Response::builder()
+            return bad_request();
-                .status(StatusCode::BAD_REQUEST)
-                .body(Either::Left(format!("Invalid upgrade request: {:?}", err)))
-                .unwrap();
         }
     };
 
@@ -73,10 +68,8 @@ pub(super) async fn ws_server_upgrade(
     );
 
     let mut response = Response::from_parts(response.into_parts().0, Either::Right(BoxBody::default()));
-    if need_cookie {
+    if need_cookie && inject_cookie(&mut response, &remote_addr).is_err() {
-        if let Err(response) = inject_cookie(&mut response, &remote_addr, Either::Left) {
+        return bad_request();
-            return response;
-        }
     }
 
     response

src/tunnel/server/server.rs

@@ -6,6 +6,8 @@ use std::fmt;
 use std::fmt::{Debug, Formatter};
 use std::future::Future;
 
+use bytes::Bytes;
+use http_body_util::combinators::BoxBody;
 use std::net::SocketAddr;
 use std::ops::Deref;
 use std::path::PathBuf;
@@ -18,7 +20,7 @@ use crate::{protocols, LocalProtocol};
 use hyper::body::Incoming;
 use hyper::server::conn::{http1, http2};
 use hyper::service::service_fn;
-use hyper::{http, Request, StatusCode, Version};
+use hyper::{http, Request, Response, StatusCode, Version};
 use hyper_util::rt::TokioExecutor;
 use once_cell::sync::Lazy;
 use parking_lot::Mutex;
@@ -35,7 +37,9 @@ use crate::tunnel::listeners::{
 };
 use crate::tunnel::server::handler_http2::http_server_upgrade;
 use crate::tunnel::server::handler_websocket::ws_server_upgrade;
-use crate::tunnel::server::utils::find_mapped_port;
+use crate::tunnel::server::utils::{
+    bad_request, extract_path_prefix, extract_tunnel_info, extract_x_forwarded_for, find_mapped_port, validate_tunnel,
+};
 use crate::tunnel::tls_reloader::TlsReloader;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio::net::TcpListener;
@@ -79,7 +83,88 @@ impl WsServer {
         }
     }
 
-    pub(super) async fn run_tunnel(
+    pub(super) async fn handle_tunnel_request(
+        &self,
+        restrictions: Arc<RestrictionsRules>,
+        restrict_path_prefix: Option<String>,
+        mut client_addr: SocketAddr,
+        req: &Request<Incoming>,
+    ) -> Result<
+        (
+            RemoteAddr,
+            Pin<Box<dyn AsyncRead + Send>>,
+            Pin<Box<dyn AsyncWrite + Send>>,
+            bool,
+        ),
+        Response<Either<String, BoxBody<Bytes, anyhow::Error>>>,
+    > {
+        match extract_x_forwarded_for(req) {
+            Ok(Some((x_forward_for, x_forward_for_str))) => {
+                info!("Request X-Forwarded-For: {:?}", x_forward_for);
+                Span::current().record("forwarded_for", x_forward_for_str);
+                client_addr.set_ip(x_forward_for);
+            }
+            Ok(_) => {}
+            Err(_err) => return Err(bad_request()),
+        };
+
+        let path_prefix = match extract_path_prefix(req) {
+            Ok(p) => p,
+            Err(_err) => return Err(bad_request()),
+        };
+
+        if let Some(restrict_path) = restrict_path_prefix {
+            if path_prefix != restrict_path {
+                warn!(
+                    "Client requested upgrade path '{}' does not match upgrade path restriction '{}' (mTLS, etc.)",
+                    path_prefix, restrict_path
+                );
+                return Err(bad_request());
+            }
+        }
+
+        let jwt = match extract_tunnel_info(req) {
+            Ok(jwt) => jwt,
+            Err(_err) => return Err(bad_request()),
+        };
+
+        Span::current().record("id", &jwt.claims.id);
+        Span::current().record("remote", format!("{}:{}", jwt.claims.r, jwt.claims.rp));
+        let remote = match RemoteAddr::try_from(jwt.claims) {
+            Ok(remote) => remote,
+            Err(err) => {
+                warn!("Rejecting connection with bad tunnel info: {} {}", err, req.uri());
+                return Err(bad_request());
+            }
+        };
+
+        let restriction = match validate_tunnel(&remote, path_prefix, &restrictions) {
+            Ok(matched_restriction) => {
+                info!("Tunnel accepted due to matched restriction: {}", matched_restriction.name);
+                matched_restriction
+            }
+            Err(_err) => return Err(bad_request()),
+        };
+
+        let req_protocol = remote.protocol.clone();
+        let inject_cookie = matches!(
+            req_protocol,
+            LocalProtocol::ReverseSocks5 { .. } | LocalProtocol::ReverseHttpProxy { .. }
+        );
+        let tunnel = match self.exec_tunnel(restriction, remote, client_addr).await {
+            Ok(ret) => ret,
+            Err(err) => {
+                warn!("Rejecting connection with bad upgrade request: {} {}", err, req.uri());
+                return Err(bad_request());
+            }
+        };
+
+        let (remote_addr, local_rx, local_tx) = tunnel;
+        info!("connected to {:?} {}:{}", req_protocol, remote_addr.host, remote_addr.port);
+        Ok((remote_addr, local_rx, local_tx, inject_cookie))
+    }
+
+    async fn exec_tunnel(
         &self,
         restriction: &RestrictionConfig,
         remote: RemoteAddr,

src/tunnel/server/utils.rs

@@ -2,6 +2,9 @@ use crate::restrictions::types::{
     AllowConfig, MatchConfig, RestrictionConfig, RestrictionsRules, ReverseTunnelConfigProtocol, TunnelConfigProtocol,
 };
 use crate::tunnel::{tunnel_to_jwt_token, JwtTunnelConfig, RemoteAddr, JWT_DECODE, JWT_HEADER_PREFIX};
+use bytes::Bytes;
+use http_body_util::combinators::BoxBody;
+use http_body_util::Either;
 use hyper::body::{Body, Incoming};
 use hyper::header::{HeaderValue, COOKIE, SEC_WEBSOCKET_PROTOCOL};
 use hyper::{http, Request, Response, StatusCode};
@@ -13,6 +16,13 @@ use tracing::{error, info, warn};
 use url::Host;
 use uuid::Uuid;
 
+pub(super) fn bad_request() -> Response<Either<String, BoxBody<Bytes, anyhow::Error>>> {
+    http::Response::builder()
+        .status(StatusCode::BAD_REQUEST)
+        .body(Either::Left("Invalid request".to_string()))
+        .unwrap()
+}
+
 /// Checks if the requested (remote) port has been mapped in the configuration to another port.
 /// If it is not mapped the original port number is returned.
 #[inline]
@@ -37,7 +47,7 @@ pub(super) fn find_mapped_port(req_port: u16, restriction: &RestrictionConfig) -
 }
 
 #[inline]
-pub(super) fn extract_x_forwarded_for(req: &Request<Incoming>) -> Result<Option<(IpAddr, &str)>, Response<String>> {
+pub(super) fn extract_x_forwarded_for(req: &Request<Incoming>) -> Result<Option<(IpAddr, &str)>, ()> {
     let Some(x_forward_for) = req.headers().get("X-Forwarded-For") else {
         return Ok(None);
     };
@@ -50,38 +60,29 @@ pub(super) fn extract_x_forwarded_for(req: &Request<Incoming>) -> Result<Option<
 }
 
 #[inline]
-pub(super) fn extract_path_prefix(req: &Request<Incoming>) -> Result<&str, Response<String>> {
+pub(super) fn extract_path_prefix(req: &Request<Incoming>) -> Result<&str, ()> {
     let path = req.uri().path();
     let min_len = min(path.len(), 1);
     if &path[0..min_len] != "/" {
         warn!("Rejecting connection with bad path prefix in upgrade request: {}", req.uri());
-        return Err(http::Response::builder()
+        return Err(());
-            .status(StatusCode::BAD_REQUEST)
-            .body("Invalid upgrade request".to_string())
-            .unwrap());
     }
 
     let Some((l, r)) = path[min_len..].split_once('/') else {
         warn!("Rejecting connection with bad upgrade request: {}", req.uri());
-        return Err(http::Response::builder()
+        return Err(());
-            .status(StatusCode::BAD_REQUEST)
-            .body("Invalid upgrade request".into())
-            .unwrap());
     };
 
     if !r.ends_with("events") {
         warn!("Rejecting connection with bad upgrade request: {}", req.uri());
-        return Err(http::Response::builder()
+        return Err(());
-            .status(StatusCode::BAD_REQUEST)
-            .body("Invalid upgrade request".into())
-            .unwrap());
     }
 
     Ok(l)
 }
 
 #[inline]
-pub(super) fn extract_tunnel_info(req: &Request<Incoming>) -> Result<TokenData<JwtTunnelConfig>, Response<String>> {
+pub(super) fn extract_tunnel_info(req: &Request<Incoming>) -> Result<TokenData<JwtTunnelConfig>, ()> {
     let jwt = req
         .headers()
         .get(SEC_WEBSOCKET_PROTOCOL)
@@ -100,10 +101,7 @@ pub(super) fn extract_tunnel_info(req: &Request<Incoming>) -> Result<TokenData<J
                 err,
                 req.headers().get(SEC_WEBSOCKET_PROTOCOL)
             );
-            return Err(http::Response::builder()
+            return Err(());
-                .status(StatusCode::BAD_REQUEST)
-                .body("Invalid upgrade request".to_string())
-                .unwrap());
         }
     };
 
@@ -115,7 +113,7 @@ pub(super) fn validate_tunnel<'a>(
     remote: &RemoteAddr,
     path_prefix: &str,
     restrictions: &'a RestrictionsRules,
-) -> Result<&'a RestrictionConfig, Response<String>> {
+) -> Result<&'a RestrictionConfig, ()> {
     for restriction in &restrictions.restrictions {
         if !restriction.r#match.iter().all(|m| match m {
             MatchConfig::Any => true,
@@ -208,26 +206,13 @@ pub(super) fn validate_tunnel<'a>(
     }
 
     warn!("Rejecting connection with not allowed destination: {:?}", remote);
-    Err(http::Response::builder()
+    Err(())
-        .status(StatusCode::BAD_REQUEST)
-        .body("Invalid upgrade request".to_string())
-        .unwrap())
 }
 
-pub(super) fn inject_cookie<B>(
+pub(super) fn inject_cookie(response: &mut http::Response<impl Body>, remote_addr: &RemoteAddr) -> Result<(), ()> {
-    response: &mut http::Response<B>,
-    remote_addr: &RemoteAddr,
-    mk_body: impl FnOnce(String) -> B,
-) -> Result<(), Response<B>>
-where
-    B: Body,
-{
     let Ok(header_val) = HeaderValue::from_str(&tunnel_to_jwt_token(Uuid::from_u128(0), remote_addr)) else {
         error!("Bad header value for reverse socks5: {} {}", remote_addr.host, remote_addr.port);
-        return Err(http::Response::builder()
+        return Err(());
-            .status(StatusCode::BAD_REQUEST)
-            .body(mk_body("Invalid upgrade request".to_string()))
-            .unwrap());
     };
     response.headers_mut().insert(COOKIE, header_val);