Add tun mode on Linux (#187)

* fix module names for classes in submodules

* TunnelInfo::Udp -> TunnelInfo::None

* add `mitmproxy_rs.tun`

Cargo.toml

@@ -56,6 +56,7 @@ socket2 = "0.5.7"
 # tokio = { path = "../tokio/tokio" }
 smoltcp = { git = 'https://github.com/smoltcp-rs/smoltcp', rev = 'ef67e7b46cabf49783053cbf68d8671ed97ff8d4' }
 boringtun = { git = 'https://github.com/cloudflare/boringtun', rev = 'e3252d9c4f4c8fc628995330f45369effd4660a1' }
+# tun2 = { path = "../rust-tun" }
 
 [target.'cfg(windows)'.dependencies.windows]
 version = "0.57.0"
@@ -82,6 +83,9 @@ cocoa = "0.26"
 objc = "0.2"
 sysinfo = "0.29.10"
 
+[target.'cfg(target_os = "linux")'.dependencies]
+tun2 = { version = "3.1.8", features = ["async"] }
+
 [dev-dependencies]
 env_logger = "0.11"
 rand = "0.8"

mitmproxy-rs/mitmproxy_rs/__init__.pyi

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import Any, Literal
 from typing import final, overload, TypeVar
-from . import certs, dns, local, process_info, udp, wireguard
+from . import certs, dns, local, process_info, tun, udp, wireguard
 
 T = TypeVar("T")
 
@@ -60,6 +60,7 @@ __all__ = [
     "dns",
     "local",
     "process_info",
+    "tun",
     "udp",
     "wireguard",
     "Stream",

mitmproxy-rs/mitmproxy_rs/tun.pyi

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import final
+from . import Stream
+
+async def create_tun_interface(
+        handle_tcp_stream: Callable[[Stream], Awaitable[None]],
+        handle_udp_stream: Callable[[Stream], Awaitable[None]],
+        tun_name: str | None = None,
+) -> TunInterface: ...
+@final
+class TunInterface:
+    def tun_name(self) -> str: ...
+    def close(self) -> None: ...
+    async def wait_closed(self) -> None: ...
+    def __repr__(self) -> str: ...

mitmproxy-rs/pytests/test_task.rs

@@ -82,7 +82,7 @@ mod tests {
                 connection_id: ConnectionId::unassigned_udp(),
                 src_addr: "127.0.0.1:51232".parse()?,
                 dst_addr: "127.0.0.1:53".parse()?,
-                tunnel_info: TunnelInfo::Udp,
+                tunnel_info: TunnelInfo::None,
                 command_tx: None,
             })
             .await

mitmproxy-rs/src/lib.rs

@@ -59,6 +59,12 @@ mod mitmproxy_rs {
         use crate::process_info::{active_executables, executable_icon, Process};
     }
 
+    #[pymodule]
+    mod tun {
+        #[pymodule_export]
+        use crate::server::{create_tun_interface, TunInterface};
+    }
+
     #[pymodule]
     mod udp {
         #[pymodule_export]

mitmproxy-rs/src/process_info.rs

@@ -6,7 +6,7 @@ use pyo3::prelude::*;
 #[cfg(any(windows, target_os = "macos"))]
 use mitmproxy::processes;
 
-#[pyclass(module = "mitmproxy_rs", frozen)]
+#[pyclass(module = "mitmproxy_rs.process_info", frozen)]
 pub struct Process(mitmproxy::processes::ProcessInfo);
 
 #[pymethods]

mitmproxy-rs/src/server/local_redirector.rs

@@ -14,7 +14,7 @@ use std::path::PathBuf;
 use crate::server::base::Server;
 use tokio::sync::mpsc;
 
-#[pyclass(module = "mitmproxy_rs")]
+#[pyclass(module = "mitmproxy_rs.local")]
 #[derive(Debug)]
 pub struct LocalRedirector {
     server: Server,

mitmproxy-rs/src/server/mod.rs

@@ -1,8 +1,10 @@
 mod base;
 mod local_redirector;
+mod tun;
 mod udp;
 mod wireguard;
 
 pub use local_redirector::{start_local_redirector, LocalRedirector};
+pub use tun::{create_tun_interface, TunInterface};
 pub use udp::{start_udp_server, UdpServer};
 pub use wireguard::{start_wireguard_server, WireGuardServer};

mitmproxy-rs/src/server/tun.rs

@@ -0,0 +1,63 @@
+use crate::server::base::Server;
+use pyo3::prelude::*;
+
+/// An open TUN interface.
+///
+/// A new tun interface can be created by calling `create_tun_interface`.
+#[pyclass(module = "mitmproxy_rs.tun")]
+#[derive(Debug)]
+pub struct TunInterface {
+    tun_name: String,
+    server: Server,
+}
+
+#[pymethods]
+impl TunInterface {
+    /// Get the tunnel interface name.
+    pub fn tun_name(&self) -> &str {
+        &self.tun_name
+    }
+
+    /// Request the interface to be closed.
+    pub fn close(&mut self) {
+        self.server.close()
+    }
+
+    /// Wait until the interface has shut down.
+    pub fn wait_closed<'p>(&self, py: Python<'p>) -> PyResult<Bound<'p, PyAny>> {
+        self.server.wait_closed(py)
+    }
+
+    pub fn __repr__(&self) -> String {
+        format!("TunInterface({})", self.tun_name)
+    }
+}
+
+/// Create a TUN interface that is configured with the given parameters:
+///
+/// - `handle_tcp_stream`: An async function that will be called for each new TCP `Stream`.
+/// - `handle_udp_stream`: An async function that will be called for each new UDP `Stream`.
+/// - `tun_name`: An optional string to specify the tunnel name. By default, tun0, ... will be used.
+///
+/// *Availability: Linux*
+#[pyfunction]
+pub fn create_tun_interface(
+    py: Python<'_>,
+    handle_tcp_stream: PyObject,
+    handle_udp_stream: PyObject,
+    tun_name: Option<String>,
+) -> PyResult<Bound<PyAny>> {
+    #[cfg(target_os = "linux")]
+    {
+        let conf = mitmproxy::packet_sources::tun::TunConf { tun_name };
+        pyo3_asyncio_0_21::tokio::future_into_py(py, async move {
+            let (server, tun_name) =
+                Server::init(conf, handle_tcp_stream, handle_udp_stream).await?;
+            Ok(TunInterface { server, tun_name })
+        })
+    }
+    #[cfg(not(target_os = "linux"))]
+    Err(pyo3::exceptions::PyNotImplementedError::new_err(
+        "TUN proxy mode is only available on Linux",
+    ))
+}

mitmproxy-rs/src/server/udp.rs

@@ -14,7 +14,7 @@ use crate::util::socketaddr_to_py;
 /// The public API is intended to be similar to the API provided by
 /// [`asyncio.Server`](https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.Server)
 /// from the Python standard library.
-#[pyclass(module = "mitmproxy_rs")]
+#[pyclass(module = "mitmproxy_rs.udp")]
 #[derive(Debug)]
 pub struct UdpServer {
     /// local address of the UDP socket

mitmproxy-rs/src/server/wireguard.rs

@@ -16,7 +16,7 @@ use crate::server::base::Server;
 /// The public API is intended to be similar to the API provided by
 /// [`asyncio.Server`](https://docs.python.org/3/library/asyncio-eventloop.html#asyncio.Server)
 /// from the Python standard library.
-#[pyclass(module = "mitmproxy_rs")]
+#[pyclass(module = "mitmproxy_rs.wireguard")]
 #[derive(Debug)]
 pub struct WireGuardServer {
     /// local address of the WireGuard UDP socket

mitmproxy-rs/src/stream.rs

@@ -194,7 +194,7 @@ impl Stream {
                 }
                 _ => (),
             },
-            TunnelInfo::Udp {} => (),
+            TunnelInfo::None {} => (),
         }
         match default {
             Some(x) => Ok(x),

mitmproxy-rs/src/udp_client.rs

@@ -50,7 +50,7 @@ pub fn open_udp_connection(
             command_tx,
             peername,
             sockname,
-            tunnel_info: TunnelInfo::Udp,
+            tunnel_info: TunnelInfo::None,
         };
 
         Ok(stream)

src/messages.rs

@@ -19,7 +19,7 @@ pub enum TunnelInfo {
         /// an unresolved remote_endpoint instead.
         remote_endpoint: Option<(String, u16)>,
     },
-    Udp,
+    None,
 }
 
 /// Events that are sent by WireGuard to the TCP stack.

src/packet_sources/mod.rs

@@ -6,6 +6,8 @@ use crate::messages::{TransportCommand, TransportEvent};
 
 #[cfg(target_os = "macos")]
 pub mod macos;
+#[cfg(target_os = "linux")]
+pub mod tun;
 pub mod udp;
 #[cfg(windows)]
 pub mod windows;

src/packet_sources/tun.rs

@@ -0,0 +1,113 @@
+use anyhow::{Context, Result};
+
+use crate::messages::{
+    NetworkCommand, NetworkEvent, SmolPacket, TransportCommand, TransportEvent, TunnelInfo,
+};
+use crate::network::{add_network_layer, MAX_PACKET_SIZE};
+use crate::packet_sources::{PacketSourceConf, PacketSourceTask};
+use tokio::sync::mpsc::{Permit, Receiver, UnboundedReceiver};
+use tokio::sync::{broadcast, mpsc::Sender};
+use tun2::AbstractDevice;
+
+pub struct TunConf {
+    pub tun_name: Option<String>,
+}
+
+impl PacketSourceConf for TunConf {
+    type Task = TunTask;
+    type Data = String;
+
+    fn name(&self) -> &'static str {
+        "TUN interface"
+    }
+
+    async fn build(
+        self,
+        transport_events_tx: Sender<TransportEvent>,
+        transport_commands_rx: UnboundedReceiver<TransportCommand>,
+        shutdown: broadcast::Receiver<()>,
+    ) -> Result<(Self::Task, Self::Data)> {
+        let mut config = tun2::Configuration::default();
+        config.mtu(MAX_PACKET_SIZE as u16);
+        config.up();
+        if let Some(tun_name) = self.tun_name {
+            config.tun_name(&tun_name);
+        }
+
+        let device = tun2::create_as_async(&config).context("Failed to create TUN device")?;
+        let tun_name = device.tun_name().context("Failed to get TUN name")?;
+
+        let (network_task_handle, net_tx, net_rx) =
+            add_network_layer(transport_events_tx, transport_commands_rx, shutdown)?;
+
+        Ok((
+            TunTask {
+                device,
+                net_tx,
+                net_rx,
+                network_task_handle,
+            },
+            tun_name,
+        ))
+    }
+}
+
+pub struct TunTask {
+    device: tun2::AsyncDevice,
+
+    net_tx: Sender<NetworkEvent>,
+    net_rx: Receiver<NetworkCommand>,
+    network_task_handle: tokio::task::JoinHandle<Result<()>>,
+}
+
+impl PacketSourceTask for TunTask {
+    async fn run(mut self) -> Result<()> {
+        let size = self.device.mtu()? as usize + tun2::PACKET_INFORMATION_LENGTH;
+        let mut buf = vec![0; size];
+
+        let mut packet_to_send = Vec::new();
+        let mut permit: Option<Permit<NetworkEvent>> = None;
+
+        // Required on macOS, but currently crashes on Linux with tokio.
+        //let (mut writer, mut reader) = self.device.split().context("failed to split device")?;
+
+        loop {
+            tokio::select! {
+                // Monitor the network task for errors or planned shutdown.
+                // This way we implicitly monitor the shutdown broadcast channel.
+                exit = &mut self.network_task_handle => break exit.context("network task panic")?.context("network task error")?,
+                // wait for transport_events_tx channel capacity...
+                Ok(p) = self.net_tx.reserve(), if permit.is_none() => {
+                    permit = Some(p);
+                },
+                // ... or process incoming packets
+                r = self.device.recv(buf.as_mut_slice()), if permit.is_some() => {
+                    let len = r.context("TUN read() failed")?;
+
+                    let Ok(packet) = SmolPacket::try_from(buf[..len].to_vec()) else {
+                        log::error!("Skipping invalid packet from tun interface: {:?}", &buf[..len]);
+                        continue;
+                    };
+                    permit.take().unwrap().send(NetworkEvent::ReceivePacket {
+                        packet,
+                        tunnel_info: TunnelInfo::None,
+                    });
+                },
+                // send_to is cancel safe, so we can use that for backpressure.
+                r = self.device.send(&packet_to_send), if !packet_to_send.is_empty() => {
+                    r.context("TUN write() failed")?;
+                    packet_to_send.clear();
+                },
+                Some(command) = self.net_rx.recv(), if packet_to_send.is_empty() => {
+                    match command {
+                        NetworkCommand::SendPacket(packet) => {
+                            packet_to_send = packet.into_inner();
+                        }
+                    }
+                }
+            }
+        }
+        log::debug!("TUN interface task shutting down.");
+        Ok(())
+    }
+}

src/packet_sources/udp.rs

@@ -131,7 +131,7 @@ impl PacketSourceTask for UdpTask {
                             dst_addr: self.local_addr,
                             payload: udp_buf[..len].to_vec(),
                         },
-                        TunnelInfo::Udp {},
+                        TunnelInfo::None {},
                         permit.take().unwrap()
                     );
                 },