99 lines
3.1 KiB
Rust
99 lines
3.1 KiB
Rust
use crate::config::Upstream;
|
|
use crate::plugins::kcp::{KcpConfig, KcpListener, KcpStream};
|
|
use crate::servers::{copy, Proxy};
|
|
use futures::future::try_join;
|
|
use log::{debug, error, warn};
|
|
use std::net::SocketAddr;
|
|
use std::sync::Arc;
|
|
use tokio::io;
|
|
use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
|
|
use tokio::net::TcpStream;
|
|
|
|
pub async fn proxy(config: Arc<Proxy>) -> Result<(), Box<dyn std::error::Error>> {
|
|
let kcp_config = KcpConfig::default();
|
|
let mut listener = KcpListener::bind(kcp_config, config.listen).await?;
|
|
let config = config.clone();
|
|
|
|
loop {
|
|
let thread_proxy = config.clone();
|
|
match listener.accept().await {
|
|
Err(err) => {
|
|
error!("Failed to accept connection: {}", err);
|
|
return Err(Box::new(err));
|
|
}
|
|
Ok((stream, peer)) => {
|
|
tokio::spawn(async move {
|
|
match accept(stream, peer, thread_proxy).await {
|
|
Ok(_) => {}
|
|
Err(err) => {
|
|
error!("Relay thread returned an error: {}", err);
|
|
}
|
|
};
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn accept(
|
|
inbound: KcpStream,
|
|
peer: SocketAddr,
|
|
proxy: Arc<Proxy>,
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
|
debug!("New connection from {:?}", peer);
|
|
|
|
let upstream_name = proxy.default_action.clone();
|
|
|
|
debug!("Upstream: {}", upstream_name);
|
|
|
|
let upstream = match proxy.upstream.get(&upstream_name) {
|
|
Some(upstream) => upstream,
|
|
None => {
|
|
warn!(
|
|
"No upstream named {:?} on server {:?}",
|
|
proxy.default, proxy.name
|
|
);
|
|
return process(inbound, proxy.upstream.get(&proxy.default_action).unwrap()).await;
|
|
// ToDo: Remove unwrap and check default option
|
|
}
|
|
};
|
|
return process(inbound, upstream).await;
|
|
}
|
|
|
|
async fn process(
|
|
mut inbound: KcpStream,
|
|
upstream: &Upstream,
|
|
) -> Result<(), Box<dyn std::error::Error>> {
|
|
match upstream {
|
|
Upstream::Ban => {
|
|
let _ = inbound.shutdown();
|
|
}
|
|
Upstream::Echo => {
|
|
let (mut ri, mut wi) = io::split(inbound);
|
|
let inbound_to_inbound = copy(&mut ri, &mut wi);
|
|
let bytes_tx = inbound_to_inbound.await;
|
|
debug!("Bytes read: {:?}", bytes_tx);
|
|
}
|
|
Upstream::Custom(custom) => match custom.protocol.as_ref() {
|
|
"tcp" => {
|
|
let outbound = TcpStream::connect(custom.addr.clone()).await?;
|
|
|
|
let (mut ri, mut wi) = io::split(inbound);
|
|
let (mut ro, mut wo) = io::split(outbound);
|
|
|
|
let inbound_to_outbound = copy(&mut ri, &mut wo);
|
|
let outbound_to_inbound = copy(&mut ro, &mut wi);
|
|
|
|
let (bytes_tx, bytes_rx) =
|
|
try_join(inbound_to_outbound, outbound_to_inbound).await?;
|
|
|
|
debug!("Bytes read: {:?} write: {:?}", bytes_tx, bytes_rx);
|
|
}
|
|
_ => {
|
|
error!("Reached unknown protocol: {:?}", custom.protocol);
|
|
}
|
|
},
|
|
};
|
|
Ok(())
|
|
}
|