port forward (#736)

* support tcp port forward
* support udp port forward
* command line option for port forward

easytier/src/gateway/tokio_smoltcp/mod.rs

@@ -20,7 +20,7 @@ use smoltcp::{
     time::{Duration, Instant},
     wire::{HardwareAddress, IpAddress, IpCidr},
 };
-pub use socket::{TcpListener, TcpStream};
+pub use socket::{TcpListener, TcpStream, UdpSocket};
 pub use socket_allocator::BufferSize;
 use tokio::sync::Notify;
 
@@ -158,6 +158,13 @@ impl Net {
         )
         .await
     }
+
+    /// This function will create a new UDP socket and attempt to bind it to the `addr` provided.
+    pub async fn udp_bind(&self, addr: SocketAddr) -> io::Result<UdpSocket> {
+        let addr = self.set_address(addr);
+        UdpSocket::new(self.reactor.clone(), addr.into()).await
+    }
+
     fn set_address(&self, mut addr: SocketAddr) -> SocketAddr {
         if addr.ip().is_unspecified() {
             addr.set_ip(match self.ip_addr.address() {

easytier/src/gateway/tokio_smoltcp/socket.rs

@@ -2,6 +2,7 @@ use super::{reactor::Reactor, socket_allocator::SocketHandle};
 use futures::future::{self, poll_fn};
 use futures::{ready, Stream};
 pub use smoltcp::socket::tcp;
+use smoltcp::socket::udp;
 use smoltcp::wire::{IpAddress, IpEndpoint};
 use std::mem::replace;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
@@ -247,3 +248,86 @@ impl AsyncWrite for TcpStream {
         Poll::Pending
     }
 }
+
+/// A UDP socket.
+pub struct UdpSocket {
+    handle: SocketHandle,
+    reactor: Arc<Reactor>,
+    local_addr: SocketAddr,
+}
+
+impl UdpSocket {
+    pub(super) async fn new(
+        reactor: Arc<Reactor>,
+        local_endpoint: IpEndpoint,
+    ) -> io::Result<UdpSocket> {
+        let handle = reactor.socket_allocator().new_udp_socket();
+        {
+            let mut socket = reactor.get_socket::<udp::Socket>(*handle);
+            socket.bind(local_endpoint).map_err(map_err)?;
+        }
+
+        let local_addr = ep2sa(&local_endpoint);
+
+        Ok(UdpSocket {
+            handle,
+            reactor,
+            local_addr,
+        })
+    }
+    /// Note that on multiple calls to a poll_* method in the send direction, only the Waker from the Context passed to the most recent call will be scheduled to receive a wakeup.
+    pub fn poll_send_to(
+        &self,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+        target: SocketAddr,
+    ) -> Poll<io::Result<usize>> {
+        let mut socket = self.reactor.get_socket::<udp::Socket>(*self.handle);
+        let target_ip: IpEndpoint = target.into();
+
+        match socket.send_slice(buf, target_ip) {
+            // the buffer is full
+            Err(udp::SendError::BufferFull) => {}
+            r => {
+                r.map_err(map_err)?;
+                self.reactor.notify();
+                return Poll::Ready(Ok(buf.len()));
+            }
+        }
+
+        socket.register_send_waker(cx.waker());
+        Poll::Pending
+    }
+    /// See note on `poll_send_to`
+    pub async fn send_to(&self, buf: &[u8], target: SocketAddr) -> io::Result<usize> {
+        poll_fn(|cx| self.poll_send_to(cx, buf, target)).await
+    }
+    /// Note that on multiple calls to a poll_* method in the recv direction, only the Waker from the Context passed to the most recent call will be scheduled to receive a wakeup.
+    pub fn poll_recv_from(
+        &self,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<(usize, SocketAddr)>> {
+        let mut socket = self.reactor.get_socket::<udp::Socket>(*self.handle);
+
+        match socket.recv_slice(buf) {
+            // the buffer is empty
+            Err(udp::RecvError::Exhausted) => {}
+            r => {
+                let (size, metadata) = r.map_err(map_err)?;
+                self.reactor.notify();
+                return Poll::Ready(Ok((size, ep2sa(&metadata.endpoint))));
+            }
+        }
+
+        socket.register_recv_waker(cx.waker());
+        Poll::Pending
+    }
+    /// See note on `poll_recv_from`
+    pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
+        poll_fn(|cx| self.poll_recv_from(cx, buf)).await
+    }
+    pub fn local_addr(&self) -> io::Result<SocketAddr> {
+        Ok(self.local_addr)
+    }
+}

easytier/src/gateway/tokio_smoltcp/socket_allocator.rs

@@ -1,7 +1,7 @@
 use parking_lot::Mutex;
 use smoltcp::{
     iface::{SocketHandle as InnerSocketHandle, SocketSet},
-    socket::tcp,
+    socket::{tcp, udp},
     time::Duration,
 };
 use std::{
@@ -14,6 +14,11 @@ use std::{
 pub struct BufferSize {
     pub tcp_rx_size: usize,
     pub tcp_tx_size: usize,
+
+    pub udp_rx_size: usize,
+    pub udp_tx_size: usize,
+    pub udp_rx_meta_size: usize,
+    pub udp_tx_meta_size: usize,
 }
 
 impl Default for BufferSize {
@@ -21,6 +26,11 @@ impl Default for BufferSize {
         BufferSize {
             tcp_rx_size: 8192,
             tcp_tx_size: 8192,
+
+            udp_rx_size: 8192,
+            udp_tx_size: 8192,
+            udp_rx_meta_size: 32,
+            udp_tx_meta_size: 32,
         }
     }
 }
@@ -59,6 +69,26 @@ impl SocketAlloctor {
 
         tcp
     }
+
+    pub fn new_udp_socket(&self) -> SocketHandle {
+        let mut set = self.sockets.lock();
+        let handle = set.add(self.alloc_udp_socket());
+        SocketHandle::new(handle, self.sockets.clone())
+    }
+
+    fn alloc_udp_socket(&self) -> udp::Socket<'static> {
+        let rx_buffer = udp::PacketBuffer::new(
+            vec![udp::PacketMetadata::EMPTY; self.buffer_size.udp_rx_meta_size],
+            vec![0; self.buffer_size.udp_rx_size],
+        );
+        let tx_buffer = udp::PacketBuffer::new(
+            vec![udp::PacketMetadata::EMPTY; self.buffer_size.udp_tx_meta_size],
+            vec![0; self.buffer_size.udp_tx_size],
+        );
+        let udp = udp::Socket::new(rx_buffer, tx_buffer);
+
+        udp
+    }
 }
 
 pub struct SocketHandle(InnerSocketHandle, SharedSocketSet);