From 006a1b0b4e2949e6658f7f1712d259db81a24d35 Mon Sep 17 00:00:00 2001
From: Aram Peres <aram.peres@wavy.fm>
Date: Fri, 15 Oct 2021 19:26:26 -0400
Subject: [PATCH] Graceful shutdown. Docs.

---
 src/main.rs           | 29 +++++++++++++++++++++++++++--
 src/port_pool.rs      |  6 ++++++
 src/virtual_device.rs |  4 ++++
 src/wg.rs             |  3 +++
 4 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 0975f45..57a13b4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -298,11 +298,20 @@ async fn virtual_tcp_interface(
     let _server_handle = socket_set.add(server_socket?);
     let client_handle = socket_set.add(client_socket?);
 
+    let mut graceful_shutdown = false;
+
     loop {
         let loop_start = smoltcp::time::Instant::now();
+        let forceful_shutdown = abort.load(Ordering::Relaxed);
 
-        if abort.load(Ordering::Relaxed) {
-            break;
+        if forceful_shutdown {
+            // Un-graceful shutdown: sends a RST packet.
+            trace!(
+                "[{}] Forcefully shutting down virtual interface",
+                virtual_port
+            );
+            let mut client_socket = socket_set.get::<TcpSocket>(client_handle);
+            client_socket.abort();
         }
 
         match virtual_interface.poll(&mut socket_set, loop_start) {
@@ -347,10 +356,26 @@ async fn virtual_tcp_interface(
                     }
                 }
             }
+            if !graceful_shutdown && !forceful_shutdown && !client_socket.is_active() {
+                // Graceful shutdown
+                client_socket.close();
+                trace!(
+                    "[{}] Gracefully shutting down virtual interface",
+                    virtual_port
+                );
+                // We don't break the loop right away so that the FIN segment can be sent in the next poll.
+                graceful_shutdown = true;
+                continue;
+            }
+        }
+
+        if graceful_shutdown || forceful_shutdown {
+            break;
         }
 
         tokio::time::sleep(Duration::from_millis(1)).await;
     }
     trace!("[{}] Virtual interface task terminated", virtual_port);
+    abort.store(true, Ordering::Relaxed);
     Ok(())
 }
diff --git a/src/port_pool.rs b/src/port_pool.rs
index 6c22e73..771c7a9 100644
--- a/src/port_pool.rs
+++ b/src/port_pool.rs
@@ -6,6 +6,8 @@ const MIN_PORT: u16 = 32768;
 const MAX_PORT: u16 = 60999;
 const PORT_RANGE: Range<u16> = MIN_PORT..MAX_PORT;
 
+/// A pool of virtual ports available.
+/// This structure is thread-safe and lock-free; you can use it safely in an `Arc`.
 pub struct PortPool {
     /// Remaining ports
     inner: lockfree::queue::Queue<u16>,
@@ -20,6 +22,7 @@ impl Default for PortPool {
 }
 
 impl PortPool {
+    /// Initializes a new pool of virtual ports.
     pub fn new() -> Self {
         let inner = lockfree::queue::Queue::default();
         PORT_RANGE.for_each(|p| inner.push(p) as ());
@@ -29,6 +32,7 @@ impl PortPool {
         }
     }
 
+    /// Requests a free port from the pool. An error is returned if none is available (exhaused max capacity).
     pub fn next(&self) -> anyhow::Result<u16> {
         let port = self
             .inner
@@ -41,11 +45,13 @@ impl PortPool {
         Ok(port)
     }
 
+    /// Releases a port back into the pool.
     pub fn release(&self, port: u16) {
         self.inner.push(port);
         self.taken.remove(&port);
     }
 
+    /// Whether the given port is in use by a virtual interface.
     pub fn is_in_use(&self, port: u16) -> bool {
         self.taken.contains(&port)
     }
diff --git a/src/virtual_device.rs b/src/virtual_device.rs
index 97aa267..9692963 100644
--- a/src/virtual_device.rs
+++ b/src/virtual_device.rs
@@ -3,9 +3,13 @@ use smoltcp::phy::{Device, DeviceCapabilities, Medium};
 use smoltcp::time::Instant;
 use std::sync::Arc;
 
+/// A virtual device that processes IP packets. IP packets received from the WireGuard endpoint
+/// are made available to this device using a broadcast channel receiver. IP packets sent from this device
+/// are asynchronously sent out to the WireGuard tunnel.
 pub struct VirtualIpDevice {
     /// Tunnel to send IP packets to.
     wg: Arc<WireGuardTunnel>,
+    /// Broadcast channel receiver for received IP packets.
     ip_broadcast_rx: tokio::sync::broadcast::Receiver<Vec<u8>>,
 }
 
diff --git a/src/wg.rs b/src/wg.rs
index 7a15555..4b88ef6 100644
--- a/src/wg.rs
+++ b/src/wg.rs
@@ -19,6 +19,9 @@ use crate::MAX_PACKET;
 /// The capacity of the broadcast channel for received IP packets.
 const BROADCAST_CAPACITY: usize = 1_000;
 
+/// A WireGuard tunnel. Encapsulates and decapsulates IP packets
+/// to be sent to and received from a remote UDP endpoint.
+/// This tunnel supports at most 1 peer IP at a time, but supports simultaneous ports.
 pub struct WireGuardTunnel {
     source_peer_ip: IpAddr,
     /// `boringtun` peer/tunnel implementation, used for crypto & WG protocol.