-
-
Notifications
You must be signed in to change notification settings - Fork 34
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
9 changed files
with
235 additions
and
89 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,151 @@ | ||
| use anyhow::Context; | ||
| use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; | ||
|
|
||
| use anyhow::Result; | ||
| use pyo3::prelude::*; | ||
| use tokio::net::{lookup_host, UdpSocket}; | ||
| use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver}; | ||
| use tokio::sync::oneshot; | ||
|
|
||
| use crate::stream::{Stream, StreamState}; | ||
| use mitmproxy::messages::{ConnectionId, TransportCommand, TunnelInfo}; | ||
| use mitmproxy::MAX_PACKET_SIZE; | ||
|
|
||
| /// Start a UDP client that is configured with the given parameters: | ||
| /// | ||
| /// - `host`: The host address. | ||
| /// - `port`: The listen port. | ||
| /// - `local_addr`: The local address to bind to. | ||
| #[pyfunction] | ||
| #[pyo3(signature = (host, port, *, local_addr = None))] | ||
| pub fn open_udp_connection( | ||
| py: Python<'_>, | ||
| host: String, | ||
| port: u16, | ||
| local_addr: Option<(String, u16)>, | ||
| ) -> PyResult<&PyAny> { | ||
| pyo3_asyncio::tokio::future_into_py(py, async move { | ||
| let socket = udp_connect(host, port, local_addr).await?; | ||
|
|
||
| let peername = socket.peer_addr()?; | ||
| let sockname = socket.local_addr()?; | ||
|
|
||
| let (command_tx, command_rx) = unbounded_channel(); | ||
|
|
||
| tokio::spawn( | ||
| UdpClientTask { | ||
| socket, | ||
| transport_commands_rx: command_rx, | ||
| } | ||
| .run(), | ||
| ); | ||
|
|
||
| let stream = Stream { | ||
| connection_id: ConnectionId::unassigned(), | ||
| state: StreamState::Open, | ||
| command_tx, | ||
| peername, | ||
| sockname, | ||
| tunnel_info: TunnelInfo::Udp, | ||
| }; | ||
|
|
||
| Ok(stream) | ||
| }) | ||
| } | ||
|
|
||
| /// Open an UDP socket from bind_to to host:port. | ||
| /// This is a bit trickier than expected because we want to support IPv4 and IPv6. | ||
| async fn udp_connect( | ||
| host: String, | ||
| port: u16, | ||
| local_addr: Option<(String, u16)>, | ||
| ) -> Result<UdpSocket> { | ||
| let addrs: Vec<SocketAddr> = lookup_host((host.as_str(), port)) | ||
| .await | ||
| .with_context(|| format!("unable to resolve hostname: {}", host))? | ||
| .collect(); | ||
|
|
||
| if let Some((host, port)) = local_addr { | ||
| let socket = UdpSocket::bind((host.as_str(), port)) | ||
| .await | ||
| .with_context(|| format!("unable to bind to ({}, {})", host, port))?; | ||
| socket | ||
| .connect(addrs.as_slice()) | ||
| .await | ||
| .context("unable to connect to remote address")?; | ||
| Ok(socket) | ||
| } else { | ||
| if let Ok(socket) = | ||
| UdpSocket::bind(SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0)).await | ||
| { | ||
| if socket.connect(addrs.as_slice()).await.is_ok() { | ||
| return Ok(socket); | ||
| } | ||
| } | ||
| let socket = UdpSocket::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)) | ||
| .await | ||
| .context("unable to bind to 127.0.0.1:0")?; | ||
| socket | ||
| .connect(addrs.as_slice()) | ||
| .await | ||
| .context("unable to connect to remote address")?; | ||
| Ok(socket) | ||
| } | ||
| } | ||
|
|
||
| #[derive(Debug)] | ||
| pub struct UdpClientTask { | ||
| socket: UdpSocket, | ||
| transport_commands_rx: UnboundedReceiver<TransportCommand>, | ||
| } | ||
|
|
||
| impl UdpClientTask { | ||
| pub async fn run(mut self) { | ||
| let mut udp_buf = [0; MAX_PACKET_SIZE]; | ||
|
|
||
| // this here isn't perfect because we block the entire transport_commands_rx channel if we | ||
| // cannot send (so we also block receiving new packets), but that's hopefully good enough. | ||
| let mut packet_needs_sending = false; | ||
| let mut packet_payload = Vec::new(); | ||
|
|
||
| let mut packet_tx: Option<oneshot::Sender<Vec<u8>>> = None; | ||
|
|
||
| loop { | ||
| tokio::select! { | ||
| // wait for transport_events_tx channel capacity... | ||
| Ok(len) = self.socket.recv(&mut udp_buf), if packet_tx.is_some() => { | ||
| packet_tx | ||
| .take() | ||
| .unwrap() | ||
| .send(udp_buf[..len].to_vec()) | ||
| .ok(); | ||
| }, | ||
| // send_to is cancel safe, so we can use that for backpressure. | ||
| _ = self.socket.send(&packet_payload), if packet_needs_sending => { | ||
| packet_needs_sending = false; | ||
| }, | ||
| Some(command) = self.transport_commands_rx.recv(), if !packet_needs_sending => { | ||
| match command { | ||
| TransportCommand::ReadData(_,_,tx) => { | ||
| packet_tx = Some(tx); | ||
| }, | ||
| TransportCommand::WriteData(_, data) => { | ||
| packet_payload = data; | ||
| packet_needs_sending = true; | ||
| }, | ||
| TransportCommand::DrainWriter(_,tx) => { | ||
| tx.send(()).ok(); | ||
| }, | ||
| TransportCommand::CloseConnection(_, half_close) => { | ||
| if !half_close { | ||
| break; | ||
| } | ||
| }, | ||
| } | ||
| } | ||
| else => break, | ||
| } | ||
| } | ||
| log::debug!("UDP client task shutting down."); | ||
| } | ||
| } |
Oops, something went wrong.