Cleanup

src/main.rs

@@ -10,6 +10,7 @@ use tokio::net::{TcpListener, TcpStream, UdpSocket};
 
 use crate::config::Config;
 use crate::port_pool::PortPool;
+use crate::wg::WireGuardTunnel;
 
 pub mod client;
 pub mod config;
@@ -23,22 +24,15 @@ pub const MAX_PACKET: usize = 65536;
 async fn main() -> anyhow::Result<()> {
     pretty_env_logger::init_custom_env("ONETUN_LOG");
     let config = Config::from_args().with_context(|| "Failed to read config")?;
-    let peer = Arc::new(wg::create_tunnel(&config)?);
     let port_pool = Arc::new(PortPool::new());
 
-    // endpoint_addr: The address of the public WireGuard endpoint; UDP.
+    let wg = WireGuardTunnel::new(&config)
-    let endpoint_addr = config.endpoint_addr;
-
-    // wireguard_udp: The UDP socket used to communicate with the public WireGuard endpoint.
-    let wireguard_udp = UdpSocket::bind("0.0.0.0:0")
         .await
-        .with_context(|| "Failed to create UDP socket for WireGuard connection")?;
+        .with_context(|| "Failed to initialize WireGuard tunnel")?;
-    let wireguard_udp = Arc::new(wireguard_udp);
+    let wg = Arc::new(wg);
 
     // Start routine task for WireGuard
-    tokio::spawn(
+    tokio::spawn(async move { Arc::clone(&wg).routine_task().await });
-        async move { wg::routine(peer.clone(), wireguard_udp.clone(), endpoint_addr).await },
-    );
 
     info!(
         "Tunnelling [{}]->[{}] (via [{}] as peer {})",

src/wg.rs

@@ -1,64 +1,88 @@
-use anyhow::Context;
-use boringtun::noise::{Tunn, TunnResult};
 use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Duration;
+
+use anyhow::Context;
+use boringtun::noise::{Tunn, TunnResult};
 use tokio::net::UdpSocket;
 
 use crate::config::Config;
 use crate::MAX_PACKET;
 
-pub fn create_tunnel(config: &Config) -> anyhow::Result<Box<Tunn>> {
+pub struct WireGuardTunnel {
-    Tunn::new(
+    /// `boringtun` peer/tunnel implementation, used for crypto & WG protocol.
-        config.private_key.clone(),
+    peer: Box<Tunn>,
-        config.endpoint_public_key.clone(),
+    /// The UDP socket for the public WireGuard endpoint to connect to.
-        None,
+    udp: UdpSocket,
-        None,
+    /// The address of the public WireGuard endpoint (UDP).
-        0,
+    endpoint: SocketAddr,
-        None,
-    )
-    .map_err(|s| anyhow::anyhow!("{}", s))
-    .with_context(|| "Failed to initialize peer")
 }
 
-/// WireGuard Routine task. Handles Handshake, keep-alive, etc.
+impl WireGuardTunnel {
-pub async fn routine(
+    /// Initialize a new WireGuard tunnel.
-    peer: Arc<Box<Tunn>>,
+    pub async fn new(config: &Config) -> anyhow::Result<Self> {
-    wireguard_udp: Arc<UdpSocket>,
+        let peer = Self::create_tunnel(&config)?;
-    endpoint_addr: SocketAddr,
+        let udp = UdpSocket::bind("0.0.0.0:0")
-) {
+            .await
-    debug!("Started WireGuard routine thread");
+            .with_context(|| "Failed to create UDP socket for WireGuard connection")?;
-    loop {
+        let endpoint = config.endpoint_addr;
-        let mut send_buf = [0u8; MAX_PACKET];
+
-        match peer.update_timers(&mut send_buf) {
+        Ok(Self {
-            TunnResult::WriteToNetwork(packet) => {
+            peer,
-                debug!(
+            udp,
-                    "Sending routine packet of {} bytes to WireGuard endpoint",
+            endpoint,
-                    packet.len()
+        })
-                );
+    }
-                match wireguard_udp.send_to(packet, endpoint_addr).await {
+
-                    Ok(_) => {}
+    /// WireGuard Routine task. Handles Handshake, keep-alive, etc.
-                    Err(e) => {
+    pub async fn routine_task(&self) -> ! {
-                        error!(
+        trace!("Starting WireGuard routine task");
-                            "Failed to send routine packet to WireGuard endpoint: {:?}",
+
-                            e
+        loop {
-                        );
+            let mut send_buf = [0u8; MAX_PACKET];
-                    }
+            match self.peer.update_timers(&mut send_buf) {
-                };
+                TunnResult::WriteToNetwork(packet) => {
-            }
+                    debug!(
-            TunnResult::Err(e) => {
+                        "Sending routine packet of {} bytes to WireGuard endpoint",
-                error!(
+                        packet.len()
-                    "Failed to prepare routine packet for WireGuard endpoint: {:?}",
+                    );
-                    e
+                    match self.udp.send_to(packet, self.endpoint).await {
-                );
+                        Ok(_) => {}
-            }
+                        Err(e) => {
-            TunnResult::Done => {
+                            error!(
-                // Sleep for a bit
+                                "Failed to send routine packet to WireGuard endpoint: {:?}",
-                tokio::time::sleep(Duration::from_millis(100)).await;
+                                e
-            }
+                            );
-            other => {
+                        }
-                warn!("Unexpected WireGuard routine task state: {:?}", other);
+                    };
+                }
+                TunnResult::Err(e) => {
+                    error!(
+                        "Failed to prepare routine packet for WireGuard endpoint: {:?}",
+                        e
+                    );
+                }
+                TunnResult::Done => {
+                    // Sleep for a bit
+                    tokio::time::sleep(Duration::from_millis(100)).await;
+                }
+                other => {
+                    warn!("Unexpected WireGuard routine task state: {:?}", other);
+                }
             }
         }
     }
+
+    fn create_tunnel(config: &Config) -> anyhow::Result<Box<Tunn>> {
+        Tunn::new(
+            config.private_key.clone(),
+            config.endpoint_public_key.clone(),
+            None,
+            None,
+            0,
+            None,
+        )
+        .map_err(|s| anyhow::anyhow!("{}", s))
+        .with_context(|| "Failed to initialize boringtun Tunn")
+    }
 }