WIP on UDP and multi-port-forward support

Cargo.lock

@@ -488,6 +488,12 @@ version = "2.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "308cc39be01b73d0d18f82a0e7b2a3df85245f84af96fdddc5d202d27e47b86a"
 
+[[package]]
+name = "minimal-lexical"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9c64630dcdd71f1a64c435f54885086a0de5d6a12d104d69b165fb7d5286d677"
+
 [[package]]
 name = "miniz_oxide"
 version = "0.4.4"
@@ -520,6 +526,17 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "nom"
+version = "7.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ffd9d26838a953b4af82cbeb9f1592c6798916983959be223a7124e992742c1"
+dependencies = [
+ "memchr",
+ "minimal-lexical",
+ "version_check",
+]
+
 [[package]]
 name = "ntapi"
 version = "0.3.6"
@@ -583,6 +600,7 @@ dependencies = [
  "futures",
  "lockfree",
  "log",
+ "nom",
  "pretty_env_logger",
  "rand",
  "smoltcp",

Cargo.toml

@@ -16,3 +16,4 @@ tokio = { version = "1", features = ["full"] }
 lockfree = "0.5.1"
 futures = "0.3.17"
 rand = "0.8.4"
+nom = "7"

src/config.rs

@@ -1,3 +1,6 @@
+use std::collections::HashSet;
+use std::convert::TryFrom;
+use std::fmt::{Display, Formatter};
 use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
 use std::sync::Arc;
 
@@ -7,8 +10,7 @@ use clap::{App, Arg};
 
 #[derive(Clone, Debug)]
 pub struct Config {
-    pub(crate) source_addr: SocketAddr,
-    pub(crate) dest_addr: SocketAddr,
+    pub(crate) port_forwards: Vec<PortForwardConfig>,
     pub(crate) private_key: Arc<X25519SecretKey>,
     pub(crate) endpoint_public_key: Arc<X25519PublicKey>,
     pub(crate) endpoint_addr: SocketAddr,
@@ -23,16 +25,12 @@ impl Config {
             .author("Aram Peres <aram.peres@gmail.com>")
             .version(env!("CARGO_PKG_VERSION"))
             .args(&[
-                Arg::with_name("SOURCE_ADDR")
-                    .required(true)
+                Arg::with_name("PORT_FORWARD")
+                    .required(false)
+                    .multiple(true)
                     .takes_value(true)
-                    .env("ONETUN_SOURCE_ADDR")
-                    .help("The source address (IP + port) to forward from. Example: 127.0.0.1:2115"),
-                Arg::with_name("DESTINATION_ADDR")
-                    .required(true)
-                    .takes_value(true)
-                    .env("ONETUN_DESTINATION_ADDR")
-                    .help("The destination address (IP + port) to forward to. The IP should be a peer registered in the Wireguard endpoint. Example: 192.168.4.2:2116"),
+                    .help("Port forward configurations. The format of each argument is [src_host:]<src_port>:<dst_host>:<dst_port>[:TCP,UDP,...]. \
+                    Environment variables of the form 'ONETUN_PORT_FORWARD_[#]' are also accepted, where [#] starts at 1."),
                 Arg::with_name("private-key")
                     .required(true)
                     .takes_value(true)
@@ -72,11 +70,40 @@ impl Config {
                     .help("Configures the log level and format.")
             ]).get_matches();
 
+        // Combine `PORT_FORWARD` arg and `ONETUN_PORT_FORWARD_#` strings
+        let mut port_forward_strings = HashSet::new();
+        matches.values_of("PORT_FORWARD").map(|values| {
+            values
+                .into_iter()
+                .map(|v| port_forward_strings.insert(v.to_string()))
+                .map(|_| ())
+        });
+        for n in 1.. {
+            if let Ok(env) = std::env::var(format!("ONETUN_PORT_FORWARD_{}", n)) {
+                port_forward_strings.insert(env);
+            } else {
+                break;
+            }
+        }
+        if port_forward_strings.is_empty() {
+            return Err(anyhow::anyhow!("No port forward configurations given."));
+        }
+
+        // Parse `PORT_FORWARD` strings into `PortForwardConfig`
+        let port_forwards: Vec<anyhow::Result<Vec<PortForwardConfig>>> = port_forward_strings
+            .into_iter()
+            .map(|s| PortForwardConfig::from_str(&s))
+            .collect();
+        let port_forwards: anyhow::Result<Vec<Vec<PortForwardConfig>>> =
+            port_forwards.into_iter().collect();
+        let port_forwards: Vec<PortForwardConfig> = port_forwards
+            .with_context(|| "Failed to parse port forward config")?
+            .into_iter()
+            .flatten()
+            .collect();
+
         Ok(Self {
-            source_addr: parse_addr(matches.value_of("SOURCE_ADDR"))
-                .with_context(|| "Invalid source address")?,
-            dest_addr: parse_addr(matches.value_of("DESTINATION_ADDR"))
-                .with_context(|| "Invalid destination address")?,
+            port_forwards,
             private_key: Arc::new(
                 parse_private_key(matches.value_of("private-key"))
                     .with_context(|| "Invalid private key")?,
@@ -137,3 +164,89 @@ fn parse_keep_alive(s: Option<&str>) -> anyhow::Result<Option<u16>> {
         Ok(None)
     }
 }
+
+#[derive(Debug, Clone, Copy)]
+pub struct PortForwardConfig {
+    /// The source IP and port where the local server will run.
+    pub source: SocketAddr,
+    /// The destination IP and port to which traffic will be forwarded.
+    pub destination: SocketAddr,
+    /// The transport protocol to use for the port (Layer 4).
+    pub protocol: PortProtocol,
+}
+
+impl PortForwardConfig {
+    /// Converts a string representation into `PortForwardConfig`.
+    ///
+    /// Sample formats:
+    ///  - `127.0.0.1:8080:192.168.4.1:8081:TCP,UDP`
+    ///  - `127.0.0.1:8080:192.168.4.1:8081:TCP`
+    ///  - `0.0.0.0:8080:192.168.4.1:8081`
+    ///  - `[::1]:8080:192.168.4.1:8081`
+    ///  - `8080:192.168.4.1:8081`
+    ///  - `8080:192.168.4.1:8081:TCP`
+    ///
+    /// Implementation Notes:
+    ///  - The format is formalized as `[src_host:]<src_port>:<dst_host>:<dst_port>[:PROTO1,PROTO2,...]`
+    ///  - `src_host` is optional and defaults to `127.0.0.1`.
+    ///  - `src_host` and `dst_host` may be specified as IPv4, IPv6, or a FQDN to be resolved by DNS.
+    ///  - IPv6 addresses must be prefixed with `[` and suffixed with `]`. Example: `[::1]`.
+    ///  - Any `u16` is accepted as `src_port` and `dst_port`
+    ///  - Specifying protocols (`PROTO1,PROTO2,...`) is optional and defaults to `TCP`. Values must be separated by commas.
+    pub fn from_str<'a>(s: &'a str) -> anyhow::Result<Vec<PortForwardConfig>> {
+        use nom::branch::alt;
+        use nom::bytes::complete::{is_not, take_until, take_while};
+        use nom::character::complete::char;
+        use nom::combinator::opt;
+        use nom::multi::separated_list0;
+        use nom::sequence::{delimited, terminated};
+        use nom::IResult;
+
+        Err(anyhow::anyhow!("TODO"))
+    }
+}
+
+impl Display for PortForwardConfig {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}:{}", self.source, self.destination, self.protocol)
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum PortProtocol {
+    Tcp,
+    Udp,
+}
+
+impl TryFrom<&str> for PortProtocol {
+    type Error = anyhow::Error;
+
+    fn try_from(value: &str) -> anyhow::Result<Self> {
+        match value.to_uppercase().as_str() {
+            "TCP" => Ok(Self::Tcp),
+            "UDP" => Ok(Self::Udp),
+            _ => Err(anyhow::anyhow!("Invalid protocol specifier: {}", value)),
+        }
+    }
+}
+
+impl Display for PortProtocol {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "{}",
+            match self {
+                Self::Tcp => "TCP",
+                Self::Udp => "UDP",
+            }
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Tests the parsing of `PortForwardConfig`.
+    fn test_parse_port_forward_config() {}
+}

src/main.rs

@@ -12,7 +12,7 @@ use smoltcp::socket::{SocketSet, TcpSocket, TcpSocketBuffer, TcpState};
 use smoltcp::wire::{IpAddress, IpCidr};
 use tokio::net::{TcpListener, TcpStream};
 
-use crate::config::Config;
+use crate::config::{Config, PortForwardConfig, PortProtocol};
 use crate::port_pool::PortPool;
 use crate::virtual_device::VirtualIpDevice;
 use crate::wg::WireGuardTunnel;
@@ -54,26 +54,63 @@ async fn main() -> anyhow::Result<()> {
         tokio::spawn(async move { ip_sink::run_ip_sink_interface(wg).await });
     }
 
+    {
+        let port_forwards = config.port_forwards;
+        let source_peer_ip = config.source_peer_ip;
+
+        futures::future::try_join_all(
+            port_forwards
+                .into_iter()
+                .map(|pf| (pf, wg.clone(), port_pool.clone()))
+                .map(|(pf, wg, port_pool)| {
+                    tokio::spawn(async move {
+                        port_forward(pf, source_peer_ip, port_pool, wg)
+                            .await
+                            .with_context(|| format!("Port-forward failed: {})", pf))
+                    })
+                }),
+        )
+        .await
+        .with_context(|| "A port-forward instance failed.")
+        .map(|_| ())
+    }
+}
+
+async fn port_forward(
+    port_forward: PortForwardConfig,
+    source_peer_ip: IpAddr,
+    port_pool: Arc<PortPool>,
+    wg: Arc<WireGuardTunnel>,
+) -> anyhow::Result<()> {
     info!(
-        "Tunnelling [{}]->[{}] (via [{}] as peer {})",
-        &config.source_addr, &config.dest_addr, &config.endpoint_addr, &config.source_peer_ip
+        "Tunnelling {} [{}]->[{}] (via [{}] as peer {})",
+        port_forward.protocol,
+        port_forward.source,
+        port_forward.destination,
+        &wg.endpoint,
+        source_peer_ip
     );
 
+    match port_forward.protocol {
+        PortProtocol::Tcp => {
             tcp_proxy_server(
-        config.source_addr,
-        config.source_peer_ip,
-        config.dest_addr,
-        port_pool.clone(),
+                port_forward.source,
+                port_forward.destination,
+                source_peer_ip,
+                port_pool,
                 wg,
             )
             .await
+        }
+        PortProtocol::Udp => Err(anyhow::anyhow!("UDP isn't supported just yet.")),
+    }
 }
 
 /// Starts the server that listens on TCP connections.
 async fn tcp_proxy_server(
     listen_addr: SocketAddr,
-    source_peer_ip: IpAddr,
     dest_addr: SocketAddr,
+    source_peer_ip: IpAddr,
     port_pool: Arc<PortPool>,
     wg: Arc<WireGuardTunnel>,
 ) -> anyhow::Result<()> {

src/wg.rs

@@ -24,7 +24,7 @@ pub struct WireGuardTunnel {
     /// The UDP socket for the public WireGuard endpoint to connect to.
     udp: UdpSocket,
     /// The address of the public WireGuard endpoint (UDP).
-    endpoint: SocketAddr,
+    pub(crate) endpoint: SocketAddr,
     /// Maps virtual ports to the corresponding IP packet dispatcher.
     virtual_port_ip_tx: lockfree::map::Map<u16, tokio::sync::mpsc::Sender<Vec<u8>>>,
     /// IP packet dispatcher for unroutable packets. `None` if not initialized.