fix nits

mitmproxy · Dec 18, 2023 · 204910c · 204910c
1 parent d72424e
commit 204910c
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 23 deletions.
diff --git a/mitmproxy-rs/src/task.rs b/mitmproxy-rs/src/task.rs
@@ -17,7 +17,7 @@ pub struct PyInteropTask {
     transport_events: mpsc::Receiver<TransportEvent>,
     py_tcp_handler: PyObject,
     py_udp_handler: PyObject,
-    sd_watcher: broadcast::Receiver<()>,
+    shutdown: broadcast::Receiver<()>,
 }
 
 impl PyInteropTask {
@@ -36,7 +36,7 @@ impl PyInteropTask {
             transport_events,
             py_tcp_handler,
             py_udp_handler,
-            sd_watcher,
+            shutdown: sd_watcher,
         }
     }
 
@@ -48,9 +48,9 @@ impl PyInteropTask {
         })?;
 
         loop {
-            tokio::select!(
+            tokio::select! {
                 // wait for graceful shutdown
-                _ = self.sd_watcher.recv() => break,
+                _ = self.shutdown.recv() => break,
                 // wait for network events
                 event = self.transport_events.recv() => {
                     let Some(event) = event else {
@@ -111,8 +111,8 @@ impl PyInteropTask {
                             };
                         },
                     }
-                },
-            );
+                }
+            };
         }
 
         log::debug!("Python interoperability task shutting down.");

diff --git a/mitmproxy-windows/redirector/src/main2.rs b/mitmproxy-windows/redirector/src/main2.rs
@@ -420,7 +420,8 @@ async fn handle_ipc(
                     }
                 }
             },
-            Some(packet) = ipc_rx.recv() => {
+            r = ipc_rx.recv() => {
+                let Some(packet) = r else { break };
 
                 packet.encode(&mut buf.as_mut_slice())?;
                 let len = packet.encoded_len();

diff --git a/src/network/tests.rs b/src/network/tests.rs
@@ -15,7 +15,6 @@ use tokio::{
 use crate::messages::{
     NetworkCommand, NetworkEvent, SmolPacket, TransportCommand, TransportEvent, TunnelInfo,
 };
-use crate::packet_sources::PacketSourceConf;
 
 use super::task::NetworkTask;
 

diff --git a/src/packet_sources/macos.rs b/src/packet_sources/macos.rs
@@ -243,10 +243,7 @@ impl ConnectionTask {
         loop {
             tokio::select! {
                 _ = self.shutdown.recv() => break,
-                packet = stream.next(), if state.packet_queue_len() < 10 => {
-                    let Some(packet) = packet else {
-                        break
-                    };
+                Some(packet) = stream.next(), if state.packet_queue_len() < 10 => {
                     let packet = ipc::UdpPacket::decode(
                         packet.context("IPC read error")?
                     ).context("invalid IPC message")?;
@@ -270,10 +267,7 @@ impl ConnectionTask {
                     }
                     state.add_packet(packet.data);
                 },
-                command = command_rx.recv() => {
-                    let Some(command) = command else {
-                        break;
-                    };
+                Some(command) = command_rx.recv() => {
                     match command {
                         TransportCommand::ReadData(_, _, tx) => {
                             state.add_reader(tx);
@@ -350,10 +344,7 @@ impl ConnectionTask {
                     self.stream.read_buf(&mut data).await.context("failed to read from socket")?;
                     tx.send(data).ok();
                 },
-                command = command_rx.recv() => {
-                    let Some(command) = command else {
-                        break;
-                    };
+                Some(command) = command_rx.recv() => {
                     match command {
                         TransportCommand::ReadData(_, n, tx) => {
                             assert!(read_tx.is_none());

diff --git a/src/packet_sources/udp.rs b/src/packet_sources/udp.rs
@@ -90,7 +90,8 @@ impl PacketSourceTask for UdpTask {
                     permit = Some(p);
                 },
                 // ... or process incoming packets
-                Ok((len, src_addr)) = self.socket.recv_from(&mut udp_buf), if py_tx_available => {
+                r = self.socket.recv_from(&mut udp_buf), if py_tx_available => {
+                    let (len, src_addr) = r.context("UDP recv() failed")?;
                     self.handler.receive_data(
                         UdpPacket {
                             src_addr,
@@ -102,7 +103,8 @@ impl PacketSourceTask for UdpTask {
                     );
                 },
                 // send_to is cancel safe, so we can use that for backpressure.
-                Ok(_) = self.socket.send_to(&packet_payload, packet_dst), if packet_needs_sending => {
+                r = self.socket.send_to(&packet_payload, packet_dst), if packet_needs_sending => {
+                    r.context("UDP send_to() failed")?;
                     packet_needs_sending = false;
                 },
                 Some(command) = self.transport_commands_rx.recv(), if !packet_needs_sending => {

diff --git a/src/packet_sources/wireguard.rs b/src/packet_sources/wireguard.rs
@@ -158,7 +158,8 @@ impl PacketSourceTask for WireGuardTask {
             tokio::select! {
                 exit = &mut self.network_task_handle => break exit.context("network task panic")?.context("network task error")?,
                 // wait for WireGuard packets incoming on the UDP socket
-                Ok((len, src_orig)) = self.socket.recv_from(&mut udp_buf) => {
+                r = self.socket.recv_from(&mut udp_buf) => {
+                    let (len, src_orig) = r.context("UDP recv() failed")?;
                     self.process_incoming_datagram(&udp_buf[..len], src_orig).await?;
                 },
                 // wait for outgoing IP packets