From 81db9cb18188e3f3343548b1ec3d572f099b7fa4 Mon Sep 17 00:00:00 2001 From: ssrlive <30760636+ssrlive@users.noreply.github.com> Date: Fri, 6 Oct 2023 21:21:33 +0800 Subject: [PATCH 1/2] rename have_data to data_size --- src/http.rs | 10 +++++----- src/socks.rs | 10 +++++----- src/tun2proxy.rs | 19 ++++++------------- 3 files changed, 16 insertions(+), 23 deletions(-) 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 6ae46ba..e7338bd 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -197,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; } @@ -395,13 +395,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); @@ -411,12 +408,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) { From 2c9ac6f3a98dec61f6759ac41353a8f69cde9764 Mon Sep 17 00:00:00 2001 From: ssrlive <30760636+ssrlive@users.noreply.github.com> Date: Fri, 6 Oct 2023 23:12:06 +0800 Subject: [PATCH 2/2] fix memory exhaustion --- src/tun2proxy.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index e7338bd..652631b 100644 --- a/src/tun2proxy.rs +++ b/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, @@ -436,6 +437,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,