From 688f693f41c799ada255a7e92c97f8b802400a07 Mon Sep 17 00:00:00 2001 From: "Y. T. Chung" Date: Thu, 3 Nov 2016 02:05:40 +0800 Subject: [PATCH] remove dummy relay struct --- src/bin/local.rs | 7 +- src/bin/server.rs | 7 +- src/lib.rs | 26 + src/relay/local.rs | 39 +- src/relay/mod.rs | 4 +- src/relay/server.rs | 37 +- src/relay/tcprelay/{http.rs => http_local.rs} | 231 ++++++++- src/relay/tcprelay/local.rs | 459 +----------------- src/relay/tcprelay/mod.rs | 3 +- src/relay/tcprelay/server.rs | 135 +++--- src/relay/tcprelay/socks5_local.rs | 227 +++++++++ src/relay/udprelay/local.rs | 71 ++- src/relay/udprelay/server.rs | 99 ++-- tests/socks5.rs | 6 +- tests/udp.rs | 6 +- 15 files changed, 695 insertions(+), 662 deletions(-) rename src/relay/tcprelay/{http.rs => http_local.rs} (62%) create mode 100644 src/relay/tcprelay/socks5_local.rs diff --git a/src/bin/local.rs b/src/bin/local.rs index 13959fa4..e9f2ed9d 100644 --- a/src/bin/local.rs +++ b/src/bin/local.rs @@ -41,8 +41,7 @@ use std::env; use env_logger::LogBuilder; use log::{LogRecord, LogLevelFilter}; -use shadowsocks::config::{self, Config, ServerConfig, ServerAddr}; -use shadowsocks::relay::RelayLocal; +use shadowsocks::{Config, ServerConfig, ServerAddr, ConfigType, run_local}; fn main() { let matches = App::new("shadowsocks") @@ -152,7 +151,7 @@ fn main() { let mut config = match matches.value_of("CONFIG") { Some(cpath) => { - match Config::load_from_file(cpath, config::ConfigType::Local) { + match Config::load_from_file(cpath, ConfigType::Local) { Ok(cfg) => { has_provided_config = true; cfg @@ -228,5 +227,5 @@ fn main() { debug!("Config: {:?}", config); - RelayLocal::run(config).unwrap(); + run_local(config).unwrap(); } diff --git a/src/bin/server.rs b/src/bin/server.rs index 5898fed6..893b7ed0 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -42,8 +42,7 @@ use clap::{App, Arg}; use env_logger::LogBuilder; use log::{LogRecord, LogLevelFilter}; -use shadowsocks::config::{self, Config, ServerConfig, ServerAddr}; -use shadowsocks::relay::RelayServer; +use shadowsocks::{Config, ServerConfig, ServerAddr, ConfigType, run_server}; fn main() { let matches = App::new("shadowsocks") @@ -157,7 +156,7 @@ fn main() { let mut has_provided_config = false; let mut config = match matches.value_of("CONFIG") { Some(cpath) => { - match Config::load_from_file(cpath, config::ConfigType::Server) { + match Config::load_from_file(cpath, ConfigType::Server) { Ok(cfg) => { has_provided_config = true; cfg @@ -224,5 +223,5 @@ fn main() { debug!("Config: {:?}", config); - RelayServer::run(config).unwrap(); + run_server(config).unwrap(); } diff --git a/src/lib.rs b/src/lib.rs index 75522eb3..318a4ddd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -66,6 +66,27 @@ //! addresses. //! //! Of course, you can also use `cargo install` to install binaries. +//! +//! ## API Usage +//! +//! Example to write a local server +//! +//! ```no_run +//! use shadowsocks::{Config, ConfigType, run_local}; +//! +//! let config = Config::load_from_file("shadowsocks.json", ConfigType::Local).unwrap(); +//! run_local(config).unwrap(); +//! ``` +//! +//! That's all! And let me show you how to run a proxy server +//! +//! ```no_run +//! use shadowsocks::{Config, ConfigType, run_server}; +//! +//! let config = Config::load_from_file("shadowsocks.json", ConfigType::Server).unwrap(); +//! run_server(config).unwrap(); +//! ``` +//! #![crate_type = "lib"] #![crate_name = "shadowsocks"] @@ -95,6 +116,11 @@ extern crate libc; /// ShadowSocks version pub const VERSION: &'static str = env!("CARGO_PKG_VERSION"); +pub use self::config::{Config, ServerConfig, ServerAddr, ClientConfig, ConfigType}; +pub use self::relay::local::run as run_local; +pub use self::relay::server::run as run_server; +pub use self::relay::tcprelay::client::Socks5Client; + pub mod config; pub mod relay; pub mod crypto; diff --git a/src/relay/local.rs b/src/relay/local.rs index bf50cd4e..2720857b 100644 --- a/src/relay/local.rs +++ b/src/relay/local.rs @@ -28,8 +28,8 @@ use tokio_core::reactor::Core; use futures::{self, Future}; -use relay::tcprelay::local::TcpRelayLocal; -use relay::udprelay::local::UdpRelayLocal; +use relay::tcprelay::local::run as run_tcp; +use relay::udprelay::local::run as run_udp; use relay::dns_resolver::DnsResolver; use relay::boxed_future; use config::Config; @@ -37,7 +37,7 @@ use config::Config; /// Relay server running under local environment. /// /// ```no_run -/// use shadowsocks::relay::RelayLocal; +/// use shadowsocks::relay::local::run; /// use shadowsocks::config::{Config, ServerConfig}; /// use shadowsocks::crypto::CipherType; /// @@ -47,28 +47,23 @@ use config::Config; /// ServerConfig::basic("127.0.0.1:8388".parse().unwrap(), /// "server-password".to_string(), /// CipherType::Aes256Cfb)]; -/// RelayLocal::run(config).unwrap(); +/// run(config).unwrap(); /// ``` -#[derive(Clone)] -pub struct RelayLocal; +pub fn run(config: Config) -> io::Result<()> { + let mut lp = try!(Core::new()); + let handle = lp.handle(); + let config = Rc::new(config); + let dns_resolver = DnsResolver::new(config.dns_cache_capacity); -impl RelayLocal { - pub fn run(config: Config) -> io::Result<()> { - let mut lp = try!(Core::new()); - let handle = lp.handle(); - let config = Rc::new(config); - let dns_resolver = DnsResolver::new(config.dns_cache_capacity); + let tcp_fut = run_tcp(config.clone(), handle.clone(), dns_resolver.clone()); - let tcp_fut = TcpRelayLocal::run(config.clone(), handle.clone(), dns_resolver.clone()); + let udp_fut = if config.enable_udp { + run_udp(config, handle, dns_resolver) + } else { + boxed_future(futures::finished(())) + }; - let udp_fut = if config.enable_udp { - UdpRelayLocal::run(config, handle, dns_resolver) - } else { - boxed_future(futures::finished(())) - }; + let join = tcp_fut.join(udp_fut).map(|_| ()); - let join = tcp_fut.join(udp_fut).map(|_| ()); - - lp.run(join) - } + lp.run(join) } diff --git a/src/relay/mod.rs b/src/relay/mod.rs index 523dee5c..f1c3db37 100644 --- a/src/relay/mod.rs +++ b/src/relay/mod.rs @@ -21,9 +21,6 @@ //! Relay server in local and server side implementations. -pub use self::local::RelayLocal; -pub use self::server::RelayServer; - use std::io; use futures::Future; @@ -36,6 +33,7 @@ mod loadbalancing; mod dns_resolver; pub mod socks5; +/// Alias for Boxed Future without Send pub type BoxIoFuture = Box>; fn boxed_future(f: F) -> Box> diff --git a/src/relay/server.rs b/src/relay/server.rs index 230a3ee3..78af4f78 100644 --- a/src/relay/server.rs +++ b/src/relay/server.rs @@ -28,8 +28,8 @@ use tokio_core::reactor::Core; use futures::{self, Future}; -use relay::udprelay::server::UdpRelayServer; -use relay::tcprelay::server::TcpRelayServer; +use relay::udprelay::server::run as run_udp; +use relay::tcprelay::server::run as run_tcp; use relay::dns_resolver::DnsResolver; use relay::boxed_future; use config::Config; @@ -37,7 +37,7 @@ use config::Config; /// Relay server running on server side. /// /// ```no_run -/// use shadowsocks::relay::RelayServer; +/// use shadowsocks::relay::server::run; /// use shadowsocks::config::{Config, ServerConfig}; /// use shadowsocks::crypto::CipherType; /// @@ -46,29 +46,24 @@ use config::Config; /// ServerConfig::basic("127.0.0.1:8388".parse().unwrap(), /// "server-password".to_string(), /// CipherType::Aes256Cfb)]; -/// RelayServer::run(config).unwrap(); +/// run(config).unwrap(); /// ``` /// -#[derive(Clone)] -pub struct RelayServer; +pub fn run(config: Config) -> io::Result<()> { + let mut lp = try!(Core::new()); -impl RelayServer { - pub fn run(config: Config) -> io::Result<()> { - let mut lp = try!(Core::new()); + let handle = lp.handle(); + let config = Rc::new(config); - let handle = lp.handle(); - let config = Rc::new(config); + let dns_resolver = DnsResolver::new(config.dns_cache_capacity); - let dns_resolver = DnsResolver::new(config.dns_cache_capacity); + let tcp_fut = run_tcp(config.clone(), handle.clone(), dns_resolver.clone()); - let tcp_fut = TcpRelayServer::run(config.clone(), handle.clone(), dns_resolver.clone()); + let udp_fut = if config.enable_udp { + run_udp(config, handle, dns_resolver) + } else { + boxed_future(futures::finished(())) + }; - let udp_fut = if config.enable_udp { - UdpRelayServer::run(config, handle, dns_resolver) - } else { - boxed_future(futures::finished(())) - }; - - lp.run(tcp_fut.join(udp_fut).map(|_| ())) - } + lp.run(tcp_fut.join(udp_fut).map(|_| ())) } diff --git a/src/relay/tcprelay/http.rs b/src/relay/tcprelay/http_local.rs similarity index 62% rename from src/relay/tcprelay/http.rs rename to src/relay/tcprelay/http_local.rs index 2562517e..c4b47240 100644 --- a/src/relay/tcprelay/http.rs +++ b/src/relay/tcprelay/http_local.rs @@ -26,6 +26,7 @@ use std::net::{SocketAddr, SocketAddrV4, SocketAddrV6}; use std::mem; use std::str; use std::fmt; +use std::rc::Rc; use hyper::uri::RequestUri; use hyper::header::{Header, HeaderFormat, Headers, ContentLength}; @@ -42,11 +43,23 @@ use url::Host; use ip::IpAddr; use futures::{self, Future, Poll}; +use futures::stream::Stream; -use tokio_core::io::{write_all, flush}; +use tokio_core::net::{TcpStream, TcpListener}; +use tokio_core::reactor::Handle; +use tokio_core::io::Io; +use tokio_core::io::{ReadHalf, WriteHalf}; +use tokio_core::io::{flush, write_all, copy}; + +use config::{Config, ServerConfig}; use relay::socks5::Address; +use relay::loadbalancing::server::RoundRobin; +use relay::loadbalancing::server::LoadBalancer; use relay::{BoxIoFuture, boxed_future}; +use relay::dns_resolver::DnsResolver; + +use super::tunnel; use super::stream::EncryptedWriter; #[derive(Debug)] @@ -435,4 +448,218 @@ pub fn should_keep_alive(req: &HttpRequest) -> bool { } None => default_keep_alive, } -} \ No newline at end of file +} + +fn handle_connect(handle: Handle, + (r, w): (ReadHalf, WriteHalf), + req: HttpRequest, + addr: Address, + remains: Vec, + svr_cfg: Rc, + dns_resolver: DnsResolver) + -> Box> { + let cloned_addr = addr.clone(); + let http_version = req.version; + let cloned_svr_cfg = svr_cfg.clone(); + + let fut = super::connect_proxy_server(&handle, svr_cfg, dns_resolver) + .and_then(move |svr_s| { + trace!("Proxy server connected"); + + // Tell the client that we are ready + let handshake_resp = format!("{} 200 Connection Established\r\n\r\n", http_version); + trace!("Sending HTTP tunnel handshake response"); + write_all(w, handshake_resp.into_bytes()) + .and_then(|(w, _)| flush(w)) + .map(|w| (svr_s, w)) + }) + .and_then(move |(svr_s, w)| { + super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr).and_then(move |(svr_r, svr_w)| { + let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); + let whalf = svr_w.and_then(move |svr_w| svr_w.write_all_encrypted(remains)) + .and_then(move |(svr_w, _)| svr_w.copy_from_encrypted(r)); + + tunnel(cloned_addr, whalf, rhalf) + }) + }); + + Box::new(fut) +} + +fn handle_http_keepalive(r: ReadHalf, svr_w: super::EncryptedHalf, req_remains: Vec) -> BoxIoFuture<()> { + let fut = HttpRequestFut::with_buf(r, req_remains).then(|res| { + match res { + Ok((r, req, remains)) => { + let should_keep_alive = should_keep_alive(&req); + trace!("Going to proxy request: {:?}", req); + trace!("Should keep alive? {}", should_keep_alive); + + let fut = proxy_request_encrypted((r, svr_w), None, req, remains) + .and_then(move |(r, svr_w, req_remains)| { + if should_keep_alive { + handle_http_keepalive(r, svr_w, req_remains) + } else { + futures::finished(()).boxed() + } + }); + Box::new(fut) as BoxIoFuture<()> + } + Err(err) => { + let fut = futures::lazy(|| { + use std::io::ErrorKind; + match err.kind() { + // It is Ok for client to close connection + ErrorKind::UnexpectedEof | ErrorKind::BrokenPipe => Ok(()), + _ => Err(err), + } + }); + Box::new(fut) as BoxIoFuture<()> + } + } + }); + Box::new(fut) +} + +fn handle_http_proxy(handle: Handle, + (r, w): (ReadHalf, WriteHalf), + client_addr: &SocketAddr, + req: HttpRequest, + addr: Address, + remains: Vec, + svr_cfg: Rc, + dns_resolver: DnsResolver) + -> Box> { + trace!("Using HTTP Proxy for {} -> {}", client_addr, addr); + + let should_keep_alive = should_keep_alive(&req); + let fut = super::connect_proxy_server(&handle, svr_cfg.clone(), dns_resolver).and_then(move |svr_s| { + trace!("Proxy server connected"); + + let cloned_addr = addr.clone(); + super::proxy_server_handshake(svr_s, svr_cfg, addr).and_then(move |(svr_r, svr_w)| { + // Just proxy anything to client + let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); + let whalf = svr_w.and_then(move |svr_w| { + // Send the first request to server + trace!("Going to proxy request: {:?}", req); + trace!("Should keep alive? {}", should_keep_alive); + proxy_request_encrypted((r, svr_w), None, req, remains).and_then(move |(r, svr_w, req_remains)| { + if should_keep_alive { + handle_http_keepalive(r, svr_w, req_remains) + } else { + futures::finished(()).boxed() + } + }) + }); + + rhalf.join(whalf) + .then(move |_| { + trace!("Relay to {} is finished", cloned_addr); + Ok(()) + }) + }) + }); + + Box::new(fut) +} + +fn handle_client(handle: &Handle, + socket: TcpStream, + _: SocketAddr, + svr_cfg: Rc, + dns_resolver: DnsResolver) + -> io::Result<()> { + let cloned_handle = handle.clone(); + let client_addr = try!(socket.peer_addr()); + let fut = futures::lazy(|| Ok(socket.split())) + .and_then(|(r, w)| { + // Process the first request to see whether client wants CONNECT tunnel or normal HTTP proxy + + HttpRequestFut::new(r).and_then(move |(r, mut req, remains)| { + trace!("Got HTTP Request, version: {}, method: {}, uri: {}", + req.version, + req.method, + req.request_uri); + + match req.get_address() { + Ok(addr) => { + req.clear_request_uri_host(); + boxed_future(futures::finished((r, w, req, addr, remains))) + } + Err(status_code) => { + error!("Invalid Uri: {}", req.request_uri); + let fut = write_response(w, req.version, status_code).then(|_| { + let err = io::Error::new(io::ErrorKind::Other, "Invalid Uri"); + Err(err) + }); + boxed_future(fut) + } + } + }) + }) + .and_then(move |(r, w, req, addr, remains)| { + match req.method.clone() { + Method::Connect => { + info!("CONNECT (Http) {}", addr); + handle_connect(cloned_handle, + (r, w), + req, + addr, + remains, + svr_cfg, + dns_resolver) + } + met => { + info!("{} (Http) {}", met, addr); + handle_http_proxy(cloned_handle, + (r, w), + &client_addr, + req, + addr, + remains, + svr_cfg, + dns_resolver) + } + } + }); + + handle.spawn(fut.then(|res| { + match res { + Ok(..) => Ok(()), + Err(err) => { + if err.kind() != io::ErrorKind::BrokenPipe { + error!("Failed to handle client: {}", err); + } + + Err(()) + } + } + })); + + Ok(()) +} + +/// TCP local server using HTTP proxy protocol +pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> Box> { + let listener = { + let local_addr = config.http_proxy.as_ref().unwrap(); + let listener = TcpListener::bind(local_addr, &handle).unwrap(); + info!("ShadowSocks HTTP Listening on {}", local_addr); + listener + }; + + let mut servers = RoundRobin::new(&*config); + let listening = listener.incoming() + .for_each(move |(socket, addr)| { + let server_cfg = servers.pick_server(); + trace!("Got connection, addr: {}", addr); + trace!("Picked proxy server: {:?}", server_cfg); + let dns_resolver = dns_resolver.clone(); + handle_client(&handle, socket, addr, server_cfg, dns_resolver) + }); + + Box::new(listening.map_err(|err| { + error!("HTTP server run failed: {}", err); + err + })) +} diff --git a/src/relay/tcprelay/local.rs b/src/relay/tcprelay/local.rs index 409a0c48..0a9f165e 100644 --- a/src/relay/tcprelay/local.rs +++ b/src/relay/tcprelay/local.rs @@ -21,462 +21,29 @@ //! TcpRelay server that running on local environment -use std::io; -use std::net::SocketAddr; use std::rc::Rc; -use futures::{self, Future}; -use futures::stream::Stream; +use futures::Future; -use tokio_core::net::{TcpStream, TcpListener}; use tokio_core::reactor::Handle; -use tokio_core::io::Io; -use tokio_core::io::{ReadHalf, WriteHalf}; -use tokio_core::io::{flush, write_all, copy}; -use hyper::method::Method; +use config::Config; -use config::{Config, ServerConfig}; - -use relay::socks5::{self, HandshakeRequest, HandshakeResponse, Address}; -use relay::socks5::{TcpRequestHeader, TcpResponseHeader}; -use relay::loadbalancing::server::RoundRobin; -use relay::loadbalancing::server::LoadBalancer; use relay::{BoxIoFuture, boxed_future}; use relay::dns_resolver::DnsResolver; -use super::http::{self, HttpRequestFut}; -use super::{tunnel, ignore_until_end}; +use super::socks5_local; +use super::http_local; -#[derive(Debug, Clone)] -struct UdpConfig { - enable_udp: bool, - client_addr: Rc, -} - -/// TCP relay local server -pub struct TcpRelayLocal; - -impl TcpRelayLocal { - /// Starts a TCP local server - pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> BoxIoFuture<()> { - let tcp_fut = Socks5RelayLocal::run(config.clone(), handle.clone(), dns_resolver.clone()); - match &config.http_proxy { - &Some(..) => { - let http_fut = HttpRelayServer::run(config, handle, dns_resolver); - Box::new(tcp_fut.join(http_fut) - .map(|_| ())) - } - &None => tcp_fut, +/// Starts a TCP local server +pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> BoxIoFuture<()> { + let tcp_fut = socks5_local::run(config.clone(), handle.clone(), dns_resolver.clone()); + match &config.http_proxy { + &Some(..) => { + let http_fut = http_local::run(config, handle, dns_resolver); + boxed_future(tcp_fut.join(http_fut) + .map(|_| ())) } - } -} - -/// Socks5 local server -pub struct Socks5RelayLocal; - -impl Socks5RelayLocal { - fn handle_socks5_connect(handle: &Handle, - (r, w): (ReadHalf, WriteHalf), - client_addr: SocketAddr, - addr: Address, - svr_cfg: Rc, - dns_resolver: DnsResolver) - -> BoxIoFuture<()> { - let cloned_addr = addr.clone(); - let cloned_svr_cfg = svr_cfg.clone(); - let fut = super::connect_proxy_server(handle, svr_cfg, dns_resolver) - .and_then(move |svr_s| { - trace!("Proxy server connected"); - - // Tell the client that we are ready - let header = TcpResponseHeader::new(socks5::Reply::Succeeded, - Address::SocketAddress(client_addr)); - trace!("Send header: {:?}", header); - - header.write_to(w) - .and_then(flush) - .map(move |w| (svr_s, w)) - }) - .and_then(move |(svr_s, w)| { - super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr).and_then(move |(svr_r, svr_w)| { - let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); - let whalf = svr_w.and_then(move |svr_w| svr_w.copy_from_encrypted(r)); - - tunnel(cloned_addr, whalf, rhalf) - }) - }); - - Box::new(fut) - } - - fn handle_client(handle: &Handle, - s: TcpStream, - conf: Rc, - dns_resolver: DnsResolver, - udp_conf: UdpConfig) - -> io::Result<()> { - let cloned_handle = handle.clone(); - let client_addr = try!(s.peer_addr()); - let cloned_client_addr = client_addr.clone(); - let fut = futures::lazy(|| Ok(s.split())) - .and_then(|(r, w)| { - // Socks5 handshakes - HandshakeRequest::read_from(r).and_then(move |(r, req)| { - trace!("Socks5 {:?}", req); - - if !req.methods.contains(&socks5::SOCKS5_AUTH_METHOD_NONE) { - let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE); - let fut = resp.write_to(w) - .then(|_| { - warn!("Currently shadowsocks-rust does not support authentication"); - Err(io::Error::new(io::ErrorKind::Other, - "Currently shadowsocks-rust does not support authentication")) - }); - boxed_future(fut) - } else { - // Reply to client - let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NONE); - trace!("Reply handshake {:?}", resp); - let fut = resp.write_to(w).and_then(flush).and_then(|w| Ok((r, w))); - boxed_future(fut) - } - }) - }) - .and_then(move |(r, w)| { - // Fetch headers - TcpRequestHeader::read_from(r).then(move |res| { - match res { - Ok((r, h)) => boxed_future(futures::finished((r, w, h))), - Err(err) => { - error!("Failed to get TcpRequestHeader: {}", err); - let fut = TcpResponseHeader::new(err.reply, Address::SocketAddress(client_addr)) - .write_to(w) - .then(|_| Err(From::from(err))); - boxed_future(fut) - } - } - }) - }) - .and_then(move |(r, w, header)| { - trace!("Socks5 {:?}", header); - - let addr = header.address; - match header.command { - socks5::Command::TcpConnect => { - info!("CONNECT {}", addr); - Socks5RelayLocal::handle_socks5_connect(&cloned_handle, - (r, w), - cloned_client_addr, - addr, - conf, - dns_resolver) - } - socks5::Command::TcpBind => { - warn!("BIND is not supported"); - let fut = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr) - .write_to(w) - .map(|_| ()); - boxed_future(fut) - } - socks5::Command::UdpAssociate => { - if udp_conf.enable_udp { - info!("UDP ASSOCIATE {}", addr); - let fut = TcpResponseHeader::new(socks5::Reply::Succeeded, - From::from((&*udp_conf.client_addr).clone())) - .write_to(w) - .and_then(flush) - .and_then(|_| { - // Hold the connection until it ends by its own - ignore_until_end(r).map(|_| ()) - }); - - boxed_future(fut) - } else { - warn!("UDP Associate is not enabled"); - let fut = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr) - .write_to(w) - .map(|_| ()); - boxed_future(fut) - } - } - } - }); - - // Runs in Tokio - handle.spawn(fut.then(|res| { - match res { - Ok(..) => Ok(()), - Err(err) => { - if err.kind() != io::ErrorKind::BrokenPipe { - error!("Failed to handle client: {}", err); - } - Err(()) - } - } - })); - - Ok(()) - } - - /// Starts a TCP local server with Socks5 proxy protocol - pub fn run(config: Rc, - handle: Handle, - dns_resolver: DnsResolver) - -> Box> { - let (listener, local_addr) = { - let local_addr = config.local.as_ref().unwrap(); - let listener = TcpListener::bind(local_addr, &handle).unwrap(); - info!("ShadowSocks TCP Listening on {}", local_addr); - (listener, local_addr.clone()) - }; - - let dns_resolver = dns_resolver.clone(); - - let udp_conf = UdpConfig { - enable_udp: config.enable_udp, - client_addr: Rc::new(local_addr), - }; - - let mut servers = RoundRobin::new(&*config); - let listening = listener.incoming() - .for_each(move |(socket, addr)| { - let server_cfg = servers.pick_server(); - trace!("Got connection, addr: {}", addr); - trace!("Picked proxy server: {:?}", server_cfg); - let dns_resolver = dns_resolver.clone(); - Socks5RelayLocal::handle_client(&handle, socket, server_cfg, dns_resolver, udp_conf.clone()) - }); - - Box::new(listening.map_err(|err| { - error!("Socks5 server run failed: {}", err); - err - })) - } -} - -/// HTTP local server -pub struct HttpRelayServer; - -impl HttpRelayServer { - fn handle_connect(handle: Handle, - (r, w): (ReadHalf, WriteHalf), - req: http::HttpRequest, - addr: Address, - remains: Vec, - svr_cfg: Rc, - dns_resolver: DnsResolver) - -> Box> { - let cloned_addr = addr.clone(); - let http_version = req.version; - let cloned_svr_cfg = svr_cfg.clone(); - - let fut = super::connect_proxy_server(&handle, svr_cfg, dns_resolver) - .and_then(move |svr_s| { - trace!("Proxy server connected"); - - // Tell the client that we are ready - let handshake_resp = format!("{} 200 Connection Established\r\n\r\n", http_version); - trace!("Sending HTTP tunnel handshake response"); - write_all(w, handshake_resp.into_bytes()) - .and_then(|(w, _)| flush(w)) - .map(|w| (svr_s, w)) - }) - .and_then(move |(svr_s, w)| { - super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr).and_then(move |(svr_r, svr_w)| { - let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); - let whalf = svr_w.and_then(move |svr_w| svr_w.write_all_encrypted(remains)) - .and_then(move |(svr_w, _)| svr_w.copy_from_encrypted(r)); - - tunnel(cloned_addr, whalf, rhalf) - }) - }); - - Box::new(fut) - } - - fn handle_http_keepalive(r: ReadHalf, - svr_w: super::EncryptedHalf, - req_remains: Vec) - -> BoxIoFuture<()> { - let fut = HttpRequestFut::with_buf(r, req_remains).then(|res| { - match res { - Ok((r, req, remains)) => { - let should_keep_alive = http::should_keep_alive(&req); - trace!("Going to proxy request: {:?}", req); - trace!("Should keep alive? {}", should_keep_alive); - - let fut = http::proxy_request_encrypted((r, svr_w), None, req, remains) - .and_then(move |(r, svr_w, req_remains)| { - if should_keep_alive { - HttpRelayServer::handle_http_keepalive(r, svr_w, req_remains) - } else { - futures::finished(()).boxed() - } - }); - Box::new(fut) as BoxIoFuture<()> - } - Err(err) => { - let fut = futures::lazy(|| { - use std::io::ErrorKind; - match err.kind() { - // It is Ok for client to close connection - ErrorKind::UnexpectedEof | ErrorKind::BrokenPipe => Ok(()), - _ => Err(err), - } - }); - Box::new(fut) as BoxIoFuture<()> - } - } - }); - Box::new(fut) - } - - fn handle_http_proxy(handle: Handle, - (r, w): (ReadHalf, WriteHalf), - client_addr: &SocketAddr, - req: http::HttpRequest, - addr: Address, - remains: Vec, - svr_cfg: Rc, - dns_resolver: DnsResolver) - -> Box> { - trace!("Using HTTP Proxy for {} -> {}", client_addr, addr); - - let should_keep_alive = http::should_keep_alive(&req); - let fut = super::connect_proxy_server(&handle, svr_cfg.clone(), dns_resolver).and_then(move |svr_s| { - trace!("Proxy server connected"); - - let cloned_addr = addr.clone(); - super::proxy_server_handshake(svr_s, svr_cfg, addr).and_then(move |(svr_r, svr_w)| { - // Just proxy anything to client - let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); - let whalf = svr_w.and_then(move |svr_w| { - // Send the first request to server - trace!("Going to proxy request: {:?}", req); - trace!("Should keep alive? {}", should_keep_alive); - http::proxy_request_encrypted((r, svr_w), None, req, remains) - .and_then(move |(r, svr_w, req_remains)| { - if should_keep_alive { - HttpRelayServer::handle_http_keepalive(r, svr_w, req_remains) - } else { - futures::finished(()).boxed() - } - }) - }); - - rhalf.join(whalf) - .then(move |_| { - trace!("Relay to {} is finished", cloned_addr); - Ok(()) - }) - }) - }); - - Box::new(fut) - } - - fn handle_client(handle: &Handle, - socket: TcpStream, - _: SocketAddr, - svr_cfg: Rc, - dns_resolver: DnsResolver) - -> io::Result<()> { - let cloned_handle = handle.clone(); - let client_addr = try!(socket.peer_addr()); - let fut = futures::lazy(|| Ok(socket.split())) - .and_then(|(r, w)| { - // Process the first request to see whether client wants CONNECT tunnel or normal HTTP proxy - - HttpRequestFut::new(r).and_then(move |(r, mut req, remains)| { - trace!("Got HTTP Request, version: {}, method: {}, uri: {}", - req.version, - req.method, - req.request_uri); - - match req.get_address() { - Ok(addr) => { - req.clear_request_uri_host(); - boxed_future(futures::finished((r, w, req, addr, remains))) - } - Err(status_code) => { - error!("Invalid Uri: {}", req.request_uri); - let fut = http::write_response(w, req.version, status_code).then(|_| { - let err = io::Error::new(io::ErrorKind::Other, "Invalid Uri"); - Err(err) - }); - boxed_future(fut) - } - } - }) - }) - .and_then(move |(r, w, req, addr, remains)| { - match req.method.clone() { - Method::Connect => { - info!("CONNECT (Http) {}", addr); - HttpRelayServer::handle_connect(cloned_handle, - (r, w), - req, - addr, - remains, - svr_cfg, - dns_resolver) - } - met => { - info!("{} (Http) {}", met, addr); - HttpRelayServer::handle_http_proxy(cloned_handle, - (r, w), - &client_addr, - req, - addr, - remains, - svr_cfg, - dns_resolver) - } - } - }); - - handle.spawn(fut.then(|res| { - match res { - Ok(..) => Ok(()), - Err(err) => { - if err.kind() != io::ErrorKind::BrokenPipe { - error!("Failed to handle client: {}", err); - } - - Err(()) - } - } - })); - - Ok(()) - } - - /// TCP local server using HTTP proxy protocol - pub fn run(config: Rc, - handle: Handle, - dns_resolver: DnsResolver) - -> Box> { - let listener = { - let local_addr = config.http_proxy.as_ref().unwrap(); - let listener = TcpListener::bind(local_addr, &handle).unwrap(); - info!("ShadowSocks HTTP Listening on {}", local_addr); - listener - }; - - let mut servers = RoundRobin::new(&*config); - let listening = listener.incoming() - .for_each(move |(socket, addr)| { - let server_cfg = servers.pick_server(); - trace!("Got connection, addr: {}", addr); - trace!("Picked proxy server: {:?}", server_cfg); - let dns_resolver = dns_resolver.clone(); - HttpRelayServer::handle_client(&handle, socket, addr, server_cfg, dns_resolver) - }); - - Box::new(listening.map_err(|err| { - error!("HTTP server run failed: {}", err); - err - })) + &None => tcp_fut, } } diff --git a/src/relay/tcprelay/mod.rs b/src/relay/tcprelay/mod.rs index c48dbe54..565495c3 100644 --- a/src/relay/tcprelay/mod.rs +++ b/src/relay/tcprelay/mod.rs @@ -47,9 +47,10 @@ use ip::IpAddr; use self::stream::{EncryptedWriter, DecryptedReader}; pub mod local; +mod socks5_local; +mod http_local; pub mod server; mod stream; -mod http; pub mod client; /// Directions in the tunnel diff --git a/src/relay/tcprelay/server.rs b/src/relay/tcprelay/server.rs index e66a29f6..74eccb45 100644 --- a/src/relay/tcprelay/server.rs +++ b/src/relay/tcprelay/server.rs @@ -44,7 +44,8 @@ use ip::IpAddr; use super::{tunnel, proxy_handshake, DecryptedHalf, EncryptedHalfFut}; -struct TcpRelayClientHandshake { +/// Context for doing handshake with client +pub struct TcpRelayClientHandshake { handle: Handle, dns_resolver: DnsResolver, s: TcpStream, @@ -53,11 +54,11 @@ struct TcpRelayClientHandshake { } impl TcpRelayClientHandshake { - fn handshake(self) -> BoxIoFuture { - let handle = self.handle; - let forbidden_ip = self.forbidden_ip; - let dns_resolver = self.dns_resolver; - let fut = proxy_handshake(self.s, self.svr_cfg).and_then(|(r_fut, w_fut)| { + /// Doing handshake with client + pub fn handshake(self) -> BoxIoFuture { + let TcpRelayClientHandshake { handle, forbidden_ip, dns_resolver, s, svr_cfg } = self; + + let fut = proxy_handshake(s, svr_cfg).and_then(|(r_fut, w_fut)| { r_fut.and_then(|r| Address::read_from(r).map_err(From::from)) .map(move |(r, addr)| { TcpRelayClientPending { @@ -74,7 +75,8 @@ impl TcpRelayClientHandshake { } } -struct TcpRelayClientPending { +/// Context for connecting remote +pub struct TcpRelayClientPending { handle: Handle, r: DecryptedHalf, addr: Address, @@ -84,6 +86,7 @@ struct TcpRelayClientPending { } impl TcpRelayClientPending { + /// Resolve Address to SocketAddr fn resolve_address(addr: Address, dns_resolver: DnsResolver) -> BoxIoFuture { match addr { Address::SocketAddress(addr) => Box::new(futures::finished(addr)), @@ -100,6 +103,8 @@ impl TcpRelayClientPending { } } + /// Resolve remote address to SocketAddr + /// Report failure if the SocketAddr is forbidden by `forbidden_ip` fn resolve_remote(dns_resolver: DnsResolver, addr: Address, forbidden_ip: Rc>) @@ -122,6 +127,7 @@ impl TcpRelayClientPending { Box::new(fut) } + /// Connect to the remote server fn connect_remote(dns_resolver: DnsResolver, handle: Handle, addr: Address, @@ -132,7 +138,8 @@ impl TcpRelayClientPending { .and_then(move |addr| TcpStream::connect(&addr, &handle))) } - fn connect(self) -> BoxIoFuture { + /// Connect to the remote server + pub fn connect(self) -> BoxIoFuture { let addr = self.addr.clone(); let client_pair = (self.r, self.w); let fut = TcpRelayClientPending::connect_remote(self.dns_resolver, self.handle, self.addr, self.forbidden_ip) @@ -147,14 +154,16 @@ impl TcpRelayClientPending { } } -struct TcpRelayClientConnected { +/// Context for extablishing tunnel +pub struct TcpRelayClientConnected { server: (ReadHalf, WriteHalf), client: (DecryptedHalf, EncryptedHalfFut), addr: Address, } impl TcpRelayClientConnected { - fn tunnel(self) -> BoxIoFuture<()> { + /// Establish tunnel + pub fn tunnel(self) -> BoxIoFuture<()> { let (svr_r, svr_w) = self.server; let (r, w_fut) = self.client; tunnel(self.addr, @@ -163,71 +172,63 @@ impl TcpRelayClientConnected { } } -/// TCP Relay backend -pub struct TcpRelayServer; +/// Runs the server +pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> Box> { + let mut fut: Option>> = None; -impl TcpRelayServer { - /// Runs the server - pub fn run(config: Rc, - handle: Handle, - dns_resolver: DnsResolver) - -> Box> { - let mut fut: Option>> = None; + let ref forbidden_ip = config.forbidden_ip; + let forbidden_ip = Rc::new(forbidden_ip.clone()); - let ref forbidden_ip = config.forbidden_ip; - let forbidden_ip = Rc::new(forbidden_ip.clone()); + for svr_cfg in &config.server { + let listener = { + let addr = svr_cfg.addr(); + let addr = addr.listen_addr(); + let listener = TcpListener::bind(addr, &handle).unwrap(); + trace!("ShadowSocks TCP Listening on {}", addr); + listener + }; - for svr_cfg in &config.server { - let listener = { - let addr = svr_cfg.addr(); - let addr = addr.listen_addr(); - let listener = TcpListener::bind(addr, &handle).unwrap(); - trace!("ShadowSocks TCP Listening on {}", addr); - listener - }; + let svr_cfg = Rc::new(svr_cfg.clone()); + let handle = handle.clone(); + let dns_resolver = dns_resolver.clone(); + let forbidden_ip = forbidden_ip.clone(); + let listening = listener.incoming() + .for_each(move |(socket, addr)| { + let server_cfg = svr_cfg.clone(); + let forbidden_ip = forbidden_ip.clone(); + let dns_resolver = dns_resolver.clone(); - let svr_cfg = Rc::new(svr_cfg.clone()); - let handle = handle.clone(); - let dns_resolver = dns_resolver.clone(); - let forbidden_ip = forbidden_ip.clone(); - let listening = listener.incoming() - .for_each(move |(socket, addr)| { - let server_cfg = svr_cfg.clone(); - let forbidden_ip = forbidden_ip.clone(); - let dns_resolver = dns_resolver.clone(); + trace!("Got connection, addr: {}", addr); + trace!("Picked proxy server: {:?}", server_cfg); - trace!("Got connection, addr: {}", addr); - trace!("Picked proxy server: {:?}", server_cfg); + let client = TcpRelayClientHandshake { + handle: handle.clone(), + dns_resolver: dns_resolver, + s: socket, + svr_cfg: server_cfg, + forbidden_ip: forbidden_ip, + }; - let client = TcpRelayClientHandshake { - handle: handle.clone(), - dns_resolver: dns_resolver, - s: socket, - svr_cfg: server_cfg, - forbidden_ip: forbidden_ip, - }; + let fut = client.handshake() + .and_then(|c| c.connect()) + .and_then(|c| c.tunnel()) + .map_err(move |err| { + error!("Failed to handle client ({}): {}", addr, err); + }); - let fut = client.handshake() - .and_then(|c| c.connect()) - .and_then(|c| c.tunnel()) - .map_err(move |err| { - error!("Failed to handle client ({}): {}", addr, err); - }); - - handle.spawn(fut); - Ok(()) - }) - .map_err(|err| { - error!("Server run failed: {}", err); - err - }); - - fut = Some(match fut.take() { - Some(fut) => Box::new(fut.join(listening).map(|_| ())) as Box>, - None => Box::new(listening) as Box>, + handle.spawn(fut); + Ok(()) }) - } + .map_err(|err| { + error!("Server run failed: {}", err); + err + }); - fut.expect("Must have at least one server") + fut = Some(match fut.take() { + Some(fut) => Box::new(fut.join(listening).map(|_| ())) as Box>, + None => Box::new(listening) as Box>, + }) } + + fut.expect("Must have at least one server") } diff --git a/src/relay/tcprelay/socks5_local.rs b/src/relay/tcprelay/socks5_local.rs new file mode 100644 index 00000000..c2a3ac6c --- /dev/null +++ b/src/relay/tcprelay/socks5_local.rs @@ -0,0 +1,227 @@ +// The MIT License (MIT) + +// Copyright (c) 2014 Y. T. CHUNG + +// Permission is hereby granted, free of charge, to any person obtaining a copy of +// this software and associated documentation files (the "Software"), to deal in +// the Software without restriction, including without limitation the rights to +// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +// the Software, and to permit persons to whom the Software is furnished to do so, +// subject to the following conditions: + +// The above copyright notice and this permission notice shall be included in all +// copies or substantial portions of the Software. + +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +//! Local server that accepts SOCKS 5 protocol + +use std::io; +use std::net::SocketAddr; +use std::rc::Rc; + +use futures::{self, Future}; +use futures::stream::Stream; + +use tokio_core::net::{TcpStream, TcpListener}; +use tokio_core::reactor::Handle; +use tokio_core::io::Io; +use tokio_core::io::{ReadHalf, WriteHalf}; +use tokio_core::io::{flush, copy}; + +use config::{Config, ServerConfig}; + +use relay::socks5::{self, HandshakeRequest, HandshakeResponse, Address}; +use relay::socks5::{TcpRequestHeader, TcpResponseHeader}; +use relay::loadbalancing::server::RoundRobin; +use relay::loadbalancing::server::LoadBalancer; +use relay::{BoxIoFuture, boxed_future}; +use relay::dns_resolver::DnsResolver; + +use super::{tunnel, ignore_until_end}; + +#[derive(Debug, Clone)] +struct UdpConfig { + enable_udp: bool, + client_addr: Rc, +} + +fn handle_socks5_connect(handle: &Handle, + (r, w): (ReadHalf, WriteHalf), + client_addr: SocketAddr, + addr: Address, + svr_cfg: Rc, + dns_resolver: DnsResolver) + -> BoxIoFuture<()> { + let cloned_addr = addr.clone(); + let cloned_svr_cfg = svr_cfg.clone(); + let fut = super::connect_proxy_server(handle, svr_cfg, dns_resolver) + .and_then(move |svr_s| { + trace!("Proxy server connected"); + + // Tell the client that we are ready + let header = TcpResponseHeader::new(socks5::Reply::Succeeded, + Address::SocketAddress(client_addr)); + trace!("Send header: {:?}", header); + + header.write_to(w) + .and_then(flush) + .map(move |w| (svr_s, w)) + }) + .and_then(move |(svr_s, w)| { + super::proxy_server_handshake(svr_s, cloned_svr_cfg, addr).and_then(move |(svr_r, svr_w)| { + let rhalf = svr_r.and_then(move |svr_r| copy(svr_r, w)); + let whalf = svr_w.and_then(move |svr_w| svr_w.copy_from_encrypted(r)); + + tunnel(cloned_addr, whalf, rhalf) + }) + }); + + Box::new(fut) +} + +fn handle_socks5_client(handle: &Handle, + s: TcpStream, + conf: Rc, + dns_resolver: DnsResolver, + udp_conf: UdpConfig) + -> io::Result<()> { + let cloned_handle = handle.clone(); + let client_addr = try!(s.peer_addr()); + let cloned_client_addr = client_addr.clone(); + let fut = futures::lazy(|| Ok(s.split())) + .and_then(|(r, w)| { + // Socks5 handshakes + HandshakeRequest::read_from(r).and_then(move |(r, req)| { + trace!("Socks5 {:?}", req); + + if !req.methods.contains(&socks5::SOCKS5_AUTH_METHOD_NONE) { + let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NOT_ACCEPTABLE); + let fut = resp.write_to(w) + .then(|_| { + warn!("Currently shadowsocks-rust does not support authentication"); + Err(io::Error::new(io::ErrorKind::Other, + "Currently shadowsocks-rust does not support authentication")) + }); + boxed_future(fut) + } else { + // Reply to client + let resp = HandshakeResponse::new(socks5::SOCKS5_AUTH_METHOD_NONE); + trace!("Reply handshake {:?}", resp); + let fut = resp.write_to(w).and_then(flush).and_then(|w| Ok((r, w))); + boxed_future(fut) + } + }) + }) + .and_then(move |(r, w)| { + // Fetch headers + TcpRequestHeader::read_from(r).then(move |res| { + match res { + Ok((r, h)) => boxed_future(futures::finished((r, w, h))), + Err(err) => { + error!("Failed to get TcpRequestHeader: {}", err); + let fut = TcpResponseHeader::new(err.reply, Address::SocketAddress(client_addr)) + .write_to(w) + .then(|_| Err(From::from(err))); + boxed_future(fut) + } + } + }) + }) + .and_then(move |(r, w, header)| { + trace!("Socks5 {:?}", header); + + let addr = header.address; + match header.command { + socks5::Command::TcpConnect => { + info!("CONNECT {}", addr); + handle_socks5_connect(&cloned_handle, + (r, w), + cloned_client_addr, + addr, + conf, + dns_resolver) + } + socks5::Command::TcpBind => { + warn!("BIND is not supported"); + let fut = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr) + .write_to(w) + .map(|_| ()); + boxed_future(fut) + } + socks5::Command::UdpAssociate => { + if udp_conf.enable_udp { + info!("UDP ASSOCIATE {}", addr); + let fut = TcpResponseHeader::new(socks5::Reply::Succeeded, + From::from((&*udp_conf.client_addr).clone())) + .write_to(w) + .and_then(flush) + .and_then(|_| { + // Hold the connection until it ends by its own + ignore_until_end(r).map(|_| ()) + }); + + boxed_future(fut) + } else { + warn!("UDP Associate is not enabled"); + let fut = TcpResponseHeader::new(socks5::Reply::CommandNotSupported, addr) + .write_to(w) + .map(|_| ()); + boxed_future(fut) + } + } + } + }); + + // Runs in Tokio + handle.spawn(fut.then(|res| { + match res { + Ok(..) => Ok(()), + Err(err) => { + if err.kind() != io::ErrorKind::BrokenPipe { + error!("Failed to handle client: {}", err); + } + Err(()) + } + } + })); + + Ok(()) +} + +/// Starts a TCP local server with Socks5 proxy protocol +pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> Box> { + let (listener, local_addr) = { + let local_addr = config.local.as_ref().unwrap(); + let listener = TcpListener::bind(local_addr, &handle).unwrap(); + info!("ShadowSocks TCP Listening on {}", local_addr); + (listener, local_addr.clone()) + }; + + let dns_resolver = dns_resolver.clone(); + + let udp_conf = UdpConfig { + enable_udp: config.enable_udp, + client_addr: Rc::new(local_addr), + }; + + let mut servers = RoundRobin::new(&*config); + let listening = listener.incoming() + .for_each(move |(socket, addr)| { + let server_cfg = servers.pick_server(); + trace!("Got connection, addr: {}", addr); + trace!("Picked proxy server: {:?}", server_cfg); + let dns_resolver = dns_resolver.clone(); + handle_socks5_client(&handle, socket, server_cfg, dns_resolver, udp_conf.clone()) + }); + + Box::new(listening.map_err(|err| { + error!("Socks5 server run failed: {}", err); + err + })) +} \ No newline at end of file diff --git a/src/relay/udprelay/local.rs b/src/relay/udprelay/local.rs index 741c740d..a4505881 100644 --- a/src/relay/udprelay/local.rs +++ b/src/relay/udprelay/local.rs @@ -310,47 +310,42 @@ impl Client { } } -/// UDP relay local server -pub struct UdpRelayLocal; +// Recursive method for handling clients +// Handle one by one +fn handle_client(client: Client) -> BoxIoFuture<()> { + let fut = client.handle_once() + .and_then(|c| handle_client(c)); + boxed_future(fut) +} -impl UdpRelayLocal { - // Recursive method for handling clients - // Handle one by one - fn handle_client(client: Client) -> BoxIoFuture<()> { - let fut = client.handle_once() - .and_then(|c| UdpRelayLocal::handle_client(c)); - boxed_future(fut) - } +fn listen(config: Rc, l: UdpSocket, dns_resolver: DnsResolver) -> BoxIoFuture<()> { + let assoc = Rc::new(RefCell::new(AssociateMap::new(MAXIMUM_ASSOCIATE_MAP_SIZE))); + let server_picker = Rc::new(RefCell::new(RoundRobin::new(&*config))); + let servers: Rc> = Rc::new(RefCell::new(ServerCache::new(config.server.len()))); - fn run_server(config: Rc, l: UdpSocket, dns_resolver: DnsResolver) -> BoxIoFuture<()> { - let assoc = Rc::new(RefCell::new(AssociateMap::new(MAXIMUM_ASSOCIATE_MAP_SIZE))); - let server_picker = Rc::new(RefCell::new(RoundRobin::new(&*config))); - let servers: Rc> = Rc::new(RefCell::new(ServerCache::new(config.server.len()))); + let c = Client { + assoc: assoc, + server_picker: server_picker, + servers: servers, + dns_resolver: dns_resolver, + socket: l, + }; - let c = Client { - assoc: assoc, - server_picker: server_picker, - servers: servers, - dns_resolver: dns_resolver, - socket: l, - }; + // Starts to handle all connections after initialization + handle_client(c) +} - // Starts to handle all connections after initialization - boxed_future(UdpRelayLocal::handle_client(c)) - } +/// Starts a UDP local server +pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> BoxIoFuture<()> { + let fut = futures::lazy(move || { + let l = { + let local_addr = config.local.as_ref().unwrap(); + info!("ShadowSocks UDP Listening on {}", local_addr); + try!(UdpSocket::bind(local_addr, &handle)) + }; + Ok((config, l)) + }) + .and_then(move |(config, l)| listen(config, l, dns_resolver)); - /// Starts a UDP local server - pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> BoxIoFuture<()> { - let fut = futures::lazy(move || { - let l = { - let local_addr = config.local.as_ref().unwrap(); - info!("ShadowSocks UDP Listening on {}", local_addr); - try!(UdpSocket::bind(local_addr, &handle)) - }; - Ok((config, l)) - }) - .and_then(move |(config, l)| UdpRelayLocal::run_server(config, l, dns_resolver)); - - boxed_future(fut) - } + boxed_future(fut) } \ No newline at end of file diff --git a/src/relay/udprelay/server.rs b/src/relay/udprelay/server.rs index d315a928..521446c5 100644 --- a/src/relay/udprelay/server.rs +++ b/src/relay/udprelay/server.rs @@ -53,19 +53,20 @@ struct Associate { type AssociateMap = LruCache; -struct Client { +/// Context for handling a UDP packet +pub struct ConnectionContext { assoc: Rc>, svr_cfg: Rc, dns_resolver: DnsResolver, socket: UdpSocket, } -impl Client { +impl ConnectionContext { /// Handles Client to Remote /// /// Extract and send the actual request body and associate remote with client - fn handle_c2s(self, buf: Vec, n: usize, src: SocketAddr) -> BoxIoFuture { - let Client { assoc, svr_cfg, dns_resolver, socket } = self; + fn handle_c2s(self, buf: Vec, n: usize, src: SocketAddr) -> BoxIoFuture { + let ConnectionContext { assoc, svr_cfg, dns_resolver, socket } = self; // Client -> Remote let fut = futures::lazy(move || { @@ -105,7 +106,7 @@ impl Client { let cloned_assoc = assoc.clone(); let cloned_addr = addr.clone(); - Client::resolve_remote_addr(addr, dns_resolver.clone()) + ConnectionContext::resolve_remote_addr(addr, dns_resolver.clone()) .and_then(move |remote_addr| { // Associate client address with remote let mut assoc = cloned_assoc.borrow_mut(); @@ -119,7 +120,7 @@ impl Client { }) .map(move |(socket, body, len)| { trace!("Sent body {} bytes, actual {} bytes", body.len(), len); - Client { + ConnectionContext { assoc: assoc, svr_cfg: svr_cfg, dns_resolver: dns_resolver, @@ -134,8 +135,12 @@ impl Client { /// Handle Remote to Client /// /// Return packet to Client with encryption - fn handle_s2c(self, Associate { address, client_addr }: Associate, buf: Vec, n: usize) -> BoxIoFuture { - let Client { assoc, svr_cfg, dns_resolver, socket } = self; + fn handle_s2c(self, + Associate { address, client_addr }: Associate, + buf: Vec, + n: usize) + -> BoxIoFuture { + let ConnectionContext { assoc, svr_cfg, dns_resolver, socket } = self; let buf_len = buf[..n].len(); info!("UDP ASSOCIATE {} <- {}, payload length {} bytes", @@ -164,7 +169,7 @@ impl Client { .and_then(move |final_buf| send_to(socket, final_buf, client_addr)) .map(|(socket, buf, len)| { trace!("Sent body {} actual {}", buf.len(), len); - Client { + ConnectionContext { assoc: assoc, svr_cfg: svr_cfg, dns_resolver: dns_resolver, @@ -176,12 +181,12 @@ impl Client { } // Handle one packet - fn handle_once(self) -> BoxIoFuture { - let Client { assoc, svr_cfg, dns_resolver, socket } = self; + fn handle_once(self) -> BoxIoFuture { + let ConnectionContext { assoc, svr_cfg, dns_resolver, socket } = self; let fut = recv_from(socket, vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE]) .and_then(move |(socket, buf, n, src)| { - let c = Client { + let c = ConnectionContext { assoc: assoc.clone(), svr_cfg: svr_cfg, dns_resolver: dns_resolver, @@ -218,43 +223,41 @@ impl Client { } } -/// UDP relay proxy server -pub struct UdpRelayServer; +fn handle_client(c: ConnectionContext) -> BoxIoFuture<()> { + let fut = c.handle_once().and_then(|c| handle_client(c)); + boxed_future(fut) +} -impl UdpRelayServer { - fn handle_client(c: Client) -> BoxIoFuture<()> { - let fut = c.handle_once().and_then(|c| UdpRelayServer::handle_client(c)); - boxed_future(fut) +fn listen(svr_cfg: Rc, handle: Handle, dns_resolver: DnsResolver) -> BoxIoFuture<()> { + let listen_addr = svr_cfg.addr().listen_addr().clone(); + info!("ShadowSocks UDP listening on {}", listen_addr); + let fut = futures::lazy(move || UdpSocket::bind(&listen_addr, &handle)).and_then(|socket| { + let c = ConnectionContext { + assoc: Rc::new(RefCell::new(AssociateMap::new(MAXIMUM_ASSOCIATE_MAP_SIZE))), + dns_resolver: dns_resolver, + svr_cfg: svr_cfg, + socket: socket, + }; + handle_client(c) + }); + boxed_future(fut) +} + +/// Starts a UDP relay server +pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> BoxIoFuture<()> { + let mut fut = None; + + for svr in &config.server { + let handle = handle.clone(); + let dns_resolver = dns_resolver.clone(); + let svr_cfg = Rc::new(svr.clone()); + + let svr_fut = listen(svr_cfg, handle, dns_resolver); + fut = match fut { + None => Some(svr_fut), + Some(fut) => Some(boxed_future(fut.join(svr_fut).map(|_| ()))), + }; } - fn run_server(svr_cfg: Rc, handle: Handle, dns_resolver: DnsResolver) -> BoxIoFuture<()> { - let listen_addr = svr_cfg.addr().listen_addr().clone(); - info!("ShadowSocks UDP listening on {}", listen_addr); - let fut = futures::lazy(move || UdpSocket::bind(&listen_addr, &handle)).and_then(|socket| { - let c = Client { - assoc: Rc::new(RefCell::new(AssociateMap::new(MAXIMUM_ASSOCIATE_MAP_SIZE))), - dns_resolver: dns_resolver, - svr_cfg: svr_cfg, - socket: socket, - }; - UdpRelayServer::handle_client(c) - }); - boxed_future(fut) - } - - /// Starts a UDP relay server - pub fn run(config: Rc, handle: Handle, dns_resolver: DnsResolver) -> BoxIoFuture<()> { - let mut fut = boxed_future(futures::finished(())); - - for svr in &config.server { - let handle = handle.clone(); - let dns_resolver = dns_resolver.clone(); - let svr_cfg = Rc::new(svr.clone()); - - let svr_fut = fut.join(UdpRelayServer::run_server(svr_cfg, handle, dns_resolver)); - fut = boxed_future(svr_fut.map(|_| ())); - } - - fut - } + fut.expect("Should have at least one server") } \ No newline at end of file diff --git a/tests/socks5.rs b/tests/socks5.rs index 4b2d3621..c540d2a5 100644 --- a/tests/socks5.rs +++ b/tests/socks5.rs @@ -20,8 +20,8 @@ use futures::Future; use shadowsocks::relay::tcprelay::client::Socks5Client; use shadowsocks::config::{Config, ServerConfig}; use shadowsocks::crypto::CipherType; -use shadowsocks::relay::{RelayLocal, RelayServer}; use shadowsocks::relay::socks5::Address; +use shadowsocks::{run_local, run_server}; const SERVER_ADDR: &'static str = "127.0.0.1:8092"; const LOCAL_ADDR: &'static str = "127.0.0.1:8290"; @@ -44,7 +44,7 @@ fn start_server(bar: Arc) { thread::spawn(move || { drop(env_logger::init()); bar.wait(); - RelayServer::run(get_config()).unwrap(); + run_server(get_config()).unwrap(); }); } @@ -52,7 +52,7 @@ fn start_local(bar: Arc) { thread::spawn(move || { drop(env_logger::init()); bar.wait(); - RelayLocal::run(get_config()).unwrap(); + run_local(get_config()).unwrap(); }); } diff --git a/tests/udp.rs b/tests/udp.rs index 82572adc..aac08658 100644 --- a/tests/udp.rs +++ b/tests/udp.rs @@ -20,8 +20,8 @@ use futures::Future; use shadowsocks::relay::tcprelay::client::Socks5Client; use shadowsocks::config::{Config, ServerConfig}; use shadowsocks::crypto::CipherType; -use shadowsocks::relay::{RelayLocal, RelayServer}; use shadowsocks::relay::socks5::{Address, UdpAssociateHeader}; +use shadowsocks::{run_local, run_server}; const SERVER_ADDR: &'static str = "127.0.0.1:8093"; const LOCAL_ADDR: &'static str = "127.0.0.1:8291"; @@ -48,7 +48,7 @@ fn start_server(bar: Arc) { thread::spawn(move || { drop(env_logger::init()); bar.wait(); - RelayServer::run(get_config()).unwrap(); + run_server(get_config()).unwrap(); }); } @@ -56,7 +56,7 @@ fn start_local(bar: Arc) { thread::spawn(move || { drop(env_logger::init()); bar.wait(); - RelayLocal::run(get_config()).unwrap(); + run_local(get_config()).unwrap(); }); }