UDP Streams

Cargo.toml

@@ -12,7 +12,7 @@ authors = [
     "Fabio Valentini <[email protected]>",
     "Maximilian Hils <[email protected]>",
 ]
-version = "0.4.0"
+version = "0.5.0"
 publish = false
 repository = "https://github.com/mitmproxy/mitmproxy-rs"
 edition = "2021"
@@ -44,6 +44,8 @@ image = "0.24.6"
 prost = "0.12.3"
 tokio-util = { version = "0.7.10", features = ["codec"] }
 futures-util = { version = "0.3.29", features = ["sink"] }
+lru_time_cache = "0.11.11"
+internet-packet = { version = "0.2.0", features = ["smoltcp"] }
 
 # [patch.crates-io]
 # tokio = { path = "../tokio/tokio" }

mitmproxy-macos/redirector/network-extension/TransparentProxyProvider.swift

@@ -157,7 +157,7 @@ class TransparentProxyProvider: NETransparentProxyProvider {
             guard let remoteEndpoint = tcp_flow.remoteEndpoint as? NWHostEndpoint else {
                 throw TransparentProxyError.noRemoteEndpoint
             }
-            log.debug("remoteEndpoint: \(String(describing: remoteEndpoint), privacy: .public)")
+            // log.debug("remoteEndpoint: \(String(describing: remoteEndpoint), privacy: .public)")
             // It would be nice if we could also include info on the local endpoint here, but that's not exposed.
             message = MitmproxyIpc_NewFlow.with {
                 $0.tcp = MitmproxyIpc_TcpFlow.with {

mitmproxy-rs/mitmproxy_rs/__init__.pyi

@@ -1,36 +1,39 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Awaitable, Callable, Any
-from typing import final
+from typing import Awaitable, Callable, Any, Literal
+from typing import final, overload, TypeVar
 
 
+T = TypeVar("T")
+
 # WireGuard
 
 async def start_wireguard_server(
     host: str,
     port: int,
     private_key: str,
     peer_public_keys: list[str],
-    handle_connection: Callable[[TcpStream], Awaitable[None]],
-    receive_datagram: Callable[[DatagramTransport, bytes, tuple[str, int], tuple[str, int]], None],
+    handle_tcp_stream: Callable[[Stream], Awaitable[None]],
+    handle_udp_stream: Callable[[Stream], Awaitable[None]],
 ) -> WireGuardServer: ...
 
 @final
 class WireGuardServer:
     def getsockname(self) -> tuple[str, int]: ...
     def close(self) -> None: ...
     async def wait_closed(self) -> None: ...
+    def __repr__(self) -> str: ...
 
 def genkey() -> str: ...
 def pubkey(private_key: str) -> str: ...
 
 
-# Windows
+# Local Redirector
 
 async def start_local_redirector(
-    handle_connection: Callable[[TcpStream], Awaitable[None]],
-    receive_datagram: Callable[[DatagramTransport, bytes, tuple[str, int], tuple[str, int]], None],
+    handle_tcp_stream: Callable[[Stream], Awaitable[None]],
+    handle_udp_stream: Callable[[Stream], Awaitable[None]],
 ) -> LocalRedirector: ...
 
 @final
@@ -41,14 +44,33 @@ class LocalRedirector:
     def close(self) -> None: ...
     async def wait_closed(self) -> None: ...
 
-# MacOS
-def add_cert(pem: str) -> None: ...
-def remove_cert() -> None: ...
+
+# UDP
+
+async def start_udp_server(
+    host: str,
+    port: int,
+    handle_udp_stream: Callable[[Stream], Awaitable[None]],
+) -> UdpServer: ...
+
+@final
+class UdpServer:
+    def getsockname(self) -> tuple[str, int]: ...
+    def close(self) -> None: ...
+    async def wait_closed(self) -> None: ...
+    def __repr__(self) -> str: ...
+
+async def open_udp_connection(
+    host: str,
+    port: int,
+    *,
+    local_addr: tuple[str, int] | None = None,
+) -> Stream: ...
 
 # TCP / UDP
 
 @final
-class TcpStream:
+class Stream:
     async def read(self, n: int) -> bytes: ...
     def write(self, data: bytes): ...
     async def drain(self) -> None: ...
@@ -58,28 +80,36 @@ class TcpStream:
     def is_closing(self) -> bool: ...
     async def wait_closed(self) -> None: ...
 
-    def get_extra_info(self, name: str, default: Any = None) -> Any: ...
-    def __repr__(self) -> str: ...
-
+    @overload
+    def get_extra_info(self, name: Literal["transport_protocol"], default: None = None) -> Literal["tcp", "udp"]: ...
+    @overload
+    def get_extra_info(self, name: Literal["transport_protocol"], default: T) -> Literal["tcp", "udp"] | T: ...
+    @overload
+    def get_extra_info(self, name: Literal["peername", "sockname", "original_src", "original_dst", "remote_endpoint"], default: None = None) -> tuple[str, int]: ...
+    @overload
+    def get_extra_info(self, name: Literal["peername", "sockname", "original_src", "original_dst", "remote_endpoint"], default: T) -> tuple[str, int] | T: ...
+    @overload
+    def get_extra_info(self, name: Literal["pid"], default: None = None) -> int: ...
+    @overload
+    def get_extra_info(self, name: Literal["pid"], default: T) -> int | T: ...
+    @overload
+    def get_extra_info(self, name: Literal["process_name"], default: None = None) -> str: ...
+    @overload
+    def get_extra_info(self, name: Literal["process_name"], default: T) -> str | T: ...
+    @overload
+    def get_extra_info(self, name: str, default: Any) -> Any: ...
 
-@final
-class DatagramTransport:
-    def sendto(self, data: bytes, addr: tuple[str, int] | None = None): ...
-    async def drain(self) -> None: ...
+    def __repr__(self) -> str: ...
 
-    def close(self): ...
-    def is_closing(self) -> bool: ...
-    async def wait_closed(self) -> None: ...
 
-    def get_extra_info(self, name: str, default: Any = None) -> Any: ...
-    def __repr__(self) -> str: ...
+# Certificate Installation
 
-    def get_protocol(self) -> DatagramTransport: ...
+def add_cert(pem: str) -> None: ...
+def remove_cert() -> None: ...
 
 
 # Process Info
 
-
 def active_executables() -> list[Process]: ...
 def executable_icon(path: Path | str) -> bytes: ...
 

mitmproxy-rs/src/lib.rs

@@ -5,11 +5,11 @@ use std::sync::RwLock;
 use once_cell::sync::Lazy;
 use pyo3::{exceptions::PyException, prelude::*};
 
-mod datagram_transport;
 mod process_info;
 mod server;
+mod stream;
 mod task;
-mod tcp_stream;
+mod udp_client;
 mod util;
 
 static LOGGER_INITIALIZED: Lazy<RwLock<bool>> = Lazy::new(|| RwLock::new(false));
@@ -50,12 +50,16 @@ pub fn mitmproxy_rs(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(server::start_local_redirector, m)?)?;
     m.add_class::<server::LocalRedirector>()?;
 
+    m.add_function(wrap_pyfunction!(server::start_udp_server, m)?)?;
+    m.add_class::<server::UdpServer>()?;
+
+    m.add_function(wrap_pyfunction!(udp_client::open_udp_connection, m)?)?;
+
     m.add_function(wrap_pyfunction!(process_info::active_executables, m)?)?;
     m.add_class::<process_info::Process>()?;
     m.add_function(wrap_pyfunction!(process_info::executable_icon, m)?)?;
 
-    m.add_class::<tcp_stream::TcpStream>()?;
-    m.add_class::<datagram_transport::DatagramTransport>()?;
+    m.add_class::<stream::Stream>()?;
 
     // Import platform-specific modules here so that missing dependencies are raising immediately.
     #[cfg(target_os = "macos")]