handle_tunnel_request: small code cleanup

more idiomatic, less code, better readability

Author: tsnoam

src/protocols/tcp/server.rs

@@ -233,7 +233,6 @@ mod tests {
     use super::*;
     use futures_util::pin_mut;
     use std::borrow::Cow;
-    use std::net::SocketAddr;
     use testcontainers::core::WaitFor;
     use testcontainers::runners::AsyncRunner;
     use testcontainers::{ContainerAsync, Image, ImageExt};

src/tunnel/server/handler_http2.rs

@@ -1,5 +1,5 @@
 use crate::restrictions::types::RestrictionsRules;
-use crate::tunnel::server::utils::{bad_request, inject_cookie};
+use crate::tunnel::server::utils::{bad_request, inject_cookie, HttpResponse};
 use crate::tunnel::server::WsServer;
 use crate::tunnel::transport;
 use crate::tunnel::transport::http2::{Http2TunnelRead, Http2TunnelWrite};
@@ -22,7 +22,7 @@ pub(super) async fn http_server_upgrade(
     restrict_path_prefix: Option<String>,
     client_addr: SocketAddr,
     mut req: Request<Incoming>,
-) -> Response<Either<String, BoxBody<Bytes, anyhow::Error>>> {
+) -> HttpResponse {
     let (remote_addr, local_rx, local_tx, need_cookie) = match server
         .handle_tunnel_request(restrictions, restrict_path_prefix, client_addr, &req)
         .await

src/tunnel/server/handler_websocket.rs

@@ -1,5 +1,5 @@
 use crate::restrictions::types::RestrictionsRules;
-use crate::tunnel::server::utils::{bad_request, inject_cookie};
+use crate::tunnel::server::utils::{bad_request, inject_cookie, HttpResponse};
 use crate::tunnel::server::WsServer;
 use crate::tunnel::transport;
 use crate::tunnel::transport::websocket::mk_websocket_tunnel;
@@ -21,7 +21,7 @@ pub(super) async fn ws_server_upgrade(
     restrict_path_prefix: Option<String>,
     client_addr: SocketAddr,
     mut req: Request<Incoming>,
-) -> Response<Either<String, BoxBody<Bytes, anyhow::Error>>> {
+) -> HttpResponse {
     if !fastwebsockets::upgrade::is_upgrade_request(&req) {
         warn!("Rejecting connection with bad upgrade request: {}", req.uri());
         return bad_request();

src/tunnel/server/server.rs

@@ -33,6 +33,7 @@ use crate::tunnel::server::handler_websocket::ws_server_upgrade;
 use crate::tunnel::server::reverse_tunnel::ReverseTunnelServer;
 use crate::tunnel::server::utils::{
     bad_request, extract_path_prefix, extract_tunnel_info, extract_x_forwarded_for, find_mapped_port, validate_tunnel,
+    HttpResponse,
 };
 use crate::tunnel::tls_reloader::TlsReloader;
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
@@ -89,68 +90,46 @@ impl WsServer {
             Pin<Box<dyn AsyncWrite + Send>>,
             bool,
         ),
-        Response<Either<String, BoxBody<Bytes, anyhow::Error>>>,
+        HttpResponse,
     > {
-        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()),
+        if let Some((x_forward_for, x_forward_for_str)) = extract_x_forwarded_for(req) {
+            info!("Request X-Forwarded-For: {x_forward_for:?}");
+            Span::current().record("forwarded_for", x_forward_for_str);
+            client_addr.set_ip(x_forward_for);
         };
 
-        let path_prefix = match extract_path_prefix(req) {
-            Ok(p) => p,
-            Err(_err) => return Err(bad_request()),
-        };
+        let path_prefix = extract_path_prefix(req)?;
 
         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
+                    "Client requested upgrade path '{path_prefix}' does not match upgrade path restriction '{restrict_path}' (mTLS, etc.)"
                 );
                 return Err(bad_request());
             }
         }
 
-        let jwt = match extract_tunnel_info(req) {
-            Ok(jwt) => jwt,
-            Err(_err) => return Err(bad_request()),
-        };
+        let jwt = extract_tunnel_info(req)?;
 
         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 remote = RemoteAddr::try_from(jwt.claims)
+            .inspect_err(|err| warn!("Rejecting connection with bad tunnel info: {err} {}", req.uri()))
+            .map_err(|_| bad_request())?;
 
-        let restriction = match validate_tunnel(&remote, path_prefix, &restrictions) {
-            Some(matched_restriction) => {
-                info!("Tunnel accepted due to matched restriction: {}", matched_restriction.name);
-                matched_restriction
-            }
-            None => {
-                warn!("Rejecting connection with not allowed destination: {:?}", remote);
-                return Err(bad_request());
-            }
-        };
+        let restriction = validate_tunnel(&remote, path_prefix, &restrictions).ok_or_else(|| {
+            warn!("Rejecting connection with not allowed destination: {remote:?}");
+            bad_request()
+        })?;
+        info!("Tunnel accepted due to matched restriction: {}", restriction.name);
 
         let req_protocol = remote.protocol.clone();
         let inject_cookie = req_protocol.is_dynamic_reverse_tunnel();
-        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 tunnel = self
+            .exec_tunnel(restriction, remote, client_addr)
+            .await
+            .inspect_err(|err| warn!("Rejecting connection with bad upgrade request: {err} {}", req.uri()))
+            .map_err(|_| bad_request())?;
 
         let (remote_addr, local_rx, local_tx) = tunnel;
         info!("connected to {:?} {}:{}", req_protocol, remote_addr.host, remote_addr.port);

src/tunnel/server/utils.rs

@@ -17,7 +17,9 @@ use tracing::{error, info, warn};
 use url::Host;
 use uuid::Uuid;
 
-pub(super) fn bad_request() -> Response<Either<String, BoxBody<Bytes, anyhow::Error>>> {
+pub type HttpResponse = Response<Either<String, BoxBody<Bytes, anyhow::Error>>>;
+
+pub(super) fn bad_request() -> HttpResponse {
     http::Response::builder()
         .status(StatusCode::BAD_REQUEST)
         .body(Either::Left("Invalid request".to_string()))
@@ -48,42 +50,40 @@ 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)>, ()> {
-    let Some(x_forward_for) = req.headers().get("X-Forwarded-For") else {
-        return Ok(None);
-    };
+pub(super) fn extract_x_forwarded_for(req: &Request<Incoming>) -> Option<(IpAddr, &str)> {
+    let x_forward_for = req.headers().get("X-Forwarded-For")?;
 
     // X-Forwarded-For: <client>, <proxy1>, <proxy2>
     let x_forward_for = x_forward_for.to_str().unwrap_or_default();
     let x_forward_for = x_forward_for.split_once(',').map(|x| x.0).unwrap_or(x_forward_for);
     let ip: Option<IpAddr> = x_forward_for.parse().ok();
-    Ok(ip.map(|ip| (ip, x_forward_for)))
+    ip.map(|ip| (ip, x_forward_for))
 }
 
 #[inline]
-pub(super) fn extract_path_prefix(req: &Request<Incoming>) -> Result<&str, ()> {
+pub(super) fn extract_path_prefix(req: &Request<Incoming>) -> Result<&str, HttpResponse> {
     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(());
+        return Err(bad_request());
     }
 
     let Some((l, r)) = path[min_len..].split_once('/') else {
         warn!("Rejecting connection with bad upgrade request: {}", req.uri());
-        return Err(());
+        return Err(bad_request());
     };
 
     if !r.ends_with("events") {
         warn!("Rejecting connection with bad upgrade request: {}", req.uri());
-        return Err(());
+        return Err(bad_request());
     }
 
     Ok(l)
 }
 
 #[inline]
-pub(super) fn extract_tunnel_info(req: &Request<Incoming>) -> Result<TokenData<JwtTunnelConfig>, ()> {
+pub(super) fn extract_tunnel_info(req: &Request<Incoming>) -> anyhow::Result<TokenData<JwtTunnelConfig>, HttpResponse> {
     let jwt = req
         .headers()
         .get(SEC_WEBSOCKET_PROTOCOL)
@@ -93,19 +93,13 @@ pub(super) fn extract_tunnel_info(req: &Request<Incoming>) -> Result<TokenData<J
         .or_else(|| req.headers().get(COOKIE).and_then(|header| header.to_str().ok()))
         .unwrap_or_default();
 
-    let jwt = match jwt_token_to_tunnel(jwt) {
-        Ok(jwt) => jwt,
-        err => {
-            warn!(
-                "error while decoding jwt for tunnel info {:?} header {:?}",
-                err,
-                req.headers().get(SEC_WEBSOCKET_PROTOCOL)
-            );
-            return Err(());
-        }
-    };
-
-    Ok(jwt)
+    jwt_token_to_tunnel(jwt).map_err(|err| {
+        warn!(
+            "error while decoding jwt for tunnel info {err:?} header {:?}",
+            req.headers().get(SEC_WEBSOCKET_PROTOCOL)
+        );
+        bad_request()
+    })
 }
 
 impl RestrictionConfig {