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 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")
+    let wg = WireGuardTunnel::new(&config)
         .await
-        .with_context(|| "Failed to create UDP socket for WireGuard connection")?;
-    let wireguard_udp = Arc::new(wireguard_udp);
+        .with_context(|| "Failed to initialize WireGuard tunnel")?;
+    let wg = Arc::new(wg);
 
     // Start routine task for WireGuard
-    tokio::spawn(
-        async move { wg::routine(peer.clone(), wireguard_udp.clone(), endpoint_addr).await },
-    );
+    tokio::spawn(async move { Arc::clone(&wg).routine_task().await });
 
     info!(
         "Tunnelling [{}]->[{}] (via [{}] as peer {})",

src/wg.rs

@@ -1,42 +1,52 @@
-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>> {
-    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 peer")
+pub struct WireGuardTunnel {
+    /// `boringtun` peer/tunnel implementation, used for crypto & WG protocol.
+    peer: Box<Tunn>,
+    /// The UDP socket for the public WireGuard endpoint to connect to.
+    udp: UdpSocket,
+    /// The address of the public WireGuard endpoint (UDP).
+    endpoint: SocketAddr,
 }
 
-/// WireGuard Routine task. Handles Handshake, keep-alive, etc.
-pub async fn routine(
-    peer: Arc<Box<Tunn>>,
-    wireguard_udp: Arc<UdpSocket>,
-    endpoint_addr: SocketAddr,
-) {
-    debug!("Started WireGuard routine thread");
+impl WireGuardTunnel {
+    /// Initialize a new WireGuard tunnel.
+    pub async fn new(config: &Config) -> anyhow::Result<Self> {
+        let peer = Self::create_tunnel(&config)?;
+        let udp = UdpSocket::bind("0.0.0.0:0")
+            .await
+            .with_context(|| "Failed to create UDP socket for WireGuard connection")?;
+        let endpoint = config.endpoint_addr;
+
+        Ok(Self {
+            peer,
+            udp,
+            endpoint,
+        })
+    }
+
+    /// WireGuard Routine task. Handles Handshake, keep-alive, etc.
+    pub async fn routine_task(&self) -> ! {
+        trace!("Starting WireGuard routine task");
+
         loop {
             let mut send_buf = [0u8; MAX_PACKET];
-        match peer.update_timers(&mut send_buf) {
+            match self.peer.update_timers(&mut send_buf) {
                 TunnResult::WriteToNetwork(packet) => {
                     debug!(
                         "Sending routine packet of {} bytes to WireGuard endpoint",
                         packet.len()
                     );
-                match wireguard_udp.send_to(packet, endpoint_addr).await {
+                    match self.udp.send_to(packet, self.endpoint).await {
                         Ok(_) => {}
                         Err(e) => {
                             error!(
@@ -61,4 +71,18 @@ pub async fn routine(
                 }
             }
         }
+    }
+
+    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")
+    }
 }