Memory exhaustion (#69)

This commit is contained in:
ssrlive 2023-10-10 14:22:33 +08:00 committed by GitHub
parent 299b51667d
commit b50cac82c0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 22 additions and 23 deletions

View file

@ -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(),
},
}
}

View file

@ -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(),
},
}
}

View file

@ -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,