1 files changed, 39 insertions, 30 deletions

diff --git a/vhost/src/net.rs b/vhost/src/net.rs
index 0c6eaad..1d3f94e 100644
--- a/vhost/src/net.rs
+++ b/vhost/src/net.rs
@@ -3,8 +3,9 @@
 // found in the LICENSE file.
 
 use libc;
-use net_util;
+use net_util::TapT;
 use std::fs::{File, OpenOptions};
+use std::marker::PhantomData;
 use std::os::unix::fs::OpenOptionsExt;
 use std::os::unix::io::{AsRawFd, RawFd};
 use virtio_sys;
@@ -19,30 +20,37 @@ static DEVICE: &'static str = "/dev/vhost-net";
 ///
 /// This provides a simple wrapper around a VHOST_NET file descriptor and
 /// methods that safely run ioctls on that file descriptor.
-pub struct Net {
+pub struct Net<T> {
     // fd must be dropped first, which will stop and tear down the
     // vhost-net worker before GuestMemory can potentially be unmapped.
     fd: File,
     mem: GuestMemory,
+    phantom: PhantomData<T>,
 }
 
-pub trait NetT {
+pub trait NetT<T: TapT>: Vhost + AsRawFd + Send + Sized {
+    /// Create a new NetT instance
+    fn new(mem: &GuestMemory) -> Result<Self>;
+
     /// Set the tap file descriptor that will serve as the VHOST_NET backend.
     /// This will start the vhost worker for the given queue.
     ///
     /// # Arguments
     /// * `queue_index` - Index of the queue to modify.
     /// * `fd` - Tap interface that will be used as the backend.
-    fn set_backend(&self, queue_index: usize, fd: &net_util::Tap) -> Result<()>;
+    fn set_backend(&self, queue_index: usize, fd: &T) -> Result<()>;
 }
 
-impl Net {
+impl<T> NetT<T> for Net<T>
+where
+    T: TapT,
+{
     /// Opens /dev/vhost-net and holds a file descriptor open for it.
     ///
     /// # Arguments
     /// * `mem` - Guest memory mapping.
-    pub fn new(mem: &GuestMemory) -> Result<Net> {
-        Ok(Net {
+    fn new(mem: &GuestMemory) -> Result<Net<T>> {
+        Ok(Net::<T> {
             fd: OpenOptions::new()
                 .read(true)
                 .write(true)
@@ -50,12 +58,11 @@ impl Net {
                 .open(DEVICE)
                 .map_err(Error::VhostOpen)?,
             mem: mem.clone(),
+            phantom: PhantomData,
         })
     }
-}
 
-impl NetT for Net {
-    fn set_backend(&self, queue_index: usize, fd: &net_util::Tap) -> Result<()> {
+    fn set_backend(&self, queue_index: usize, fd: &T) -> Result<()> {
         let vring_file = virtio_sys::vhost_vring_file {
             index: queue_index as u32,
             fd: fd.as_raw_fd(),
@@ -72,64 +79,66 @@ impl NetT for Net {
     }
 }
 
-impl Vhost for Net {
+impl<T> Vhost for Net<T> {
     fn mem(&self) -> &GuestMemory {
         &self.mem
     }
 }
 
-impl AsRawFd for Net {
+impl<T> AsRawFd for Net<T> {
     fn as_raw_fd(&self) -> RawFd {
         self.fd.as_raw_fd()
     }
 }
 
-#[cfg(test)]
-pub mod tests {
+pub mod fakes {
     use super::*;
     use std::fs::OpenOptions;
     use std::fs::remove_file;
 
     const TMP_FILE: &str = "/tmp/crosvm_vhost_test_file";
 
-    pub struct FakeNet {
+    pub struct FakeNet<T> {
         fd: File,
         mem: GuestMemory,
+        phantom: PhantomData<T>,
+    }
+
+    impl<T> Drop for FakeNet<T> {
+        fn drop(&mut self) {
+            let _ = remove_file(TMP_FILE);
+        }
     }
 
-    impl FakeNet {
-        pub fn new(mem: &GuestMemory) -> Result<FakeNet> {
-            Ok(FakeNet {
+    impl<T> NetT<T> for FakeNet<T>
+    where
+        T: TapT,
+    {
+        fn new(mem: &GuestMemory) -> Result<FakeNet<T>> {
+            Ok(FakeNet::<T> {
                 fd: OpenOptions::new()
                     .read(true)
                     .append(true)
                     .create(true)
                     .open(TMP_FILE)
                     .unwrap(),
-                mem: mem.clone()
+                mem: mem.clone(),
+                phantom: PhantomData,
             })
         }
-    }
-
-    impl Drop for FakeNet {
-        fn drop(&mut self) {
-            let _ = remove_file(TMP_FILE);
-        }
-    }
 
-    impl NetT for FakeNet {
-        fn set_backend(&self, _queue_index: usize, _fd: &net_util::Tap) -> Result<()> {
+        fn set_backend(&self, _queue_index: usize, _fd: &T) -> Result<()> {
             Ok(())
         }
     }
 
-    impl Vhost for FakeNet {
+    impl<T> Vhost for FakeNet<T> {
         fn mem(&self) -> &GuestMemory {
             &self.mem
         }
     }
 
-    impl AsRawFd for FakeNet {
+    impl<T> AsRawFd for FakeNet<T> {
         fn as_raw_fd(&self) -> RawFd {
             self.fd.as_raw_fd()
         }