extract dns logic to separate functions

src/tun2proxy.rs

@@ -757,11 +757,9 @@ impl<'a> TunToProxy<'a> {
         Ok(())
     }
 
-    fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> {
+    fn receive_udp_packet_and_write_to_client(&mut self, info: &ConnectionInfo) -> Result<()> {
-        if let Some(info) = self.find_info_by_udp_token(event.token()) {
-            let info = info.clone();
         let err = "udp connection state not found";
-            let state = self.connection_map.get_mut(&info).ok_or(err)?;
+        let state = self.connection_map.get_mut(info).ok_or(err)?;
         state.expiry = Some(Self::udp_associate_timeout());
         let mut to_send: LinkedList<Vec<u8>> = LinkedList::new();
         if let Some(udp_socket) = state.udp_socket.as_ref() {
@@ -789,8 +787,27 @@ impl<'a> TunToProxy<'a> {
         while let Some(packet) = to_send.pop_front() {
             self.send_udp_packet_to_client(src, info.src, &packet)?;
         }
+        Ok(())
+    }
 
-            return Ok(());
+    fn comsume_cached_udp_packets(&mut self, info: &ConnectionInfo) -> Result<()> {
+        // Try to send the first UDP packets to remote SOCKS5 server for UDP associate session
+        if let Some(state) = self.connection_map.get_mut(info) {
+            if let Some(udp_socket) = state.udp_socket.as_ref() {
+                if let Some(addr) = state.tcp_proxy_handler.get_udp_associate() {
+                    // Consume udp_data_cache data
+                    while let Some(buf) = state.udp_data_cache.pop_front() {
+                        udp_socket.send_to(&buf, addr)?;
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> {
+        if let Some(info) = self.find_info_by_udp_token(event.token()) {
+            return self.receive_udp_packet_and_write_to_client(&info.clone());
         }
 
         let conn_info = match self.find_info_by_token(event.token()) {
@@ -872,17 +889,7 @@ impl<'a> TunToProxy<'a> {
                 // server.
                 self.write_to_server(&conn_info)?;
 
-                // Try to send the first UDP packet to remote SOCKS5 server for UDP associate session
+                self.comsume_cached_udp_packets(&conn_info)?;
-                if let Some(state) = self.connection_map.get_mut(&conn_info) {
-                    if let Some(udp_socket) = state.udp_socket.as_ref() {
-                        if let Some(addr) = state.tcp_proxy_handler.get_udp_associate() {
-                            // Consume udp_data_cache data
-                            while let Some(buf) = state.udp_data_cache.pop_front() {
-                                udp_socket.send_to(&buf, addr)?;
-                            }
-                        }
-                    }
-                }
             }
 
             if event.is_writable() {