Work on Up/Down operation on WireGuard device

3 files changed, 57 insertions, 46 deletions

diff --git a/src/wireguard/router/device.rs b/src/wireguard/router/device.rs
index 7c3b0a1..a5028e1 100644
--- a/src/wireguard/router/device.rs
+++ b/src/wireguard/router/device.rs
@@ -27,6 +27,8 @@ use super::route::get_route;
 use super::super::{bind, tun, Endpoint, KeyPair};
 
 pub struct DeviceInner<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> {
+    pub enabled: AtomicBool,
+
     // inbound writer (TUN)
     pub inbound: T,
 
@@ -91,6 +93,7 @@ impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> Device<E, C,
         // allocate shared device state
         let inner = DeviceInner {
             inbound: tun,
+            enabled: AtomicBool::new(true),
             outbound: RwLock::new(None),
             queues: Mutex::new(Vec::with_capacity(num_workers)),
             queue_next: AtomicUsize::new(0),
@@ -114,6 +117,16 @@ impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> Device<E, C,
         }
     }
 
+    /// Brings the router down.
+    /// When the router is brought down it:
+    /// - Prevents transmission of outbound messages.
+    /// - Erases all key state (key-wheels) of all peers
+    pub fn down(&self) {}
+
+    /// Brints the router up
+    /// When the router is brought up it enables the transmission of outbound messages.
+    pub fn up(&self) {}
+
     /// A new secret key has been set for the device.
     /// According to WireGuard semantics, this should cause all "sending" keys to be discarded.
     pub fn new_sk(&self) {}
diff --git a/src/wireguard/router/peer.rs b/src/wireguard/router/peer.rs
index 5522a3e..7527a60 100644
--- a/src/wireguard/router/peer.rs
+++ b/src/wireguard/router/peer.rs
@@ -1,5 +1,6 @@
 use std::mem;
 use std::net::{IpAddr, SocketAddr};
+use std::ops::Deref;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::mpsc::{sync_channel, SyncSender};
@@ -11,7 +12,6 @@ use log::debug;
 use spin::Mutex;
 use treebitmap::address::Address;
 use treebitmap::IpLookupTable;
-use zerocopy::LayoutVerified;
 
 use super::super::constants::*;
 use super::super::{bind, tun, Endpoint, KeyPair};
@@ -55,6 +55,14 @@ pub struct Peer<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> {
     thread_inbound: Option<thread::JoinHandle<()>>,
 }
 
+impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> Deref for Peer<E, C, T, B> {
+    type Target = Arc<PeerInner<E, C, T, B>>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.state
+    }
+}
+
 fn treebit_list<A, R, E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>>(
     peer: &Arc<PeerInner<E, C, T, B>>,
     table: &spin::RwLock<IpLookupTable<A, Arc<PeerInner<E, C, T, B>>>>,
@@ -199,7 +207,7 @@ pub fn new_peer<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>>(
     let thread_inbound = {
         let peer = peer.clone();
         let device = device.clone();
-        thread::spawn(move || worker_outbound(device, peer, out_rx))
+        thread::spawn(move || worker_outbound(peer, out_rx))
     };
 
     // spawn inbound thread
@@ -217,6 +225,36 @@ pub fn new_peer<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>>(
 }
 
 impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> PeerInner<E, C, T, B> {
+    /// Send a raw message to the peer (used for handshake messages)
+    ///
+    /// # Arguments
+    ///
+    /// - `msg`, message body to send to peer
+    ///
+    /// # Returns
+    ///
+    /// Unit if packet was sent, or an error indicating why sending failed
+    pub fn send(&self, msg: &[u8]) -> Result<(), RouterError> {
+        debug!("peer.send");
+
+        // check if device is enabled
+        if !self.device.enabled.load(Ordering::Acquire) {
+            return Ok(());
+        }
+
+        // send to endpoint (if known)
+        match self.endpoint.lock().as_ref() {
+            Some(endpoint) => self
+                .device
+                .outbound
+                .read()
+                .as_ref()
+                .ok_or(RouterError::SendError)
+                .and_then(|w| w.write(msg, endpoint).map_err(|_| RouterError::SendError)),
+            None => Err(RouterError::NoEndpoint),
+        }
+    }
+
     fn send_staged(&self) -> bool {
         debug!("peer.send_staged");
         let mut sent = false;
@@ -498,7 +536,7 @@ impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> Peer<E, C, T
 
     pub fn send_keepalive(&self) -> bool {
         debug!("peer.send_keepalive");
-        self.state.send_raw(vec![0u8; SIZE_MESSAGE_PREFIX])
+        self.send_raw(vec![0u8; SIZE_MESSAGE_PREFIX])
     }
 
     /// Map a subnet to the peer
@@ -565,30 +603,6 @@ impl<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>> Peer<E, C, T
         treebit_remove(self, &self.state.device.ipv6);
     }
 
-    /// Send a raw message to the peer (used for handshake messages)
-    ///
-    /// # Arguments
-    ///
-    /// - `msg`, message body to send to peer
-    ///
-    /// # Returns
-    ///
-    /// Unit if packet was sent, or an error indicating why sending failed
-    pub fn send(&self, msg: &[u8]) -> Result<(), RouterError> {
-        debug!("peer.send");
-        let inner = &self.state;
-        match inner.endpoint.lock().as_ref() {
-            Some(endpoint) => inner
-                .device
-                .outbound
-                .read()
-                .as_ref()
-                .ok_or(RouterError::SendError)
-                .and_then(|w| w.write(msg, endpoint).map_err(|_| RouterError::SendError)),
-            None => Err(RouterError::NoEndpoint),
-        }
-    }
-
     pub fn clear_src(&self) {
         (*self.state.endpoint.lock())
             .as_mut()
diff --git a/src/wireguard/router/workers.rs b/src/wireguard/router/workers.rs
index 08c2db9..5482cee 100644
--- a/src/wireguard/router/workers.rs
+++ b/src/wireguard/router/workers.rs
@@ -141,8 +141,7 @@ pub fn worker_inbound<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer
 /* TODO: Replace with run-queue
  */
 pub fn worker_outbound<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Writer<E>>(
-    device: Arc<DeviceInner<E, C, T, B>>, // related device
-    peer: Arc<PeerInner<E, C, T, B>>,     // related peer
+    peer: Arc<PeerInner<E, C, T, B>>,
     receiver: Receiver<JobOutbound>,
 ) {
     loop {
@@ -160,23 +159,8 @@ pub fn worker_outbound<E: Endpoint, C: Callbacks, T: tun::Writer, B: bind::Write
             .map(|buf| {
                 debug!("outbound worker: job complete");
 
-                // write to UDP bind
-                let xmit = if let Some(dst) = peer.endpoint.lock().as_ref() {
-                    let send: &Option<B> = &*device.outbound.read();
-                    if let Some(writer) = send.as_ref() {
-                        match writer.write(&buf.msg[..], dst) {
-                            Err(e) => {
-                                debug!("failed to send outbound packet: {:?}", e);
-                                false
-                            }
-                            Ok(_) => true,
-                        }
-                    } else {
-                        false
-                    }
-                } else {
-                    false
-                };
+                // send to peer
+                let xmit = peer.send(&buf.msg[..]).is_ok();
 
                 // trigger callback
                 C::send(&peer.opaque, buf.msg.len(), xmit, &buf.keypair, buf.counter);