diff --git a/src/http.rs b/src/http.rs index 4d3e5bb..7ebd02e 100644 --- a/src/http.rs +++ b/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(), }, } } diff --git a/src/socks.rs b/src/socks.rs index 7c62aa9..c7b60aa 100644 --- a/src/socks.rs +++ b/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(), }, } } diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index a13c83c..318b46c 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -173,6 +173,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, @@ -196,7 +197,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; } @@ -406,13 +407,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::(state.smoltcp_handle); @@ -422,12 +420,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) { @@ -454,6 +448,11 @@ impl<'a> TunToProxy<'a> { let socket = self.sockets.get_mut::(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,