feat: Add connection pool to speed up creation of tunnel

src/tunnel/client.rs

@@ -0,0 +1,129 @@
+use super::{JwtTunnelConfig, MaybeTlsStream, JWT_KEY};
+use crate::{LocalProtocol, LocalToRemote, WsClientConfig};
+use anyhow::{anyhow, Context};
+
+use fastwebsockets::WebSocket;
+use hyper::header::{AUTHORIZATION, SEC_WEBSOCKET_VERSION, UPGRADE};
+use hyper::header::{CONNECTION, HOST, SEC_WEBSOCKET_KEY};
+use hyper::upgrade::Upgraded;
+use hyper::{Body, Request};
+use std::future::Future;
+use std::ops::{Deref, DerefMut};
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::sync::oneshot;
+use tracing::log::debug;
+use tracing::{Instrument, Span};
+use uuid::Uuid;
+
+struct SpawnExecutor;
+
+impl<Fut> hyper::rt::Executor<Fut> for SpawnExecutor
+where
+    Fut: Future + Send + 'static,
+    Fut::Output: Send + 'static,
+{
+    fn execute(&self, fut: Fut) {
+        tokio::task::spawn(fut);
+    }
+}
+
+pub async fn connect(
+    request_id: Uuid,
+    client_cfg: &WsClientConfig,
+    tunnel_cfg: &LocalToRemote,
+) -> anyhow::Result<WebSocket<Upgraded>> {
+    let mut tcp_stream = match client_cfg.cnx_pool().get().await {
+        Ok(tcp_stream) => tcp_stream,
+        Err(err) => Err(anyhow!(
+            "failed to get a connection to the server from the pool: {err:?}"
+        ))?,
+    };
+
+    let data = JwtTunnelConfig {
+        id: request_id.to_string(),
+        p: match tunnel_cfg.local_protocol {
+            LocalProtocol::Tcp => LocalProtocol::Tcp,
+            LocalProtocol::Udp { .. } => tunnel_cfg.local_protocol,
+            LocalProtocol::Stdio => LocalProtocol::Tcp,
+            LocalProtocol::Socks5 => LocalProtocol::Tcp,
+        },
+        r: tunnel_cfg.remote.0.to_string(),
+        rp: tunnel_cfg.remote.1,
+    };
+    let (alg, secret) = JWT_KEY.deref();
+    let mut req = Request::builder()
+        .method("GET")
+        .uri(format!(
+            "/{}/events?bearer={}",
+            &client_cfg.http_upgrade_path_prefix,
+            jsonwebtoken::encode(alg, &data, secret).unwrap_or_default(),
+        ))
+        .header(HOST, client_cfg.remote_addr.0.to_string())
+        .header(UPGRADE, "websocket")
+        .header(CONNECTION, "upgrade")
+        .header(SEC_WEBSOCKET_KEY, fastwebsockets::handshake::generate_key())
+        .header(SEC_WEBSOCKET_VERSION, "13")
+        .version(hyper::Version::HTTP_11);
+
+    for (k, v) in &client_cfg.http_headers {
+        req = req.header(k.clone(), v.clone());
+    }
+    if let Some(auth) = &client_cfg.http_upgrade_credentials {
+        req = req.header(AUTHORIZATION, auth.clone());
+    }
+
+    let req = req.body(Body::empty()).with_context(|| {
+        format!(
+            "failed to build HTTP request to contact the server {:?}",
+            client_cfg.remote_addr
+        )
+    })?;
+    debug!("with HTTP upgrade request {:?}", req);
+    let ws_handshake = match tcp_stream.deref_mut() {
+        MaybeTlsStream::Plain(cnx) => {
+            fastwebsockets::handshake::client(&SpawnExecutor, req, cnx.take().unwrap()).await
+        }
+        MaybeTlsStream::Tls(cnx) => {
+            fastwebsockets::handshake::client(&SpawnExecutor, req, cnx.take().unwrap()).await
+        }
+    };
+
+    let (ws, _) = ws_handshake.with_context(|| {
+        format!(
+            "failed to do websocket handshake with the server {:?}",
+            client_cfg.remote_addr
+        )
+    })?;
+
+    Ok(ws)
+}
+
+pub async fn connect_to_server<R, W>(
+    request_id: Uuid,
+    client_cfg: &WsClientConfig,
+    remote_cfg: &LocalToRemote,
+    duplex_stream: (R, W),
+) -> anyhow::Result<()>
+where
+    R: AsyncRead + Send + 'static,
+    W: AsyncWrite + Send + 'static,
+{
+    let mut ws = connect(request_id, client_cfg, remote_cfg).await?;
+    ws.set_auto_apply_mask(client_cfg.websocket_mask_frame);
+
+    let (ws_rx, ws_tx) = ws.split(tokio::io::split);
+    let (local_rx, local_tx) = duplex_stream;
+    let (close_tx, close_rx) = oneshot::channel::<()>();
+
+    // Forward local tx to websocket tx
+    let ping_frequency = client_cfg.websocket_ping_frequency;
+    tokio::spawn(
+        super::io::propagate_read(local_rx, ws_tx, close_tx, ping_frequency)
+            .instrument(Span::current()),
+    );
+
+    // Forward websocket rx to local rx
+    let _ = super::io::propagate_write(local_tx, ws_rx, close_rx).await;
+
+    Ok(())
+}

src/tunnel/io.rs

@@ -0,0 +1,123 @@
+use fastwebsockets::{Frame, OpCode, Payload, WebSocketError, WebSocketRead, WebSocketWrite};
+use futures_util::pin_mut;
+use hyper::upgrade::Upgraded;
+use std::time::Duration;
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};
+use tokio::select;
+use tokio::sync::oneshot;
+use tokio::time::timeout;
+use tracing::log::debug;
+use tracing::{error, info, trace, warn};
+
+pub(super) async fn propagate_read(
+    local_rx: impl AsyncRead,
+    mut ws_tx: WebSocketWrite<WriteHalf<Upgraded>>,
+    mut close_tx: oneshot::Sender<()>,
+    ping_frequency: Duration,
+) -> Result<(), WebSocketError> {
+    let _guard = scopeguard::guard((), |_| {
+        info!("Closing local tx ==> websocket tx tunnel");
+    });
+
+    let mut buffer = vec![0u8; 8 * 1024];
+    pin_mut!(local_rx);
+    loop {
+        let read = select! {
+            biased;
+
+            read_len = local_rx.read(buffer.as_mut_slice()) => read_len,
+
+            _ = close_tx.closed() => break,
+
+            _ = timeout(ping_frequency, futures_util::future::pending::<()>()) => {
+                debug!("sending ping to keep websocket connection alive");
+                ws_tx.write_frame(Frame::new(true, OpCode::Ping, None, Payload::Borrowed(&[]))).await?;
+                continue;
+            }
+        };
+
+        let read_len = match read {
+            Ok(read_len) if read_len > 0 => read_len,
+            Ok(_) => break,
+            Err(err) => {
+                warn!(
+                    "error while reading incoming bytes from local tx tunnel {}",
+                    err
+                );
+                break;
+            }
+        };
+
+        trace!("read {} bytes", read_len);
+        match ws_tx
+            .write_frame(Frame::binary(Payload::Borrowed(&buffer[..read_len])))
+            .await
+        {
+            Ok(_) => {}
+            Err(err) => {
+                warn!("error while writing to websocket tx tunnel {}", err);
+                break;
+            }
+        }
+
+        if read_len == buffer.len() {
+            buffer.resize(read_len * 2, 0);
+        }
+    }
+
+    let _ = ws_tx.write_frame(Frame::close(1000, &[])).await;
+
+    Ok(())
+}
+
+pub(super) async fn propagate_write(
+    local_tx: impl AsyncWrite,
+    mut ws_rx: WebSocketRead<ReadHalf<Upgraded>>,
+    mut close_rx: oneshot::Receiver<()>,
+) -> Result<(), WebSocketError> {
+    let _guard = scopeguard::guard((), |_| {
+        info!("Closing local rx <== websocket rx tunnel");
+    });
+    let mut x = |x: Frame<'_>| {
+        debug!("frame {:?} {:?}", x.opcode, x.payload);
+        futures_util::future::ready(anyhow::Ok(()))
+    };
+
+    pin_mut!(local_tx);
+    loop {
+        let ret = select! {
+            biased;
+            ret = ws_rx.read_frame(&mut x) => ret,
+
+            _ = &mut close_rx => break,
+        };
+
+        let msg = match ret {
+            Ok(msg) => msg,
+            Err(err) => {
+                error!("error while reading from websocket rx {}", err);
+                break;
+            }
+        };
+
+        trace!("receive ws frame {:?} {:?}", msg.opcode, msg.payload);
+        let ret = match msg.opcode {
+            OpCode::Continuation | OpCode::Text | OpCode::Binary => {
+                local_tx.write_all(msg.payload.as_ref()).await
+            }
+            OpCode::Close => break,
+            OpCode::Ping => Ok(()),
+            OpCode::Pong => Ok(()),
+        };
+
+        match ret {
+            Ok(_) => {}
+            Err(err) => {
+                error!("error while writing bytes to local for rx tunnel {}", err);
+                break;
+            }
+        }
+    }
+
+    Ok(())
+}

src/tunnel/mod.rs

@@ -0,0 +1,83 @@
+pub mod client;
+mod io;
+pub mod server;
+
+use crate::{tcp, tls, LocalProtocol, WsClientConfig};
+use async_trait::async_trait;
+use bb8::ManageConnection;
+use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, Validation};
+use once_cell::sync::Lazy;
+use serde::{Deserialize, Serialize};
+use std::collections::HashSet;
+use tokio::net::TcpStream;
+use tokio_rustls::client::TlsStream;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct JwtTunnelConfig {
+    pub id: String,
+    pub p: LocalProtocol,
+    pub r: String,
+    pub rp: u16,
+}
+
+static JWT_SECRET: &[u8; 15] = b"champignonfrais";
+static JWT_KEY: Lazy<(Header, EncodingKey)> = Lazy::new(|| {
+    (
+        Header::new(Algorithm::HS256),
+        EncodingKey::from_secret(JWT_SECRET),
+    )
+});
+
+static JWT_DECODE: Lazy<(Validation, DecodingKey)> = Lazy::new(|| {
+    let mut validation = Validation::new(Algorithm::HS256);
+    validation.required_spec_claims = HashSet::with_capacity(0);
+    (validation, DecodingKey::from_secret(JWT_SECRET))
+});
+
+pub enum MaybeTlsStream {
+    Plain(Option<TcpStream>),
+    Tls(Option<TlsStream<TcpStream>>),
+}
+
+impl MaybeTlsStream {
+    pub fn is_used(&self) -> bool {
+        match self {
+            MaybeTlsStream::Plain(Some(_)) | MaybeTlsStream::Tls(Some(_)) => false,
+            MaybeTlsStream::Plain(None) | MaybeTlsStream::Tls(None) => true,
+        }
+    }
+}
+
+#[async_trait]
+impl ManageConnection for WsClientConfig {
+    type Connection = MaybeTlsStream;
+    type Error = anyhow::Error;
+
+    async fn connect(&self) -> Result<Self::Connection, Self::Error> {
+        let (host, port) = &self.remote_addr;
+        let so_mark = &self.socket_so_mark;
+        let timeout = self.timeout_connect;
+
+        let tcp_stream = if let Some(http_proxy) = &self.http_proxy {
+            tcp::connect_with_http_proxy(http_proxy, host, *port, so_mark, timeout).await?
+        } else {
+            tcp::connect(host, *port, so_mark, timeout).await?
+        };
+
+        match &self.tls {
+            None => Ok(MaybeTlsStream::Plain(Some(tcp_stream))),
+            Some(tls_cfg) => {
+                let tls_stream = tls::connect(self, tls_cfg, tcp_stream).await?;
+                Ok(MaybeTlsStream::Tls(Some(tls_stream)))
+            }
+        }
+    }
+
+    async fn is_valid(&self, _conn: &mut Self::Connection) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
+    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
+        conn.is_used()
+    }
+}

src/tunnel/server.rs

@@ -0,0 +1,263 @@
+use std::cmp::min;
+use std::ops::{Deref, Not};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use super::{JwtTunnelConfig, JWT_DECODE};
+use crate::udp::MyUdpSocket;
+use crate::{tcp, tls, LocalProtocol, WsServerConfig};
+use hyper::server::conn::Http;
+use hyper::service::service_fn;
+use hyper::{http, Body, Request, Response, StatusCode};
+use jsonwebtoken::TokenData;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio::net::{TcpListener, UdpSocket};
+use tokio::sync::oneshot;
+use tracing::{error, info, instrument, span, warn, Instrument, Level, Span};
+use url::Host;
+
+async fn from_query(
+    server_config: &WsServerConfig,
+    query: &str,
+) -> anyhow::Result<(
+    LocalProtocol,
+    Host,
+    u16,
+    Pin<Box<dyn AsyncRead + Send>>,
+    Pin<Box<dyn AsyncWrite + Send>>,
+)> {
+    let jwt: TokenData<JwtTunnelConfig> = match query.split_once('=') {
+        Some(("bearer", jwt)) => {
+            let (validation, decode_key) = JWT_DECODE.deref();
+            match jsonwebtoken::decode(jwt, decode_key, validation) {
+                Ok(jwt) => jwt,
+                err => {
+                    error!("error while decoding jwt for tunnel info {:?}", err);
+                    return Err(anyhow::anyhow!("Invalid upgrade request"));
+                }
+            }
+        }
+        _err => return Err(anyhow::anyhow!("Invalid upgrade request")),
+    };
+
+    Span::current().record("id", jwt.claims.id);
+    Span::current().record("remote", format!("{}:{}", jwt.claims.r, jwt.claims.rp));
+    if let Some(allowed_dests) = &server_config.restrict_to {
+        let requested_dest = format!("{}:{}", jwt.claims.r, jwt.claims.rp);
+        if allowed_dests
+            .iter()
+            .any(|dest| dest == &requested_dest)
+            .not()
+        {
+            warn!(
+                "Rejecting connection with not allowed destination: {}",
+                requested_dest
+            );
+            return Err(anyhow::anyhow!("Invalid upgrade request"));
+        }
+    }
+
+    match jwt.claims.p {
+        LocalProtocol::Udp { .. } => {
+            let host = Host::parse(&jwt.claims.r)?;
+            let cnx = Arc::new(UdpSocket::bind("[::]:0").await?);
+            cnx.connect((host.to_string(), jwt.claims.rp)).await?;
+            Ok((
+                LocalProtocol::Udp { timeout: None },
+                host,
+                jwt.claims.rp,
+                Box::pin(MyUdpSocket::new(cnx.clone())),
+                Box::pin(MyUdpSocket::new(cnx)),
+            ))
+        }
+        LocalProtocol::Tcp { .. } => {
+            let host = Host::parse(&jwt.claims.r)?;
+            let port = jwt.claims.rp;
+            let (rx, tx) = tcp::connect(
+                &host,
+                port,
+                &server_config.socket_so_mark,
+                Duration::from_secs(10),
+            )
+            .await?
+            .into_split();
+
+            Ok((jwt.claims.p, host, port, Box::pin(rx), Box::pin(tx)))
+        }
+        _ => Err(anyhow::anyhow!("Invalid upgrade request")),
+    }
+}
+
+async fn server_upgrade(
+    server_config: Arc<WsServerConfig>,
+    mut req: Request<Body>,
+) -> Result<Response<Body>, anyhow::Error> {
+    if let Some(x) = req.headers().get("X-Forwarded-For") {
+        info!("Request X-Forwarded-For: {:?}", x);
+        Span::current().record("forwarded_for", x.to_str().unwrap_or_default());
+    }
+
+    if !req.uri().path().ends_with("/events") {
+        warn!(
+            "Rejecting connection with bad upgrade request: {}",
+            req.uri()
+        );
+        return Ok(http::Response::builder()
+            .status(StatusCode::BAD_REQUEST)
+            .body(Body::from("Invalid upgrade request"))
+            .unwrap_or_default());
+    }
+
+    if let Some(path_prefix) = &server_config.restrict_http_upgrade_path_prefix {
+        let path = req.uri().path();
+        let min_len = min(path.len(), 1);
+        let max_len = min(path.len(), path_prefix.len() + 1);
+        if &path[0..min_len] != "/"
+            || &path[min_len..max_len] != path_prefix.as_str()
+            || !path[max_len..].starts_with('/')
+        {
+            warn!(
+                "Rejecting connection with bad path prefix in upgrade request: {}",
+                req.uri()
+            );
+            return Ok(http::Response::builder()
+                .status(StatusCode::BAD_REQUEST)
+                .body(Body::from("Invalid upgrade request"))
+                .unwrap_or_default());
+        }
+    }
+
+    let (protocol, dest, port, local_rx, local_tx) =
+        match from_query(&server_config, req.uri().query().unwrap_or_default()).await {
+            Ok(ret) => ret,
+            Err(err) => {
+                warn!(
+                    "Rejecting connection with bad upgrade request: {} {}",
+                    err,
+                    req.uri()
+                );
+                return Ok(http::Response::builder()
+                    .status(StatusCode::BAD_REQUEST)
+                    .body(Body::from(format!("Invalid upgrade request: {:?}", err)))
+                    .unwrap_or_default());
+            }
+        };
+
+    info!("connected to {:?} {:?} {:?}", protocol, dest, port);
+    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 Ok(http::Response::builder()
+                .status(StatusCode::BAD_REQUEST)
+                .body(Body::from(format!("Invalid upgrade request: {:?}", err)))
+                .unwrap_or_default());
+        }
+    };
+
+    tokio::spawn(
+        async move {
+            let (ws_rx, mut ws_tx) = match fut.await {
+                Ok(ws) => ws.split(tokio::io::split),
+                Err(err) => {
+                    error!("Error during http upgrade request: {:?}", err);
+                    return;
+                }
+            };
+            let (close_tx, close_rx) = oneshot::channel::<()>();
+            let ping_frequency = server_config
+                .websocket_ping_frequency
+                .unwrap_or(Duration::MAX);
+            ws_tx.set_auto_apply_mask(server_config.websocket_mask_frame);
+
+            tokio::task::spawn(
+                super::io::propagate_write(local_tx, ws_rx, close_rx).instrument(Span::current()),
+            );
+
+            let _ = super::io::propagate_read(local_rx, ws_tx, close_tx, ping_frequency).await;
+        }
+        .instrument(Span::current()),
+    );
+
+    Ok(response)
+}
+
+#[instrument(name="tunnel", level="info", skip_all, fields(id=tracing::field::Empty, remote=tracing::field::Empty, peer=tracing::field::Empty, forwarded_for=tracing::field::Empty))]
+pub async fn run_server(server_config: Arc<WsServerConfig>) -> anyhow::Result<()> {
+    info!(
+        "Starting wstunnel server listening on {}",
+        server_config.bind
+    );
+
+    let config = server_config.clone();
+    let upgrade_fn = move |req: Request<Body>| server_upgrade(config.clone(), req);
+
+    let listener = TcpListener::bind(&server_config.bind).await?;
+    let tls_acceptor = if let Some(tls) = &server_config.tls {
+        Some(tls::tls_acceptor(tls, Some(vec![b"http/1.1".to_vec()]))?)
+    } else {
+        None
+    };
+
+    loop {
+        let (stream, peer_addr) = listener.accept().await?;
+        let _ = stream.set_nodelay(true);
+
+        let span = span!(
+            Level::INFO,
+            "tunnel",
+            id = tracing::field::Empty,
+            remote = tracing::field::Empty,
+            peer = peer_addr.to_string(),
+            forwarded_for = tracing::field::Empty
+        );
+
+        info!("Accepting connection");
+        let upgrade_fn = upgrade_fn.clone();
+        // TLS
+        if let Some(tls_acceptor) = &tls_acceptor {
+            let tls_acceptor = tls_acceptor.clone();
+            let fut = async move {
+                info!("Doing TLS handshake");
+                let tls_stream = match tls_acceptor.accept(stream).await {
+                    Ok(tls_stream) => tls_stream,
+                    Err(err) => {
+                        error!("error while accepting TLS connection {}", err);
+                        return;
+                    }
+                };
+                let conn_fut = Http::new()
+                    .http1_only(true)
+                    .serve_connection(tls_stream, service_fn(upgrade_fn))
+                    .with_upgrades();
+
+                if let Err(e) = conn_fut.await {
+                    error!("Error while upgrading cnx to websocket: {:?}", e);
+                }
+            }
+            .instrument(span);
+
+            tokio::spawn(fut);
+            // Normal
+        } else {
+            let conn_fut = Http::new()
+                .http1_only(true)
+                .serve_connection(stream, service_fn(upgrade_fn))
+                .with_upgrades();
+
+            let fut = async move {
+                if let Err(e) = conn_fut.await {
+                    error!("Error while upgrading cnx to weboscket: {:?}", e);
+                }
+            }
+            .instrument(span);
+
+            tokio::spawn(fut);
+        };
+    }
+}