luckyimnotanigger/tun2proxy
1
0
Fork 0
mirror of https://github.com/tun2proxy/tun2proxy.git synced 2025-04-22 23:09:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

extract dns logic to separate functions

Browse source
This commit is contained in:
ssrlive 2023-08-22 10:44:46 +08:00
parent 10ade80488
commit d42d3a8287
1 changed files with 48 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

89
src/tun2proxy.rs
View file

@ -757,40 +757,57 @@ impl<'a> TunToProxy<'a> {
Ok(()) Ok(())
} }
fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> { fn receive_udp_packet_and_write_to_client(&mut self, info: &ConnectionInfo) -> Result<()> {
if let Some(info) = self.find_info_by_udp_token(event.token()) { let err = "udp connection state not found";
let info = info.clone(); let state = self.connection_map.get_mut(info).ok_or(err)?;
let err = "udp connection state not found"; state.expiry = Some(Self::udp_associate_timeout());
let state = self.connection_map.get_mut(&info).ok_or(err)?; let mut to_send: LinkedList<Vec<u8>> = LinkedList::new();
state.expiry = Some(Self::udp_associate_timeout()); if let Some(udp_socket) = state.udp_socket.as_ref() {
let mut to_send: LinkedList<Vec<u8>> = LinkedList::new(); let mut buf = [0; 1 << 16];
// Receive UDP packet from remote SOCKS5 server
while let Ok((packet_size, _svr_addr)) = udp_socket.recv_from(&mut buf) {
let buf = buf[..packet_size].to_vec();
let header = UdpHeader::retrieve_from_stream(&mut &buf[..])?;
let buf = if info.dst.port() == DNS_PORT {
let mut message = dns::parse_data_to_dns_message(&buf[header.len()..], false)?;
dns::remove_ipv6_entries(&mut message); // TODO: Configurable
message.to_vec()?
} else {
buf[header.len()..].to_vec()
};
// Escape the borrow checker madness
to_send.push_back(buf);
}
}
// Write to client
let src = state.udp_origin_dst.ok_or("udp address")?;
while let Some(packet) = to_send.pop_front() {
self.send_udp_packet_to_client(src, info.src, &packet)?;
}
Ok(())
}
fn comsume_cached_udp_packets(&mut self, info: &ConnectionInfo) -> Result<()> {
// Try to send the first UDP packets to remote SOCKS5 server for UDP associate session
if let Some(state) = self.connection_map.get_mut(info) {
if let Some(udp_socket) = state.udp_socket.as_ref() { if let Some(udp_socket) = state.udp_socket.as_ref() {
let mut buf = [0; 1 << 16]; if let Some(addr) = state.tcp_proxy_handler.get_udp_associate() {
// Receive UDP packet from remote SOCKS5 server // Consume udp_data_cache data
while let Ok((packet_size, _svr_addr)) = udp_socket.recv_from(&mut buf) { while let Some(buf) = state.udp_data_cache.pop_front() {
let buf = buf[..packet_size].to_vec(); udp_socket.send_to(&buf, addr)?;
let header = UdpHeader::retrieve_from_stream(&mut &buf[..])?; }
let buf = if info.dst.port() == DNS_PORT {
let mut message = dns::parse_data_to_dns_message(&buf[header.len()..], false)?;
dns::remove_ipv6_entries(&mut message); // TODO: Configurable
message.to_vec()?
} else {
buf[header.len()..].to_vec()
};
// Escape the borrow checker madness
to_send.push_back(buf);
} }
} }
}
Ok(())
}
// Write to client fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> {
let src = state.udp_origin_dst.ok_or("udp address")?; if let Some(info) = self.find_info_by_udp_token(event.token()) {
while let Some(packet) = to_send.pop_front() { return self.receive_udp_packet_and_write_to_client(&info.clone());
self.send_udp_packet_to_client(src, info.src, &packet)?;
}
return Ok(());
} }
let conn_info = match self.find_info_by_token(event.token()) { let conn_info = match self.find_info_by_token(event.token()) {
@ -872,17 +889,7 @@ impl<'a> TunToProxy<'a> {
// server. // server.
self.write_to_server(&conn_info)?; self.write_to_server(&conn_info)?;
// Try to send the first UDP packet to remote SOCKS5 server for UDP associate session self.comsume_cached_udp_packets(&conn_info)?;
if let Some(state) = self.connection_map.get_mut(&conn_info) {
if let Some(udp_socket) = state.udp_socket.as_ref() {
if let Some(addr) = state.tcp_proxy_handler.get_udp_associate() {
// Consume udp_data_cache data
while let Some(buf) = state.udp_data_cache.pop_front() {
udp_socket.send_to(&buf, addr)?;
}
}
}
}
} }
if event.is_writable() { if event.is_writable() {