Memory exhaustion (#69)

src/http.rs

@@ -366,15 +366,15 @@ impl ProxyHandler for HttpConnection {
         self.state == HttpState::Established
     }
 
-    fn have_data(&mut self, dir: Direction) -> bool {
+    fn data_len(&self, dir: Direction) -> usize {
         match dir {
             Direction::Incoming(incoming) => match incoming {
-                IncomingDirection::FromServer => !self.server_inbuf.is_empty(),
-                IncomingDirection::FromClient => !self.client_inbuf.is_empty() || !self.data_buf.is_empty(),
+                IncomingDirection::FromServer => self.server_inbuf.len(),
+                IncomingDirection::FromClient => self.client_inbuf.len().max(self.data_buf.len()),
             },
             Direction::Outgoing(outgoing) => match outgoing {
-                OutgoingDirection::ToServer => !self.server_outbuf.is_empty(),
-                OutgoingDirection::ToClient => !self.client_outbuf.is_empty(),
+                OutgoingDirection::ToServer => self.server_outbuf.len(),
+                OutgoingDirection::ToClient => self.client_outbuf.len(),
             },
         }
     }

src/socks.rs

@@ -314,15 +314,15 @@ impl ProxyHandler for SocksProxyImpl {
         self.state == SocksState::Established
     }
 
-    fn have_data(&mut self, dir: Direction) -> bool {
+    fn data_len(&self, dir: Direction) -> usize {
         match dir {
             Direction::Incoming(incoming) => match incoming {
-                IncomingDirection::FromServer => !self.server_inbuf.is_empty(),
-                IncomingDirection::FromClient => !self.client_inbuf.is_empty() || !self.data_buf.is_empty(),
+                IncomingDirection::FromServer => self.server_inbuf.len(),
+                IncomingDirection::FromClient => self.client_inbuf.len().max(self.data_buf.len()),
             },
             Direction::Outgoing(outgoing) => match outgoing {
-                OutgoingDirection::ToServer => !self.server_outbuf.is_empty(),
-                OutgoingDirection::ToClient => !self.client_outbuf.is_empty(),
+                OutgoingDirection::ToServer => self.server_outbuf.len(),
+                OutgoingDirection::ToClient => self.client_outbuf.len(),
             },
         }
     }

src/tun2proxy.rs

@@ -174,6 +174,7 @@ const CLIENT_WRITE_CLOSED: u8 = 2;
 
 const UDP_ASSO_TIMEOUT: u64 = 10; // seconds
 const DNS_PORT: u16 = 53;
+const IP_PACKAGE_MAX_SIZE: usize = 0xFFFF;
 
 struct ConnectionState {
     smoltcp_handle: SocketHandle,
@@ -197,7 +198,7 @@ pub(crate) trait ProxyHandler {
     fn consume_data(&mut self, dir: OutgoingDirection, size: usize);
     fn peek_data(&mut self, dir: OutgoingDirection) -> OutgoingDataEvent;
     fn connection_established(&self) -> bool;
-    fn have_data(&mut self, dir: Direction) -> bool;
+    fn data_len(&self, dir: Direction) -> usize;
     fn reset_connection(&self) -> bool;
     fn get_udp_associate(&self) -> Option<SocketAddr>;
 }
@@ -395,13 +396,10 @@ impl<'a> TunToProxy<'a> {
             None => return Ok(()),
         };
         let mut closed_ends = 0;
+        let handler = state.proxy_handler.as_ref();
         if (state.close_state & SERVER_WRITE_CLOSED) == SERVER_WRITE_CLOSED
-            && !state
-                .proxy_handler
-                .have_data(Direction::Incoming(IncomingDirection::FromServer))
-            && !state
-                .proxy_handler
-                .have_data(Direction::Outgoing(OutgoingDirection::ToClient))
+            && handler.data_len(Direction::Incoming(IncomingDirection::FromServer)) == 0
+            && handler.data_len(Direction::Outgoing(OutgoingDirection::ToClient)) == 0
         {
             // Close tun interface
             let socket = self.sockets.get_mut::<tcp::Socket>(state.smoltcp_handle);
@@ -411,12 +409,8 @@ impl<'a> TunToProxy<'a> {
         }
 
         if (state.close_state & CLIENT_WRITE_CLOSED) == CLIENT_WRITE_CLOSED
-            && !state
-                .proxy_handler
-                .have_data(Direction::Incoming(IncomingDirection::FromClient))
-            && !state
-                .proxy_handler
-                .have_data(Direction::Outgoing(OutgoingDirection::ToServer))
+            && handler.data_len(Direction::Incoming(IncomingDirection::FromClient)) == 0
+            && handler.data_len(Direction::Outgoing(OutgoingDirection::ToServer)) == 0
         {
             // Close remote server
             if let Err(err) = state.mio_stream.shutdown(Shutdown::Write) {
@@ -443,6 +437,11 @@ impl<'a> TunToProxy<'a> {
             let socket = self.sockets.get_mut::<tcp::Socket>(state.smoltcp_handle);
             let mut error = Ok(());
             while socket.can_recv() && error.is_ok() {
+                let dir = Direction::Outgoing(OutgoingDirection::ToServer);
+                if state.proxy_handler.data_len(dir) >= IP_PACKAGE_MAX_SIZE {
+                    break;
+                }
+
                 socket.recv(|data| {
                     let event = IncomingDataEvent {
                         direction: IncomingDirection::FromClient,