├── .github └── workflows │ └── rust.yml ├── .gitignore ├── Cargo.toml ├── LICENSE ├── README.md ├── rustfmt.toml └── src ├── lib.rs ├── opt.rs ├── sys └── mod.rs ├── yamux-plugin-local.rs └── yamux-plugin-server.rs /.github/workflows/rust.yml: -------------------------------------------------------------------------------- 1 | name: Rust 2 | 3 | on: 4 | push: 5 | branches: [ "main" ] 6 | pull_request: 7 | branches: [ "main" ] 8 | 9 | env: 10 | CARGO_TERM_COLOR: always 11 | 12 | jobs: 13 | build: 14 | 15 | runs-on: ubuntu-latest 16 | 17 | steps: 18 | - uses: actions/checkout@v3 19 | - name: Build 20 | run: cargo build --verbose 21 | - name: Run tests 22 | run: cargo test --verbose 23 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | # will have compiled files and executables 3 | /target/ 4 | 5 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 6 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 7 | Cargo.lock 8 | 9 | # These are backup files generated by rustfmt 10 | **/*.rs.bk 11 | 12 | 13 | # Added by cargo 14 | 15 | /target 16 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "yamux-plugin" 3 | version = "0.3.1" 4 | edition = "2021" 5 | 6 | [[bin]] 7 | name = "yamux-plugin-local" 8 | path = "src/yamux-plugin-local.rs" 9 | 10 | [[bin]] 11 | name = "yamux-plugin-server" 12 | path = "src/yamux-plugin-server.rs" 13 | 14 | [features] 15 | mimalloc = ["dep:mimalloc"] 16 | 17 | [dependencies] 18 | tokio = { version = "1.21.2", features = [ 19 | "net", 20 | "rt", 21 | "rt-multi-thread", 22 | "macros", 23 | "io-util", 24 | ] } 25 | tokio-yamux = "0.3.7" 26 | tracing-subscriber = { version = "0.3", features = [ 27 | "std", 28 | "fmt", 29 | "env-filter", 30 | "time", 31 | "local-time", 32 | ] } 33 | time = "0.3" 34 | futures = "0.3.24" 35 | log = "0.4.17" 36 | serde_urlencoded = "0.7" 37 | serde = { version = "1.0", features = ["derive"] } 38 | cfg-if = "1.0" 39 | libc = { version = "0.2", features = ["extra_traits"] } 40 | once_cell = "1.17" 41 | lru_time_cache = "0.11" 42 | shadowsocks = { git = "https://github.com/shadowsocks/shadowsocks-rust.git", default-features = false, features = [ 43 | "hickory-dns", 44 | ] } 45 | mimalloc = { version = "0.1.41", optional = true } 46 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 ty 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # shadowsocks-yamux-plugin 2 | 3 | A shadowsocks SIP003 (SIP003u) Plugin with connection multiplexor in YAMUX protocol 4 | 5 | ```plain 6 | ClientA ----+ 7 | | N Connections 8 | ClientB ----+---- sslocal ---- yamux-plugin-local 9 | | | 10 | ClientC ----+ | 11 | | M (TCP) Connections 12 | RemoteA ----+ | 13 | | | 14 | RemoteB ----+---- ssserver ---- yamux-plugin-server 15 | | N Connections 16 | RemoteC ----+ 17 | ``` 18 | 19 | `yamux-plugin` could mulplex `N` TCP / UDP connections into `M` TCP tunnels, which `N >= M`. 20 | 21 | ## Build 22 | 23 | ```bash 24 | cargo build --release 25 | ``` 26 | 27 | ## Plugin Options 28 | 29 | ```json 30 | { 31 | "plugin_opts": "outbound_fwmark=100&outbound_user_cookie=100&outbound_bind_interface=eth1&outbound_bind_addr=1.2.3.4" 32 | } 33 | ``` 34 | 35 | * `outbound_fwmark`: Linux (or Android) sockopt `SO_MARK` 36 | * `outbound_user_cookie`: FreeBSD sockopt `SO_USER_COOKIE` 37 | * `outbound_bind_interface`: Socket binds to interface, Linux `SO_BINDTODEVICE`, macOS `IP_BOUND_IF`, Windows `IP_UNICAST_IF` 38 | * `outbound_bind_addr`: Socket binds to IP 39 | * `udp_timeout`: UDP tunnel timeout (default 5 minutes) 40 | * `tcp_keep_alive`: TCP socket keep-alive time (default 15 seconds) 41 | * `tcp_fast_open`: TCP Fast Open 42 | * `mptcp`: Multipath-TCP 43 | * `ipv6_first`: Connect IPv6 first (default true) 44 | -------------------------------------------------------------------------------- /rustfmt.toml: -------------------------------------------------------------------------------- 1 | edition = "2021" 2 | max_width = 120 3 | #indent_style = "Visual" 4 | #fn_call_width = 120 5 | reorder_imports = true 6 | reorder_modules = true 7 | #reorder_imports_in_group = true 8 | #reorder_imported_names = true 9 | condense_wildcard_suffixes = true 10 | #fn_args_layout = "Visual" 11 | #fn_call_style = "Visual" 12 | #chain_indent = "Visual" 13 | normalize_comments = true 14 | use_try_shorthand = true 15 | reorder_impl_items = true 16 | #use_small_heuristics = "Max" 17 | imports_layout = "HorizontalVertical" 18 | imports_granularity = "Crate" 19 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | //! shadowsocks yamux SIP003 plugin 2 | 3 | pub mod opt; 4 | mod sys; 5 | 6 | pub use self::opt::PluginOpts; 7 | #[cfg(all(unix, not(target_os = "android")))] 8 | pub use sys::adjust_nofile; 9 | 10 | /// Default UDP timeout duration 11 | pub const UDP_DEFAULT_TIMEOUT_SEC: u64 = 5 * 60; 12 | -------------------------------------------------------------------------------- /src/opt.rs: -------------------------------------------------------------------------------- 1 | //! shadowsocks yamux plugin options 2 | 3 | use std::{ 4 | net::{IpAddr, SocketAddr}, 5 | time::Duration, 6 | }; 7 | 8 | use serde::{Deserialize, Serialize}; 9 | use serde_urlencoded::{self, de::Error as DeError, ser::Error as SerError}; 10 | use shadowsocks::net::{AcceptOpts, ConnectOpts}; 11 | 12 | #[derive(Debug, Default, Serialize, Deserialize)] 13 | pub struct PluginOpts { 14 | /// Set `SO_MARK` socket option for outbound sockets 15 | #[cfg(any(target_os = "linux", target_os = "android"))] 16 | pub outbound_fwmark: Option, 17 | /// Set `SO_USER_COOKIE` socket option for outbound sockets 18 | #[cfg(target_os = "freebsd")] 19 | pub outbound_user_cookie: Option, 20 | /// Set `SO_BINDTODEVICE` (Linux), `IP_BOUND_IF` (BSD), `IP_UNICAST_IF` (Windows) socket option for outbound sockets 21 | pub outbound_bind_interface: Option, 22 | /// Outbound sockets will `bind` to this address 23 | pub outbound_bind_addr: Option, 24 | /// UDP tunnel timeout (in seconds) 25 | pub udp_timeout: Option, 26 | /// TCP Keep Alive 27 | pub tcp_keep_alive: Option, 28 | /// TCP Fast Open 29 | pub tcp_fast_open: Option, 30 | /// MPTCP 31 | pub mptcp: Option, 32 | /// IPv6 First 33 | pub ipv6_first: Option, 34 | } 35 | 36 | impl PluginOpts { 37 | pub fn from_str(opt: &str) -> Result { 38 | serde_urlencoded::from_str(opt) 39 | } 40 | 41 | pub fn to_string(&self) -> Result { 42 | serde_urlencoded::to_string(self) 43 | } 44 | 45 | pub fn as_connect_opts(&self) -> ConnectOpts { 46 | let mut connect_opts = ConnectOpts::default(); 47 | 48 | #[cfg(any(target_os = "linux", target_os = "android"))] 49 | if let Some(outbound_fwmark) = self.outbound_fwmark { 50 | connect_opts.fwmark = Some(outbound_fwmark); 51 | } 52 | 53 | #[cfg(target_os = "freebsd")] 54 | if let Some(outbound_user_cookie) = self.outbound_user_cookie { 55 | connect_opts.user_cookie = Some(outbound_user_cookie); 56 | } 57 | 58 | connect_opts.bind_interface = self.outbound_bind_interface.clone(); 59 | connect_opts.bind_local_addr = self.outbound_bind_addr.map(|ip| SocketAddr::new(ip, 0)); 60 | 61 | connect_opts.tcp.keepalive = self.tcp_keep_alive.map(|sec| Duration::from_secs(sec)); 62 | connect_opts.tcp.fastopen = self.tcp_fast_open.unwrap_or(false); 63 | connect_opts.tcp.mptcp = self.mptcp.unwrap_or(false); 64 | 65 | connect_opts 66 | } 67 | 68 | pub fn as_accept_opts(&self) -> AcceptOpts { 69 | let mut accept_opts = AcceptOpts::default(); 70 | 71 | accept_opts.tcp.keepalive = self.tcp_keep_alive.map(|sec| Duration::from_secs(sec)); 72 | accept_opts.tcp.fastopen = self.tcp_fast_open.unwrap_or(false); 73 | accept_opts.tcp.mptcp = self.mptcp.unwrap_or(false); 74 | 75 | accept_opts 76 | } 77 | } 78 | -------------------------------------------------------------------------------- /src/sys/mod.rs: -------------------------------------------------------------------------------- 1 | /// Some systems set an artificially low soft limit on open file count, for compatibility 2 | /// with code that uses select and its hard-coded maximum file descriptor 3 | /// (limited by the size of fd_set). 4 | /// 5 | /// Tokio (Mio) doesn't use select. 6 | /// 7 | /// http://0pointer.net/blog/file-descriptor-limits.html 8 | /// https://github.com/golang/go/issues/46279 9 | #[cfg(all(unix, not(target_os = "android")))] 10 | pub fn adjust_nofile() { 11 | use log::{debug, trace}; 12 | use std::{io::Error, mem}; 13 | 14 | unsafe { 15 | let mut lim: libc::rlimit = mem::zeroed(); 16 | let ret = libc::getrlimit(libc::RLIMIT_NOFILE, &mut lim); 17 | if ret < 0 { 18 | debug!("getrlimit NOFILE failed, {}", Error::last_os_error()); 19 | return; 20 | } 21 | 22 | if lim.rlim_cur != lim.rlim_max { 23 | trace!("rlimit NOFILE {:?} require adjustion", lim); 24 | lim.rlim_cur = lim.rlim_max; 25 | 26 | // On older macOS, setrlimit with rlim_cur = infinity will fail. 27 | #[cfg(any(target_os = "macos", target_os = "ios", target_os = "watchos", target_os = "tvos"))] 28 | { 29 | use std::ptr; 30 | 31 | extern "C" { 32 | fn sysctlbyname( 33 | name: *const libc::c_char, 34 | oldp: *mut libc::c_void, 35 | oldlenp: *mut libc::size_t, 36 | newp: *mut libc::c_void, 37 | newlen: libc::size_t, 38 | ) -> libc::c_int; 39 | } 40 | 41 | // CTL_KERN 42 | // 43 | // Name Type Changeable 44 | // kern.maxfiles int32_t yes 45 | // kern.maxfilesperproc int32_t yes 46 | 47 | let name = b"kern.maxfilesperproc\0"; 48 | let mut nfile: i32 = 0; 49 | let mut nfile_len: libc::size_t = mem::size_of_val(&nfile); 50 | 51 | let ret = sysctlbyname( 52 | name.as_ptr() as *const _, 53 | &mut nfile as *mut _ as *mut _, 54 | &mut nfile_len, 55 | ptr::null_mut(), 56 | 0, 57 | ); 58 | 59 | if ret < 0 { 60 | debug!("sysctlbyname kern.maxfilesperproc failed, {}", Error::last_os_error()); 61 | } else { 62 | lim.rlim_cur = nfile as libc::rlim_t; 63 | } 64 | } 65 | 66 | let ret = libc::setrlimit(libc::RLIMIT_NOFILE, &lim); 67 | if ret < 0 { 68 | debug!("setrlimit NOFILE {:?} failed, {}", lim, Error::last_os_error()); 69 | } else { 70 | debug!("rlimit NOFILE adjusted {:?}", lim); 71 | } 72 | } 73 | } 74 | } 75 | -------------------------------------------------------------------------------- /src/yamux-plugin-local.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | cell::RefCell, 3 | collections::LinkedList, 4 | env, 5 | io::{self, Cursor, ErrorKind, IsTerminal}, 6 | net::SocketAddr, 7 | sync::Arc, 8 | time::Duration, 9 | }; 10 | 11 | use futures::StreamExt; 12 | use log::{error, info, trace, warn}; 13 | use lru_time_cache::LruCache; 14 | #[cfg(feature = "mimalloc")] 15 | use mimalloc::MiMalloc; 16 | use once_cell::sync::OnceCell; 17 | use shadowsocks::{ 18 | config::ServerType, 19 | context::Context, 20 | dns_resolver::DnsResolver, 21 | lookup_then, 22 | lookup_then_connect, 23 | net::{TcpListener, TcpStream, UdpSocket}, 24 | relay::tcprelay::utils::copy_bidirectional, 25 | }; 26 | use time::UtcOffset; 27 | use tokio::{ 28 | io::{AsyncReadExt, AsyncWriteExt, WriteHalf}, 29 | sync::Mutex, 30 | }; 31 | use tokio_yamux::{Config, Control, Error, Session, StreamHandle}; 32 | use tracing_subscriber::{filter::EnvFilter, fmt::time::OffsetTime, FmtSubscriber}; 33 | 34 | use yamux_plugin::PluginOpts; 35 | 36 | #[cfg(feature = "mimalloc")] 37 | #[global_allocator] 38 | static GLOBAL: MiMalloc = MiMalloc; 39 | 40 | #[derive(Debug, Clone, Copy, Eq, PartialEq)] 41 | enum ConnectionType { 42 | Tcp, 43 | Udp, 44 | } 45 | 46 | async fn get_yamux_stream( 47 | connection_type: ConnectionType, 48 | context: &Context, 49 | remote_host: &str, 50 | remote_port: u16, 51 | plugin_opts: &PluginOpts, 52 | ) -> io::Result { 53 | thread_local! { 54 | static TCP_YAMUX_SESSION_LIST: RefCell> = RefCell::new(LinkedList::new()); 55 | static UDP_YAMUX_SESSION_LIST: RefCell> = RefCell::new(LinkedList::new()); 56 | } 57 | 58 | let session_list = match connection_type { 59 | ConnectionType::Tcp => &TCP_YAMUX_SESSION_LIST, 60 | ConnectionType::Udp => &UDP_YAMUX_SESSION_LIST, 61 | }; 62 | 63 | const YAMUX_CONNECT_RETRY_COUNT: usize = 3; 64 | 65 | let mut connect_tried_count = 0; 66 | 67 | let custom_config = Config { 68 | max_stream_window_size: (100 * 1024 * 1024), 69 | ..Config::default() 70 | }; 71 | 72 | let yamux_stream = loop { 73 | connect_tried_count += 1; 74 | 75 | if connect_tried_count > YAMUX_CONNECT_RETRY_COUNT { 76 | return Err(io::Error::new(ErrorKind::Other, "failed to connect remote")); 77 | } 78 | 79 | let control_opt = session_list.with(|list| list.borrow_mut().pop_front()); 80 | 81 | if let Some(mut control) = control_opt { 82 | match control.open_stream().await { 83 | Ok(s) => { 84 | trace!("yamux opened stream {:?}", s); 85 | session_list.with(|list| list.borrow_mut().push_back(control)); 86 | break s; 87 | } 88 | Err(Error::StreamsExhausted) => { 89 | trace!("yamux connection stream id exhaused"); 90 | session_list.with(|list| list.borrow_mut().push_back(control)); 91 | } 92 | Err(Error::SessionShutdown) => { 93 | trace!("yamux connection already shutdown"); 94 | continue; 95 | } 96 | Err(err) => { 97 | error!("yamux connection open stream failed, error: {}", err); 98 | continue; 99 | } 100 | } 101 | } 102 | 103 | let connect_opts = plugin_opts.as_connect_opts(); 104 | let remote_stream_result = lookup_then_connect!(context, remote_host, remote_port, |addr| { 105 | TcpStream::connect_with_opts(&addr, &connect_opts).await 106 | }); 107 | 108 | let remote_stream = match remote_stream_result { 109 | Ok((_, s)) => { 110 | trace!( 111 | "connected tcp host {}:{}, opts: {:?}", 112 | remote_host, 113 | remote_port, 114 | plugin_opts 115 | ); 116 | s 117 | } 118 | Err(err) => { 119 | error!( 120 | "failed to connect to remote {}:{}, error: {}", 121 | remote_host, remote_port, err 122 | ); 123 | continue; 124 | } 125 | }; 126 | 127 | let mut yamux_session = Session::new_client(remote_stream, custom_config); 128 | let yamux_control = yamux_session.control(); 129 | 130 | tokio::spawn(async move { 131 | loop { 132 | match yamux_session.next().await { 133 | Some(Ok(..)) => {} 134 | Some(Err(e)) => { 135 | error!("yamux connection aborted with connection error: {}", e); 136 | break; 137 | } 138 | None => { 139 | trace!("yamux client session closed"); 140 | break; 141 | } 142 | } 143 | } 144 | }); 145 | 146 | session_list.with(|list| list.borrow_mut().push_front(yamux_control)); 147 | }; 148 | 149 | Ok(yamux_stream) 150 | } 151 | 152 | #[inline] 153 | async fn get_tcp_yamux_stream( 154 | context: &Context, 155 | remote_host: &str, 156 | remote_port: u16, 157 | plugin_opts: &PluginOpts, 158 | ) -> io::Result { 159 | get_yamux_stream(ConnectionType::Tcp, context, remote_host, remote_port, plugin_opts).await 160 | } 161 | 162 | #[inline] 163 | async fn get_udp_yamux_stream( 164 | context: &Context, 165 | remote_host: &str, 166 | remote_port: u16, 167 | plugin_opts: &PluginOpts, 168 | ) -> io::Result { 169 | get_yamux_stream(ConnectionType::Udp, context, remote_host, remote_port, plugin_opts).await 170 | } 171 | 172 | async fn start_tcp( 173 | context: &Context, 174 | local_host: &str, 175 | local_port: u16, 176 | remote_host: &str, 177 | remote_port: u16, 178 | plugin_opts: &PluginOpts, 179 | ) -> io::Result<()> { 180 | let accept_opts = plugin_opts.as_accept_opts(); 181 | 182 | let (_, listener) = lookup_then!(context, local_host, local_port, |addr| { 183 | TcpListener::bind_with_opts(&addr, accept_opts.clone()).await 184 | })?; 185 | info!( 186 | "yamux-plugin TCP listening on {}:{}, remote {}:{}", 187 | local_host, local_port, remote_host, remote_port 188 | ); 189 | 190 | loop { 191 | let (mut stream, peer_addr) = match listener.accept().await { 192 | Ok(s) => s, 193 | Err(err) => { 194 | error!("TcpListener::accept failed, error: {}", err); 195 | tokio::time::sleep(Duration::from_secs(1)).await; 196 | continue; 197 | } 198 | }; 199 | 200 | trace!("accepted TCP (shadowsocks) client {}", peer_addr); 201 | 202 | let mut yamux_stream = match get_tcp_yamux_stream(context, remote_host, remote_port, plugin_opts).await { 203 | Ok(s) => s, 204 | Err(err) => { 205 | error!( 206 | "failed to get a valid yamux stream to {}:{}, error: {}", 207 | remote_host, remote_port, err 208 | ); 209 | continue; 210 | } 211 | }; 212 | 213 | tokio::spawn(async move { 214 | let _ = copy_bidirectional(&mut stream, &mut yamux_stream).await; 215 | }); 216 | } 217 | } 218 | 219 | async fn start_udp( 220 | context: &Context, 221 | local_host: &str, 222 | local_port: u16, 223 | remote_host: &str, 224 | remote_port: u16, 225 | plugin_opts: &PluginOpts, 226 | ) -> io::Result<()> { 227 | let accept_opts = plugin_opts.as_accept_opts(); 228 | 229 | let (_, listener) = lookup_then!(context, local_host, local_port, |addr| { 230 | UdpSocket::listen_with_opts(&addr, accept_opts.clone()).await 231 | })?; 232 | 233 | info!( 234 | "yamux-plugin UDP listening on {}:{}, remote {}:{}", 235 | local_host, local_port, remote_host, remote_port 236 | ); 237 | 238 | let listener = Arc::new(listener); 239 | let mut buffer = [0u8; 65535]; 240 | 241 | loop { 242 | let (n, peer_addr) = match listener.recv_from(&mut buffer).await { 243 | Ok(s) => s, 244 | Err(err) => { 245 | error!("UdpSocket::recv_from failed, error: {}", err); 246 | tokio::time::sleep(Duration::from_secs(1)).await; 247 | continue; 248 | } 249 | }; 250 | 251 | trace!("received UDP packet {} bytes from {}", n, peer_addr); 252 | 253 | struct UnderlyingStream { 254 | handle: WriteHalf, 255 | } 256 | 257 | static UDP_TUNNEL_MAP: OnceCell>> = OnceCell::new(); 258 | 259 | let mut tunnel_map = UDP_TUNNEL_MAP 260 | .get_or_init(|| { 261 | let timeout = 262 | Duration::from_secs(plugin_opts.udp_timeout.unwrap_or(yamux_plugin::UDP_DEFAULT_TIMEOUT_SEC)); 263 | Mutex::new(LruCache::with_expiry_duration(timeout)) 264 | }) 265 | .lock() 266 | .await; 267 | 268 | let (opt_yamux_stream, expired_streams) = tunnel_map.notify_get_mut(&peer_addr); 269 | 270 | // Close expired streams gracefully 271 | // Create a new task, don't block the main loop 272 | if !expired_streams.is_empty() { 273 | tokio::spawn(async move { 274 | for (_, mut stream) in expired_streams { 275 | if let Err(err) = stream.handle.shutdown().await { 276 | warn!("UDP tunnel expired. closing with FIN failed with error: {}", err); 277 | } 278 | } 279 | }); 280 | } 281 | 282 | let yamux_stream = match opt_yamux_stream { 283 | Some(s) => s, 284 | // Create a new YAMUX stream. slow-path 285 | None => { 286 | let new_stream = match get_udp_yamux_stream(context, remote_host, remote_port, plugin_opts).await { 287 | Ok(s) => s, 288 | Err(err) => { 289 | error!( 290 | "failed to get valid yamux stream to {}:{}, error: {}", 291 | remote_host, remote_port, err 292 | ); 293 | continue; 294 | } 295 | }; 296 | 297 | let (mut rx, tx) = tokio::io::split(new_stream); 298 | 299 | let listener = listener.clone(); 300 | tokio::spawn(async move { 301 | let mut buffer = Vec::new(); 302 | 303 | loop { 304 | // [LENGTH 8-bytes][PACKET .. LENGTH bytes] 305 | let length = match rx.read_u64().await { 306 | Ok(n) => n, 307 | Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => { 308 | break; 309 | } 310 | Err(err) => { 311 | error!("UDP tunnel for {} ended with error: {}", peer_addr, err); 312 | break; 313 | } 314 | }; 315 | 316 | if length > usize::MAX as u64 { 317 | error!( 318 | "UDP tunnel received packet length {} > usize::MAX {}", 319 | length, 320 | usize::MAX 321 | ); 322 | break; 323 | } 324 | 325 | let length = length as usize; 326 | 327 | if buffer.len() < length { 328 | buffer.resize(length, 0); 329 | } 330 | 331 | if let Err(err) = rx.read_exact(&mut buffer[0..length]).await { 332 | error!("UDP tunnel for {} read with error: {}", peer_addr, err); 333 | break; 334 | } 335 | 336 | match listener.send_to(&buffer[0..length], peer_addr).await { 337 | Ok(n) => { 338 | trace!( 339 | "UDP tunnel sent back {} bytes (expected {} bytes) to {}", 340 | n, 341 | length, 342 | peer_addr 343 | ); 344 | } 345 | Err(err) => { 346 | error!( 347 | "UDP tunnel send back {} bytes to {} failed with error: {}", 348 | length, peer_addr, err 349 | ); 350 | } 351 | } 352 | } 353 | }); 354 | 355 | let new_stream = UnderlyingStream { handle: tx }; 356 | 357 | tunnel_map.insert(peer_addr, new_stream); 358 | tunnel_map.get_mut(&peer_addr).unwrap() 359 | } 360 | }; 361 | 362 | // [LENGTH 8-bytes][PACKET .. LENGTH bytes] 363 | let result: io::Result<()> = async { 364 | let buffer_len = 8 + n; 365 | 366 | let mut packet_buffer = vec![0u8; buffer_len]; 367 | let mut packet_buffer_cursor = Cursor::new(&mut packet_buffer); 368 | packet_buffer_cursor.write_u64(n as u64).await?; 369 | packet_buffer_cursor.write_all(&buffer[..n]).await?; 370 | 371 | yamux_stream.handle.write_all(&packet_buffer).await?; 372 | 373 | Ok(()) 374 | } 375 | .await; 376 | 377 | if let Err(err) = result { 378 | error!("UDP tunnel send packet from {} failed with error: {}", peer_addr, err); 379 | tunnel_map.remove(&peer_addr); 380 | } 381 | } 382 | } 383 | 384 | #[tokio::main] 385 | async fn main() -> io::Result<()> { 386 | let mut builder = FmtSubscriber::builder() 387 | .with_level(true) 388 | .with_timer(match OffsetTime::local_rfc_3339() { 389 | Ok(t) => t, 390 | Err(..) => { 391 | // Reinit with UTC time 392 | OffsetTime::new(UtcOffset::UTC, time::format_description::well_known::Rfc3339) 393 | } 394 | }) 395 | .with_env_filter(EnvFilter::from_default_env()) 396 | .compact(); 397 | 398 | // NOTE: ansi is enabled by default. 399 | // Could be disabled by `NO_COLOR` environment variable. 400 | // https://no-color.org/ 401 | if !std::io::stdout().is_terminal() { 402 | builder = builder.with_ansi(false); 403 | } 404 | 405 | builder.init(); 406 | 407 | #[cfg(all(unix, not(target_os = "android")))] 408 | yamux_plugin::adjust_nofile(); 409 | 410 | let remote_host = env::var("SS_REMOTE_HOST").expect("require SS_REMOTE_HOST"); 411 | let remote_port = env::var("SS_REMOTE_PORT").expect("require SS_REMOTE_PORT"); 412 | let local_host = env::var("SS_LOCAL_HOST").expect("require SS_LOCAL_HOST"); 413 | let local_port = env::var("SS_LOCAL_PORT").expect("require SS_LOCAL_PORT"); 414 | 415 | let remote_port = remote_port.parse::().expect("SS_REMOTE_PORT must be a valid port"); 416 | let local_port = local_port.parse::().expect("SS_LOCAL_PORT must be a valid port"); 417 | 418 | let mut plugin_opts = PluginOpts::default(); 419 | if let Ok(opts) = env::var("SS_PLUGIN_OPTIONS") { 420 | plugin_opts = PluginOpts::from_str(&opts).expect("unrecognized SS_PLUGIN_OPTIONS"); 421 | } 422 | 423 | let connect_opts = plugin_opts.as_connect_opts(); 424 | let dns_resolver = Arc::new(DnsResolver::hickory_dns_system_resolver(None, connect_opts).await?); 425 | 426 | let mut context = Context::new(ServerType::Local); 427 | context.set_dns_resolver(dns_resolver); 428 | context.set_ipv6_first(plugin_opts.ipv6_first.unwrap_or(true)); 429 | 430 | let tcp_fut = start_tcp( 431 | &context, 432 | &local_host, 433 | local_port, 434 | &remote_host, 435 | remote_port, 436 | &plugin_opts, 437 | ); 438 | 439 | let udp_remote_port = if remote_port == u16::MAX { 440 | remote_port - 1 441 | } else { 442 | remote_port + 1 443 | }; 444 | 445 | let udp_fut = start_udp( 446 | &context, 447 | &local_host, 448 | local_port, 449 | &remote_host, 450 | udp_remote_port, 451 | &plugin_opts, 452 | ); 453 | 454 | info!( 455 | "yamux-plugin listening on {}:{}, remote {}:{} (udp: {})", 456 | local_host, local_port, remote_host, remote_port, udp_remote_port 457 | ); 458 | 459 | tokio::pin!(tcp_fut); 460 | tokio::pin!(udp_fut); 461 | 462 | loop { 463 | let tcp_fut = tcp_fut.as_mut(); 464 | let udp_fut = udp_fut.as_mut(); 465 | 466 | tokio::select! { 467 | result = tcp_fut => { 468 | error!("TCP service ended with result {:?}", result); 469 | return Err(io::Error::new(io::ErrorKind::Other, "TCP service exited unexpectly")); 470 | } 471 | result = udp_fut => { 472 | error!("UDP service ended with result {:?}", result); 473 | return Err(io::Error::new(io::ErrorKind::Other, "UDP service exited unexpectly")); 474 | } 475 | } 476 | } 477 | } 478 | -------------------------------------------------------------------------------- /src/yamux-plugin-server.rs: -------------------------------------------------------------------------------- 1 | use std::{ 2 | env, 3 | io::{self, Cursor, ErrorKind, IsTerminal}, 4 | sync::Arc, 5 | time::Duration, 6 | }; 7 | 8 | use futures::StreamExt; 9 | use log::{debug, error, info, trace, warn}; 10 | #[cfg(feature = "mimalloc")] 11 | use mimalloc::MiMalloc; 12 | use shadowsocks::{ 13 | config::ServerType, 14 | context::{Context, SharedContext}, 15 | dns_resolver::DnsResolver, 16 | lookup_then, 17 | lookup_then_connect, 18 | net::{TcpListener, TcpStream, UdpSocket}, 19 | relay::tcprelay::utils::copy_bidirectional, 20 | }; 21 | use time::UtcOffset; 22 | use tokio::{ 23 | io::{AsyncReadExt, AsyncWriteExt}, 24 | net::TcpStream as TokioTcpStream, 25 | time::Instant, 26 | }; 27 | use tokio_yamux::{Config, Session, StreamHandle}; 28 | use tracing_subscriber::{filter::EnvFilter, fmt::time::OffsetTime, FmtSubscriber}; 29 | 30 | use yamux_plugin::PluginOpts; 31 | 32 | #[cfg(feature = "mimalloc")] 33 | #[global_allocator] 34 | static GLOBAL: MiMalloc = MiMalloc; 35 | 36 | #[derive(Debug, Copy, Clone, Eq, PartialEq)] 37 | enum ConnectionType { 38 | Tcp, 39 | Udp, 40 | } 41 | 42 | async fn handle_tcp_connection( 43 | context: &Context, 44 | local_host: &str, 45 | local_port: u16, 46 | plugin_opts: &PluginOpts, 47 | mut yamux_stream: StreamHandle, 48 | ) -> io::Result<()> { 49 | let connect_opts = plugin_opts.as_connect_opts(); 50 | 51 | let local_stream_result = lookup_then_connect!(context, local_host, local_port, |addr| { 52 | TcpStream::connect_with_opts(&addr, &connect_opts).await 53 | }); 54 | 55 | let mut local_stream = match local_stream_result { 56 | Ok((_, s)) => { 57 | trace!( 58 | "connected tcp host {}:{}, opts: {:?}", 59 | local_host, 60 | local_port, 61 | plugin_opts 62 | ); 63 | s 64 | } 65 | Err(err) => { 66 | error!( 67 | "failed to connect to local {}:{}, error: {}", 68 | local_host, local_port, err 69 | ); 70 | return Err(err); 71 | } 72 | }; 73 | 74 | let _ = copy_bidirectional(&mut yamux_stream, &mut local_stream).await; 75 | Ok(()) 76 | } 77 | 78 | async fn handle_udp_connection( 79 | context: &Context, 80 | local_host: &str, 81 | local_port: u16, 82 | plugin_opts: &PluginOpts, 83 | mut yamux_stream: StreamHandle, 84 | ) -> io::Result<()> { 85 | let connect_opts = plugin_opts.as_connect_opts(); 86 | 87 | let (_, socket) = lookup_then!(context, local_host, local_port, |addr| { 88 | UdpSocket::connect_with_opts(&addr, &connect_opts).await 89 | })?; 90 | 91 | let mut udp_recv_buffer = [0u8; 65535]; 92 | let mut yamux_recv_buffer = Vec::new(); 93 | let timeout = Duration::from_secs(plugin_opts.udp_timeout.unwrap_or(yamux_plugin::UDP_DEFAULT_TIMEOUT_SEC)); 94 | let timer = tokio::time::sleep(timeout); 95 | tokio::pin!(timer); 96 | 97 | loop { 98 | let mut timer = timer.as_mut(); 99 | 100 | tokio::select! { 101 | _ = timer.as_mut() => { 102 | debug!("UDP tunnel timed out"); 103 | break; 104 | } 105 | 106 | udp_result = socket.recv(&mut udp_recv_buffer) => { 107 | let n = match udp_result { 108 | Ok(n) => n, 109 | Err(err) => { 110 | error!("UDP tunnel recv failed, error: {}", err); 111 | return Err(err); 112 | } 113 | }; 114 | timer.reset(Instant::now() + timeout); 115 | 116 | // [LENGTH 8-bytes][PACKET .. LENGTH bytes] 117 | let mut buffer = vec![0u8; 8 + n]; 118 | let mut buffer_cursor = Cursor::new(&mut buffer); 119 | 120 | buffer_cursor.write_u64(n as u64).await?; 121 | buffer_cursor.write_all(&udp_recv_buffer[..n]).await?; 122 | 123 | yamux_stream.write_all(&mut buffer).await?; 124 | } 125 | 126 | yamux_result = yamux_stream.read_u64() => { 127 | let length = match yamux_result { 128 | Ok(n) => n, 129 | Err(ref err) if err.kind() == ErrorKind::UnexpectedEof => { 130 | break; 131 | } 132 | Err(err) => { 133 | error!("UDP tunnel ended with error: {}", err); 134 | return Err(err); 135 | } 136 | }; 137 | 138 | if length > usize::MAX as u64 { 139 | error!( 140 | "UDP tunnel received packet length {} > usize::MAX {}", 141 | length, 142 | usize::MAX 143 | ); 144 | return Err(io::Error::new(ErrorKind::Other, "UDP tunnel received packet too large")); 145 | } 146 | 147 | timer.reset(Instant::now() + timeout); 148 | 149 | let length = length as usize; 150 | 151 | if yamux_recv_buffer.len() < length { 152 | yamux_recv_buffer.resize(length, 0); 153 | } 154 | 155 | yamux_stream.read_exact(&mut yamux_recv_buffer[0..length]).await?; 156 | 157 | match socket.send(&yamux_recv_buffer[0..length]).await { 158 | Ok(n) => { 159 | trace!("UDP tunnel sent back {} bytes (expected {} bytes)", n, length,); 160 | } 161 | Err(err) => { 162 | error!("UDP tunnel send back {} bytes failed with error: {}", length, err); 163 | } 164 | } 165 | } 166 | } 167 | } 168 | 169 | // shutdown() will call StreamHandle::close(), which will send Fin 170 | if let Err(err) = yamux_stream.shutdown().await { 171 | warn!("UDP tunnel shutdown() with FIN gracefully failed, error: {}", err); 172 | } 173 | 174 | Ok(()) 175 | } 176 | 177 | async fn handle_tcp_session( 178 | context: SharedContext, 179 | local_host: Arc, 180 | local_port: u16, 181 | plugin_opts: Arc, 182 | mut yamux_session: Session, 183 | connection_type: ConnectionType, 184 | ) -> io::Result<()> { 185 | loop { 186 | let yamux_stream = match yamux_session.next().await { 187 | Some(Ok(s)) => s, 188 | Some(Err(err)) => { 189 | error!("yamux session accept failed, error: {}", err); 190 | break; 191 | } 192 | None => break, 193 | }; 194 | 195 | trace!("yamux session accepted new stream. {:?}", yamux_stream); 196 | 197 | let local_host = local_host.clone(); 198 | let plugin_opts = plugin_opts.clone(); 199 | let context = context.clone(); 200 | tokio::spawn(async move { 201 | match connection_type { 202 | ConnectionType::Tcp => { 203 | handle_tcp_connection(&context, &local_host, local_port, &plugin_opts, yamux_stream).await 204 | } 205 | ConnectionType::Udp => { 206 | handle_udp_connection(&context, &local_host, local_port, &plugin_opts, yamux_stream).await 207 | } 208 | } 209 | }); 210 | } 211 | 212 | Ok(()) 213 | } 214 | 215 | #[tokio::main] 216 | async fn main() -> io::Result<()> { 217 | let mut builder = FmtSubscriber::builder() 218 | .with_timer(match OffsetTime::local_rfc_3339() { 219 | Ok(t) => t, 220 | Err(..) => { 221 | // Reinit with UTC time 222 | OffsetTime::new(UtcOffset::UTC, time::format_description::well_known::Rfc3339) 223 | } 224 | }) 225 | .compact() 226 | .with_env_filter(EnvFilter::from_default_env()); 227 | 228 | // NOTE: ansi is enabled by default. 229 | // Could be disabled by `NO_COLOR` environment variable. 230 | // https://no-color.org/ 231 | if !std::io::stdout().is_terminal() { 232 | builder = builder.with_ansi(false); 233 | } 234 | 235 | builder.init(); 236 | 237 | #[cfg(all(unix, not(target_os = "android")))] 238 | yamux_plugin::adjust_nofile(); 239 | 240 | let remote_host = env::var("SS_REMOTE_HOST").expect("require SS_REMOTE_HOST"); 241 | let remote_port = env::var("SS_REMOTE_PORT").expect("require SS_REMOTE_PORT"); 242 | let local_host = env::var("SS_LOCAL_HOST").expect("require SS_LOCAL_HOST"); 243 | let local_port = env::var("SS_LOCAL_PORT").expect("require SS_LOCAL_PORT"); 244 | 245 | let remote_port = remote_port.parse::().expect("SS_REMOTE_PORT must be a valid port"); 246 | let local_port = local_port.parse::().expect("SS_LOCAL_PORT must be a valid port"); 247 | 248 | let mut plugin_opts = PluginOpts::default(); 249 | if let Ok(opts) = env::var("SS_PLUGIN_OPTIONS") { 250 | plugin_opts = PluginOpts::from_str(&opts).expect("unrecognized SS_PLUGIN_OPTIONS"); 251 | } 252 | 253 | let connect_opts = plugin_opts.as_connect_opts(); 254 | let accept_opts = plugin_opts.as_accept_opts(); 255 | let dns_resolver = Arc::new(DnsResolver::hickory_dns_system_resolver(None, connect_opts).await?); 256 | 257 | let mut context = Context::new(ServerType::Server); 258 | context.set_dns_resolver(dns_resolver); 259 | context.set_ipv6_first(plugin_opts.ipv6_first.unwrap_or(true)); 260 | 261 | let (_, tcp_listener) = lookup_then!(context, &remote_host, remote_port, |addr| { 262 | TcpListener::bind_with_opts(&addr, accept_opts.clone()).await 263 | })?; 264 | 265 | let udp_remote_port = if remote_port == u16::MAX { 266 | remote_port - 1 267 | } else { 268 | remote_port + 1 269 | }; 270 | 271 | let (_, udp_listener) = lookup_then!(context, &remote_host, udp_remote_port, |addr| { 272 | TcpListener::bind_with_opts(&addr, accept_opts.clone()).await 273 | })?; 274 | 275 | info!( 276 | "yamux-plugin listening on {}:{} (udp: {}), local {}:{}", 277 | remote_host, remote_port, udp_remote_port, local_host, local_port 278 | ); 279 | 280 | let local_host = Arc::new(local_host); 281 | let plugin_opts = Arc::new(plugin_opts); 282 | let context = Arc::new(context); 283 | 284 | let custom_config = Config { 285 | max_stream_window_size: (100 * 1024 * 1024), 286 | ..Config::default() 287 | }; 288 | 289 | let tcp_fut = { 290 | let local_host = local_host.clone(); 291 | let plugin_opts = plugin_opts.clone(); 292 | let context = context.clone(); 293 | 294 | async move { 295 | loop { 296 | let (stream, peer_addr) = match tcp_listener.accept().await { 297 | Ok(s) => s, 298 | Err(err) => { 299 | error!("TcpListener::accept failed, error: {}", err); 300 | tokio::time::sleep(Duration::from_secs(1)).await; 301 | continue; 302 | } 303 | }; 304 | 305 | trace!("accepted TCP (shadowsocks) client {}", peer_addr); 306 | 307 | let local_host = local_host.clone(); 308 | let plugin_opts = plugin_opts.clone(); 309 | let context = context.clone(); 310 | let yamux_session = Session::new_server(stream, custom_config); 311 | tokio::spawn(handle_tcp_session( 312 | context, 313 | local_host, 314 | local_port, 315 | plugin_opts, 316 | yamux_session, 317 | ConnectionType::Tcp, 318 | )); 319 | } 320 | } 321 | }; 322 | 323 | let udp_fut = async move { 324 | loop { 325 | let (stream, peer_addr) = match udp_listener.accept().await { 326 | Ok(s) => s, 327 | Err(err) => { 328 | error!("TcpListener::accept failed, error: {}", err); 329 | tokio::time::sleep(Duration::from_secs(1)).await; 330 | continue; 331 | } 332 | }; 333 | 334 | trace!("accepted UDP (shadowsocks) client {}", peer_addr); 335 | 336 | let local_host = local_host.clone(); 337 | let plugin_opts = plugin_opts.clone(); 338 | let context = context.clone(); 339 | let yamux_session = Session::new_server(stream, custom_config); 340 | tokio::spawn(handle_tcp_session( 341 | context, 342 | local_host, 343 | local_port, 344 | plugin_opts, 345 | yamux_session, 346 | ConnectionType::Udp, 347 | )); 348 | } 349 | }; 350 | 351 | tokio::pin!(tcp_fut); 352 | tokio::pin!(udp_fut); 353 | 354 | loop { 355 | tokio::select! { 356 | _ = &mut tcp_fut => {}, 357 | _ = &mut udp_fut => {}, 358 | } 359 | } 360 | } 361 | --------------------------------------------------------------------------------