From 538050415f522230d699b0616df40e04e5a2e343 Mon Sep 17 00:00:00 2001 From: X1r0z Date: Sat, 21 Dec 2024 15:44:18 +0800 Subject: [PATCH] refactor: move command parsing progress to lib.rs 1. move command parsing progress to lib.rs 2. rename some variables --- src/forward.rs | 28 +++++------ src/lib.rs | 119 +++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 127 ++----------------------------------------------- src/proxy.rs | 20 ++++---- 4 files changed, 146 insertions(+), 148 deletions(-) diff --git a/src/forward.rs b/src/forward.rs index d026800..1609661 100644 --- a/src/forward.rs +++ b/src/forward.rs @@ -11,24 +11,24 @@ pub struct Forward { remote_addrs: Vec, socket: Option, udp: bool, - local_ssl_opts: Vec, - remote_ssl_opts: Vec, + local_opts: Vec, + remote_opts: Vec, } impl Forward { pub fn new( local_addrs: Vec, remote_addrs: Vec, - local_ssl_opts: Vec, - remote_ssl_opts: Vec, + local_opts: Vec, + remote_opts: Vec, socket: Option, udp: bool, ) -> Self { Self { local_addrs, remote_addrs, - local_ssl_opts, - remote_ssl_opts, + local_opts, + remote_opts, socket, udp, } @@ -81,12 +81,12 @@ impl Forward { info!("Bind to {} success", listener1.local_addr()?); info!("Bind to {} success", listener1.local_addr()?); - let acceptor1 = Arc::new(match self.local_ssl_opts[0] { + let acceptor1 = Arc::new(match self.local_opts[0] { true => Some(crypto::get_tls_acceptor(&self.local_addrs[0])), false => None, }); - let acceptor2 = Arc::new(match self.local_ssl_opts[1] { + let acceptor2 = Arc::new(match self.local_opts[1] { true => Some(crypto::get_tls_acceptor(&self.local_addrs[1])), false => None, }); @@ -116,12 +116,12 @@ impl Forward { let listener = TcpListener::bind(&self.local_addrs[0]).await?; info!("Bind to {} success", listener.local_addr()?); - let acceptor = Arc::new(match self.local_ssl_opts[0] { + let acceptor = Arc::new(match self.local_opts[0] { true => Some(crypto::get_tls_acceptor(&self.local_addrs[0])), false => None, }); - let connector = Arc::new(match self.remote_ssl_opts[0] { + let connector = Arc::new(match self.remote_opts[0] { true => Some(crypto::get_tls_connector()), false => None, }); @@ -152,12 +152,12 @@ impl Forward { } async fn remote_to_remote_tcp(&self) -> Result<()> { - let connector1 = Arc::new(match self.remote_ssl_opts[0] { + let connector1 = Arc::new(match self.remote_opts[0] { true => Some(crypto::get_tls_connector()), false => None, }); - let connector2 = Arc::new(match self.remote_ssl_opts[1] { + let connector2 = Arc::new(match self.remote_opts[1] { true => Some(crypto::get_tls_connector()), false => None, }); @@ -193,7 +193,7 @@ impl Forward { let local_listener = TcpListener::bind(&self.local_addrs[0]).await?; info!("Bind to {} success", local_listener.local_addr()?); - let acceptor = Arc::new(match self.local_ssl_opts[0] { + let acceptor = Arc::new(match self.local_opts[0] { true => Some(crypto::get_tls_acceptor(&self.local_addrs[0])), false => None, }); @@ -221,7 +221,7 @@ impl Forward { async fn socket_to_remote_tcp(&self) -> Result<()> { let socket_path = self.socket.as_ref().unwrap(); - let connector = Arc::new(match self.remote_ssl_opts[0] { + let connector = Arc::new(match self.remote_opts[0] { true => Some(crypto::get_tls_connector()), false => None, }); diff --git a/src/lib.rs b/src/lib.rs index 0d8e338..942105c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,125 @@ +use std::io::Result; + +use clap::{Parser, Subcommand}; +use forward::Forward; +use proxy::Proxy; +use tracing::info; + pub mod crypto; pub mod forward; pub mod proxy; pub mod socks; pub mod tcp; pub mod udp; + +#[derive(Parser)] +#[command(author, version, about = "Rsproxy: Port-Forwarding and Proxy Tool")] +pub struct Cli { + #[command(subcommand)] + command: Commands, +} + +#[derive(Subcommand)] +pub enum Commands { + /// Port forwarding mode + Fwd { + /// Local listen address, format: [+][IP:]PORT + #[arg(short, long)] + local: Vec, + + /// Remote connect address, format: [+]IP:PORT + #[arg(short, long)] + remote: Vec, + + /// Unix domain socket path + #[arg(short, long)] + socket: Option, + + /// Enable UDP forward mode + #[arg(short, long)] + udp: bool, + }, + + /// Socks proxy mode + Socks { + /// Local listen address, format: [+][IP:]PORT + #[arg(short, long)] + local: Vec, + + /// Reverse server address, format: [+]IP:PORT + #[arg(short, long)] + remote: Option, + }, +} + +pub async fn run(cli: Cli) -> Result<()> { + match cli.command { + Commands::Fwd { + mut local, + mut remote, + socket, + udp, + } => { + info!("Starting forward mode"); + + if udp { + info!("Using UDP protocol"); + } else { + info!("Using TCP protocol"); + } + + let mut local_opts = Vec::new(); + let mut remote_opts = Vec::new(); + + load_opts(&mut local, &mut local_opts); + load_opts(&mut remote, &mut remote_opts); + + format_addrs(&mut local); + + let forward = Forward::new(local, remote, local_opts, remote_opts, socket, udp); + forward.start().await?; + } + Commands::Socks { mut local, remote } => { + info!("Starting proxy mode"); + + let mut local_opts = Vec::new(); + let mut remote_opt = false; + + load_opts(&mut local, &mut local_opts); + + let remote = match remote { + Some(remote) => Some(if remote.starts_with('+') { + remote_opt = true; + remote.replace("+", "") + } else { + remote + }), + None => None, + }; + + let proxy = Proxy::new(local, remote, local_opts, remote_opt); + proxy.start().await?; + } + } + + Ok(()) +} + +pub fn load_opts(addrs: &mut Vec, opts: &mut Vec) { + for addr in addrs { + if addr.starts_with('+') { + *addr = addr.replace("+", ""); + opts.push(true); + } else { + opts.push(false); + } + } +} + +pub fn format_addrs(addrs: &mut Vec) { + for addr in addrs { + if !addr.contains(":") { + *addr = "0.0.0.0:".to_string() + addr; + } + } +} diff --git a/src/main.rs b/src/main.rs index ecdac5a..093848f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,132 +1,11 @@ +use clap::Parser; +use rsproxy::Cli; use std::io::Result; -use clap::{Parser, Subcommand}; -use rsproxy::{forward::Forward, proxy::Proxy}; -use tracing::info; - -#[derive(Parser)] -#[command(author, version, about = "Rsproxy: Port-Forwarding and Proxy Tool")] -struct Cli { - #[command(subcommand)] - command: Commands, -} - -#[derive(Subcommand)] -enum Commands { - /// Port forwarding mode - Fwd { - /// Local listen address, format: [+][IP:]PORT - #[arg(short, long)] - local: Vec, - - /// Remote connect address, format: [+]IP:PORT - #[arg(short, long)] - remote: Vec, - - /// Unix domain socket path - #[arg(short, long)] - socket: Option, - - /// Enable UDP forward mode - #[arg(short, long)] - udp: bool, - }, - - /// Socks proxy mode - Socks { - /// Local listen address, format: [+][IP:]PORT - #[arg(short, long)] - local: Vec, - - /// Reverse server address, format: [+]IP:PORT - #[arg(short, long)] - remote: Option, - }, -} #[tokio::main] async fn main() -> Result<()> { tracing_subscriber::fmt::init(); - let cli = Cli::parse(); - match cli.command { - Commands::Fwd { - mut local, - mut remote, - socket, - udp, - } => { - info!("Starting forward mode"); - - if udp { - info!("Using UDP protocol"); - } else { - info!("Using TCP protocol"); - } - - let mut local_ssl_opts = Vec::new(); - let mut remote_ssl_opts = Vec::new(); - - for addr in &mut local { - if addr.starts_with('+') { - *addr = addr.replace("+", ""); - local_ssl_opts.push(true); - } else { - local_ssl_opts.push(false); - } - - if !addr.contains(":") { - *addr = "0.0.0.0:".to_string() + addr; - } - } - - for addr in &mut remote { - if addr.starts_with('+') { - *addr = addr.replace("+", ""); - remote_ssl_opts.push(true); - } else { - remote_ssl_opts.push(false); - } - } - - let forward = Forward::new(local, remote, local_ssl_opts, remote_ssl_opts, socket, udp); - forward.start().await?; - } - Commands::Socks { mut local, remote } => { - info!("Starting proxy mode"); - - let mut local_ssl_opts = Vec::new(); - let mut remote_ssl_opt = false; - - for addr in &mut local { - if addr.starts_with('+') { - *addr = addr.replace("+", ""); - local_ssl_opts.push(true); - } else { - local_ssl_opts.push(false); - } - - if !addr.contains(":") { - *addr = "0.0.0.0:".to_string() + addr; - } - } - - let remote = match remote { - Some(remote) => { - if remote.starts_with('+') { - remote_ssl_opt = true; - Some(remote.replace("+", "")) - } else { - Some(remote) - } - } - None => None, - }; - - let proxy = Proxy::new(local, remote, local_ssl_opts, remote_ssl_opt); - proxy.start().await?; - } - } - - Ok(()) + rsproxy::run(cli).await } diff --git a/src/proxy.rs b/src/proxy.rs index 787015b..cdef030 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -13,22 +13,22 @@ use crate::{ pub struct Proxy { local_addrs: Vec, remote_addr: Option, - local_ssl_opts: Vec, - remote_ssl_opt: bool, + local_opts: Vec, + remote_opt: bool, } impl Proxy { pub fn new( local_addrs: Vec, remote_addr: Option, - local_ssl_opts: Vec, - remote_ssl_opt: bool, + local_opts: Vec, + remote_opt: bool, ) -> Self { Self { local_addrs, remote_addr, - local_ssl_opts, - remote_ssl_opt, + local_opts, + remote_opt, } } @@ -47,7 +47,7 @@ impl Proxy { let listener = TcpListener::bind(&self.local_addrs[0]).await?; info!("Start socks server on {}", listener.local_addr()?); - let acceptor = Arc::new(match self.local_ssl_opts[0] { + let acceptor = Arc::new(match self.local_opts[0] { true => Some(crypto::get_tls_acceptor(&self.local_addrs[0])), false => None, }); @@ -76,7 +76,7 @@ impl Proxy { pub async fn socks_reverse_client(&self) -> Result<()> { let remote_addr = self.remote_addr.clone().unwrap(); - let connector = Arc::new(match self.remote_ssl_opt { + let connector = Arc::new(match self.remote_opt { true => Some(crypto::get_tls_connector()), false => None, }); @@ -120,12 +120,12 @@ impl Proxy { info!("Bind to {} success", control_listener.local_addr()?); info!("Bind to {} success", proxy_listener.local_addr()?); - let control_acceptor = Arc::new(match self.local_ssl_opts[0] { + let control_acceptor = Arc::new(match self.local_opts[0] { true => Some(crypto::get_tls_acceptor(&self.local_addrs[0])), false => None, }); - let proxy_acceptor = Arc::new(match self.local_ssl_opts[1] { + let proxy_acceptor = Arc::new(match self.local_opts[1] { true => Some(crypto::get_tls_acceptor(&self.local_addrs[1])), false => None, });