luckyimnotanigger/tun2proxy
1
0
Fork 0
mirror of https://github.com/tun2proxy/tun2proxy.git synced 2025-04-20 13:59:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

test code

Browse source
This commit is contained in:
ssrlive 2023-11-04 14:34:47 +08:00
parent fe85ecd15c
commit 9396db4a52
1 changed files with 20 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
src/tun2proxy.rs
View file

@ -189,6 +189,7 @@ struct ConnectionState {
udp_token: Option<Token>, udp_token: Option<Token>,
udp_data_cache: LinkedList<Vec<u8>>, udp_data_cache: LinkedList<Vec<u8>>,
dns_over_tcp_expiry: Option<::std::time::Instant>, dns_over_tcp_expiry: Option<::std::time::Instant>,
is_tcp_closed: bool,
} }
pub(crate) trait ProxyHandler { pub(crate) trait ProxyHandler {
@ -580,7 +581,7 @@ impl<'a> TunToProxy<'a> {
state.dns_over_tcp_expiry = Some(Self::common_udp_life_timeout()); state.dns_over_tcp_expiry = Some(Self::common_udp_life_timeout());
let mut vecbuf = vec![]; let mut vecbuf = vec![];
Self::read_data_from_tcp_stream(&mut state.mio_stream, |data| { Self::read_data_from_tcp_stream(&mut state.mio_stream, &mut state.is_tcp_closed, |data| {
vecbuf.extend_from_slice(data); vecbuf.extend_from_slice(data);
Ok(()) Ok(())
})?; })?;
@ -835,6 +836,7 @@ impl<'a> TunToProxy<'a> {
origin_dst: dst, origin_dst: dst,
udp_data_cache: LinkedList::new(), udp_data_cache: LinkedList::new(),
dns_over_tcp_expiry: None, dns_over_tcp_expiry: None,
is_tcp_closed: false,
}; };
Ok(state) Ok(state)
} }
@ -852,13 +854,25 @@ impl<'a> TunToProxy<'a> {
false false
} }
fn clearup_expired_udp_associate(&mut self) -> Result<()> { fn tcp_is_closed(&self, info: &ConnectionInfo) -> bool {
if let Some(state) = self.connection_map.get(info) {
return state.is_tcp_closed;
}
false
}
fn clearup_expired_connection(&mut self) -> Result<()> {
let keys = self.connection_map.keys().cloned().collect::<Vec<_>>(); let keys = self.connection_map.keys().cloned().collect::<Vec<_>>();
for key in keys { for key in keys {
if self.udp_associate_timeout_expired(&key) { if self.udp_associate_timeout_expired(&key) {
log::trace!("UDP associate timeout: {}", key); log::trace!("UDP associate timeout: {}", key);
self.remove_connection(&key)?; self.remove_connection(&key)?;
} }
if self.tcp_is_closed(&key) {
log::trace!("TCP closed: {}", key);
self.remove_connection(&key)?;
}
} }
Ok(()) Ok(())
} }
@ -1061,7 +1075,7 @@ impl<'a> TunToProxy<'a> {
// TODO: Move this reading process to its own function. // TODO: Move this reading process to its own function.
let mut vecbuf = vec![]; let mut vecbuf = vec![];
Self::read_data_from_tcp_stream(&mut state.mio_stream, |data| { Self::read_data_from_tcp_stream(&mut state.mio_stream, &mut state.is_tcp_closed, |data| {
vecbuf.extend_from_slice(data); vecbuf.extend_from_slice(data);
Ok(()) Ok(())
})?; })?;
@ -1130,7 +1144,7 @@ impl<'a> TunToProxy<'a> {
Ok(()) Ok(())
} }
fn read_data_from_tcp_stream<F>(stream: &mut TcpStream, mut callback: F) -> Result<()> fn read_data_from_tcp_stream<F>(stream: &mut TcpStream, is_closed: &mut bool, mut callback: F) -> Result<()>
where where
F: FnMut(&mut [u8]) -> Result<()>, F: FnMut(&mut [u8]) -> Result<()>,
{ {
@ -1139,6 +1153,7 @@ impl<'a> TunToProxy<'a> {
match stream.read(&mut tmp) { match stream.read(&mut tmp) {
Ok(0) => { Ok(0) => {
// The tcp connection closed // The tcp connection closed
*is_closed = true;
break; break;
} }
Ok(read_result) => { Ok(read_result) => {
@ -1219,7 +1234,7 @@ impl<'a> TunToProxy<'a> {
} }
} }
self.send_to_smoltcp()?; self.send_to_smoltcp()?;
self.clearup_expired_udp_associate()?; self.clearup_expired_connection()?;
self.clearup_expired_dns_over_tcp()?; self.clearup_expired_dns_over_tcp()?;
} }
} }