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

remove lots of unwrap()

Browse source
This commit is contained in:
ssrlive 2023-03-24 14:31:22 +08:00
parent ad388f897a
commit 1d1e50c36a
3 changed files with 164 additions and 150 deletions
Download patch file Download diff file Expand all files Collapse all files

44
src/error.rs
View file

@ -9,33 +9,49 @@ pub fn s2e(s: &str) -> Error {
impl From<std::io::Error> for Error { impl From<std::io::Error> for Error {
fn from(err: std::io::Error) -> Self { fn from(err: std::io::Error) -> Self {
Self { From::<String>::from(err.to_string())
message: err.to_string(),
}
} }
} }
impl From<std::net::AddrParseError> for Error { impl From<std::net::AddrParseError> for Error {
fn from(err: std::net::AddrParseError) -> Self { fn from(err: std::net::AddrParseError) -> Self {
Self { From::<String>::from(err.to_string())
message: err.to_string(),
}
} }
} }
impl From<smoltcp::iface::RouteTableFull> for Error { impl From<smoltcp::iface::RouteTableFull> for Error {
fn from(err: smoltcp::iface::RouteTableFull) -> Self { fn from(err: smoltcp::iface::RouteTableFull) -> Self {
Self { From::<String>::from(format!("{err:?}"))
message: format!("{err:?}"), }
} }
impl From<smoltcp::socket::tcp::RecvError> for Error {
fn from(err: smoltcp::socket::tcp::RecvError) -> Self {
From::<String>::from(format!("{err:?}"))
}
}
impl From<smoltcp::socket::tcp::ListenError> for Error {
fn from(err: smoltcp::socket::tcp::ListenError) -> Self {
From::<String>::from(format!("{err:?}"))
}
}
impl From<smoltcp::socket::udp::BindError> for Error {
fn from(err: smoltcp::socket::udp::BindError) -> Self {
From::<String>::from(format!("{err:?}"))
}
}
impl From<smoltcp::socket::tcp::SendError> for Error {
fn from(err: smoltcp::socket::tcp::SendError) -> Self {
From::<String>::from(format!("{err:?}"))
} }
} }
impl From<&str> for Error { impl From<&str> for Error {
fn from(err: &str) -> Self { fn from(err: &str) -> Self {
Self { From::<String>::from(err.to_string())
message: err.to_string(),
}
} }
} }
@ -47,9 +63,7 @@ impl From<String> for Error {
impl From<&String> for Error { impl From<&String> for Error {
fn from(err: &String) -> Self { fn from(err: &String) -> Self {
Self { From::<String>::from(err.to_string())
message: err.to_string(),
}
} }
} }

255
src/tun2proxy.rs
View file

@ -9,7 +9,7 @@ use mio::unix::SourceFd;
use mio::{Events, Interest, Poll, Token}; use mio::{Events, Interest, Poll, Token};
use smoltcp::iface::{Config, Interface, SocketHandle, SocketSet}; use smoltcp::iface::{Config, Interface, SocketHandle, SocketSet};
use smoltcp::phy::{Device, Medium, RxToken, TunTapInterface, TxToken}; use smoltcp::phy::{Device, Medium, RxToken, TunTapInterface, TxToken};
use smoltcp::socket::tcp; use smoltcp::socket::{tcp, udp};
use smoltcp::time::Instant; use smoltcp::time::Instant;
use smoltcp::wire::{IpCidr, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket}; use smoltcp::wire::{IpCidr, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket};
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
@ -312,14 +312,14 @@ impl<'a> TunToProxy<'a> {
tun, tun,
poll, poll,
iface, iface,
connections: Default::default(), connections: HashMap::default(),
next_token: 2, next_token: 2,
token_to_connection: Default::default(), token_to_connection: HashMap::default(),
connection_managers: Default::default(), connection_managers: Vec::default(),
sockets: SocketSet::new([]), sockets: SocketSet::new([]),
device: virt, device: virt,
options, options,
write_sockets: Default::default(), write_sockets: HashSet::default(),
}; };
Ok(tun) Ok(tun)
} }
@ -328,7 +328,7 @@ impl<'a> TunToProxy<'a> {
self.connection_managers.push(manager); self.connection_managers.push(manager);
} }
fn expect_smoltcp_send(&mut self) { fn expect_smoltcp_send(&mut self) -> Result<(), Error> {
self.iface self.iface
.poll(Instant::now(), &mut self.device, &mut self.sockets); .poll(Instant::now(), &mut self.device, &mut self.sockets);
@ -338,22 +338,22 @@ impl<'a> TunToProxy<'a> {
// TODO: Actual write. Replace. // TODO: Actual write. Replace.
self.tun self.tun
.transmit(Instant::now()) .transmit(Instant::now())
.unwrap() .ok_or("tx token not available")?
.consume(slice.len(), |buf| { .consume(slice.len(), |buf| {
buf[..].clone_from_slice(slice); buf[..].clone_from_slice(slice);
}); });
} }
Ok(())
} }
fn remove_connection(&mut self, connection: &Connection) { fn remove_connection(&mut self, connection: &Connection) -> Result<(), Error> {
let mut connection_state = self.connections.remove(connection).unwrap(); let e = "connection not exist";
let token = &connection_state.token; let mut conn = self.connections.remove(connection).ok_or(e)?;
let token = &conn.token;
self.token_to_connection.remove(token); self.token_to_connection.remove(token);
self.poll self.poll.registry().deregister(&mut conn.mio_stream)?;
.registry()
.deregister(&mut connection_state.mio_stream)
.unwrap();
info!("CLOSE {}", connection); info!("CLOSE {}", connection);
Ok(())
} }
fn get_connection_manager(&self, connection: &Connection) -> Option<Rc<dyn ConnectionManager>> { fn get_connection_manager(&self, connection: &Connection) -> Option<Rc<dyn ConnectionManager>> {
@ -365,23 +365,21 @@ impl<'a> TunToProxy<'a> {
None None
} }
fn tunsocket_read_and_forward(&mut self, connection: &Connection) { fn tunsocket_read_and_forward(&mut self, connection: &Connection) -> Result<(), Error> {
if let Some(state) = self.connections.get_mut(connection) { if let Some(state) = self.connections.get_mut(connection) {
let closed = { let closed = {
let socket = self.sockets.get_mut::<tcp::Socket>(state.smoltcp_handle); let socket = self.sockets.get_mut::<tcp::Socket>(state.smoltcp_handle);
let mut error = Ok(()); let mut error = Ok(());
while socket.can_recv() && error.is_ok() { while socket.can_recv() && error.is_ok() {
socket socket.recv(|data| {
.recv(|data| { let event = IncomingDataEvent {
let event = IncomingDataEvent { direction: IncomingDirection::FromClient,
direction: IncomingDirection::FromClient, buffer: data,
buffer: data, };
}; error = state.handler.push_data(event);
error = state.handler.push_data(event);
(data.len(), ()) (data.len(), ())
}) })?;
.unwrap();
} }
match error { match error {
@ -394,24 +392,26 @@ impl<'a> TunToProxy<'a> {
}; };
// Expect ACKs etc. from smoltcp sockets. // Expect ACKs etc. from smoltcp sockets.
self.expect_smoltcp_send(); self.expect_smoltcp_send()?;
if closed { if closed {
let connection_state = self.connections.get_mut(connection).unwrap(); let e = "connection not exist";
connection_state.mio_stream.shutdown(Both).unwrap(); let connection_state = self.connections.get_mut(connection).ok_or(e)?;
self.remove_connection(connection); connection_state.mio_stream.shutdown(Both)?;
self.remove_connection(connection)?;
} }
} }
Ok(())
} }
fn receive_tun(&mut self, frame: &mut [u8]) { fn receive_tun(&mut self, frame: &mut [u8]) -> Result<(), Error> {
if let Some((connection, first_packet, _payload_offset, _payload_size)) = if let Some((connection, first_packet, _payload_offset, _payload_size)) =
connection_tuple(frame) connection_tuple(frame)
{ {
let resolved_conn = match &self.options.virtdns { let resolved_conn = match &self.options.virtdns {
None => connection.clone(), None => connection.clone(),
Some(virt_dns) => { Some(virt_dns) => {
let ip = SocketAddr::try_from(connection.dst.clone()).unwrap().ip(); let ip = SocketAddr::try_from(connection.dst.clone())?.ip();
match virt_dns.ip_to_name(&ip) { match virt_dns.ip_to_name(&ip) {
None => connection.clone(), None => connection.clone(),
Some(name) => connection.to_named(name.clone()), Some(name) => connection.to_named(name.clone()),
@ -421,9 +421,9 @@ impl<'a> TunToProxy<'a> {
if resolved_conn.proto == IpProtocol::Tcp.into() { if resolved_conn.proto == IpProtocol::Tcp.into() {
let cm = self.get_connection_manager(&resolved_conn); let cm = self.get_connection_manager(&resolved_conn);
if cm.is_none() { if cm.is_none() {
return; return Ok(());
} }
let server = cm.unwrap().get_server(); let server = cm.ok_or("no connect manager")?.get_server();
if first_packet { if first_packet {
for manager in self.connection_managers.iter_mut() { for manager in self.connection_managers.iter_mut() {
if let Some(handler) = if let Some(handler) =
@ -434,11 +434,11 @@ impl<'a> TunToProxy<'a> {
tcp::SocketBuffer::new(vec![0; 4096]), tcp::SocketBuffer::new(vec![0; 4096]),
); );
socket.set_ack_delay(None); socket.set_ack_delay(None);
let dst = SocketAddr::try_from(connection.dst).unwrap(); let dst = SocketAddr::try_from(connection.dst)?;
socket.listen(dst).unwrap(); socket.listen(dst)?;
let handle = self.sockets.add(socket); let handle = self.sockets.add(socket);
let client = TcpStream::connect(server).unwrap(); let client = TcpStream::connect(server)?;
let token = Token(self.next_token); let token = Token(self.next_token);
self.next_token += 1; self.next_token += 1;
@ -453,14 +453,11 @@ impl<'a> TunToProxy<'a> {
self.token_to_connection self.token_to_connection
.insert(token, resolved_conn.clone()); .insert(token, resolved_conn.clone());
self.poll self.poll.registry().register(
.registry() &mut state.mio_stream,
.register( token,
&mut state.mio_stream, Interest::READABLE | Interest::WRITABLE,
token, )?;
Interest::READABLE | Interest::WRITABLE,
)
.unwrap();
self.connections.insert(resolved_conn.clone(), state); self.connections.insert(resolved_conn.clone(), state);
@ -469,7 +466,7 @@ impl<'a> TunToProxy<'a> {
} }
} }
} else if !self.connections.contains_key(&resolved_conn) { } else if !self.connections.contains_key(&resolved_conn) {
return; return Ok(());
} }
// Inject the packet to advance the smoltcp socket state // Inject the packet to advance the smoltcp socket state
@ -478,10 +475,10 @@ impl<'a> TunToProxy<'a> {
// Having advanced the socket state, we expect the socket to ACK // Having advanced the socket state, we expect the socket to ACK
// Exfiltrate the response packets generated by the socket and inject them // Exfiltrate the response packets generated by the socket and inject them
// into the tunnel interface. // into the tunnel interface.
self.expect_smoltcp_send(); self.expect_smoltcp_send()?;
// Read from the smoltcp socket and push the data to the connection handler. // Read from the smoltcp socket and push the data to the connection handler.
self.tunsocket_read_and_forward(&resolved_conn); self.tunsocket_read_and_forward(&resolved_conn)?;
// The connection handler builds up the connection or encapsulates the data. // The connection handler builds up the connection or encapsulates the data.
// Therefore, we now expect it to write data to the server. // Therefore, we now expect it to write data to the server.
@ -491,28 +488,25 @@ impl<'a> TunToProxy<'a> {
if let Some(virtual_dns) = &mut self.options.virtdns { if let Some(virtual_dns) = &mut self.options.virtdns {
let payload = &frame[_payload_offset.._payload_offset + _payload_size]; let payload = &frame[_payload_offset.._payload_offset + _payload_size];
if let Some(response) = virtual_dns.receive_query(payload) { if let Some(response) = virtual_dns.receive_query(payload) {
let rx_buffer = smoltcp::socket::udp::PacketBuffer::new( let rx_buffer =
vec![smoltcp::socket::udp::PacketMetadata::EMPTY], udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
vec![0; 4096], let tx_buffer =
); udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
let tx_buffer = smoltcp::socket::udp::PacketBuffer::new( let mut socket = udp::Socket::new(rx_buffer, tx_buffer);
vec![smoltcp::socket::udp::PacketMetadata::EMPTY], let dst = SocketAddr::try_from(connection.dst)?;
vec![0; 4096], socket.bind(dst)?;
);
let mut socket = smoltcp::socket::udp::Socket::new(rx_buffer, tx_buffer);
let dst = SocketAddr::try_from(connection.dst).unwrap();
socket.bind(dst).unwrap();
socket socket
.send_slice(response.as_slice(), resolved_conn.src.into()) .send_slice(response.as_slice(), resolved_conn.src.into())
.expect("failed to send DNS response"); .expect("failed to send DNS response");
let handle = self.sockets.add(socket); let handle = self.sockets.add(socket);
self.expect_smoltcp_send(); self.expect_smoltcp_send()?;
self.sockets.remove(handle); self.sockets.remove(handle);
} }
} }
// Otherwise, UDP is not yet supported. // Otherwise, UDP is not yet supported.
} }
} }
Ok(())
} }
fn write_to_server(&mut self, connection: &Connection) { fn write_to_server(&mut self, connection: &Connection) {
@ -538,7 +532,7 @@ impl<'a> TunToProxy<'a> {
} }
} }
fn write_to_client(&mut self, token: Token, connection: &Connection) { fn write_to_client(&mut self, token: Token, connection: &Connection) -> Result<(), Error> {
loop { loop {
if let Some(state) = self.connections.get_mut(connection) { if let Some(state) = self.connections.get_mut(connection) {
let socket_state = state.smoltcp_socket_state; let socket_state = state.smoltcp_socket_state;
@ -549,11 +543,11 @@ impl<'a> TunToProxy<'a> {
{ {
let socket = self.sockets.get_mut::<tcp::Socket>(socket_handle); let socket = self.sockets.get_mut::<tcp::Socket>(socket_handle);
if socket.may_send() { if socket.may_send() {
consumed = socket.send_slice(event.buffer).unwrap(); consumed = socket.send_slice(event.buffer)?;
state state
.handler .handler
.consume_data(OutgoingDirection::ToClient, consumed); .consume_data(OutgoingDirection::ToClient, consumed);
self.expect_smoltcp_send(); self.expect_smoltcp_send()?;
if consumed < buflen { if consumed < buflen {
self.write_sockets.insert(token); self.write_sockets.insert(token);
break; break;
@ -570,97 +564,102 @@ impl<'a> TunToProxy<'a> {
let socket = self.sockets.get_mut::<tcp::Socket>(socket_handle); let socket = self.sockets.get_mut::<tcp::Socket>(socket_handle);
if socket_state & WRITE_CLOSED != 0 && consumed == buflen { if socket_state & WRITE_CLOSED != 0 && consumed == buflen {
socket.close(); socket.close();
self.expect_smoltcp_send(); self.expect_smoltcp_send()?;
self.write_sockets.remove(&token); self.write_sockets.remove(&token);
self.remove_connection(connection); self.remove_connection(connection)?;
break; break;
} }
} }
} }
Ok(())
} }
fn tun_event(&mut self, event: &Event) { fn tun_event(&mut self, event: &Event) -> Result<(), Error> {
if event.is_readable() { if event.is_readable() {
while let Some((rx_token, _)) = self.tun.receive(Instant::now()) { while let Some((rx_token, _)) = self.tun.receive(Instant::now()) {
rx_token.consume(|frame| { rx_token.consume(|frame| self.receive_tun(frame))?;
self.receive_tun(frame);
});
} }
} }
Ok(())
} }
fn send_to_smoltcp(&mut self) { fn send_to_smoltcp(&mut self) -> Result<(), Error> {
let cloned = self.write_sockets.clone(); let cloned = self.write_sockets.clone();
for token in cloned.iter() { for token in cloned.iter() {
if let Some(connection) = self.token_to_connection.get(token) { if let Some(connection) = self.token_to_connection.get(token) {
self.write_to_client(*token, &connection.clone()); self.write_to_client(*token, &connection.clone())?;
} }
} }
Ok(())
} }
fn mio_socket_event(&mut self, event: &Event) { fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> {
if let Some(conn_ref) = self.token_to_connection.get(&event.token()) { let e = "connection not found";
let connection = conn_ref.clone(); let conn_ref = self.token_to_connection.get(&event.token());
if event.is_readable() || event.is_read_closed() { if conn_ref.is_none() {
{ return Ok(());
let state = self.connections.get_mut(&connection).unwrap(); }
let connection = conn_ref.ok_or(e)?.clone();
if event.is_readable() || event.is_read_closed() {
{
let state = self.connections.get_mut(&connection).ok_or(e)?;
// TODO: Move this reading process to its own function. // TODO: Move this reading process to its own function.
let mut vecbuf = Vec::<u8>::new(); let mut vecbuf = Vec::<u8>::new();
let read_result = state.mio_stream.read_to_end(&mut vecbuf); let read_result = state.mio_stream.read_to_end(&mut vecbuf);
let read = match read_result { let read = match read_result {
Ok(read_result) => read_result, Ok(read_result) => read_result,
Err(error) => { Err(error) => {
if error.kind() != std::io::ErrorKind::WouldBlock { if error.kind() != std::io::ErrorKind::WouldBlock {
error!("READ from proxy: {}", error); error!("READ from proxy: {}", error);
}
vecbuf.len()
} }
}; vecbuf.len()
}
};
if read == 0 { if read == 0 {
{ {
let socket = self.sockets.get_mut::<tcp::Socket>( let socket = self.sockets.get_mut::<tcp::Socket>(
self.connections.get(&connection).unwrap().smoltcp_handle, self.connections.get(&connection).ok_or(e)?.smoltcp_handle,
); );
socket.close(); socket.close();
}
self.expect_smoltcp_send();
self.remove_connection(&connection.clone());
return;
}
let data = vecbuf.as_slice();
let data_event = IncomingDataEvent {
direction: IncomingDirection::FromServer,
buffer: &data[0..read],
};
if let Err(error) = state.handler.push_data(data_event) {
state.mio_stream.shutdown(Both).unwrap();
{
let socket = self.sockets.get_mut::<tcp::Socket>(
self.connections.get(&connection).unwrap().smoltcp_handle,
);
socket.close();
}
self.expect_smoltcp_send();
log::error! {"{error}"}
self.remove_connection(&connection.clone());
return;
}
if event.is_read_closed() {
state.smoltcp_socket_state |= WRITE_CLOSED;
} }
self.expect_smoltcp_send()?;
self.remove_connection(&connection.clone())?;
return Ok(());
} }
// We have read from the proxy server and pushed the data to the connection handler. let data = vecbuf.as_slice();
// Thus, expect data to be processed (e.g. decapsulated) and forwarded to the client. let data_event = IncomingDataEvent {
self.write_to_client(event.token(), &connection); direction: IncomingDirection::FromServer,
} buffer: &data[0..read],
if event.is_writable() { };
self.write_to_server(&connection); if let Err(error) = state.handler.push_data(data_event) {
state.mio_stream.shutdown(Both)?;
{
let socket = self.sockets.get_mut::<tcp::Socket>(
self.connections.get(&connection).ok_or(e)?.smoltcp_handle,
);
socket.close();
}
self.expect_smoltcp_send()?;
log::error! {"{error}"}
self.remove_connection(&connection.clone())?;
return Ok(());
}
if event.is_read_closed() {
state.smoltcp_socket_state |= WRITE_CLOSED;
}
} }
// We have read from the proxy server and pushed the data to the connection handler.
// Thus, expect data to be processed (e.g. decapsulated) and forwarded to the client.
self.write_to_client(event.token(), &connection)?;
} }
if event.is_writable() {
self.write_to_server(&connection);
}
Ok(())
} }
fn udp_event(&mut self, _event: &Event) {} fn udp_event(&mut self, _event: &Event) {}
@ -672,12 +671,12 @@ impl<'a> TunToProxy<'a> {
self.poll.poll(&mut events, None)?; self.poll.poll(&mut events, None)?;
for event in events.iter() { for event in events.iter() {
match event.token() { match event.token() {
TCP_TOKEN => self.tun_event(event), TCP_TOKEN => self.tun_event(event)?,
UDP_TOKEN => self.udp_event(event), UDP_TOKEN => self.udp_event(event),
_ => self.mio_socket_event(event), _ => self.mio_socket_event(event)?,
} }
} }
self.send_to_smoltcp(); self.send_to_smoltcp()?;
} }
} }
} }

15
src/virtdns.rs
View file

@ -133,7 +133,7 @@ impl VirtualDns {
Some(response) Some(response)
} }
fn increment_ip(addr: IpAddr) -> IpAddr { fn increment_ip(addr: IpAddr) -> Option<IpAddr> {
let mut ip_bytes = match addr as IpAddr { let mut ip_bytes = match addr as IpAddr {
IpAddr::V4(ip) => Vec::<u8>::from(ip.octets()), IpAddr::V4(ip) => Vec::<u8>::from(ip.octets()),
IpAddr::V6(ip) => Vec::<u8>::from(ip.octets()), IpAddr::V6(ip) => Vec::<u8>::from(ip.octets()),
@ -151,13 +151,14 @@ impl VirtualDns {
ip_bytes[i] = 0; ip_bytes[i] = 0;
} }
} }
if addr.is_ipv4() { let addr = if addr.is_ipv4() {
let bytes: [u8; 4] = ip_bytes.as_slice().try_into().unwrap(); let bytes: [u8; 4] = ip_bytes.as_slice().try_into().ok()?;
IpAddr::V4(Ipv4Addr::from(bytes)) IpAddr::V4(Ipv4Addr::from(bytes))
} else { } else {
let bytes: [u8; 16] = ip_bytes.as_slice().try_into().unwrap(); let bytes: [u8; 16] = ip_bytes.as_slice().try_into().ok()?;
IpAddr::V6(Ipv6Addr::from(bytes)) IpAddr::V6(Ipv6Addr::from(bytes))
} };
Some(addr)
} }
pub fn ip_to_name(&self, addr: &IpAddr) -> Option<&String> { pub fn ip_to_name(&self, addr: &IpAddr) -> Option<&String> {
@ -168,7 +169,7 @@ impl VirtualDns {
let now = Instant::now(); let now = Instant::now();
while let Some((ip, expiry)) = self.expiry.front() { while let Some((ip, expiry)) = self.expiry.front() {
if now > *expiry { if now > *expiry {
let name = self.ip_to_name.remove(ip).unwrap(); let name = self.ip_to_name.remove(ip)?;
self.name_to_ip.remove(&name); self.name_to_ip.remove(&name);
self.expiry.pop_front(); self.expiry.pop_front();
} else { } else {
@ -194,7 +195,7 @@ impl VirtualDns {
)); ));
return Some(self.next_addr); return Some(self.next_addr);
} }
self.next_addr = Self::increment_ip(self.next_addr); self.next_addr = Self::increment_ip(self.next_addr)?;
if self.next_addr == self.broadcast_addr { if self.next_addr == self.broadcast_addr {
// Wrap around. // Wrap around.
self.next_addr = self.network_addr; self.next_addr = self.network_addr;