mirror of
https://github.com/tun2proxy/tun2proxy.git
synced 2025-04-22 14:59:09 +00:00
Merge improved error handling from ssrlive
This commit is contained in:
commit
5930b9b3b1
7 changed files with 211 additions and 203 deletions
50
src/error.rs
50
src/error.rs
|
@ -3,23 +3,51 @@ pub struct Error {
|
|||
message: String,
|
||||
}
|
||||
|
||||
pub fn s2e(s: &str) -> Error {
|
||||
Error::from(s)
|
||||
}
|
||||
|
||||
impl From<std::io::Error> for Error {
|
||||
fn from(err: std::io::Error) -> Self {
|
||||
Self {
|
||||
message: err.to_string(),
|
||||
From::<String>::from(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 {
|
||||
fn from(err: &str) -> Self {
|
||||
Self {
|
||||
message: err.to_string(),
|
||||
}
|
||||
From::<String>::from(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
|
@ -31,9 +59,7 @@ impl From<String> for Error {
|
|||
|
||||
impl From<&String> for Error {
|
||||
fn from(err: &String) -> Self {
|
||||
Self {
|
||||
message: err.to_string(),
|
||||
}
|
||||
From::<String>::from(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
|
|
15
src/http.rs
15
src/http.rs
|
@ -62,19 +62,16 @@ impl HttpConnection {
|
|||
}
|
||||
|
||||
fn state_change(&mut self) -> Result<(), Error> {
|
||||
let http_len = "HTTP/1.1 200".len();
|
||||
match self.state {
|
||||
HttpState::ExpectStatusCode if self.server_inbuf.len() >= "HTTP/1.1 200 ".len() => {
|
||||
let status_line: Vec<u8> = self
|
||||
.server_inbuf
|
||||
.range(0.."HTTP/1.1 200 ".len())
|
||||
.copied()
|
||||
.collect();
|
||||
HttpState::ExpectStatusCode if self.server_inbuf.len() > http_len => {
|
||||
let status_line: Vec<u8> =
|
||||
self.server_inbuf.range(0..http_len + 1).copied().collect();
|
||||
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"
|
||||
|| self.server_inbuf["HTTP/1.1 200 ".len() - 1] != b' '
|
||||
|| self.server_inbuf[http_len] != b' '
|
||||
{
|
||||
let status_str =
|
||||
String::from_utf8_lossy(&status_line.as_slice()[0.."HTTP/1.1 200".len()]);
|
||||
let status_str = String::from_utf8_lossy(&status_line.as_slice()[0..http_len]);
|
||||
let e =
|
||||
format!("Expected success status code. Server replied with {status_str}.");
|
||||
return Err(e.into());
|
||||
|
|
22
src/lib.rs
22
src/lib.rs
|
@ -1,4 +1,4 @@
|
|||
use crate::error::{s2e, Error};
|
||||
use crate::error::Error;
|
||||
use crate::tun2proxy::{Credentials, Options};
|
||||
use crate::{http::HttpManager, socks5::Socks5Manager, tun2proxy::TunToProxy};
|
||||
use std::net::{SocketAddr, ToSocketAddrs};
|
||||
|
@ -20,21 +20,21 @@ pub struct Proxy {
|
|||
impl Proxy {
|
||||
pub fn from_url(s: &str) -> Result<Proxy, Error> {
|
||||
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 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 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_str(port.to_string().as_str());
|
||||
|
||||
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 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() {
|
||||
None
|
||||
|
@ -51,7 +51,7 @@ impl Proxy {
|
|||
"http" => Some(ProxyType::Http),
|
||||
_ => None,
|
||||
}
|
||||
.ok_or(s2e(&format!("`{scheme}` is an invalid proxy type")))?;
|
||||
.ok_or(Error::from(&format!("`{scheme}` is an invalid proxy type")))?;
|
||||
|
||||
Ok(Proxy {
|
||||
proxy_type,
|
||||
|
@ -76,8 +76,8 @@ impl std::fmt::Display for ProxyType {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn main_entry(tun: &str, proxy: Proxy, options: Options) {
|
||||
let mut ttp = TunToProxy::new(tun, options);
|
||||
pub fn main_entry(tun: &str, proxy: Proxy, options: Options) -> Result<(), Error> {
|
||||
let mut ttp = TunToProxy::new(tun, options)?;
|
||||
match proxy.proxy_type {
|
||||
ProxyType::Socks5 => {
|
||||
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));
|
||||
}
|
||||
}
|
||||
if let Err(e) = ttp.run() {
|
||||
log::error!("{e}");
|
||||
}
|
||||
ttp.run()
|
||||
}
|
||||
|
|
|
@ -47,5 +47,7 @@ fn main() {
|
|||
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}");
|
||||
}
|
||||
}
|
||||
|
|
202
src/tun2proxy.rs
202
src/tun2proxy.rs
|
@ -9,7 +9,7 @@ use mio::unix::SourceFd;
|
|||
use mio::{Events, Interest, Poll, Token};
|
||||
use smoltcp::iface::{Config, Interface, SocketHandle, SocketSet};
|
||||
use smoltcp::phy::{Device, Medium, RxToken, TunTapInterface, TxToken};
|
||||
use smoltcp::socket::tcp;
|
||||
use smoltcp::socket::{tcp, udp};
|
||||
use smoltcp::time::Instant;
|
||||
use smoltcp::wire::{IpCidr, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
@ -17,7 +17,7 @@ use std::convert::{From, TryFrom};
|
|||
use std::fmt::{Display, Formatter};
|
||||
use std::io::{Read, Write};
|
||||
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::rc::Rc;
|
||||
use std::str::FromStr;
|
||||
|
@ -286,64 +286,49 @@ pub(crate) struct TunToProxy<'a> {
|
|||
}
|
||||
|
||||
impl<'a> TunToProxy<'a> {
|
||||
pub(crate) fn new(interface: &str, options: Options) -> Self {
|
||||
let tun = TunTapInterface::new(interface, Medium::Ip).unwrap();
|
||||
let poll = Poll::new().unwrap();
|
||||
poll.registry()
|
||||
.register(
|
||||
pub(crate) fn new(interface: &str, options: Options) -> Result<Self, Error> {
|
||||
let tun = TunTapInterface::new(interface, Medium::Ip)?;
|
||||
let poll = Poll::new()?;
|
||||
poll.registry().register(
|
||||
&mut SourceFd(&tun.as_raw_fd()),
|
||||
TCP_TOKEN,
|
||||
Interest::READABLE,
|
||||
)
|
||||
.unwrap();
|
||||
)?;
|
||||
|
||||
let config = Config::new();
|
||||
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);
|
||||
iface.update_ip_addrs(|ip_addrs| {
|
||||
ip_addrs
|
||||
.push(IpCidr::new(
|
||||
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()
|
||||
ip_addrs.push(IpCidr::new(gateway4.into(), 0)).unwrap();
|
||||
ip_addrs.push(IpCidr::new(gateway6.into(), 0)).unwrap()
|
||||
});
|
||||
iface
|
||||
.routes_mut()
|
||||
.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.routes_mut().add_default_ipv4_route(gateway4.into())?;
|
||||
iface.routes_mut().add_default_ipv6_route(gateway6.into())?;
|
||||
iface.set_any_ip(true);
|
||||
|
||||
Self {
|
||||
let tun = Self {
|
||||
tun,
|
||||
poll,
|
||||
iface,
|
||||
connections: Default::default(),
|
||||
connections: HashMap::default(),
|
||||
next_token: 2,
|
||||
token_to_connection: Default::default(),
|
||||
connection_managers: Default::default(),
|
||||
token_to_connection: HashMap::default(),
|
||||
connection_managers: Vec::default(),
|
||||
sockets: SocketSet::new([]),
|
||||
device: virt,
|
||||
options,
|
||||
write_sockets: Default::default(),
|
||||
}
|
||||
write_sockets: HashSet::default(),
|
||||
};
|
||||
Ok(tun)
|
||||
}
|
||||
|
||||
pub(crate) fn add_connection_manager(&mut self, manager: Rc<dyn ConnectionManager>) {
|
||||
self.connection_managers.push(manager);
|
||||
}
|
||||
|
||||
fn expect_smoltcp_send(&mut self) {
|
||||
fn expect_smoltcp_send(&mut self) -> Result<(), Error> {
|
||||
self.iface
|
||||
.poll(Instant::now(), &mut self.device, &mut self.sockets);
|
||||
|
||||
|
@ -353,22 +338,22 @@ impl<'a> TunToProxy<'a> {
|
|||
// TODO: Actual write. Replace.
|
||||
self.tun
|
||||
.transmit(Instant::now())
|
||||
.unwrap()
|
||||
.ok_or("tx token not available")?
|
||||
.consume(slice.len(), |buf| {
|
||||
buf[..].clone_from_slice(slice);
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn remove_connection(&mut self, connection: &Connection) {
|
||||
let mut connection_state = self.connections.remove(connection).unwrap();
|
||||
let token = &connection_state.token;
|
||||
fn remove_connection(&mut self, connection: &Connection) -> Result<(), Error> {
|
||||
let e = "connection not exist";
|
||||
let mut conn = self.connections.remove(connection).ok_or(e)?;
|
||||
let token = &conn.token;
|
||||
self.token_to_connection.remove(token);
|
||||
self.poll
|
||||
.registry()
|
||||
.deregister(&mut connection_state.mio_stream)
|
||||
.unwrap();
|
||||
self.poll.registry().deregister(&mut conn.mio_stream)?;
|
||||
info!("CLOSE {}", connection);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_connection_manager(&self, connection: &Connection) -> Option<Rc<dyn ConnectionManager>> {
|
||||
|
@ -380,14 +365,13 @@ impl<'a> TunToProxy<'a> {
|
|||
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) {
|
||||
let closed = {
|
||||
let socket = self.sockets.get_mut::<tcp::Socket>(state.smoltcp_handle);
|
||||
let mut error = Ok(());
|
||||
while socket.can_recv() && error.is_ok() {
|
||||
socket
|
||||
.recv(|data| {
|
||||
socket.recv(|data| {
|
||||
let event = IncomingDataEvent {
|
||||
direction: IncomingDirection::FromClient,
|
||||
buffer: data,
|
||||
|
@ -395,8 +379,7 @@ impl<'a> TunToProxy<'a> {
|
|||
error = state.handler.push_data(event);
|
||||
|
||||
(data.len(), ())
|
||||
})
|
||||
.unwrap();
|
||||
})?;
|
||||
}
|
||||
|
||||
match error {
|
||||
|
@ -409,24 +392,26 @@ impl<'a> TunToProxy<'a> {
|
|||
};
|
||||
|
||||
// Expect ACKs etc. from smoltcp sockets.
|
||||
self.expect_smoltcp_send();
|
||||
self.expect_smoltcp_send()?;
|
||||
|
||||
if closed {
|
||||
let connection_state = self.connections.get_mut(connection).unwrap();
|
||||
connection_state.mio_stream.shutdown(Both).unwrap();
|
||||
self.remove_connection(connection);
|
||||
let e = "connection not exist";
|
||||
let connection_state = self.connections.get_mut(connection).ok_or(e)?;
|
||||
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)) =
|
||||
connection_tuple(frame)
|
||||
{
|
||||
let resolved_conn = match &mut self.options.virtdns {
|
||||
None => connection.clone(),
|
||||
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);
|
||||
match virt_dns.resolve_ip(&ip) {
|
||||
None => connection.clone(),
|
||||
|
@ -437,9 +422,9 @@ impl<'a> TunToProxy<'a> {
|
|||
if resolved_conn.proto == IpProtocol::Tcp.into() {
|
||||
let cm = self.get_connection_manager(&resolved_conn);
|
||||
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 {
|
||||
for manager in self.connection_managers.iter_mut() {
|
||||
if let Some(handler) =
|
||||
|
@ -450,11 +435,11 @@ impl<'a> TunToProxy<'a> {
|
|||
tcp::SocketBuffer::new(vec![0; 4096]),
|
||||
);
|
||||
socket.set_ack_delay(None);
|
||||
let dst = SocketAddr::try_from(connection.dst).unwrap();
|
||||
socket.listen(dst).unwrap();
|
||||
let dst = SocketAddr::try_from(connection.dst)?;
|
||||
socket.listen(dst)?;
|
||||
let handle = self.sockets.add(socket);
|
||||
|
||||
let client = TcpStream::connect(server).unwrap();
|
||||
let client = TcpStream::connect(server)?;
|
||||
|
||||
let token = Token(self.next_token);
|
||||
self.next_token += 1;
|
||||
|
@ -469,14 +454,11 @@ impl<'a> TunToProxy<'a> {
|
|||
|
||||
self.token_to_connection
|
||||
.insert(token, resolved_conn.clone());
|
||||
self.poll
|
||||
.registry()
|
||||
.register(
|
||||
self.poll.registry().register(
|
||||
&mut state.mio_stream,
|
||||
token,
|
||||
Interest::READABLE | Interest::WRITABLE,
|
||||
)
|
||||
.unwrap();
|
||||
)?;
|
||||
|
||||
self.connections.insert(resolved_conn.clone(), state);
|
||||
|
||||
|
@ -485,7 +467,7 @@ impl<'a> TunToProxy<'a> {
|
|||
}
|
||||
}
|
||||
} else if !self.connections.contains_key(&resolved_conn) {
|
||||
return;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// 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
|
||||
// Exfiltrate the response packets generated by the socket and inject them
|
||||
// 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.
|
||||
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.
|
||||
// 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 {
|
||||
let payload = &frame[_payload_offset.._payload_offset + _payload_size];
|
||||
if let Some(response) = virtual_dns.receive_query(payload) {
|
||||
let rx_buffer = smoltcp::socket::udp::PacketBuffer::new(
|
||||
vec![smoltcp::socket::udp::PacketMetadata::EMPTY],
|
||||
vec![0; 4096],
|
||||
);
|
||||
let tx_buffer = smoltcp::socket::udp::PacketBuffer::new(
|
||||
vec![smoltcp::socket::udp::PacketMetadata::EMPTY],
|
||||
vec![0; 4096],
|
||||
);
|
||||
let mut socket = smoltcp::socket::udp::Socket::new(rx_buffer, tx_buffer);
|
||||
let dst = SocketAddr::try_from(connection.dst).unwrap();
|
||||
socket.bind(dst).unwrap();
|
||||
let rx_buffer =
|
||||
udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
|
||||
let tx_buffer =
|
||||
udp::PacketBuffer::new(vec![udp::PacketMetadata::EMPTY], vec![0; 4096]);
|
||||
let mut socket = udp::Socket::new(rx_buffer, tx_buffer);
|
||||
let dst = SocketAddr::try_from(connection.dst)?;
|
||||
socket.bind(dst)?;
|
||||
socket
|
||||
.send_slice(response.as_slice(), resolved_conn.src.into())
|
||||
.expect("failed to send DNS response");
|
||||
let handle = self.sockets.add(socket);
|
||||
self.expect_smoltcp_send();
|
||||
self.expect_smoltcp_send()?;
|
||||
self.sockets.remove(handle);
|
||||
}
|
||||
}
|
||||
// Otherwise, UDP is not yet supported.
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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 {
|
||||
if let Some(state) = self.connections.get_mut(connection) {
|
||||
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.
|
||||
virtdns.touch_ip(&IpAddr::from(socket.local_endpoint().unwrap().addr));
|
||||
}
|
||||
consumed = socket.send_slice(event.buffer).unwrap();
|
||||
consumed = socket.send_slice(event.buffer)?;
|
||||
state
|
||||
.handler
|
||||
.consume_data(OutgoingDirection::ToClient, consumed);
|
||||
self.expect_smoltcp_send();
|
||||
self.expect_smoltcp_send()?;
|
||||
if consumed < buflen {
|
||||
self.write_sockets.insert(token);
|
||||
break;
|
||||
|
@ -590,40 +569,45 @@ impl<'a> TunToProxy<'a> {
|
|||
let socket = self.sockets.get_mut::<tcp::Socket>(socket_handle);
|
||||
if socket_state & WRITE_CLOSED != 0 && consumed == buflen {
|
||||
socket.close();
|
||||
self.expect_smoltcp_send();
|
||||
self.expect_smoltcp_send()?;
|
||||
self.write_sockets.remove(&token);
|
||||
self.remove_connection(connection);
|
||||
self.remove_connection(connection)?;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn tun_event(&mut self, event: &Event) {
|
||||
fn tun_event(&mut self, event: &Event) -> Result<(), Error> {
|
||||
if event.is_readable() {
|
||||
while let Some((rx_token, _)) = self.tun.receive(Instant::now()) {
|
||||
rx_token.consume(|frame| {
|
||||
self.receive_tun(frame);
|
||||
});
|
||||
rx_token.consume(|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();
|
||||
for token in cloned.iter() {
|
||||
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) {
|
||||
if let Some(conn_ref) = self.token_to_connection.get(&event.token()) {
|
||||
let connection = conn_ref.clone();
|
||||
fn mio_socket_event(&mut self, event: &Event) -> Result<(), Error> {
|
||||
let e = "connection not found";
|
||||
let conn_ref = self.token_to_connection.get(&event.token());
|
||||
if conn_ref.is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
let connection = conn_ref.ok_or(e)?.clone();
|
||||
if event.is_readable() || event.is_read_closed() {
|
||||
{
|
||||
let state = self.connections.get_mut(&connection).unwrap();
|
||||
let state = self.connections.get_mut(&connection).ok_or(e)?;
|
||||
|
||||
// TODO: Move this reading process to its own function.
|
||||
let mut vecbuf = Vec::<u8>::new();
|
||||
|
@ -641,13 +625,13 @@ impl<'a> TunToProxy<'a> {
|
|||
if read == 0 {
|
||||
{
|
||||
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();
|
||||
}
|
||||
self.expect_smoltcp_send();
|
||||
self.remove_connection(&connection.clone());
|
||||
return;
|
||||
self.expect_smoltcp_send()?;
|
||||
self.remove_connection(&connection.clone())?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let data = vecbuf.as_slice();
|
||||
|
@ -656,17 +640,17 @@ impl<'a> TunToProxy<'a> {
|
|||
buffer: &data[0..read],
|
||||
};
|
||||
if let Err(error) = state.handler.push_data(data_event) {
|
||||
state.mio_stream.shutdown(Both).unwrap();
|
||||
state.mio_stream.shutdown(Both)?;
|
||||
{
|
||||
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();
|
||||
}
|
||||
self.expect_smoltcp_send();
|
||||
self.expect_smoltcp_send()?;
|
||||
log::error! {"{error}"}
|
||||
self.remove_connection(&connection.clone());
|
||||
return;
|
||||
self.remove_connection(&connection.clone())?;
|
||||
return Ok(());
|
||||
}
|
||||
if event.is_read_closed() {
|
||||
state.smoltcp_socket_state |= WRITE_CLOSED;
|
||||
|
@ -675,12 +659,12 @@ impl<'a> TunToProxy<'a> {
|
|||
|
||||
// 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);
|
||||
self.write_to_client(event.token(), &connection)?;
|
||||
}
|
||||
if event.is_writable() {
|
||||
self.write_to_server(&connection);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn udp_event(&mut self, _event: &Event) {}
|
||||
|
@ -692,12 +676,12 @@ impl<'a> TunToProxy<'a> {
|
|||
self.poll.poll(&mut events, None)?;
|
||||
for event in events.iter() {
|
||||
match event.token() {
|
||||
TCP_TOKEN => self.tun_event(event),
|
||||
TCP_TOKEN => self.tun_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()?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -137,7 +137,7 @@ impl VirtualDns {
|
|||
Some(response)
|
||||
}
|
||||
|
||||
fn increment_ip(addr: IpAddr) -> IpAddr {
|
||||
fn increment_ip(addr: IpAddr) -> Option<IpAddr> {
|
||||
let mut ip_bytes = match addr as IpAddr {
|
||||
IpAddr::V4(ip) => Vec::<u8>::from(ip.octets()),
|
||||
IpAddr::V6(ip) => Vec::<u8>::from(ip.octets()),
|
||||
|
@ -155,13 +155,14 @@ impl VirtualDns {
|
|||
ip_bytes[i] = 0;
|
||||
}
|
||||
}
|
||||
if addr.is_ipv4() {
|
||||
let bytes: [u8; 4] = ip_bytes.as_slice().try_into().unwrap();
|
||||
let addr = if addr.is_ipv4() {
|
||||
let bytes: [u8; 4] = ip_bytes.as_slice().try_into().ok()?;
|
||||
IpAddr::V4(Ipv4Addr::from(bytes))
|
||||
} 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))
|
||||
}
|
||||
};
|
||||
Some(addr)
|
||||
}
|
||||
|
||||
// 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);
|
||||
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 {
|
||||
// Wrap around.
|
||||
self.next_addr = self.network_addr;
|
||||
|
|
|
@ -146,7 +146,7 @@ mod tests {
|
|||
}
|
||||
Ok(Fork::Child) => {
|
||||
prctl::set_death_signal(signal::SIGKILL as isize).unwrap(); // 9 == SIGKILL
|
||||
main_entry(
|
||||
let _ = main_entry(
|
||||
TUN_TEST_DEVICE,
|
||||
test.proxy,
|
||||
Options::new().with_virtual_dns(),
|
||||
|
|
Loading…
Add table
Reference in a new issue