luckyimnotanigger/tun2proxy
1
0
Fork 0
mirror of https://github.com/tun2proxy/tun2proxy.git synced 2025-05-14 05:52:48 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

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_managers: Vec<Rc<dyn ConnectionManager>>,
next_token: usize,
token_to_info: HashMap<Token, ConnectionInfo>,
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::<tcp::Socket>(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::<tcp::Socket>(socket_handle);
if let Some(handle) = state.smoltcp_handle {
// Close tun interface
let socket = self.sockets.get_mut::<tcp::Socket>(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::<tcp::Socket>(socket_handle);
let socket = self.sockets.get_mut::<tcp::Socket>(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,