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 self.state == HttpState::Established
} }
fn have_data(&mut self, dir: Direction) -> bool { fn data_len(&self, dir: Direction) -> usize {
match dir { match dir {
Direction::Incoming(incoming) => match incoming { Direction::Incoming(incoming) => match incoming {
IncomingDirection::FromServer => !self.server_inbuf.is_empty(), IncomingDirection::FromServer => self.server_inbuf.len(),
IncomingDirection::FromClient => !self.client_inbuf.is_empty() || !self.data_buf.is_empty(), IncomingDirection::FromClient => self.client_inbuf.len().max(self.data_buf.len()),
}, },
Direction::Outgoing(outgoing) => match outgoing { Direction::Outgoing(outgoing) => match outgoing {
OutgoingDirection::ToServer => !self.server_outbuf.is_empty(), OutgoingDirection::ToServer => self.server_outbuf.len(),
OutgoingDirection::ToClient => !self.client_outbuf.is_empty(), OutgoingDirection::ToClient => self.client_outbuf.len(),
}, },
} }
} }

View file

@ -314,15 +314,15 @@ impl ProxyHandler for SocksProxyImpl {
self.state == SocksState::Established self.state == SocksState::Established
} }
fn have_data(&mut self, dir: Direction) -> bool { fn data_len(&self, dir: Direction) -> usize {
match dir { match dir {
Direction::Incoming(incoming) => match incoming { Direction::Incoming(incoming) => match incoming {
IncomingDirection::FromServer => !self.server_inbuf.is_empty(), IncomingDirection::FromServer => self.server_inbuf.len(),
IncomingDirection::FromClient => !self.client_inbuf.is_empty() || !self.data_buf.is_empty(), IncomingDirection::FromClient => self.client_inbuf.len().max(self.data_buf.len()),
}, },
Direction::Outgoing(outgoing) => match outgoing { Direction::Outgoing(outgoing) => match outgoing {
OutgoingDirection::ToServer => !self.server_outbuf.is_empty(), OutgoingDirection::ToServer => self.server_outbuf.len(),
OutgoingDirection::ToClient => !self.client_outbuf.is_empty(), 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 UDP_ASSO_TIMEOUT: u64 = 10; // seconds
const DNS_PORT: u16 = 53; const DNS_PORT: u16 = 53;
const IP_PACKAGE_MAX_SIZE: usize = 0xFFFF;
struct ConnectionState { struct ConnectionState {
smoltcp_handle: SocketHandle, smoltcp_handle: SocketHandle,
@ -197,7 +198,7 @@ pub(crate) trait ProxyHandler {
fn consume_data(&mut self, dir: OutgoingDirection, size: usize); fn consume_data(&mut self, dir: OutgoingDirection, size: usize);
fn peek_data(&mut self, dir: OutgoingDirection) -> OutgoingDataEvent; fn peek_data(&mut self, dir: OutgoingDirection) -> OutgoingDataEvent;
fn connection_established(&self) -> bool; 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 reset_connection(&self) -> bool;
fn get_udp_associate(&self) -> Option<SocketAddr>; fn get_udp_associate(&self) -> Option<SocketAddr>;
} }
@ -395,13 +396,10 @@ impl<'a> TunToProxy<'a> {
None => return Ok(()), None => return Ok(()),
}; };
let mut closed_ends = 0; let mut closed_ends = 0;
let handler = state.proxy_handler.as_ref();
if (state.close_state & SERVER_WRITE_CLOSED) == SERVER_WRITE_CLOSED if (state.close_state & SERVER_WRITE_CLOSED) == SERVER_WRITE_CLOSED
&& !state && handler.data_len(Direction::Incoming(IncomingDirection::FromServer)) == 0
.proxy_handler && handler.data_len(Direction::Outgoing(OutgoingDirection::ToClient)) == 0
.have_data(Direction::Incoming(IncomingDirection::FromServer))
&& !state
.proxy_handler
.have_data(Direction::Outgoing(OutgoingDirection::ToClient))
{ {
// Close tun interface // Close tun interface
let socket = self.sockets.get_mut::<tcp::Socket>(state.smoltcp_handle); 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 if (state.close_state & CLIENT_WRITE_CLOSED) == CLIENT_WRITE_CLOSED
&& !state && handler.data_len(Direction::Incoming(IncomingDirection::FromClient)) == 0
.proxy_handler && handler.data_len(Direction::Outgoing(OutgoingDirection::ToServer)) == 0
.have_data(Direction::Incoming(IncomingDirection::FromClient))
&& !state
.proxy_handler
.have_data(Direction::Outgoing(OutgoingDirection::ToServer))
{ {
// Close remote server // Close remote server
if let Err(err) = state.mio_stream.shutdown(Shutdown::Write) { 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 socket = self.sockets.get_mut::<tcp::Socket>(state.smoltcp_handle);
let mut error = Ok(()); let mut error = Ok(());
while socket.can_recv() && error.is_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| { socket.recv(|data| {
let event = IncomingDataEvent { let event = IncomingDataEvent {
direction: IncomingDirection::FromClient, direction: IncomingDirection::FromClient,