Improve proxy error handling

This commit is contained in:
timvisee 2021-11-07 23:08:45 +01:00
parent 277f63d850
commit fe9f5dc936
No known key found for this signature in database
GPG Key ID: B8DB720BC383E172

View File

@ -4,6 +4,7 @@ pub(crate) mod protocol;
pub(crate) mod server; pub(crate) mod server;
pub(crate) mod types; pub(crate) mod types;
use std::error::Error;
use std::sync::Arc; use std::sync::Arc;
use bytes::BytesMut; use bytes::BytesMut;
@ -56,27 +57,22 @@ async fn main() -> Result<(), ()> {
// Proxy all incomming connections // Proxy all incomming connections
while let Ok((inbound, _)) = listener.accept().await { while let Ok((inbound, _)) = listener.accept().await {
let client = Client::default(); let client = Client::default();
// eprintln!("Client connected");
if !server_state.online() { if !server_state.online() {
// When server is not online, spawn a status server // When server is not online, spawn a status server
let transfer = status_server(client, inbound, server_state.clone()).map(|r| { let transfer = status_server(client, inbound, server_state.clone()).map(|r| {
if let Err(e) = r { if let Err(err) = r {
println!("Failed to proxy: {:?}", e); println!("Failed to serve status: {:?}", err);
} }
// eprintln!("Client disconnected");
}); });
tokio::spawn(transfer); tokio::spawn(transfer);
} else { } else {
// When server is online, proxy all // When server is online, proxy all
let transfer = proxy(inbound, ADDRESS_PROXY.to_string()).map(|r| { let transfer = proxy(inbound, ADDRESS_PROXY.to_string()).map(|r| {
if let Err(e) = r { if let Err(err) = r {
println!("Failed to proxy: {:?}", e); println!("Failed to proxy: {:?}", err);
} }
// eprintln!("Client disconnected");
}); });
tokio::spawn(transfer); tokio::spawn(transfer);
@ -265,22 +261,22 @@ async fn status_server(
/// Proxy the inbound stream to a target address. /// Proxy the inbound stream to a target address.
// TODO: do not drop error here, return Box<dyn Error> // TODO: do not drop error here, return Box<dyn Error>
async fn proxy(mut inbound: TcpStream, addr_target: String) -> Result<(), ()> { async fn proxy(mut inbound: TcpStream, addr_target: String) -> Result<(), Box<dyn Error>> {
let mut outbound = TcpStream::connect(addr_target).await.map_err(|_| ())?; // Set up connection to server
// TODO: on connect fail, ping server and redirect to status_server if offline // TODO: on connect fail, ping server and redirect to status_server if offline
let mut outbound = TcpStream::connect(addr_target).await?;
let (mut ri, mut wi) = inbound.split(); let (mut ri, mut wi) = inbound.split();
let (mut ro, mut wo) = outbound.split(); let (mut ro, mut wo) = outbound.split();
let client_to_server = async { let client_to_server = async {
io::copy(&mut ri, &mut wo).await.map_err(|_| ())?; io::copy(&mut ri, &mut wo).await?;
wo.shutdown().await.map_err(|_| ()) wo.shutdown().await
}; };
let server_to_client = async { let server_to_client = async {
io::copy(&mut ro, &mut wi).await.map_err(|_| ())?; io::copy(&mut ro, &mut wi).await?;
wi.shutdown().await.map_err(|_| ()) wi.shutdown().await
}; };
tokio::try_join!(client_to_server, server_to_client)?; tokio::try_join!(client_to_server, server_to_client)?;