chore(phantun) bump fake-tcp dependency to v0.3.0 and release

`v0.3.1`

.github/workflows/release.yml

@@ -31,7 +31,7 @@ jobs:
           - mipsel-unknown-linux-musl
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - uses: actions-rs/toolchain@v1
         with:
           toolchain: stable

.github/workflows/rust.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - uses: actions-rs/toolchain@v1
       with:
         toolchain: stable

LICENSE-APACHE

@@ -186,7 +186,7 @@ APPENDIX: How to apply the Apache License to your work.
    same "printed page" as the copyright notice for easier
    identification within third-party archives.
 
-Copyright 2014-2021 The Rust Project Developers
+Copyright 2021-2022 Datong Sun (dndx@idndx.com)
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

LICENSE-MIT

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2014-2021 The Rust Project Developers
+Copyright (c) 2021-2022 Datong Sun (dndx@idndx.com)
 
 Permission is hereby granted, free of charge, to any
 person obtaining a copy of this software and associated

README.md

@@ -24,6 +24,7 @@ Table of Contents
 * [MTU overhead](#mtu-overhead)
     * [MTU calculation for WireGuard](#mtu-calculation-for-wireguard)
 * [Version compatibility](#version-compatibility)
+* [Documentations](#documentations)
 * [Performance](#performance)
 * [Future plans](#future-plans)
 * [Compariation to udp2raw](#compariation-to-udp2raw)
@@ -31,7 +32,7 @@ Table of Contents
 
 # Latest release
 
-[v0.2.4](https://github.com/dndx/phantun/releases/tag/v0.2.4)
+[v0.3.0](https://github.com/dndx/phantun/releases/tag/v0.3.0)
 
 # Overview
 
@@ -71,7 +72,7 @@ NIC address and Phantun's TUN interface address.
 You may customize the name of Tun interface created by Phantun and the assigned addresses. Please
 run the executable with `-h` options to see how to change them.
 
-Another way to help understand this network topology:
+Another way to help understand this network topology (please see the diagram above for an illustration of this topology):
 
 Phantun Client is like a machine with private IP address (`192.168.200.2`) behind a router.
 In order for it to reach the Internet, you will need to SNAT the private IP address before it's traffic
@@ -176,6 +177,8 @@ sudo setcap cap_net_admin=+pe phantun_client
 
 **Note:** Run Phantun executable with `-h` option to see full detailed options.
 
+[Back to TOC](#table-of-contents)
+
 ### Server
 
 Note: `4567` is the TCP port Phantun should listen on and must corresponds to the DNAT
@@ -252,15 +255,25 @@ of Server/Client of Phantun on both ends to ensure maximum compatibility.
 
 [Back to TOC](#table-of-contents)
 
+# Documentations
+
+For users who wish to use `fake-tcp` library inside their own project, refer to the documentations for the library at:
+[https://docs.rs/fake-tcp](https://docs.rs/fake-tcp).
+
+[Back to TOC](#table-of-contents)
+
 # Performance
 
-Performance was tested on AWS t3.xlarge instance with 4 vCPUs and 5 Gb/s NIC. WireGuard was used
-for tunneling TCP/UDP traffic between two test instances and MTU has been tuned to avoid fragmentation.
+Performance was tested on 2 AWS `t4g.xlarge` instances with 4 vCPUs and 5 Gb/s NIC over LAN. `nftables` was used to redirect
+UDP stream of `iperf3` to go through the Phantun/udp2raw tunnel between two test instances and MTU has been tuned to avoid fragmentation.
 
-|                 | WireGuard   | WireGuard + Phantun | WireGuard + udp2raw (cipher-mode=none auth-mode=none disable-anti-replay) |
-|-----------------|-------------|---------------------|---------------------------------------------------------------------------|
-| iperf3 -c IP -R | 1.56 Gbit/s | 540 Mbit/s          | 369 Mbit/s                                                                |
-| iperf3 -c IP    | 1.71 Gbit/s | 519 Mbit/s          | 312 Mbit/s                                                                |
+Test command: `iperf3 -c <IP> -p <PORT> -R -u -l 1400 -b 1000m -t 30 -P 5`
+
+| Mode                                                          | Speed          | Overall CPU Usage        |
+|---------------------------------------------------------------|----------------|--------------------------|
+| Direct connection                                             | 3.35 Gbits/sec | 25% (1 core at 100%)     |
+| Phantun                                                       | 2.03 Gbits/sec | 95% (all cores utilized) |
+| udp2raw (cipher-mode=none auth-mode=none disable-anti-replay) | 876 Mbits/sec  | 50% (2 cores at 100%)    |
 
 [Back to TOC](#table-of-contents)
 
@@ -300,7 +313,7 @@ Here is a quick overview of comparison between those two to help you choose:
 
 # License
 
-Copyright 2021 Datong Sun <dndx@idndx.com>
+Copyright 2021-2022 Datong Sun (dndx@idndx.com)
 
 Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 [https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0)> or the MIT license

fake-tcp/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "fake-tcp"
-version = "0.2.3"
+version = "0.3.0"
 edition = "2021"
 authors = ["Datong Sun <dndx@idndx.com>"]
 license = "MIT OR Apache-2.0"
@@ -16,7 +16,7 @@ benchmark = []
 
 [dependencies]
 bytes = "1"
-pnet = "0.28"
+pnet = "0.29"
 tokio = { version = "1.14", features = ["full"] }
 rand = { version = "0.8", features = ["small_rng"] }
 log = "0.4"

fake-tcp/LICENSE-APACHE

@@ -186,7 +186,7 @@ APPENDIX: How to apply the Apache License to your work.
    same "printed page" as the copyright notice for easier
    identification within third-party archives.
 
-Copyright 2014-2021 The Rust Project Developers
+Copyright 2021-2022 Datong Sun (dndx@idndx.com)
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

fake-tcp/LICENSE-MIT

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2014-2021 The Rust Project Developers
+Copyright (c) 2021-2022 Datong Sun (dndx@idndx.com)
 
 Permission is hereby granted, free of charge, to any
 person obtaining a copy of this software and associated

fake-tcp/src/lib.rs

@@ -1,3 +1,43 @@
+//! A minimum, userspace TCP based datagram stack
+//!
+//! # Overview
+//!
+//! `fake-tcp` is a reusable library that implements a minimum TCP stack in
+//! user space using the Tun interface. It allows programs to send datagrams
+//! as if they are part of a TCP connection. `fake-tcp` has been tested to
+//! be able to pass through a variety of NAT and stateful firewalls while
+//! fully preserves certain desirable behavior such as out of order delivery
+//! and no congestion/flow controls.
+//!
+//! # Core Concepts
+//!
+//! The core of the `fake-tcp` crate compose of two structures. [`Stack`] and
+//! [`Socket`].
+//!
+//! ## [`Stack`]
+//!
+//! [`Stack`] represents a virtual TCP stack that operates at
+//! Layer 3. It is responsible for:
+//!
+//! * TCP active and passive open and handshake
+//! * `RST` handling
+//! * Interact with the Tun interface at Layer 3
+//! * Distribute incoming datagrams to corresponding [`Socket`]
+//!
+//! ## [`Socket`]
+//!
+//! [`Socket`] represents a TCP connection. It registers the identifying
+//! tuple `(src_ip, src_port, dest_ip, dest_port)` inside the [`Stack`] so
+//! so that incoming packets can be distributed to the right [`Socket`] with
+//! using a channel. It is also what the client should use for
+//! sending/receiving datagrams.
+//!
+//! # Examples
+//!
+//! Please see [`client.rs`](https://github.com/dndx/phantun/blob/main/phantun/src/bin/client.rs)
+//! and [`server.rs`](https://github.com/dndx/phantun/blob/main/phantun/src/bin/server.rs) files
+//! from the `phantun` crate for how to use this library in client/server mode, respectively.
+
 #![cfg_attr(feature = "benchmark", feature(test))]
 
 pub mod packet;
@@ -21,9 +61,10 @@ use tokio_tun::Tun;
 const TIMEOUT: time::Duration = time::Duration::from_secs(1);
 const RETRIES: usize = 6;
 const MPSC_BUFFER_LEN: usize = 512;
+const MAX_UNACKED_LEN: u32 = 128 * 1024 * 1024; // 128MB
 
 #[derive(Hash, Eq, PartialEq, Clone, Debug)]
-pub struct AddrTuple {
+struct AddrTuple {
     local_addr: SocketAddrV4,
     remote_addr: SocketAddrV4,
 }
@@ -66,9 +107,17 @@ pub struct Socket {
     remote_addr: SocketAddrV4,
     seq: AtomicU32,
     ack: AtomicU32,
+    last_ack: AtomicU32,
     state: State,
 }
 
+/// A socket that represents a unique TCP connection between a server and client.
+///
+/// The `Socket` object itself satisfies `Sync` and `Send`, which means it can
+/// be safely called within an async future.
+///
+/// To close a TCP connection that is no longer needed, simply drop this object
+/// out of scope.
 impl Socket {
     fn new(
         shared: Arc<Shared>,
@@ -89,6 +138,7 @@ impl Socket {
                 remote_addr,
                 seq: AtomicU32::new(0),
                 ack: AtomicU32::new(ack.unwrap_or(0)),
+                last_ack: AtomicU32::new(ack.unwrap_or(0)),
                 state,
             },
             incoming_tx,
@@ -96,32 +146,44 @@ impl Socket {
     }
 
     fn build_tcp_packet(&self, flags: u16, payload: Option<&[u8]>) -> Bytes {
+        let ack = self.ack.load(Ordering::Relaxed);
+        self.last_ack.store(ack, Ordering::Relaxed);
+
         build_tcp_packet(
             self.local_addr,
             self.remote_addr,
             self.seq.load(Ordering::Relaxed),
-            self.ack.load(Ordering::Relaxed),
+            ack,
             flags,
             payload,
         )
     }
 
+    /// Sends a datagram to the other end.
+    ///
+    /// This method takes `&self`, and it can be called safely by multiple threads
+    /// at the same time.
+    ///
+    /// A return of `None` means the Tun socket returned an error
+    /// and this socket must be closed.
     pub async fn send(&self, payload: &[u8]) -> Option<()> {
         match self.state {
             State::Established => {
                 let buf = self.build_tcp_packet(tcp::TcpFlags::ACK, Some(payload));
                 self.seq.fetch_add(payload.len() as u32, Ordering::Relaxed);
-
-                tokio::select! {
-                    res = self.tun.send(&buf) => {
-                        res.ok().and(Some(()))
-                    },
-                }
+                self.tun.send(&buf).await.ok().and(Some(()))
             }
             _ => unreachable!(),
         }
     }
 
+    /// Attempt to receive a datagram from the other end.
+    ///
+    /// This method takes `&self`, and it can be called safely by multiple threads
+    /// at the same time.
+    ///
+    /// A return of `None` means the TCP connection is broken
+    /// and this socket must be closed.
     pub async fn recv(&self, buf: &mut [u8]) -> Option<usize> {
         match self.state {
             State::Established => {
@@ -136,8 +198,18 @@ impl Socket {
 
                     let payload = tcp_packet.payload();
 
-                    self.ack
-                        .store(tcp_packet.get_sequence().wrapping_add(1), Ordering::Relaxed);
+                    let new_ack = tcp_packet.get_sequence().wrapping_add(payload.len() as u32);
+                    let last_ask = self.last_ack.load(Ordering::Relaxed);
+                    self.ack.store(new_ack, Ordering::Relaxed);
+
+                    if new_ack.overflowing_sub(last_ask).0 > MAX_UNACKED_LEN {
+                        let buf = self.build_tcp_packet(tcp::TcpFlags::ACK, None);
+                        if let Err(e) = self.tun.try_send(&buf) {
+                            // This should not really happen as we have not sent anything for
+                            // quite some time...
+                            info!("Connection {} unable to send idling ACK back: {}", self, e)
+                        }
+                    }
 
                     buf[..payload.len()].copy_from_slice(payload);
 
@@ -246,6 +318,7 @@ impl Socket {
 }
 
 impl Drop for Socket {
+    /// Drop the socket and close the TCP connection
     fn drop(&mut self) {
         let tuple = AddrTuple::new(self.local_addr, self.remote_addr);
         // dissociates ourself from the dispatch map
@@ -253,7 +326,14 @@ impl Drop for Socket {
         // purge cache
         self.shared.tuples_purge.send(tuple).unwrap();
 
-        let buf = self.build_tcp_packet(tcp::TcpFlags::RST, None);
+        let buf = build_tcp_packet(
+            self.local_addr,
+            self.remote_addr,
+            self.seq.load(Ordering::Relaxed),
+            0,
+            tcp::TcpFlags::RST,
+            None,
+        );
         if let Err(e) = self.tun.try_send(&buf) {
             warn!("Unable to send RST to remote end: {}", e);
         }
@@ -263,6 +343,7 @@ impl Drop for Socket {
 }
 
 impl fmt::Display for Socket {
+    /// User-friendly string representation of the socket
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(
             f,
@@ -272,7 +353,12 @@ impl fmt::Display for Socket {
     }
 }
 
+/// A userspace TCP state machine
 impl Stack {
+    /// Create a new stack, `tun` is an array of [`Tun`](tokio_tun::Tun).
+    /// When more than one [`Tun`](tokio_tun::Tun) object is passed in, same amount
+    /// of reader will be spawned later. This allows user to utilize the performance
+    /// benefit of Multiqueue Tun support on machines with SMP.
     pub fn new(tun: Vec<Tun>) -> Stack {
         let tun: Vec<Arc<Tun>> = tun.into_iter().map(Arc::new).collect();
         let (ready_tx, ready_rx) = mpsc::channel(MPSC_BUFFER_LEN);
@@ -301,14 +387,18 @@ impl Stack {
         }
     }
 
+    /// Listens for incoming connections on the given `port`.
     pub fn listen(&mut self, port: u16) {
         assert!(self.shared.listening.write().unwrap().insert(port));
     }
 
+    /// Accepts an incoming connection.
     pub async fn accept(&mut self) -> Socket {
         self.ready.recv().await.unwrap()
     }
 
+    /// Connects to the remote end. `None` returned means
+    /// the connection attempt failed.
     pub async fn connect(&mut self, addr: SocketAddrV4) -> Option<Socket> {
         let mut rng = SmallRng::from_entropy();
         let local_port: u16 = rng.gen_range(1024..65535);
@@ -414,8 +504,8 @@ impl Stack {
                                 local_addr,
                                 remote_addr,
                                 0,
-                                tcp_packet.get_sequence() + 1,
-                                tcp::TcpFlags::RST,
+                                tcp_packet.get_sequence() + tcp_packet.payload().len() as u32 + 1, // +1 because of SYN flag set
+                                tcp::TcpFlags::RST | tcp::TcpFlags::ACK,
                                 None,
                             );
                             shared.tun[0].try_send(&buf).unwrap();
@@ -426,8 +516,8 @@ impl Stack {
                             local_addr,
                             remote_addr,
                             tcp_packet.get_acknowledgement(),
-                            0,
-                            tcp::TcpFlags::RST,
+                            tcp_packet.get_sequence() + tcp_packet.payload().len() as u32,
+                            tcp::TcpFlags::RST | tcp::TcpFlags::ACK,
                             None,
                         );
                         shared.tun[0].try_send(&buf).unwrap();
@@ -436,7 +526,7 @@ impl Stack {
                 tuple = tuples_purge.recv() => {
                     let tuple = tuple.unwrap();
                     tuples.remove(&tuple);
-                    trace!("Removed cached tuple");
+                    trace!("Removed cached tuple: {:?}", tuple);
                 }
             }
         }

phantun/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "phantun"
-version = "0.2.5"
+version = "0.3.1"
 edition = "2021"
 authors = ["Datong Sun <dndx@idndx.com>"]
 license = "MIT OR Apache-2.0"
@@ -13,8 +13,9 @@ Layer 3 & Layer 4 (NAPT) firewalls/NATs.
 [dependencies]
 clap = { version = "3.0", features = ["cargo"] }
 socket2 = { version = "0.4", features = ["all"] }
-fake-tcp = "0.2"
+fake-tcp = { path = "../fake-tcp", version = "0.3" }
 tokio = { version = "1.14", features = ["full"] }
+tokio-util = "0.7"
 log = "0.4"
 pretty_env_logger = "0.4"
 tokio-tun = "0.5"

phantun/LICENSE-APACHE

@@ -186,7 +186,7 @@ APPENDIX: How to apply the Apache License to your work.
    same "printed page" as the copyright notice for easier
    identification within third-party archives.
 
-Copyright 2014-2021 The Rust Project Developers
+Copyright 2021-2022 Datong Sun (dndx@idndx.com)
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.

phantun/LICENSE-MIT

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2014-2021 The Rust Project Developers
+Copyright (c) 2021-2022 Datong Sun (dndx@idndx.com)
 
 Permission is hereby granted, free of charge, to any
 person obtaining a copy of this software and associated

phantun/src/bin/client.rs

@@ -1,44 +1,23 @@
-use clap::{crate_version, App, Arg};
+use clap::{crate_version, Arg, Command};
 use fake_tcp::packet::MAX_PACKET_LEN;
 use fake_tcp::{Socket, Stack};
 use log::{debug, error, info};
+use phantun::utils::new_udp_reuseport;
 use std::collections::HashMap;
-use std::convert::TryInto;
 use std::net::{Ipv4Addr, SocketAddr};
 use std::sync::Arc;
-use std::time::Duration;
-use tokio::net::UdpSocket;
-use tokio::sync::RwLock;
+use tokio::sync::{Notify, RwLock};
 use tokio::time;
 use tokio_tun::TunBuilder;
+use tokio_util::sync::CancellationToken;
 
-const UDP_TTL: Duration = Duration::from_secs(180);
-
-fn new_udp_reuseport(addr: SocketAddr) -> UdpSocket {
-    let udp_sock = socket2::Socket::new(
-        if addr.is_ipv4() {
-            socket2::Domain::IPV4
-        } else {
-            socket2::Domain::IPV6
-        },
-        socket2::Type::DGRAM,
-        None,
-    )
-    .unwrap();
-    udp_sock.set_reuse_port(true).unwrap();
-    // from tokio-rs/mio/blob/master/src/sys/unix/net.rs
-    udp_sock.set_cloexec(true).unwrap();
-    udp_sock.set_nonblocking(true).unwrap();
-    udp_sock.bind(&socket2::SockAddr::from(addr)).unwrap();
-    let udp_sock: std::net::UdpSocket = udp_sock.into();
-    udp_sock.try_into().unwrap()
-}
+use phantun::UDP_TTL;
 
 #[tokio::main]
 async fn main() {
     pretty_env_logger::init();
 
-    let matches = App::new("Phantun Client")
+    let matches = Command::new("Phantun Client")
         .version(crate_version!())
         .author("Datong Sun (github.com/dndx)")
         .arg(
@@ -119,6 +98,9 @@ async fn main() {
         .parse()
         .expect("bad peer address for Tun interface");
 
+    let num_cpus = num_cpus::get();
+    info!("{} cores available", num_cpus);
+
     let tun = TunBuilder::new()
         .name(matches.value_of("tun").unwrap()) // if name is empty, then it is set by kernel.
         .tap(false) // false (default): TUN, true: TAP.
@@ -126,7 +108,7 @@ async fn main() {
         .up() // or set it up manually using `sudo ip link set <tun-name> up`.
         .address(tun_local)
         .destination(tun_peer)
-        .try_build_mq(num_cpus::get())
+        .try_build_mq(num_cpus)
         .unwrap();
 
     info!("Created TUN device {}", tun[0].name());
@@ -168,52 +150,85 @@ async fn main() {
                     assert!(connections.write().await.insert(addr, sock.clone()).is_none());
                     debug!("inserted fake TCP socket into connection table");
 
-                    let connections = connections.clone();
-
                     // spawn "fastpath" UDP socket and task, this will offload main task
                     // from forwarding UDP packets
-                    tokio::spawn(async move {
-                        let mut buf_udp = [0u8; MAX_PACKET_LEN];
-                        let mut buf_tcp = [0u8; MAX_PACKET_LEN];
-                        let udp_sock = new_udp_reuseport(local_addr);
-                        udp_sock.connect(addr).await.unwrap();
 
+                    let packet_received = Arc::new(Notify::new());
+                    let quit = CancellationToken::new();
+
+                    for i in 0..num_cpus {
+                        let sock = sock.clone();
+                        let quit = quit.clone();
+                        let packet_received = packet_received.clone();
+
+                        tokio::spawn(async move {
+                            let mut buf_udp = [0u8; MAX_PACKET_LEN];
+                            let mut buf_tcp = [0u8; MAX_PACKET_LEN];
+                            let udp_sock = new_udp_reuseport(local_addr);
+                            udp_sock.connect(addr).await.unwrap();
+
+                            loop {
+                                tokio::select! {
+                                    Ok(size) = udp_sock.recv(&mut buf_udp) => {
+                                        if sock.send(&buf_udp[..size]).await.is_none() {
+                                            debug!("removed fake TCP socket from connections table");
+                                            quit.cancel();
+                                            return;
+                                        }
+
+                                        packet_received.notify_one();
+                                    },
+                                    res = sock.recv(&mut buf_tcp) => {
+                                        match res {
+                                            Some(size) => {
+                                                if size > 0 {
+                                                    if let Err(e) = udp_sock.send(&buf_tcp[..size]).await {
+                                                        error!("Unable to send UDP packet to {}: {}, closing connection", e, addr);
+                                                        quit.cancel();
+                                                        return;
+                                                    }
+                                                }
+                                            },
+                                            None => {
+                                                debug!("removed fake TCP socket from connections table");
+                                                quit.cancel();
+                                                return;
+                                            },
+                                        }
+
+                                        packet_received.notify_one();
+                                    },
+                                    _ = quit.cancelled() => {
+                                        debug!("worker {} terminated", i);
+                                        return;
+                                    },
+                                };
+                            }
+                        });
+                    }
+
+                    let connections = connections.clone();
+                    tokio::spawn(async move {
                         loop {
                             let read_timeout = time::sleep(UDP_TTL);
+                            let packet_received_fut = packet_received.notified();
 
                             tokio::select! {
-                                Ok(size) = udp_sock.recv(&mut buf_udp) => {
-                                    if sock.send(&buf_udp[..size]).await.is_none() {
-                                        connections.write().await.remove(&addr);
-                                        debug!("removed fake TCP socket from connections table");
-                                        return;
-                                    }
-                                },
-                                res = sock.recv(&mut buf_tcp) => {
-                                    match res {
-                                        Some(size) => {
-                                            if size > 0 {
-                                                if let Err(e) = udp_sock.send(&buf_tcp[..size]).await {
-                                                    connections.write().await.remove(&addr);
-                                                    error!("Unable to send UDP packet to {}: {}, closing connection", e, addr);
-                                                    return;
-                                                }
-                                            }
-                                        },
-                                        None => {
-                                            connections.write().await.remove(&addr);
-                                            debug!("removed fake TCP socket from connections table");
-                                            return;
-                                        },
-                                    }
-                                },
                                 _ = read_timeout => {
                                     info!("No traffic seen in the last {:?}, closing connection", UDP_TTL);
                                     connections.write().await.remove(&addr);
                                     debug!("removed fake TCP socket from connections table");
+
+                                    quit.cancel();
                                     return;
-                                }
-                            };
+                                },
+                                _ = quit.cancelled() => {
+                                    connections.write().await.remove(&addr);
+                                    debug!("removed fake TCP socket from connections table");
+                                    return;
+                                },
+                                _ = packet_received_fut => {},
+                            }
                         }
                     });
                 },

phantun/src/bin/server.rs

@@ -1,18 +1,23 @@
-use clap::{crate_version, App, Arg};
+use clap::{crate_version, Arg, Command};
 use fake_tcp::packet::MAX_PACKET_LEN;
 use fake_tcp::Stack;
-use log::{error, info};
+use log::{debug, error, info};
+use phantun::utils::new_udp_reuseport;
 use std::net::Ipv4Addr;
+use std::sync::Arc;
 use tokio::net::UdpSocket;
-use tokio::time::{self, Duration};
+use tokio::sync::Notify;
+use tokio::time;
 use tokio_tun::TunBuilder;
-const UDP_TTL: Duration = Duration::from_secs(180);
+use tokio_util::sync::CancellationToken;
+
+use phantun::UDP_TTL;
 
 #[tokio::main]
 async fn main() {
     pretty_env_logger::init();
 
-    let matches = App::new("Phantun Server")
+    let matches = Command::new("Phantun Server")
         .version(crate_version!())
         .author("Datong Sun (github.com/dndx)")
         .arg(
@@ -88,6 +93,9 @@ async fn main() {
         .parse()
         .expect("bad peer address for Tun interface");
 
+    let num_cpus = num_cpus::get();
+    info!("{} cores available", num_cpus);
+
     let tun = TunBuilder::new()
         .name(matches.value_of("tun").unwrap()) // if name is empty, then it is set by kernel.
         .tap(false) // false (default): TUN, true: TAP.
@@ -95,7 +103,7 @@ async fn main() {
         .up() // or set it up manually using `sudo ip link set <tun-name> up`.
         .address(tun_local)
         .destination(tun_peer)
-        .try_build_mq(num_cpus::get())
+        .try_build_mq(num_cpus)
         .unwrap();
 
     info!("Created TUN device {}", tun[0].name());
@@ -110,46 +118,82 @@ async fn main() {
         let mut buf_tcp = [0u8; MAX_PACKET_LEN];
 
         loop {
-            let sock = stack.accept().await;
+            let sock = Arc::new(stack.accept().await);
             info!("New connection: {}", sock);
 
-            tokio::spawn(async move {
-                let udp_sock = UdpSocket::bind(if remote_addr.is_ipv4() {
-                    "0.0.0.0:0"
-                } else {
-                    "[::]:0"
-                })
-                .await
-                .unwrap();
-                udp_sock.connect(remote_addr).await.unwrap();
+            let packet_received = Arc::new(Notify::new());
+            let quit = CancellationToken::new();
+            let udp_sock = UdpSocket::bind(if remote_addr.is_ipv4() {
+                "0.0.0.0:0"
+            } else {
+                "[::]:0"
+            })
+            .await
+            .unwrap();
+            let local_addr = udp_sock.local_addr().unwrap();
+            drop(udp_sock);
 
+            for i in 0..num_cpus {
+                let sock = sock.clone();
+                let quit = quit.clone();
+                let packet_received = packet_received.clone();
+                let udp_sock = new_udp_reuseport(local_addr);
+
+                tokio::spawn(async move {
+                    udp_sock.connect(remote_addr).await.unwrap();
+
+                    loop {
+                        tokio::select! {
+                            Ok(size) = udp_sock.recv(&mut buf_udp) => {
+                                if sock.send(&buf_udp[..size]).await.is_none() {
+                                    quit.cancel();
+                                    return;
+                                }
+
+                                packet_received.notify_one();
+                            },
+                            res = sock.recv(&mut buf_tcp) => {
+                                match res {
+                                    Some(size) => {
+                                        if size > 0 {
+                                            if let Err(e) = udp_sock.send(&buf_tcp[..size]).await {
+                                                error!("Unable to send UDP packet to {}: {}, closing connection", e, remote_addr);
+                                                quit.cancel();
+                                                return;
+                                            }
+                                        }
+                                    },
+                                    None => {
+                                        quit.cancel();
+                                        return;
+                                    },
+                                }
+
+                                packet_received.notify_one();
+                            },
+                            _ = quit.cancelled() => {
+                                debug!("worker {} terminated", i);
+                                return;
+                            },
+                        };
+                    }
+                });
+            }
+
+            tokio::spawn(async move {
                 loop {
                     let read_timeout = time::sleep(UDP_TTL);
+                    let packet_received_fut = packet_received.notified();
 
                     tokio::select! {
-                        Ok(size) = udp_sock.recv(&mut buf_udp) => {
-                            if sock.send(&buf_udp[..size]).await.is_none() {
-                                return;
-                            }
-                        },
-                        res = sock.recv(&mut buf_tcp) => {
-                            match res {
-                                Some(size) => {
-                                    if size > 0 {
-                                        if let Err(e) = udp_sock.send(&buf_tcp[..size]).await {
-                                            error!("Unable to send UDP packet to {}: {}, closing connection", e, remote_addr);
-                                            return;
-                                        }
-                                    }
-                                },
-                                None => { return; },
-                            }
-                        },
                         _ = read_timeout => {
                             info!("No traffic seen in the last {:?}, closing connection", UDP_TTL);
+
+                            quit.cancel();
                             return;
-                        }
-                    };
+                        },
+                        _ = packet_received_fut => {},
+                    }
                 }
             });
         }

phantun/src/lib.rs

@@ -0,0 +1,5 @@
+use std::time::Duration;
+
+pub mod utils;
+
+pub const UDP_TTL: Duration = Duration::from_secs(180);

phantun/src/utils.rs

@@ -0,0 +1,22 @@
+use std::net::SocketAddr;
+use tokio::net::UdpSocket;
+
+pub fn new_udp_reuseport(local_addr: SocketAddr) -> UdpSocket {
+    let udp_sock = socket2::Socket::new(
+        if local_addr.is_ipv4() {
+            socket2::Domain::IPV4
+        } else {
+            socket2::Domain::IPV6
+        },
+        socket2::Type::DGRAM,
+        None,
+    )
+    .unwrap();
+    udp_sock.set_reuse_port(true).unwrap();
+    // from tokio-rs/mio/blob/master/src/sys/unix/net.rs
+    udp_sock.set_cloexec(true).unwrap();
+    udp_sock.set_nonblocking(true).unwrap();
+    udp_sock.bind(&socket2::SockAddr::from(local_addr)).unwrap();
+    let udp_sock: std::net::UdpSocket = udp_sock.into();
+    udp_sock.try_into().unwrap()
+}