use std::sync::atomic::{AtomicUsize, Ordering};
use anyhow::Result;
use quinn::Connection;
use tracing::{debug, error, info, info_span, Instrument};
use crate::common::quic::create_server_endpoint;
use crate::common::remote::{Protocol, RemoteRequest};
use crate::common::socks::tunnel_socks_client;
use crate::common::tcp::{tunnel_tcp_client, tunnel_tcp_server};
use crate::common::tunnel::server_receive_remote_request;
use crate::common::udp::{tunnel_udp_client, tunnel_udp_server};
use crate::ServerConfig;
pub fn run(config: ServerConfig) -> Result<()> {
tokio::runtime::Runtime::new()?.block_on(run_async(config))
}
pub async fn run_async(config: ServerConfig) -> Result<()> {
let endpoint = create_server_endpoint(config.host, config.port, &config.tls)?;
info!("Listening on {}", endpoint.local_addr()?);
let session_counter = AtomicUsize::new(0);
while let Some(conn) = endpoint.accept().await {
let session_id = session_counter.fetch_add(1, Ordering::Relaxed) + 1;
let span = info_span!("session", id = session_id, remote = %conn.remote_address());
let fut = handle_client_connection(conn, config.allow_reverse);
tokio::spawn(
async move {
info!("client connected");
if let Err(e) = fut.await {
error!("connection failed: {}", e)
}
}
.instrument(span),
);
}
Ok(())
}
async fn handle_client_connection(conn: quinn::Incoming, allow_reverse: bool) -> Result<()> {
let connection = conn.await?;
loop {
let quic_connection = connection.clone();
let stream = quic_connection.accept_bi().await;
let stream = match stream {
Err(quinn::ConnectionError::ApplicationClosed { .. }) => {
debug!("client disconnected");
return Ok(());
}
Err(e) => {
error!("stream error: {}", e);
return Err(e.into());
}
Ok(s) => s,
};
let fut = handle_remote_stream(quic_connection, stream, allow_reverse);
tokio::spawn(async move {
if let Err(e) = fut.await {
error!("failed: {}", e);
}
});
}
}
async fn handle_remote_stream(
quic_connection: Connection,
(mut send, mut recv): (quinn::SendStream, quinn::RecvStream),
allow_reverse: bool,
) -> Result<()> {
let request = server_receive_remote_request(&mut send, &mut recv, allow_reverse).await?;
let remote_display = request.to_string();
async {
info!("tunnel established");
match request {
RemoteRequest {
local_host: _,
local_port: _,
remote_host: ref remote_host_ref,
remote_port: 0,
reversed: true,
protocol: Protocol::Tcp,
} if remote_host_ref == "socks" => {
tunnel_socks_client(quic_connection, request).await?;
}
RemoteRequest {
local_host: _,
local_port: _,
remote_host: _,
remote_port: _,
reversed: false,
protocol: Protocol::Tcp,
} => {
tunnel_tcp_server(recv, send, request).await?;
}
RemoteRequest {
local_host: _,
local_port: _,
remote_host: _,
remote_port: _,
reversed: true,
protocol: Protocol::Tcp,
} => {
tunnel_tcp_client(quic_connection, request).await?;
}
RemoteRequest {
local_host: _,
local_port: _,
remote_host: _,
remote_port: _,
reversed: false,
protocol: Protocol::Udp,
} => {
tunnel_udp_server(recv, send, request).await?;
}
RemoteRequest {
local_host: _,
local_port: _,
remote_host: _,
remote_port: _,
reversed: true,
protocol: Protocol::Udp,
} => {
tunnel_udp_client(quic_connection, request).await?;
}
}
Ok(())
}
.instrument(info_span!("tunnel", remote = %remote_display))
.await
}