support no tun mode (#141)

easytier/src/gateway/tokio_smoltcp/reactor.rs

@@ -0,0 +1,163 @@
+use super::{
+    device::{BufferDevice, Packet},
+    socket_allocator::{BufferSize, SocketAlloctor},
+};
+use futures::{stream::iter, FutureExt, SinkExt, StreamExt};
+use parking_lot::{MappedMutexGuard, Mutex, MutexGuard};
+use smoltcp::{
+    iface::{Context, Interface, SocketHandle},
+    socket::{AnySocket, Socket},
+    time::{Duration, Instant},
+};
+use std::{collections::VecDeque, future::Future, io, sync::Arc};
+use tokio::{pin, select, sync::Notify, time::sleep};
+
+pub(crate) type BufferInterface = Arc<Mutex<Interface>>;
+const MAX_BURST_SIZE: usize = 100;
+
+pub(crate) struct Reactor {
+    notify: Arc<Notify>,
+    iface: BufferInterface,
+    socket_allocator: SocketAlloctor,
+}
+
+async fn receive(
+    async_iface: &mut impl super::device::AsyncDevice,
+    recv_buf: &mut VecDeque<Packet>,
+) -> io::Result<()> {
+    if let Some(packet) = async_iface.next().await {
+        recv_buf.push_back(packet?);
+    }
+    Ok(())
+}
+
+async fn run(
+    mut async_iface: impl super::device::AsyncDevice,
+    iface: BufferInterface,
+    mut device: BufferDevice,
+    socket_allocator: SocketAlloctor,
+    notify: Arc<Notify>,
+    stopper: Arc<Notify>,
+) -> io::Result<()> {
+    let default_timeout = Duration::from_secs(60);
+    let timer = sleep(default_timeout.into());
+    let max_burst_size = async_iface
+        .capabilities()
+        .max_burst_size
+        .unwrap_or(MAX_BURST_SIZE);
+    let mut recv_buf = VecDeque::with_capacity(max_burst_size);
+    pin!(timer);
+
+    loop {
+        let packets = device.take_send_queue();
+
+        async_iface
+            .send_all(&mut iter(packets).map(|p| Ok(p)))
+            .await?;
+
+        if recv_buf.is_empty() && device.need_wait() {
+            let start = Instant::now();
+            let deadline = {
+                iface
+                    .lock()
+                    .poll_delay(start, &socket_allocator.sockets().lock())
+                    .unwrap_or(default_timeout)
+            };
+
+            timer
+                .as_mut()
+                .reset(tokio::time::Instant::now() + deadline.into());
+            select! {
+                _ = &mut timer => {},
+                _ = receive(&mut async_iface,&mut recv_buf) => {}
+                _ = notify.notified() => {}
+                _ = stopper.notified() => break,
+            };
+
+            while let (true, Some(Ok(p))) = (
+                recv_buf.len() < max_burst_size,
+                async_iface.next().now_or_never().flatten(),
+            ) {
+                recv_buf.push_back(p);
+            }
+        }
+
+        let mut iface = iface.lock();
+
+        device.push_recv_queue(recv_buf.drain(..device.avaliable_recv_queue().min(recv_buf.len())));
+
+        iface.poll(
+            Instant::now(),
+            &mut device,
+            &mut socket_allocator.sockets().lock(),
+        );
+    }
+
+    Ok(())
+}
+
+impl Reactor {
+    pub fn new(
+        async_device: impl super::device::AsyncDevice,
+        iface: Interface,
+        device: BufferDevice,
+        buffer_size: BufferSize,
+        stopper: Arc<Notify>,
+    ) -> (Self, impl Future<Output = io::Result<()>> + Send) {
+        let iface = Arc::new(Mutex::new(iface));
+        let notify = Arc::new(Notify::new());
+        let socket_allocator = SocketAlloctor::new(buffer_size);
+        let fut = run(
+            async_device,
+            iface.clone(),
+            device,
+            socket_allocator.clone(),
+            notify.clone(),
+            stopper,
+        );
+
+        (
+            Reactor {
+                notify,
+                iface: iface.clone(),
+                socket_allocator,
+            },
+            fut,
+        )
+    }
+    pub fn get_socket<T: AnySocket<'static>>(
+        &self,
+        handle: SocketHandle,
+    ) -> MappedMutexGuard<'_, T> {
+        MutexGuard::map(
+            self.socket_allocator.sockets().lock(),
+            |sockets: &mut smoltcp::iface::SocketSet<'_>| sockets.get_mut::<T>(handle),
+        )
+    }
+    pub fn context(&self) -> MappedMutexGuard<'_, Context> {
+        MutexGuard::map(self.iface.lock(), |iface| iface.context())
+    }
+    pub fn socket_allocator(&self) -> &SocketAlloctor {
+        &self.socket_allocator
+    }
+    pub fn notify(&self) {
+        self.notify.notify_waiters();
+    }
+    pub fn iface(&self) -> &BufferInterface {
+        &self.iface
+    }
+}
+
+impl Drop for Reactor {
+    fn drop(&mut self) {
+        for (_, socket) in self.socket_allocator.sockets().lock().iter_mut() {
+            match socket {
+                Socket::Tcp(tcp) => tcp.close(),
+                Socket::Raw(_) => {}
+                Socket::Udp(udp) => udp.close(),
+                #[allow(unreachable_patterns)]
+                _ => {}
+            }
+        }
+    }
+}