fixed #1242 - added streamlocal-forward support (remote UNIX socket forwarding)

Eugeny · Eugeny · commit 14e3f9d1a801 · 2025-02-09T01:15:14.000+01:00
diff --git a/warpgate-core/src/recordings/traffic.rs b/warpgate-core/src/recordings/traffic.rs
@@ -17,11 +17,16 @@ pub struct TrafficRecorder {
 }
 
 #[derive(Debug)]
-pub struct TrafficConnectionParams {
-    pub src_addr: Ipv4Addr,
-    pub src_port: u16,
-    pub dst_addr: Ipv4Addr,
-    pub dst_port: u16,
+pub enum TrafficConnectionParams {
+    Tcp {
+        src_addr: Ipv4Addr,
+        src_port: u16,
+        dst_addr: Ipv4Addr,
+        dst_port: u16,
+    },
+    Socket {
+        socket_path: String,
+    },
 }
 
 impl TrafficRecorder {
@@ -136,41 +141,59 @@ impl ConnectionRecorder {
     where
         F: FnOnce(packet::ip::v4::Builder) -> Result<Bytes>,
     {
-        f(packet::ip::v4::Builder::default()
-            .protocol(packet::ip::Protocol::Tcp)?
-            .source(self.params.src_addr)?
-            .destination(self.params.dst_addr)?)
+        match self.params {
+            TrafficConnectionParams::Socket { .. } => f(packet::ip::v4::Builder::default()
+                .protocol(packet::ip::Protocol::Tcp)?
+                .source(Ipv4Addr::UNSPECIFIED)?
+                .destination(Ipv4Addr::BROADCAST)?),
+            TrafficConnectionParams::Tcp {
+                src_addr, dst_addr, ..
+            } => f(packet::ip::v4::Builder::default()
+                .protocol(packet::ip::Protocol::Tcp)?
+                .source(src_addr)?
+                .destination(dst_addr)?),
+        }
     }
 
     fn ip_packet_rx<F>(&self, f: F) -> Result<Bytes>
     where
         F: FnOnce(packet::ip::v4::Builder) -> Result<Bytes>,
     {
-        f(packet::ip::v4::Builder::default()
-            .protocol(packet::ip::Protocol::Tcp)?
-            .source(self.params.dst_addr)?
-            .destination(self.params.src_addr)?)
+        match self.params {
+            TrafficConnectionParams::Socket { .. } => f(packet::ip::v4::Builder::default()
+                .protocol(packet::ip::Protocol::Tcp)?
+                .source(Ipv4Addr::BROADCAST)?
+                .destination(Ipv4Addr::UNSPECIFIED)?),
+            TrafficConnectionParams::Tcp {
+                src_addr, dst_addr, ..
+            } => f(packet::ip::v4::Builder::default()
+                .protocol(packet::ip::Protocol::Tcp)?
+                .source(dst_addr)?
+                .destination(src_addr)?),
+        }
     }
 
     fn tcp_packet_tx<F>(&self, f: F) -> Result<Bytes>
     where
         F: FnOnce(packet::tcp::Builder) -> Result<Bytes>,
     {
-        self.ip_packet_tx(|b| {
-            f(b.tcp()?
-                .source(self.params.src_port)?
-                .destination(self.params.dst_port)?)
+        self.ip_packet_tx(|b| match self.params {
+            TrafficConnectionParams::Socket { .. } => f(b.tcp()?.source(0)?.destination(0)?),
+            TrafficConnectionParams::Tcp {
+                src_port, dst_port, ..
+            } => f(b.tcp()?.source(src_port)?.destination(dst_port)?),
         })
     }
 
     fn tcp_packet_rx<F>(&self, f: F) -> Result<Bytes>
     where
         F: FnOnce(packet::tcp::Builder) -> Result<Bytes>,
     {
-        self.ip_packet_rx(|b| {
-            f(b.tcp()?
-                .source(self.params.dst_port)?
-                .destination(self.params.src_port)?)
+        self.ip_packet_rx(|b| match self.params {
+            TrafficConnectionParams::Socket { .. } => f(b.tcp()?.source(0)?.destination(0)?),
+            TrafficConnectionParams::Tcp {
+                src_port, dst_port, ..
+            } => f(b.tcp()?.source(dst_port)?.destination(src_port)?),
         })
     }
 
diff --git a/warpgate-protocol-ssh/src/client/handler.rs b/warpgate-protocol-ssh/src/client/handler.rs
@@ -8,13 +8,14 @@ use warpgate_common::{SessionId, TargetSSHOptions};
 use warpgate_core::Services;
 
 use crate::known_hosts::{KnownHostValidationResult, KnownHosts};
-use crate::{ConnectionError, ForwardedTcpIpParams};
+use crate::{ConnectionError, ForwardedStreamlocalParams, ForwardedTcpIpParams};
 
 #[derive(Debug)]
 pub enum ClientHandlerEvent {
     HostKeyReceived(PublicKey),
     HostKeyUnknown(PublicKey, oneshot::Sender<bool>),
     ForwardedTcpIp(Channel<Msg>, ForwardedTcpIpParams),
+    ForwardedStreamlocal(Channel<Msg>, ForwardedStreamlocalParams),
     X11(Channel<Msg>, String, u32),
     Disconnect,
 }
@@ -146,6 +147,20 @@ impl russh::client::Handler for ClientHandler {
         ));
         Ok(())
     }
+
+    async fn server_channel_open_forwarded_streamlocal(
+        &mut self,
+        channel: Channel<Msg>,
+        socket_path: &str,
+        _session: &mut Session,
+    ) -> Result<(), Self::Error> {
+        let socket_path = socket_path.to_string();
+        let _ = self.event_tx.send(ClientHandlerEvent::ForwardedStreamlocal(
+            channel,
+            ForwardedStreamlocalParams { socket_path },
+        ));
+        Ok(())
+    }
 }
 
 impl Drop for ClientHandler {
diff --git a/warpgate-protocol-ssh/src/client/mod.rs b/warpgate-protocol-ssh/src/client/mod.rs
@@ -29,7 +29,7 @@ use warpgate_core::Services;
 use self::handler::ClientHandlerEvent;
 use super::{ChannelOperation, DirectTCPIPParams};
 use crate::client::handler::ClientHandlerError;
-use crate::{load_all_usable_private_keys, ForwardedTcpIpParams};
+use crate::{load_all_usable_private_keys, ForwardedStreamlocalParams, ForwardedTcpIpParams};
 
 #[derive(Debug, thiserror::Error)]
 pub enum ConnectionError {
@@ -91,6 +91,7 @@ pub enum RCEvent {
     HostKeyReceived(PublicKey),
     HostKeyUnknown(PublicKey, oneshot::Sender<bool>),
     ForwardedTcpIp(Uuid, ForwardedTcpIpParams),
+    ForwardedStreamlocal(Uuid, ForwardedStreamlocalParams),
     X11(Uuid, String, u32),
 }
 
@@ -102,6 +103,8 @@ pub enum RCCommand {
     Channel(Uuid, ChannelOperation),
     ForwardTCPIP(String, u32),
     CancelTCPIPForward(String, u32),
+    StreamlocalForward(String),
+    CancelStreamlocalForward(String),
     Disconnect,
 }
 
@@ -126,6 +129,7 @@ pub struct RemoteClient {
     channel_pipes: Arc<Mutex<HashMap<Uuid, UnboundedSender<ChannelOperation>>>>,
     pending_ops: Vec<(Uuid, ChannelOperation)>,
     pending_forwards: Vec<(String, u32)>,
+    pending_streamlocal_forwards: Vec<String>,
     state: RCState,
     abort_rx: UnboundedReceiver<()>,
     inner_event_rx: UnboundedReceiver<InnerEvent>,
@@ -155,6 +159,7 @@ impl RemoteClient {
             channel_pipes: Arc::new(Mutex::new(HashMap::new())),
             pending_ops: vec![],
             pending_forwards: vec![],
+            pending_streamlocal_forwards: vec![],
             state: RCState::NotInitialized,
             inner_event_rx,
             inner_event_tx: inner_event_tx.clone(),
@@ -309,6 +314,11 @@ impl RemoteClient {
                         let id = self.setup_server_initiated_channel(channel).await?;
                         let _ = self.tx.send(RCEvent::ForwardedTcpIp(id, params));
                     }
+                    ClientHandlerEvent::ForwardedStreamlocal(channel, params) => {
+                        info!("New forwarded socket connection: {params:?}");
+                        let id = self.setup_server_initiated_channel(channel).await?;
+                        let _ = self.tx.send(RCEvent::ForwardedStreamlocal(id, params));
+                    }
                     ClientHandlerEvent::X11(channel, originator_address, originator_port) => {
                         info!("New X11 connection from {originator_address}:{originator_port:?}");
                         let id = self.setup_server_initiated_channel(channel).await?;
@@ -355,10 +365,19 @@ impl RemoteClient {
                     for (id, op) in ops {
                         self.apply_channel_op(id, op).await?;
                     }
+
                     let forwards = self.pending_forwards.drain(..).collect::<Vec<_>>();
                     for (address, port) in forwards {
                         self.tcpip_forward(address, port).await?;
                     }
+
+                    let forwards = self
+                        .pending_streamlocal_forwards
+                        .drain(..)
+                        .collect::<Vec<_>>();
+                    for socket_path in forwards {
+                        self.streamlocal_forward(socket_path).await?;
+                    }
                 }
                 Err(e) => {
                     debug!("Connect error: {}", e);
@@ -376,6 +395,12 @@ impl RemoteClient {
             RCCommand::CancelTCPIPForward(address, port) => {
                 self.cancel_tcpip_forward(address, port).await?;
             }
+            RCCommand::StreamlocalForward(socket_path) => {
+                self.streamlocal_forward(socket_path).await?;
+            }
+            RCCommand::CancelStreamlocalForward(socket_path) => {
+                self.cancel_streamlocal_forward(socket_path).await?;
+            }
             RCCommand::Disconnect => {
                 self.disconnect().await;
                 return Ok(true);
@@ -635,6 +660,30 @@ impl RemoteClient {
         Ok(())
     }
 
+    async fn streamlocal_forward(&mut self, socket_path: String) -> Result<(), SshClientError> {
+        if let Some(session) = &self.session {
+            let mut session = session.lock().await;
+            session.streamlocal_forward(socket_path).await?;
+        } else {
+            self.pending_streamlocal_forwards.push(socket_path);
+        }
+        Ok(())
+    }
+
+    async fn cancel_streamlocal_forward(
+        &mut self,
+        socket_path: String,
+    ) -> Result<(), SshClientError> {
+        if let Some(session) = &self.session {
+            let session = session.lock().await;
+            session.cancel_streamlocal_forward(socket_path).await?;
+        } else {
+            self.pending_streamlocal_forwards
+                .retain(|x| x != &socket_path);
+        }
+        Ok(())
+    }
+
     async fn disconnect(&mut self) {
         if let Some(session) = &mut self.session {
             let _ = session
diff --git a/warpgate-protocol-ssh/src/common.rs b/warpgate-protocol-ssh/src/common.rs
@@ -38,6 +38,11 @@ pub struct ForwardedTcpIpParams {
     pub originator_port: u32,
 }
 
+#[derive(Clone, Debug)]
+pub struct ForwardedStreamlocalParams {
+    pub socket_path: String,
+}
+
 #[derive(Clone, Debug)]
 pub struct X11Request {
     pub single_conection: bool,
diff --git a/warpgate-protocol-ssh/src/server/russh_handler.rs b/warpgate-protocol-ssh/src/server/russh_handler.rs
@@ -47,6 +47,8 @@ pub enum ServerHandlerEvent {
     X11Request(ServerChannelId, X11Request, oneshot::Sender<()>),
     TcpIpForward(String, u32, oneshot::Sender<bool>),
     CancelTcpIpForward(String, u32, oneshot::Sender<bool>),
+    StreamlocalForward(String, oneshot::Sender<bool>),
+    CancelStreamlocalForward(String, oneshot::Sender<bool>),
     Disconnect,
 }
 
@@ -477,6 +479,43 @@ impl russh::server::Handler for ServerHandler {
         }
         Ok(allowed)
     }
+
+    async fn streamlocal_forward(
+        &mut self,
+        socket_path: &str,
+        session: &mut Session,
+    ) -> Result<bool, Self::Error> {
+        let socket_path = socket_path.to_string();
+        let (tx, rx) = oneshot::channel();
+        self.send_event(ServerHandlerEvent::StreamlocalForward(socket_path, tx))?;
+        let allowed = rx.await.unwrap_or(false);
+        if allowed {
+            session.request_success()
+        } else {
+            session.request_failure()
+        }
+        Ok(allowed)
+    }
+
+    async fn cancel_streamlocal_forward(
+        &mut self,
+        socket_path: &str,
+        session: &mut Session,
+    ) -> Result<bool, Self::Error> {
+        let socket_path = socket_path.to_string();
+        let (tx, rx) = oneshot::channel();
+        self.send_event(ServerHandlerEvent::CancelStreamlocalForward(
+            socket_path,
+            tx,
+        ))?;
+        let allowed = rx.await.unwrap_or(false);
+        if allowed {
+            session.request_success()
+        } else {
+            session.request_failure()
+        }
+        Ok(allowed)
+    }
 }
 
 impl Drop for ServerHandler {
diff --git a/warpgate-protocol-ssh/src/server/session.rs b/warpgate-protocol-ssh/src/server/session.rs
