From bfeb8749aaf05b7136e1574c67d997fdbc8eb5bf Mon Sep 17 00:00:00 2001 From: "Y. T. Chung" Date: Mon, 3 Nov 2014 16:25:52 +0800 Subject: [PATCH] fixed #4, add command line option to enable udp relay --- README.md | 4 ++-- src/bin/local.rs | 3 +++ src/bin/server.rs | 8 ++++++++ src/relay/local.rs | 24 ++++++++++++++++++++---- src/relay/server.rs | 27 ++++++++++++++++++++------- src/relay/tcprelay/local.rs | 15 +++++++++++---- 6 files changed, 64 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index ef420a10..de0f0ea9 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,7 @@ Create a shadowsocks' configuration file. Example "local_port": 1080, "password": "mypassword", "timeout": 300, - "method": "aes-256-cfb", + "method": "aes-256-cfb" } ``` @@ -71,7 +71,7 @@ In shadowsocks-rust, we also have a extended configuration file format, which is } ], "local_port": 8388, - "local_address": "127.0.0.1", + "local_address": "127.0.0.1" } ``` diff --git a/src/bin/local.rs b/src/bin/local.rs index 3a934767..69cc3f94 100644 --- a/src/bin/local.rs +++ b/src/bin/local.rs @@ -48,6 +48,7 @@ fn main() { let opts = [ optflag("v", "version", "print version"), optflag("h", "help", "print this message"), + optflag("u", "enable-udp", "enable UDP relay"), optopt("c", "config", "specify config file", "config.json"), optopt("s", "server-addr", "server address", ""), optopt("b", "local-addr", "local address, listen only to this address if specified", ""), @@ -119,6 +120,8 @@ fn main() { config.local = Some(local) } + config.enable_udp = matches.opt_present("u"); + info!("ShadowSocks {}", shadowsocks::VERSION); debug!("Config: {}", config) diff --git a/src/bin/server.rs b/src/bin/server.rs index 06dac7df..62125e8e 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -49,6 +49,7 @@ fn main() { let opts = [ optflag("v", "version", "print version"), optflag("h", "help", "print this message"), + optflag("u", "enable-udp", "enable UDP relay"), optopt("c", "config", "specify config file", "config.json"), optopt("s", "server-addr", "server address", ""), optopt("b", "local-addr", "local address, listen only to this address if specified", ""), @@ -120,6 +121,13 @@ fn main() { config.local = Some(local) } + config.enable_udp = matches.opt_present("u"); + + if !cfg!(feature = "enable-udp") && config.enable_udp { + error!("Please compile shadowsocks with --cfg feature=\"enable-udp\""); + panic!("UDP relay is disabled"); + } + info!("ShadowSocks {}", shadowsocks::VERSION); debug!("Config: {}", config) diff --git a/src/relay/local.rs b/src/relay/local.rs index c3e771f4..a1c540c3 100644 --- a/src/relay/local.rs +++ b/src/relay/local.rs @@ -59,6 +59,7 @@ use config::Config; /// ``` #[deriving(Clone)] pub struct RelayLocal { + enable_udp: bool, tcprelay: TcpRelayLocal, #[cfg(feature = "enable-udp")] udprelay: UdpRelayLocal, @@ -72,6 +73,7 @@ impl RelayLocal { RelayLocal { tcprelay: tcprelay, udprelay: udprelay, + enable_udp: config.enable_udp, } } @@ -87,20 +89,34 @@ impl RelayLocal { impl Relay for RelayLocal { #[cfg(not(feature = "enable-udp"))] fn run(&self) { + if self.enable_udp { + warn!("UDP relay feature is disabled, recompile with feature=\"enable-udp\" to enable this feature"); + } let tcprelay = self.tcprelay.clone(); let tcp_future = try_future(proc() tcprelay.run()); + info!("Enabled TCP relay"); + drop(tcp_future.unwrap()); } #[cfg(feature = "enable-udp")] fn run(&self) { + let mut futures = Vec::with_capacity(2); + let tcprelay = self.tcprelay.clone(); let tcp_future = try_future(proc() tcprelay.run()); + futures.push(tcp_future); + info!("Enabled TCP relay"); - let udprelay = self.udprelay.clone(); - let udp_future = try_future(proc() udprelay.run()); + if self.enable_udp { + let udprelay = self.udprelay.clone(); + let udp_future = try_future(proc() udprelay.run()); + futures.push(udp_future); + info!("Enabled UDP relay"); + } - drop(tcp_future.unwrap()); - drop(udp_future.unwrap()); + for fut in futures.into_iter() { + drop(fut.unwrap()); + } } } diff --git a/src/relay/server.rs b/src/relay/server.rs index d96ab523..50b783a0 100644 --- a/src/relay/server.rs +++ b/src/relay/server.rs @@ -54,9 +54,9 @@ use config::Config; /// RelayServer::new(config).run(); /// ``` /// -/// It should be noticed that `config.server` has to be `Some(SingleServer)` here. #[deriving(Clone)] pub struct RelayServer { + enable_udp: bool, tcprelay: TcpRelayServer, #[cfg(feature = "enable-udp")] udprelay: UdpRelayServer, @@ -70,6 +70,7 @@ impl RelayServer { RelayServer { tcprelay: tcprelay, udprelay: udprelay, + enable_udp: config.enable_udp, } } @@ -85,21 +86,33 @@ impl RelayServer { impl Relay for RelayServer { #[cfg(feature = "enable-udp")] fn run(&self) { + let mut futures = Vec::with_capacity(2); + let tcprelay = self.tcprelay.clone(); - let udprelay = self.udprelay.clone(); + futures.push(try_future(proc() tcprelay.run())); + info!("Enabled TCP relay"); - let tcp_future = try_future(proc() tcprelay.run()); - let udp_future = try_future(proc() udprelay.run()); + if self.enable_udp { + let udprelay = self.udprelay.clone(); + let udp_future = try_future(proc() udprelay.run()); + futures.push(udp_future); + info!("Enabled UDP relay"); + } - drop(tcp_future.unwrap()); - drop(udp_future.unwrap()); + for fut in futures.into_iter() { + drop(fut.unwrap()); + } } #[cfg(not(feature = "enable-udp"))] fn run(&self) { - let tcprelay = self.tcprelay.clone(); + if self.enable_udp { + warn!("UDP relay feature is disabled, recompile with feature=\"enable-udp\" to enable this feature"); + } + let tcprelay = self.tcprelay.clone(); let tcp_future = try_future(proc() tcprelay.run()); + info!("Enabled TCP relay"); drop(tcp_future.unwrap()); } diff --git a/src/relay/tcprelay/local.rs b/src/relay/tcprelay/local.rs index 01cfd7d4..ec550240 100644 --- a/src/relay/tcprelay/local.rs +++ b/src/relay/tcprelay/local.rs @@ -107,11 +107,16 @@ impl TcpRelayLocal { let sockname = stream.socket_name().ok().expect("Failed to get socket name"); write_addr(&SocketAddress(sockname), stream).unwrap(); + + // TODO: record this client's information for udprelay local server to validate + // whether the client has already authenticated } fn handle_client(mut stream: TcpStream, server_addr: SocketAddr, - password: String, encrypt_method: String) { + password: String, + encrypt_method: String, + enable_udp: bool) { TcpRelayLocal::do_handshake(&mut stream); let raw_header_part1 = stream.read_exact(3).ok().expect("Failed to read header"); @@ -221,10 +226,10 @@ impl TcpRelayLocal { SOCKS5_CMD_UDP_ASSOCIATE => { let sockname = stream.peer_name().unwrap(); info!("{} requests for UDP ASSOCIATE", sockname); - if cfg!(feature = "enable-udp") { + if cfg!(feature = "enable-udp") && enable_udp { TcpRelayLocal::handle_udp_associate_local(&mut stream); } else { - warn!("UDP ASSOCIATE is not supported"); + warn!("UDP ASSOCIATE is disabled"); send_error_reply(&mut stream, SOCKS5_REPLY_COMMAND_NOT_SUPPORTED); } }, @@ -264,7 +269,9 @@ impl Relay for TcpRelayLocal { spawn(proc() TcpRelayLocal::handle_client(stream, server_addr, - password, encrypt_method)); + password, + encrypt_method, + self.config.enable_udp)); } } }