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

Merge improved error handling from ssrlive

Browse source
This commit is contained in:
B. Blechschmidt 2023-03-24 15:31:05 +01:00
commit 5930b9b3b1
7 changed files with 211 additions and 203 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
src/error.rs
View file

@ -3,23 +3,51 @@ pub struct Error {
message: String, message: String,
} }
pub fn s2e(s: &str) -> Error {
Error::from(s)
}
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 {
fn from(err: std::net::AddrParseError) -> Self {
From::<String>::from(err.to_string())
}
}
impl From<smoltcp::iface::RouteTableFull> for Error {
fn from(err: smoltcp::iface::RouteTableFull) -> Self {
From::<String>::from(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(),
}
} }
} }
@ -31,9 +59,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(),
}
} }
} }

15
src/http.rs
View file

@ -62,19 +62,16 @@ impl HttpConnection {
} }
fn state_change(&mut self) -> Result<(), Error> { fn state_change(&mut self) -> Result<(), Error> {
let http_len = "HTTP/1.1 200".len();
match self.state { match self.state {
HttpState::ExpectStatusCode if self.server_inbuf.len() >= "HTTP/1.1 200 ".len() => { HttpState::ExpectStatusCode if self.server_inbuf.len() > http_len => {
let status_line: Vec<u8> = self let status_line: Vec<u8> =
.server_inbuf self.server_inbuf.range(0..http_len + 1).copied().collect();
.range(0.."HTTP/1.1 200 ".len())
.copied()
.collect();
let slice = &status_line.as_slice()[0.."HTTP/1.1 2".len()]; let slice = &status_line.as_slice()[0.."HTTP/1.1 2".len()];
if slice != b"HTTP/1.1 2" && slice != b"HTTP/1.0 2" if slice != b"HTTP/1.1 2" && slice != b"HTTP/1.0 2"
|| self.server_inbuf["HTTP/1.1 200 ".len() - 1] != b' ' || self.server_inbuf[http_len] != b' '
{ {
let status_str = let status_str = String::from_utf8_lossy(&status_line.as_slice()[0..http_len]);
String::from_utf8_lossy(&status_line.as_slice()[0.."HTTP/1.1 200".len()]);
let e = let e =
format!("Expected success status code. Server replied with {status_str}."); format!("Expected success status code. Server replied with {status_str}.");
return Err(e.into()); return Err(e.into());

22
src/lib.rs
View file

@ -1,4 +1,4 @@
use crate::error::{s2e, Error}; use crate::error::Error;
use crate::tun2proxy::{Credentials, Options}; use crate::tun2proxy::{Credentials, Options};
use crate::{http::HttpManager, socks5::Socks5Manager, tun2proxy::TunToProxy}; use crate::{http::HttpManager, socks5::Socks5Manager, tun2proxy::TunToProxy};
use std::net::{SocketAddr, ToSocketAddrs}; use std::net::{SocketAddr, ToSocketAddrs};
@ -20,21 +20,21 @@ pub struct Proxy {
impl Proxy { impl Proxy {
pub fn from_url(s: &str) -> Result<Proxy, Error> { pub fn from_url(s: &str) -> Result<Proxy, Error> {
let e = format!("`{s}` is not a valid proxy URL"); let e = format!("`{s}` is not a valid proxy URL");
let url = url::Url::parse(s).map_err(|_| s2e(&e))?; let url = url::Url::parse(s).map_err(|_| Error::from(&e))?;
let e = format!("`{s}` does not contain a host"); let e = format!("`{s}` does not contain a host");
let host = url.host_str().ok_or(s2e(&e))?; let host = url.host_str().ok_or(Error::from(e))?;
let mut url_host = String::from(host); let mut url_host = String::from(host);
let e = format!("`{s}` does not contain a port"); let e = format!("`{s}` does not contain a port");
let port = url.port().ok_or(s2e(&e))?; let port = url.port().ok_or(Error::from(&e))?;
url_host.push(':'); url_host.push(':');
url_host.push_str(port.to_string().as_str()); url_host.push_str(port.to_string().as_str());
let e = format!("`{host}` could not be resolved"); let e = format!("`{host}` could not be resolved");
let mut addr_iter = url_host.to_socket_addrs().map_err(|_| s2e(&e))?; let mut addr_iter = url_host.to_socket_addrs().map_err(|_| Error::from(&e))?;
let e = format!("`{host}` does not resolve to a usable IP address"); let e = format!("`{host}` does not resolve to a usable IP address");
let addr = addr_iter.next().ok_or(s2e(&e))?; let addr = addr_iter.next().ok_or(Error::from(&e))?;
let credentials = if url.username() == "" && url.password().is_none() { let credentials = if url.username() == "" && url.password().is_none() {
None None
@ -51,7 +51,7 @@ impl Proxy {
"http" => Some(ProxyType::Http), "http" => Some(ProxyType::Http),
_ => None, _ => None,
} }
.ok_or(s2e(&format!("`{scheme}` is an invalid proxy type")))?; .ok_or(Error::from(&format!("`{scheme}` is an invalid proxy type")))?;
Ok(Proxy { Ok(Proxy {
proxy_type, proxy_type,
@ -76,8 +76,8 @@ impl std::fmt::Display for ProxyType {
} }
} }
pub fn main_entry(tun: &str, proxy: Proxy, options: Options) { pub fn main_entry(tun: &str, proxy: Proxy, options: Options) -> Result<(), Error> {
let mut ttp = TunToProxy::new(tun, options); let mut ttp = TunToProxy::new(tun, options)?;
match proxy.proxy_type { match proxy.proxy_type {
ProxyType::Socks5 => { ProxyType::Socks5 => {
ttp.add_connection_manager(Socks5Manager::new(proxy.addr, proxy.credentials)); ttp.add_connection_manager(Socks5Manager::new(proxy.addr, proxy.credentials));
@ -86,7 +86,5 @@ pub fn main_entry(tun: &str, proxy: Proxy, options: Options) {
ttp.add_connection_manager(HttpManager::new(proxy.addr, proxy.credentials)); ttp.add_connection_manager(HttpManager::new(proxy.addr, proxy.credentials));
} }
} }
if let Err(e) = ttp.run() { ttp.run()
log::error!("{e}");
}
} }

4
src/main.rs
View file

@ -47,5 +47,7 @@ fn main() {
options = options.with_virtual_dns(); options = options.with_virtual_dns();
} }
main_entry(&args.tun, args.proxy, options); if let Err(e) = main_entry(&args.tun, args.proxy, options) {
log::error!("{e}");
}
} }

306
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};
@ -17,7 +17,7 @@ use std::convert::{From, TryFrom};
use std::fmt::{Display, Formatter}; use std::fmt::{Display, Formatter};
use std::io::{Read, Write}; use std::io::{Read, Write};
use std::net::Shutdown::Both; use std::net::Shutdown::Both;
use std::net::{IpAddr, SocketAddr}; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::os::unix::io::AsRawFd; use std::os::unix::io::AsRawFd;
use std::rc::Rc; use std::rc::Rc;
use std::str::FromStr; use std::str::FromStr;
@ -286,64 +286,49 @@ pub(crate) struct TunToProxy<'a> {
} }
impl<'a> TunToProxy<'a> { impl<'a> TunToProxy<'a> {
pub(crate) fn new(interface: &str, options: Options) -> Self { pub(crate) fn new(interface: &str, options: Options) -> Result<Self, Error> {
let tun = TunTapInterface::new(interface, Medium::Ip).unwrap(); let tun = TunTapInterface::new(interface, Medium::Ip)?;
let poll = Poll::new().unwrap(); let poll = Poll::new()?;
poll.registry() poll.registry().register(
.register( &mut SourceFd(&tun.as_raw_fd()),
&mut SourceFd(&tun.as_raw_fd()), TCP_TOKEN,
TCP_TOKEN, Interest::READABLE,
Interest::READABLE, )?;
)
.unwrap();
let config = Config::new(); let config = Config::new();
let mut virt = VirtualTunDevice::new(tun.capabilities()); let mut virt = VirtualTunDevice::new(tun.capabilities());
let gateway4: Ipv4Addr = Ipv4Addr::from_str("0.0.0.1")?;
let gateway6: Ipv6Addr = Ipv6Addr::from_str("::1")?;
let mut iface = Interface::new(config, &mut virt); let mut iface = Interface::new(config, &mut virt);
iface.update_ip_addrs(|ip_addrs| { iface.update_ip_addrs(|ip_addrs| {
ip_addrs ip_addrs.push(IpCidr::new(gateway4.into(), 0)).unwrap();
.push(IpCidr::new( ip_addrs.push(IpCidr::new(gateway6.into(), 0)).unwrap()
std::net::Ipv4Addr::from_str("0.0.0.1").unwrap().into(),
0,
))
.unwrap();
ip_addrs
.push(IpCidr::new(
std::net::Ipv6Addr::from_str("::1").unwrap().into(),
0,
))
.unwrap()
}); });
iface iface.routes_mut().add_default_ipv4_route(gateway4.into())?;
.routes_mut() iface.routes_mut().add_default_ipv6_route(gateway6.into())?;
.add_default_ipv4_route(std::net::Ipv4Addr::from_str("0.0.0.1").unwrap().into())
.unwrap();
iface
.routes_mut()
.add_default_ipv6_route(std::net::Ipv6Addr::from_str("::1").unwrap().into())
.unwrap();
iface.set_any_ip(true); iface.set_any_ip(true);
Self { let tun = Self {
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)
} }
pub(crate) fn add_connection_manager(&mut self, manager: Rc<dyn ConnectionManager>) { pub(crate) fn add_connection_manager(&mut self, manager: Rc<dyn ConnectionManager>) {
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);
@ -353,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>> {
@ -380,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 {
@ -409,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 &mut self.options.virtdns { let resolved_conn = match &mut 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();
virt_dns.touch_ip(&ip); virt_dns.touch_ip(&ip);
match virt_dns.resolve_ip(&ip) { match virt_dns.resolve_ip(&ip) {
None => connection.clone(), None => connection.clone(),
@ -437,9 +422,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) =
@ -450,11 +435,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;
@ -469,14 +454,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);
@ -485,7 +467,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
@ -494,10 +476,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.
@ -507,28 +489,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) {
@ -554,7 +533,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;
@ -569,11 +548,11 @@ impl<'a> TunToProxy<'a> {
// Unwrapping is fine because every smoltcp socket is bound to an. // Unwrapping is fine because every smoltcp socket is bound to an.
virtdns.touch_ip(&IpAddr::from(socket.local_endpoint().unwrap().addr)); virtdns.touch_ip(&IpAddr::from(socket.local_endpoint().unwrap().addr));
} }
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;
@ -590,97 +569,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) {}
@ -692,12 +676,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()?;
} }
} }
} }

13
src/virtdns.rs
View file

@ -137,7 +137,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()),
@ -155,13 +155,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)
} }
// This is to be called whenever we receive or send a packet on the socket // This is to be called whenever we receive or send a packet on the socket
@ -226,7 +227,7 @@ impl VirtualDns {
self.name_to_ip.insert(name, self.next_addr); self.name_to_ip.insert(name, self.next_addr);
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;

2
tests/proxy.rs
View file

@ -146,7 +146,7 @@ mod tests {
} }
Ok(Fork::Child) => { Ok(Fork::Child) => {
prctl::set_death_signal(signal::SIGKILL as isize).unwrap(); // 9 == SIGKILL prctl::set_death_signal(signal::SIGKILL as isize).unwrap(); // 9 == SIGKILL
main_entry( let _ = main_entry(
TUN_TEST_DEVICE, TUN_TEST_DEVICE,
test.proxy, test.proxy,
Options::new().with_virtual_dns(), Options::new().with_virtual_dns(),