Support TCP port exposure via "--expose" command-line argument (#70)

lib/proxy/exposed_port.rs

@@ -0,0 +1,47 @@
+use anyhow::{anyhow, Context, Error};
+use std::str::FromStr;
+
+#[derive(Debug, Clone, Copy, Default)]
+pub struct ExposedPort {
+    pub external_port: u16,
+    pub internal_port: u16,
+}
+
+impl FromStr for ExposedPort {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let splits: Vec<&str> = s.split(':').collect();
+
+        match splits.len() {
+            2 => Ok(ExposedPort {
+                external_port: splits[0]
+                    .parse()
+                    .context(format!("invalid external port {:?}", splits[0]))?,
+                internal_port: splits[1]
+                    .parse()
+                    .context(format!("invalid internal port {:?}", splits[1]))?,
+            }),
+            _ => Err(anyhow!(
+                "invalid exposed port specification {:?}, the format should be EXTERNAL:INTERNAL",
+                s
+            )),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::proxy::exposed_port::ExposedPort;
+
+    #[test]
+    fn exposed_port() {
+        assert_eq!(
+            ExposedPort {
+                external_port: 2222,
+                internal_port: 22
+            },
+            "2222:22".parse().unwrap()
+        );
+    }
+}

lib/proxy/mod.rs

@@ -1,4 +1,6 @@
+mod exposed_port;
 mod host;
+mod port_forwarder;
 mod udp_packet_helper;
 mod vm;
 
@@ -8,8 +10,10 @@ use crate::host::NetType;
 use crate::poller::Poller;
 use crate::vm::VM;
 use anyhow::Result;
+pub use exposed_port::ExposedPort;
 use ipnet::Ipv4Net;
 use mac_address::MacAddress;
+use port_forwarder::PortForwarder;
 use prefix_trie::{Prefix, PrefixSet};
 use smoltcp::wire::EthernetFrame;
 use std::io::ErrorKind;
@@ -23,6 +27,7 @@ pub struct Proxy<'proxy> {
     dhcp_snooper: DhcpSnooper,
     allow: PrefixSet<Ipv4Net>,
     enobufs_encountered: bool,
+    port_forwarder: PortForwarder,
 }
 
 impl Proxy<'_> {
@@ -31,6 +36,7 @@ impl Proxy<'_> {
         vm_mac_address: MacAddress,
         vm_net_type: NetType,
         allow: PrefixSet<Ipv4Net>,
+        exposed_ports: Vec<ExposedPort>,
     ) -> Result<Proxy<'proxy>> {
         let vm = VM::new(vm_fd)?;
         let host = Host::new(vm_net_type, !allow.contains(&Ipv4Net::zero()))?;
@@ -44,6 +50,7 @@ impl Proxy<'_> {
             dhcp_snooper: Default::default(),
             allow,
             enobufs_encountered: false,
+            port_forwarder: PortForwarder::new(exposed_ports),
         })
     }
 
@@ -68,6 +75,12 @@ impl Proxy<'_> {
                 return Ok(());
             }
 
+            // Timeout
+            if !vm_readable && !host_readable && !interrupt {
+                self.port_forwarder
+                    .tick(&mut self.host, self.dhcp_snooper.lease());
+            }
+
             self.poller.rearm()?;
         }
     }

lib/proxy/port_forwarder.rs

@@ -0,0 +1,99 @@
+use crate::dhcp_snooper::Lease;
+use crate::host::Host;
+use crate::proxy::exposed_port::ExposedPort;
+use anyhow::Result;
+use log::error;
+use std::net::Ipv4Addr;
+
+#[derive(Default)]
+pub struct PortForwarder {
+    port_forwardings: Vec<PortForwarding>,
+    failed: bool,
+}
+
+#[derive(Debug, Clone, Copy, Default)]
+struct PortForwarding {
+    exposed_port: ExposedPort,
+    forwarding_to_addr: Option<Ipv4Addr>,
+}
+
+impl PortForwarder {
+    pub fn new(exposed_ports: Vec<ExposedPort>) -> PortForwarder {
+        let port_forwardings = exposed_ports
+            .into_iter()
+            .map(|exposed_port| PortForwarding {
+                exposed_port,
+                ..Default::default()
+            })
+            .collect();
+
+        PortForwarder {
+            port_forwardings,
+            ..Default::default()
+        }
+    }
+
+    pub fn tick(&mut self, host: &mut Host, lease: &Option<Lease>) {
+        if self.failed {
+            return;
+        }
+
+        if let Err(err) = self.tick_inner(host, lease) {
+            error!("port-forwarding failed: {}", err);
+
+            self.failed = true;
+        }
+    }
+
+    fn tick_inner(&mut self, host: &mut Host, lease: &Option<Lease>) -> Result<()> {
+        if let Some(lease) = lease {
+            // Lease exists, but is not valid, remove all port forwardings
+            if !lease.valid() {
+                self.remove_all_port_forwardings(host)?;
+
+                return Ok(());
+            }
+
+            // Lease exists and is valid, install/re-install port forwardings
+            for port_forwarding in &mut self.port_forwardings {
+                if let Some(installed_addr) = port_forwarding.forwarding_to_addr {
+                    // Port forwarding already installed, perhaps it's outdated?
+                    if installed_addr == lease.address() {
+                        // Nope, the port forwarding is up to date
+                        continue;
+                    }
+
+                    // Remove port forwarding since the lease address had changed
+                    host.port_forwarding_remove_rule(port_forwarding.exposed_port.external_port)?;
+                    port_forwarding.forwarding_to_addr = None;
+                }
+
+                // Install new port forwarding
+                host.port_forwarding_add_rule(
+                    port_forwarding.exposed_port.external_port,
+                    lease.address(),
+                    port_forwarding.exposed_port.internal_port,
+                )?;
+                port_forwarding.forwarding_to_addr = Some(lease.address());
+            }
+        } else {
+            // Lease does not exist, remove all port forwardings
+            self.remove_all_port_forwardings(host)?;
+        }
+
+        Ok(())
+    }
+
+    fn remove_all_port_forwardings(&mut self, host: &mut Host) -> Result<()> {
+        for port_forwarding in &mut self.port_forwardings {
+            if port_forwarding.forwarding_to_addr.is_none() {
+                continue;
+            }
+
+            host.port_forwarding_remove_rule(port_forwarding.exposed_port.external_port)?;
+            port_forwarding.forwarding_to_addr = None;
+        }
+
+        Ok(())
+    }
+}