feat(dgw): ngrok tunnel support

devolutions-gateway/src/config.rs

@@ -685,6 +685,9 @@ pub mod dto {
 
         /// Ignore KDC address provided by KDC token, and use this one instead
         pub override_kdc: Option<TargetAddr>,
+
+        #[serde(default)]
+        pub enable_ngrok: bool,
     }
 
     /// Manual Default trait implementation just to make sure default values are deliberates
@@ -695,6 +698,7 @@ pub mod dto {
                 dump_tokens: false,
                 disable_token_validation: false,
                 override_kdc: None,
+                enable_ngrok: false,
             }
         }
     }

devolutions-gateway/src/generic_client.rs

@@ -2,30 +2,32 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 
 use anyhow::Context as _;
-use tokio::io::AsyncWriteExt as _;
-use tokio::net::TcpStream;
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt as _};
 use typed_builder::TypedBuilder;
 
 use crate::config::Conf;
 use crate::proxy::Proxy;
-use crate::rdp_pcb::{extract_association_claims, read_preconnection_pdu};
+use crate::rdp_pcb::{extract_association_claims, read_pcb};
 use crate::session::{ConnectionModeDetails, SessionInfo, SessionManagerHandle};
 use crate::subscriber::SubscriberSender;
 use crate::token::{ConnectionMode, CurrentJrl, TokenCache};
 use crate::utils;
 
 #[derive(TypedBuilder)]
-pub struct GenericClient {
+pub struct GenericClient<S> {
     conf: Arc<Conf>,
     token_cache: Arc<TokenCache>,
     jrl: Arc<CurrentJrl>,
     client_addr: SocketAddr,
-    client_stream: TcpStream,
+    client_stream: S,
     sessions: SessionManagerHandle,
     subscriber_tx: SubscriberSender,
 }
 
-impl GenericClient {
+impl<S> GenericClient<S>
+where
+    S: AsyncWrite + AsyncRead + Unpin,
+{
     pub async fn serve(self) -> anyhow::Result<()> {
         let Self {
             conf,
@@ -37,7 +39,18 @@ impl GenericClient {
             subscriber_tx,
         } = self;
 
-        let (pdu, mut leftover_bytes) = read_preconnection_pdu(&mut client_stream).await?;
+        let timeout = tokio::time::sleep(tokio::time::Duration::from_secs(10));
+        let read_pcb_fut = read_pcb(&mut client_stream);
+
+        let (pdu, mut leftover_bytes) = tokio::select! {
+            () = timeout => {
+                anyhow::bail!("timed out at preconnection blob reception");
+            }
+            result = read_pcb_fut => {
+                result?
+            }
+        };
+
         let source_ip = client_addr.ip();
         let association_claims = extract_association_claims(&pdu, source_ip, &conf, &token_cache, &jrl)?;
 

devolutions-gateway/src/lib.rs

@@ -19,6 +19,7 @@ pub mod jrec;
 pub mod listener;
 pub mod log;
 pub mod middleware;
+pub mod ngrok;
 pub mod plugin_manager;
 pub mod proxy;
 pub mod rdp_extension;

devolutions-gateway/src/listener.rs

@@ -207,7 +207,7 @@ async fn handle_https_peer(
     handle_http_peer(tls_stream, state, peer_addr).await
 }
 
-async fn handle_http_peer<I>(io: I, state: DgwState, peer_addr: SocketAddr) -> anyhow::Result<()>
+pub(crate) async fn handle_http_peer<I>(io: I, state: DgwState, peer_addr: SocketAddr) -> anyhow::Result<()>
 where
     I: AsyncRead + AsyncWrite + Send + Unpin + 'static,
 {

devolutions-gateway/src/ngrok.rs

@@ -0,0 +1,136 @@
+use anyhow::Context as _;
+use futures::TryStreamExt as _;
+use ngrok::config::TunnelBuilder as _;
+use ngrok::tunnel::UrlTunnel as _;
+use tracing::Instrument as _;
+
+use crate::generic_client::GenericClient;
+use crate::DgwState;
+
+#[derive(Clone)]
+pub struct NgrokSession {
+    inner: ngrok::Session,
+    state: DgwState,
+}
+
+impl NgrokSession {
+    pub async fn connect(state: DgwState) -> anyhow::Result<Self> {
+        info!("Connecting to ngrok service");
+
+        let session = ngrok::Session::builder()
+            // Read the token from the NGROK_AUTHTOKEN environment variable
+            .authtoken_from_env()
+            // Connect the ngrok session
+            .connect()
+            .await
+            .context("connect to ngrok service")?;
+
+        debug!("ngrok session connected");
+
+        Ok(Self { inner: session, state })
+    }
+
+    pub async fn run_tcp_endpoint(&self) -> anyhow::Result<()> {
+        info!("Start ngrok TCP tunnel…");
+
+        let conf = self.state.conf_handle.get_conf();
+
+        // Start a tunnel with an HTTP edge
+        let mut tunnel = self
+            .inner
+            .tcp_endpoint()
+            .forwards_to(conf.hostname.clone())
+            .metadata("Devolutions Gateway Tunnel")
+            .listen()
+            .await
+            .context("TCP tunnel listen")?;
+
+        info!(url = tunnel.url(), "Bound TCP ngrok tunnel");
+
+        loop {
+            match tunnel.try_next().await {
+                Ok(Some(conn)) => {
+                    let state = self.state.clone();
+                    let peer_addr = conn.remote_addr();
+
+                    let fut = async move {
+                        if let Err(e) = GenericClient::builder()
+                            .conf(state.conf_handle.get_conf())
+                            .client_addr(peer_addr)
+                            .client_stream(conn)
+                            .token_cache(state.token_cache)
+                            .jrl(state.jrl)
+                            .sessions(state.sessions)
+                            .subscriber_tx(state.subscriber_tx)
+                            .build()
+                            .serve()
+                            .instrument(info_span!("generic-client"))
+                            .await
+                        {
+                            error!(error = format!("{e:#}"), "handle_tcp_peer failed");
+                        }
+                    }
+                    .instrument(info_span!("ngrok_tcp", client = %peer_addr));
+
+                    tokio::spawn(fut);
+                }
+                Ok(None) => {
+                    info!(url = tunnel.url(), "Tunnel closed");
+                    break;
+                }
+                Err(error) => {
+                    error!(url = tunnel.url(), %error, "Failed to accept connection");
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    pub async fn run_http_endpoint(&self) -> anyhow::Result<()> {
+        info!("Start ngrok HTTP tunnel…");
+
+        let conf = self.state.conf_handle.get_conf();
+
+        // Start a tunnel with an HTTP edge
+        let mut tunnel = self
+            .inner
+            .http_endpoint()
+            .forwards_to(conf.hostname.clone())
+            .metadata("Devolutions Gateway Tunnel")
+            .listen()
+            .await
+            .context("HTTP tunnel listen")?;
+
+        info!(url = tunnel.url(), "Bound HTTP ngrok tunnel");
+
+        loop {
+            match tunnel.try_next().await {
+                Ok(Some(conn)) => {
+                    let state = self.state.clone();
+                    let peer_addr = conn.remote_addr();
+
+                    let fut = async move {
+                        if let Err(e) = crate::listener::handle_http_peer(conn, state, peer_addr).await {
+                            error!(error = format!("{e:#}"), "handle_http_peer failed");
+                        }
+                    }
+                    .instrument(info_span!("ngrok_http", client = %peer_addr));
+
+                    tokio::spawn(fut);
+                }
+                Ok(None) => {
+                    info!(url = tunnel.url(), "Tunnel closed");
+                    break;
+                }
+                Err(error) => {
+                    error!(url = tunnel.url(), %error, "Failed to accept connection");
+                }
+            }
+        }
+
+        Ok(())
+    }
+}
+
+// fn handle_http_tunnel(mut tunnel: impl UrlTunnel, ) ->

devolutions-gateway/src/rdp_pcb.rs

@@ -4,8 +4,7 @@ use std::net::IpAddr;
 use anyhow::Context as _;
 use bytes::BytesMut;
 use ironrdp_pdu::pcb::PreconnectionBlob;
-use tokio::io::AsyncReadExt;
-use tokio::net::TcpStream;
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
 
 use crate::config::Conf;
 use crate::token::{AccessTokenClaims, AssociationTokenClaims, CurrentJrl, TokenCache, TokenValidator};
@@ -48,7 +47,7 @@ pub fn extract_association_claims(
     }
 }
 
-pub fn decode_preconnection_pdu(buf: &[u8]) -> Result<Option<PreconnectionBlob>, io::Error> {
+pub fn decode_pcb(buf: &[u8]) -> Result<Option<PreconnectionBlob>, io::Error> {
     match ironrdp_pdu::decode::<PreconnectionBlob>(buf) {
         Ok(preconnection_pdu) => Ok(Some(preconnection_pdu)),
         Err(ironrdp_pdu::Error::NotEnoughBytes { .. }) => Ok(None),
@@ -57,7 +56,7 @@ pub fn decode_preconnection_pdu(buf: &[u8]) -> Result<Option<PreconnectionBlob>,
 }
 
 /// Returns the decoded preconnection PDU and leftover bytes
-pub async fn read_preconnection_pdu(stream: &mut TcpStream) -> io::Result<(PreconnectionBlob, BytesMut)> {
+pub async fn read_pcb(mut stream: impl AsyncRead + AsyncWrite + Unpin) -> io::Result<(PreconnectionBlob, BytesMut)> {
     let mut buf = BytesMut::with_capacity(1024);
 
     loop {
@@ -70,7 +69,7 @@ pub async fn read_preconnection_pdu(stream: &mut TcpStream) -> io::Result<(Preco
             ));
         }
 
-        if let Some(pdu) = decode_preconnection_pdu(&buf)? {
+        if let Some(pdu) = decode_pcb(&buf)? {
             let leftover_bytes = buf.split_off(ironrdp_pdu::size(&pdu));
             return Ok((pdu, leftover_bytes));
         }

devolutions-gateway/src/service.rs

@@ -70,7 +70,7 @@ impl GatewayService {
         let config = self.conf_handle.clone();
 
         // create_futures needs to be run in the runtime in order to bind the sockets.
-        let futures = runtime.block_on(async { create_futures(config) })?;
+        let futures = runtime.block_on(create_futures(config))?;
 
         let join_all = futures::future::join_all(futures);
 
@@ -107,7 +107,7 @@ impl GatewayService {
 // TODO: when benchmarking facility is ready, use Handle instead of pinned futures
 type VecOfFuturesType = Vec<Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + 'static>>>;
 
-fn create_futures(conf_handle: ConfHandle) -> anyhow::Result<VecOfFuturesType> {
+async fn create_futures(conf_handle: ConfHandle) -> anyhow::Result<VecOfFuturesType> {
     let conf = conf_handle.get_conf();
 
     let token_cache = devolutions_gateway::token::new_token_cache().pipe(Arc::new);
@@ -133,12 +133,27 @@ fn create_futures(conf_handle: ConfHandle) -> anyhow::Result<VecOfFuturesType> {
                 .with_context(|| format!("Failed to initialize {}", listener.internal_url))
         })
         .collect::<anyhow::Result<Vec<GatewayListener>>>()
-        .context("Failed to bind a listener")?;
+        .context("failed to bind listener")?;
 
     for listener in listeners {
         futures.push(Box::pin(listener.run()))
     }
 
+    if conf.debug.enable_ngrok && std::env::var("NGROK_AUTHTOKEN").is_ok() {
+        let session = devolutions_gateway::ngrok::NgrokSession::connect(state)
+            .await
+            .context("couldn’t create ngrok session")?;
+
+        let tcp_fut = {
+            let session = session.clone();
+            async move { session.run_tcp_endpoint().await }
+        };
+        futures.push(Box::pin(tcp_fut));
+
+        let http_fut = async move { session.run_http_endpoint().await };
+        futures.push(Box::pin(http_fut));
+    }
+
     futures.push(Box::pin(async {
         devolutions_gateway::token::cleanup_task(token_cache).await;
         Ok(())