diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs index 33743b5..e54e7cd 100644 --- a/src/tun2proxy.rs +++ b/src/tun2proxy.rs @@ -209,7 +209,6 @@ pub struct TunToProxy<'a> { connection_map: HashMap, connection_managers: Vec>, next_token: usize, - token_to_info: HashMap, sockets: SocketSet<'a>, device: VirtualTunDevice, options: Options, @@ -256,7 +255,6 @@ impl<'a> TunToProxy<'a> { iface, connection_map: HashMap::default(), next_token: usize::from(EXIT_TOKEN) + 1, - token_to_info: HashMap::default(), connection_managers: Vec::default(), sockets: SocketSet::new([]), device, @@ -296,18 +294,27 @@ impl<'a> TunToProxy<'a> { Ok(()) } + fn find_info_by_token(&self, token: Token) -> Option<&ConnectionInfo> { + self.connection_map + .iter() + .find_map(|(info, state)| if state.token == token { Some(info) } else { None }) + } + + /// Destroy connection state machine fn remove_connection(&mut self, info: &ConnectionInfo) -> Result<(), Error> { - if let Some(mut conn) = self.connection_map.remove(info) { - _ = conn.mio_stream.shutdown(Shutdown::Both); - if let Some(handle) = conn.smoltcp_handle { + if let Some(mut state) = self.connection_map.remove(info) { + _ = state.mio_stream.shutdown(Shutdown::Both); + if let Some(handle) = state.smoltcp_handle { let socket = self.sockets.get_mut::(handle); socket.close(); self.sockets.remove(handle); } + + // FIXME: Does this line should be moved up to the beginning of this function? self.expect_smoltcp_send()?; - let token = &conn.token; - self.token_to_info.remove(token); - _ = self.poll.registry().deregister(&mut conn.mio_stream); + + _ = self.poll.registry().deregister(&mut state.mio_stream); + log::info!("Close {}", info); } Ok(()) @@ -322,10 +329,11 @@ impl<'a> TunToProxy<'a> { None } + /// Scan connection state machine and check if any connection should be closed. fn check_change_close_state(&mut self, info: &ConnectionInfo) -> Result<(), Error> { let state = match self.connection_map.get_mut(info) { - None => return Ok(()), Some(state) => state, + None => return Ok(()), }; let mut closed_ends = 0; if (state.close_state & SERVER_WRITE_CLOSED) == SERVER_WRITE_CLOSED @@ -336,8 +344,9 @@ impl<'a> TunToProxy<'a> { .tcp_proxy_handler .have_data(Direction::Outgoing(OutgoingDirection::ToClient)) { - if let Some(socket_handle) = state.smoltcp_handle { - let socket = self.sockets.get_mut::(socket_handle); + if let Some(handle) = state.smoltcp_handle { + // Close tun interface + let socket = self.sockets.get_mut::(handle); socket.close(); } closed_ends += 1; @@ -351,17 +360,20 @@ impl<'a> TunToProxy<'a> { .tcp_proxy_handler .have_data(Direction::Outgoing(OutgoingDirection::ToServer)) { + // Close remote server _ = state.mio_stream.shutdown(Shutdown::Write); closed_ends += 1; } if closed_ends == 2 { + // Close connection state machine self.remove_connection(info)?; } Ok(()) } fn tunsocket_read_and_forward(&mut self, info: &ConnectionInfo) -> Result<(), Error> { + // 1. Read data from tun and write to proxy handler (remote server). // Scope for mutable borrow of self. { let state = match self.connection_map.get_mut(info) { @@ -393,10 +405,10 @@ impl<'a> TunToProxy<'a> { // need to send data. state.close_state |= CLIENT_WRITE_CLOSED; } - - // Expect ACKs etc. from smoltcp sockets. - self.expect_smoltcp_send()?; } + // 2. Write data from proxy handler (remote server) to tun. + // Expect ACKs etc. from smoltcp sockets. + self.expect_smoltcp_send()?; self.check_change_close_state(info)?; @@ -554,7 +566,6 @@ impl<'a> TunToProxy<'a> { }; self.connection_map.insert(connection_info.clone(), state); - self.token_to_info.insert(token, connection_info.clone()); Ok(()) } @@ -593,7 +604,7 @@ impl<'a> TunToProxy<'a> { fn write_to_client(&mut self, token: Token, info: &ConnectionInfo) -> Result<(), Error> { while let Some(state) = self.connection_map.get_mut(info) { - let socket_handle = match state.smoltcp_handle { + let handle = match state.smoltcp_handle { Some(handle) => handle, None => break, }; @@ -601,7 +612,7 @@ impl<'a> TunToProxy<'a> { let buflen = event.buffer.len(); let consumed; { - let socket = self.sockets.get_mut::(socket_handle); + let socket = self.sockets.get_mut::(handle); if socket.may_send() { if let Some(virtual_dns) = &mut self.options.virtual_dns { // Unwrapping is fine because every smoltcp socket is bound to an. @@ -641,11 +652,10 @@ impl<'a> TunToProxy<'a> { } fn send_to_smoltcp(&mut self) -> Result<(), Error> { - let cloned = self.write_sockets.clone(); - for token in cloned.iter() { - if let Some(connection) = self.token_to_info.get(token) { + for token in self.write_sockets.clone().into_iter() { + if let Some(connection) = self.find_info_by_token(token) { let connection = connection.clone(); - if let Err(error) = self.write_to_client(*token, &connection) { + if let Err(error) = self.write_to_client(token, &connection) { self.remove_connection(&connection)?; log::error!("Write to client: {}: ", error); } @@ -656,7 +666,7 @@ impl<'a> TunToProxy<'a> { fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> { let e = "connection not found"; - let conn_info = match self.token_to_info.get(&event.token()) { + let conn_info = match self.find_info_by_token(event.token()) { Some(conn_info) => conn_info.clone(), None => { // We may have closed the connection in an earlier iteration over the poll events,