Improved API and unit test

2 files changed, 83 insertions, 22 deletions

diff --git a/src/device.rs b/src/device.rs
index bd9103b..c797c79 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -1,8 +1,9 @@
 //! Tunnel device handling
-use std::path::PathBuf;
-use std::io::{Read, Write};
+use std::env::temp_dir;
 use std::fs::{File, OpenOptions};
+use std::io::{Read, Write};
 use std::os::unix::io::AsRawFd;
+use std::path::PathBuf;
 
 use bindgen::*;
 use error::WgResult;
@@ -11,17 +12,27 @@ use error::WgResult;
 /// A certain device
 pub struct Device {
     /// The interface name
-    pub name: String,
+    name: String,
 
     /// The tunnel device file descriptor
-    pub fd: File,
+    fd: File,
+
+    /// The full path to the file
+    path: PathBuf,
+
+    /// Dummy indicator
+    is_dummy: bool,
+
+    /// A read/write counter
+    rw_count: u64,
 }
 
 impl Device {
     /// Create a new tunneling `Device`
     pub fn new(name: &str) -> WgResult<Self> {
         // Get a file descriptor to the operating system
-        let fd = OpenOptions::new().read(true).write(true).open("/dev/net/tun")?;
+        let path = "/dev/net/tun";
+        let fd = OpenOptions::new().read(true).write(true).open(path)?;
 
         // Get the default interface options
         let mut ifr = ifreq::new();
@@ -49,28 +60,43 @@ impl Device {
         Ok(Device {
             name: name.to_owned(),
             fd: fd,
+            path: PathBuf::from(path),
+            is_dummy: false,
+            rw_count: 0,
         })
     }
 
     /// Create a dummy device for testing
-    pub fn dummy(name: &str) -> WgResult<Self> {
+    pub fn dummy() -> WgResult<Self> {
+        let name = "wg";
+        let path = temp_dir().join(name);
         let fd = OpenOptions::new().read(true)
             .write(true)
             .create(true)
-            .open(PathBuf::from("/tmp").join(name))?;
+            .open(&path)?;
         Ok(Device {
             name: name.to_owned(),
             fd: fd,
+            path: path,
+            is_dummy: true,
+            rw_count: 0,
         })
     }
 
     /// Reads a frame from the device, returns the number of bytes read
     pub fn read(&mut self, mut buffer: &mut [u8]) -> WgResult<usize> {
+        // Increment the read/write count
+        self.increment_rw_count();
+
+        // Read from the file descriptor
         Ok(self.fd.read(&mut buffer)?)
     }
 
     /// Write a frame to the device
     pub fn write(&mut self, data: &[u8]) -> WgResult<usize> {
+        // Increment the read/write count
+        self.increment_rw_count();
+
         // Write the data
         let size = self.fd.write(data)?;
 
@@ -80,8 +106,28 @@ impl Device {
         Ok(size)
     }
 
+    /// Increment the read/write count
+    fn increment_rw_count(&mut self) {
+        self.rw_count = self.rw_count.saturating_add(1);
+    }
+
     /// Flush the device
     pub fn flush(&mut self) -> WgResult<()> {
         Ok(self.fd.flush()?)
     }
+
+    /// Returns `true` if the device is a dummy
+    pub fn is_dummy(&self) -> bool {
+        self.is_dummy
+    }
+
+    /// Returns the read/write counter of the device
+    pub fn get_rw_count(&self) -> u64 {
+        self.rw_count
+    }
+
+    /// Returns a reference to the internal file descriptor
+    pub fn get_fd(&self) -> &File {
+        &self.fd
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 987fbcd..0baac2b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@
 //! ## Fast, modern and secure VPN tunnel
 //!
 //! Target of this project is to have a user space Rust implementation of `WireGuard`.
+#![deny(missing_docs)]
 
 #[macro_use]
 extern crate log;
@@ -13,10 +14,10 @@ extern crate mowl;
 
 #[macro_use]
 pub mod error;
-pub mod device;
+mod device;
 mod bindgen;
 
-use device::Device;
+pub use device::Device;
 use error::WgResult;
 
 use std::io;
@@ -24,7 +25,7 @@ use std::net::SocketAddr;
 use std::str::FromStr;
 
 use log::LogLevel;
-use futures::{Future, Poll};
+use futures::{Async, Future, Poll};
 use tokio_core::net::UdpSocket;
 use tokio_core::reactor::{Core, Handle};
 
@@ -73,6 +74,16 @@ impl WireGuard {
         mowl::init_with_level(level)?;
         Ok(self)
     }
+
+    /// Returns the socket address of the server
+    pub fn get_addr(&self) -> WgResult<SocketAddr> {
+        Ok(self.tunnel.server.local_addr()?)
+    }
+
+    /// Returns a reference to the tunneling device
+    pub fn get_device(&self) -> &Device {
+        &self.tunnel.device
+    }
 }
 
 #[derive(Debug)]
@@ -98,11 +109,10 @@ impl WireGuardFuture {
     /// Creates a new `WireGuardFuture`
     pub fn new(handle: &Handle, addr: &str, dummy: bool) -> WgResult<Self> {
         // Create a tunneling device
-        let device_name = "wg";
         let device = if dummy {
-            Device::dummy(device_name)?
+            Device::dummy()?
         } else {
-            Device::new(device_name)?
+            Device::new("wg")?
         };
 
         // Create a server for the tunnel
@@ -134,10 +144,10 @@ impl Future for WireGuardFuture {
                 // Set to `None` if transmission is done
                 self.send_to_device = None;
 
-                debug!("Wrote {}/{} bytes from {} to tunnel device",
-                       bytes_written,
-                       length,
-                       peer);
+                info!("Wrote {}/{} bytes from {} to tunnel device",
+                      bytes_written,
+                      length,
+                      peer);
             }
 
             // Process message from the tunneling device
@@ -148,10 +158,10 @@ impl Future for WireGuardFuture {
                 // Set to `None` if transmission is done
                 self.send_to_client = None;
 
-                debug!("Wrote {}/{} bytes from the server to {}",
-                       bytes_written,
-                       length,
-                       peer);
+                info!("Wrote {}/{} bytes from the server to {}",
+                      bytes_written,
+                      length,
+                      peer);
             }
 
 
@@ -160,7 +170,12 @@ impl Future for WireGuardFuture {
 
             // If `send_to_client` is `None` we can receive the next message from the tunnel device
             // self.send_to_client = Some(try_nb!(self.device.read(&mut self.buffer)));
-            // debug!("Read {} bytes from tunnel device", self.send_to_client.0);
+            // info!("Read {} bytes from tunnel device", self.send_to_client.0);
+
+            // Stop the future when running in test mode
+            if self.device.is_dummy() && self.device.get_rw_count() > 2 {
+                return Ok(Async::Ready(()));
+            }
         }
     }
 }