luckyimnotanigger/tun2proxy
1
0
Fork 0
mirror of https://github.com/tun2proxy/tun2proxy.git synced 2025-05-15 06:22:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

token_to_info removed

Browse source
This commit is contained in:
ssrlive 2023-08-11 01:13:22 +08:00
parent 489d5fec00
commit 57851f029e
1 changed files with 32 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

54
src/tun2proxy.rs
View file

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