Prevent connection error bubbling from terminating the app

src/tun2proxy.rs

@@ -384,6 +384,8 @@ impl<'a> TunToProxy<'a> {
                     }
                 }
             };
+            let dst = connection.dst;
+            (|| -> Result<(), Error> {
                 if resolved_conn.proto == IpProtocol::Tcp {
                     let cm = self.get_connection_manager(&resolved_conn);
                     if cm.is_none() {
@@ -400,7 +402,7 @@ impl<'a> TunToProxy<'a> {
                                     tcp::SocketBuffer::new(vec![0; 4096]),
                                 );
                                 socket.set_ack_delay(None);
-                            let dst = SocketAddr::try_from(connection.dst)?;
+                                let dst = SocketAddr::try_from(dst)?;
                                 socket.listen(dst)?;
                                 let handle = self.sockets.add(socket);
 
@@ -448,17 +450,21 @@ impl<'a> TunToProxy<'a> {
 
                     // The connection handler builds up the connection or encapsulates the data.
                     // Therefore, we now expect it to write data to the server.
-                self.write_to_server(&resolved_conn);
+                    self.write_to_server(&resolved_conn)?;
                 } else if resolved_conn.proto == IpProtocol::Udp && resolved_conn.dst.port == 53 {
                     if let Some(virtual_dns) = &mut self.options.virtdns {
                         let payload = &frame[_payload_offset.._payload_offset + _payload_size];
                         if let Some(response) = virtual_dns.receive_query(payload) {
-                        let rx_buffer =
-                            udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
-                        let tx_buffer =
-                            udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
+                            let rx_buffer = udp::PacketBuffer::new(
+                                vec![udp::PacketMetadata::EMPTY],
+                                vec![0; 4096],
+                            );
+                            let tx_buffer = udp::PacketBuffer::new(
+                                vec![udp::PacketMetadata::EMPTY],
+                                vec![0; 4096],
+                            );
                             let mut socket = udp::Socket::new(rx_buffer, tx_buffer);
-                        let dst = SocketAddr::try_from(connection.dst)?;
+                            let dst = SocketAddr::try_from(dst)?;
                             socket.bind(dst)?;
                             socket
                                 .send_slice(response.as_slice(), resolved_conn.src.into())
@@ -470,15 +476,21 @@ impl<'a> TunToProxy<'a> {
                     }
                     // Otherwise, UDP is not yet supported.
                 }
+                Ok(())
+            })()
+            .or_else(|error| {
+                log::error! {"{error}"}
+                Ok::<(), Error>(())
+            })?;
         }
         Ok(())
     }
 
-    fn write_to_server(&mut self, connection: &Connection) {
+    fn write_to_server(&mut self, connection: &Connection) -> Result<(), Error> {
         if let Some(state) = self.connections.get_mut(connection) {
             let event = state.handler.peek_data(OutgoingDirection::ToServer);
             if event.buffer.is_empty() {
-                return;
+                return Ok(());
             }
             let result = state.mio_stream.write(event.buffer);
             match result {
@@ -488,13 +500,12 @@ impl<'a> TunToProxy<'a> {
                         .consume_data(OutgoingDirection::ToServer, consumed);
                 }
                 Err(error) if error.kind() != std::io::ErrorKind::WouldBlock => {
-                    panic!("Error: {:?}", error);
-                }
-                _ => {
-                    // println!("{:?}", result);
+                    return Err(error.into());
                 }
+                _ => {}
             }
         }
+        Ok(())
     }
 
     fn write_to_client(&mut self, token: Token, connection: &Connection) -> Result<(), Error> {
@@ -556,7 +567,11 @@ impl<'a> TunToProxy<'a> {
         let cloned = self.write_sockets.clone();
         for token in cloned.iter() {
             if let Some(connection) = self.token_to_connection.get(token) {
-                self.write_to_client(*token, &connection.clone())?;
+                let connection = connection.clone();
+                if let Err(error) = self.write_to_client(*token, &connection) {
+                    self.remove_connection(&connection)?;
+                    log::error!("Write to client: {}: ", error);
+                }
             }
         }
         Ok(())
@@ -565,10 +580,9 @@ impl<'a> TunToProxy<'a> {
     fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> {
         let e = "connection not found";
         let conn_ref = self.token_to_connection.get(&event.token());
-        if conn_ref.is_none() {
-            return Ok(());
-        }
         let connection = conn_ref.ok_or(e)?.clone();
+
+        (|| -> Result<(), Error> {
             if event.is_readable() || event.is_read_closed() {
                 {
                     let state = self.connections.get_mut(&connection).ok_or(e)?;
@@ -580,7 +594,7 @@ impl<'a> TunToProxy<'a> {
                         Ok(read_result) => read_result,
                         Err(error) => {
                             if error.kind() != std::io::ErrorKind::WouldBlock {
-                            error!("READ from proxy: {}", error);
+                                error!("Read from proxy: {}", error);
                             }
                             vecbuf.len()
                         }
@@ -626,9 +640,15 @@ impl<'a> TunToProxy<'a> {
                 self.write_to_client(event.token(), &connection)?;
             }
             if event.is_writable() {
-            self.write_to_server(&connection);
+                self.write_to_server(&connection)?;
             }
             Ok(())
+        })()
+        .or_else(|error| {
+            self.remove_connection(&connection)?;
+            log::error! {"{error}"}
+            Ok(())
+        })
     }
 
     fn udp_event(&mut self, _event: &Event) {}