fixed #4, add command line option to enable udp relay

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"
 }
 ```
 

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)

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)

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());
+        }
     }
 }

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());
     }

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));
         }
     }
 }