Merge improved error handling from ssrlive

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

View file

@ -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())
}
}

View file

@ -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());

View file

@ -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()
}

View file

@ -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}");
}
}

View file

@ -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()?;
}
}
}

View file

@ -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;

View file

@ -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(),