re-format code

This commit is contained in:
ssrlive 2023-08-23 09:28:11 +08:00
parent cdbed3ed9b
commit 72a00af0ed
2 changed files with 31 additions and 45 deletions

View file

@ -78,7 +78,6 @@ pub fn extract_ipaddr_from_dns_message(message: &Message) -> Result<IpAddr, Stri
if let Some(cname) = cname {
return Err(cname);
}
Err(format!("{:?}", message.answers()))
}

View file

@ -1,3 +1,13 @@
use crate::{dns, error::Error, error::Result, virtdevice::VirtualTunDevice, NetworkInterface, Options};
use mio::{event::Event, net::TcpStream, net::UdpSocket, unix::SourceFd, Events, Interest, Poll, Token};
use smoltcp::{
iface::{Config, Interface, SocketHandle, SocketSet},
phy::{Device, Medium, RxToken, TunTapInterface, TxToken},
socket::{tcp, tcp::State, udp, udp::UdpMetadata},
time::Instant,
wire::{IpCidr, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket, UDP_HEADER_LEN},
};
use socks5_impl::protocol::{Address, StreamOperation, UdpHeader, UserKey};
use std::collections::LinkedList;
use std::{
collections::{HashMap, HashSet},
@ -9,18 +19,6 @@ use std::{
str::FromStr,
};
use mio::{event::Event, net::TcpStream, net::UdpSocket, unix::SourceFd, Events, Interest, Poll, Token};
use smoltcp::{
iface::{Config, Interface, SocketHandle, SocketSet},
phy::{Device, Medium, RxToken, TunTapInterface, TxToken},
socket::{tcp, tcp::State, udp, udp::UdpMetadata},
time::Instant,
wire::{IpCidr, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket, UDP_HEADER_LEN},
};
use socks5_impl::protocol::{Address, StreamOperation, UdpHeader, UserKey};
use crate::{dns, error::Error, error::Result, virtdevice::VirtualTunDevice, NetworkInterface, Options};
#[derive(Hash, Clone, Eq, PartialEq, PartialOrd, Ord, Debug)]
pub(crate) struct ConnectionInfo {
pub(crate) src: SocketAddr,
@ -169,11 +167,10 @@ fn connection_tuple(frame: &[u8]) -> Result<(ConnectionInfo, bool, usize, usize)
const SERVER_WRITE_CLOSED: u8 = 1;
const CLIENT_WRITE_CLOSED: u8 = 2;
const UDP_ASSO_TIMEOUT: u64 = 10;
// seconds
const UDP_ASSO_TIMEOUT: u64 = 10; // seconds
const DNS_PORT: u16 = 53;
struct TcpConnectState {
struct ConnectionState {
smoltcp_handle: Option<SocketHandle>,
mio_stream: TcpStream,
token: Token,
@ -187,7 +184,6 @@ struct TcpConnectState {
udp_origin_dst: Option<SocketAddr>,
udp_data_cache: LinkedList<Vec<u8>>,
udp_over_tcp_expiry: Option<::std::time::Instant>,
udp_over_tcp_data_cache: LinkedList<Vec<u8>>,
is_tcp_dns: bool,
}
@ -216,7 +212,7 @@ pub struct TunToProxy<'a> {
tun: TunTapInterface,
poll: Poll,
iface: Interface,
connection_map: HashMap<ConnectionInfo, TcpConnectState>,
connection_map: HashMap<ConnectionInfo, ConnectionState>,
connection_manager: Option<Rc<dyn ConnectionManager>>,
next_token: usize,
sockets: SocketSet<'a>,
@ -440,7 +436,7 @@ impl<'a> TunToProxy<'a> {
Ok(())
}
fn update_mio_socket_interest(poll: &mut Poll, state: &mut TcpConnectState) -> Result<()> {
fn update_mio_socket_interest(poll: &mut Poll, state: &mut ConnectionState) -> Result<()> {
// Maybe we did not listen for any events before. Therefore, just swallow the error.
if let Err(err) = poll.registry().deregister(&mut state.mio_stream) {
log::trace!("{}", err);
@ -478,7 +474,7 @@ impl<'a> TunToProxy<'a> {
Ok(connection_info)
}
fn process_incoming_udp_packets_dns_over_tcp(
fn process_incoming_dns_over_tcp_packets(
&mut self,
manager: &Rc<dyn ConnectionManager>,
original_info: &ConnectionInfo,
@ -502,6 +498,7 @@ impl<'a> TunToProxy<'a> {
state.udp_origin_dst = Some(SocketAddr::try_from(original_info.dst.clone())?);
self.connection_map.insert(info.clone(), state);
// TODO: Move this 3 lines to the function end?
self.expect_smoltcp_send()?;
self.tunsocket_read_and_forward(info)?;
self.write_to_server(info)?;
@ -527,15 +524,6 @@ impl<'a> TunToProxy<'a> {
Ok(())
}
fn consume_cached_dns_over_tcp_packets(&mut self, info: &ConnectionInfo) -> Result<()> {
if let Some(state) = self.connection_map.get_mut(info) {
while let Some(buf) = state.udp_over_tcp_data_cache.pop_front() {
_ = state.mio_stream.write(&buf)?;
}
}
Ok(())
}
fn receive_dns_over_tcp_packet_and_write_to_client(&mut self, info: &ConnectionInfo) -> Result<()> {
let err = "udp connection state not found";
let state = self.connection_map.get_mut(info).ok_or(err)?;
@ -580,18 +568,19 @@ impl<'a> TunToProxy<'a> {
}
let data = buf[2..len + 2].to_vec();
let message = dns::parse_data_to_dns_message(&data, false)?;
let mut message = dns::parse_data_to_dns_message(&data, false)?;
let name = dns::extract_domain_from_dns_message(&message)?;
let ip = dns::extract_ipaddr_from_dns_message(&message);
log::info!("DNS over TCP query result: {} -> {:?}", name, ip);
if let (Ok(name), Ok(ip)) = (
dns::extract_domain_from_dns_message(&message),
dns::extract_ipaddr_from_dns_message(&message),
) {
log::info!("DNS over TCP ======== {} -> {}", name, ip);
}
state
.tcp_proxy_handler
.consume_data(OutgoingDirection::ToClient, len + 2);
to_send.push_back(data);
dns::remove_ipv6_entries(&mut message); // TODO: Configurable
to_send.push_back(message.to_vec()?);
if len + 2 == buf.len() {
break;
}
@ -615,7 +604,7 @@ impl<'a> TunToProxy<'a> {
false
}
fn clearup_expired_udp_over_tcp(&mut self) -> Result<()> {
fn clearup_expired_dns_over_tcp(&mut self) -> Result<()> {
let keys = self.connection_map.keys().cloned().collect::<Vec<_>>();
for key in keys {
if self.udp_over_tcp_timeout_expired(&key) {
@ -724,7 +713,7 @@ impl<'a> TunToProxy<'a> {
} else {
// Another UDP packet
if self.options.dns_over_tcp && origin_dst.port() == DNS_PORT {
self.process_incoming_udp_packets_dns_over_tcp(&manager, &info, origin_dst, payload)?;
self.process_incoming_dns_over_tcp_packets(&manager, &info, origin_dst, payload)?;
} else {
self.process_incoming_udp_packets(&manager, &info, origin_dst, payload)?;
}
@ -746,7 +735,7 @@ impl<'a> TunToProxy<'a> {
dst: SocketAddr,
tcp_proxy_handler: Box<dyn TcpProxy>,
udp_associate: bool,
) -> Result<TcpConnectState> {
) -> Result<ConnectionState> {
let mut socket = tcp::Socket::new(
tcp::SocketBuffer::new(vec![0; 1024 * 128]),
tcp::SocketBuffer::new(vec![0; 1024 * 128]),
@ -775,7 +764,7 @@ impl<'a> TunToProxy<'a> {
} else {
(None, None)
};
let state = TcpConnectState {
let state = ConnectionState {
smoltcp_handle: Some(handle),
mio_stream: client,
token,
@ -789,7 +778,6 @@ impl<'a> TunToProxy<'a> {
udp_origin_dst: None,
udp_data_cache: LinkedList::new(),
udp_over_tcp_expiry: None,
udp_over_tcp_data_cache: LinkedList::new(),
is_tcp_dns: false,
};
Ok(state)
@ -1068,7 +1056,6 @@ impl<'a> TunToProxy<'a> {
// server.
self.write_to_server(&conn_info)?;
self.consume_cached_dns_over_tcp_packets(&conn_info)?;
self.consume_cached_udp_packets(&conn_info)?;
}
@ -1106,7 +1093,7 @@ impl<'a> TunToProxy<'a> {
}
self.send_to_smoltcp()?;
self.clearup_expired_udp_associate()?;
self.clearup_expired_udp_over_tcp()?;
self.clearup_expired_dns_over_tcp()?;
}
}