tempfile: Unify the two tempdir implementations

Looks like we ended up with two totally different tempdir
implementations: one from CL:520706 and the other from CL:1409705.

This CL consolidates them into one implementation.

BUG=chromium:974059
TEST=tempfile: cargo test
TEST=crosvm: cargo check --all-features
TEST=devices: cargo check --tests
TEST=sys_util: cargo check --tests
TEST=local kokoro
TEST=./build_test

Cq-Depend: chromium:1574668
Change-Id: Id70e963c9986ed2fc5f160819c4a7f9f16092b3b
Signed-off-by: Daniel Verkamp <dverkamp@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1573227
Tested-by: kokoro <noreply+kokoro@google.com>
Legacy-Commit-Queue: Commit Bot <commit-bot@chromium.org>

9 files changed, 96 insertions, 133 deletions

diff --git a/Cargo.lock b/Cargo.lock
index beaf37b..8387f0e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -162,6 +162,7 @@ dependencies = [
  "resources 0.1.0",
  "sync 0.1.0",
  "sys_util 0.1.0",
+ "tempfile 3.0.7",
  "tpm2 0.1.0",
  "usb_util 0.1.0",
  "vhost 0.1.0",
@@ -473,6 +474,7 @@ dependencies = [
  "poll_token_derive 0.1.0",
  "sync 0.1.0",
  "syscall_defines 0.1.0",
+ "tempfile 3.0.7",
 ]
 
 [[package]]
diff --git a/devices/Cargo.toml b/devices/Cargo.toml
index 62b570e..8cfe8dd 100644
--- a/devices/Cargo.toml
+++ b/devices/Cargo.toml
@@ -37,3 +37,6 @@ usb_util = { path = "../usb_util" }
 vhost = { path = "../vhost" }
 virtio_sys = { path = "../virtio_sys" }
 vm_control = { path = "../vm_control" }
+
+[dev-dependencies]
+tempfile = { path = "../tempfile" }
diff --git a/devices/src/virtio/block.rs b/devices/src/virtio/block.rs
index 59cc51a..f638a0e 100644
--- a/devices/src/virtio/block.rs
+++ b/devices/src/virtio/block.rs
@@ -1035,15 +1035,14 @@ impl<T: 'static + AsRawFd + DiskFile + Send> VirtioDevice for Block<T> {
 #[cfg(test)]
 mod tests {
     use std::fs::{File, OpenOptions};
-    use std::path::PathBuf;
-    use sys_util::TempDir;
+    use tempfile::TempDir;
 
     use super::*;
 
     #[test]
     fn read_size() {
-        let tempdir = TempDir::new("/tmp/block_read_test").unwrap();
-        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        let tempdir = TempDir::new().unwrap();
+        let mut path = tempdir.path().to_owned();
         path.push("disk_image");
         let f = File::create(&path).unwrap();
         f.set_len(0x1000).unwrap();
@@ -1061,8 +1060,8 @@ mod tests {
 
     #[test]
     fn read_features() {
-        let tempdir = TempDir::new("/tmp/block_read_test").unwrap();
-        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        let tempdir = TempDir::new().unwrap();
+        let mut path = tempdir.path().to_owned();
         path.push("disk_image");
 
         // read-write block device
@@ -1086,8 +1085,8 @@ mod tests {
 
     #[test]
     fn read_last_sector() {
-        let tempdir = TempDir::new("/tmp/block_read_test").unwrap();
-        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        let tempdir = TempDir::new().unwrap();
+        let mut path = tempdir.path().to_owned();
         path.push("disk_image");
         let mut f = OpenOptions::new()
             .read(true)
@@ -1129,8 +1128,8 @@ mod tests {
 
     #[test]
     fn read_beyond_last_sector() {
-        let tempdir = TempDir::new("/tmp/block_read_test").unwrap();
-        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        let tempdir = TempDir::new().unwrap();
+        let mut path = tempdir.path().to_owned();
         path.push("disk_image");
         let mut f = OpenOptions::new()
             .read(true)
diff --git a/sys_util/Cargo.toml b/sys_util/Cargo.toml
index 00e39fa..beba374 100644
--- a/sys_util/Cargo.toml
+++ b/sys_util/Cargo.toml
@@ -11,5 +11,6 @@ libc = "*"
 poll_token_derive = { version = "*", path = "poll_token_derive" }
 sync = { path = "../sync" } # provided by ebuild
 syscall_defines = { path = "../syscall_defines" } # provided by ebuild
+tempfile = { path = "../tempfile" } # provided by ebuild
 
 [workspace]
diff --git a/sys_util/src/lib.rs b/sys_util/src/lib.rs
index 9e2e623..51c3898 100644
--- a/sys_util/src/lib.rs
+++ b/sys_util/src/lib.rs
@@ -33,7 +33,6 @@ pub mod signal;
 mod signalfd;
 mod sock_ctrl_msg;
 mod struct_util;
-mod tempdir;
 mod terminal;
 mod timerfd;
 mod write_zeroes;
@@ -60,7 +59,6 @@ pub use crate::signal::*;
 pub use crate::signalfd::*;
 pub use crate::sock_ctrl_msg::*;
 pub use crate::struct_util::*;
-pub use crate::tempdir::*;
 pub use crate::terminal::*;
 pub use crate::timerfd::*;
 pub use poll_token_derive::*;
diff --git a/sys_util/src/seek_hole.rs b/sys_util/src/seek_hole.rs
index 8215a8d..8cec06c 100644
--- a/sys_util/src/seek_hole.rs
+++ b/sys_util/src/seek_hole.rs
@@ -54,10 +54,9 @@ impl SeekHole for File {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::TempDir;
     use std::fs::File;
     use std::io::{Seek, SeekFrom, Write};
-    use std::path::PathBuf;
+    use tempfile::TempDir;
 
     fn seek_cur(file: &mut File) -> u64 {
         file.seek(SeekFrom::Current(0)).unwrap()
@@ -65,8 +64,8 @@ mod tests {
 
     #[test]
     fn seek_data() {
-        let tempdir = TempDir::new("/tmp/seek_data_test").unwrap();
-        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        let tempdir = TempDir::new().unwrap();
+        let mut path = tempdir.path().to_owned();
         path.push("test_file");
         let mut file = File::create(&path).unwrap();
 
@@ -112,8 +111,8 @@ mod tests {
 
     #[test]
     fn seek_hole() {
-        let tempdir = TempDir::new("/tmp/seek_hole_test").unwrap();
-        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        let tempdir = TempDir::new().unwrap();
+        let mut path = tempdir.path().to_owned();
         path.push("test_file");
         let mut file = File::create(&path).unwrap();
 
diff --git a/sys_util/src/tempdir.rs b/sys_util/src/tempdir.rs
deleted file mode 100644
index 045f245..0000000
--- a/sys_util/src/tempdir.rs
+++ /dev/null
@@ -1,102 +0,0 @@
-// Copyright 2017 The Chromium OS Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style license that can be
-// found in the LICENSE file.
-
-use std::ffi::CString;
-use std::ffi::OsStr;
-use std::ffi::OsString;
-use std::fs;
-use std::os::unix::ffi::OsStringExt;
-use std::path::Path;
-use std::path::PathBuf;
-
-use libc;
-
-use crate::{errno_result, Result};
-
-/// Create and remove a temporary directory.  The directory will be maintained for the lifetime of
-/// the `TempDir` object.
-pub struct TempDir {
-    path: Option<PathBuf>,
-}
-
-impl TempDir {
-    /// Creates a new tempory directory.
-    /// The directory will be removed when the object goes out of scope.
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// # use std::path::Path;
-    /// # use std::path::PathBuf;
-    /// # use sys_util::TempDir;
-    /// # fn test_create_temp_dir() -> Result<(), ()> {
-    ///       let t = TempDir::new("/tmp/testdir").map_err(|_| ())?;
-    ///       assert!(t.as_path().unwrap().exists());
-    /// #     Ok(())
-    /// # }
-    /// ```
-    pub fn new<P: AsRef<OsStr>>(prefix: P) -> Result<TempDir> {
-        let mut dir_string = prefix.as_ref().to_os_string();
-        dir_string.push("XXXXXX");
-        // unwrap this result as the internal bytes can't have a null with a valid path.
-        let dir_name = CString::new(dir_string.into_vec()).unwrap();
-        let mut dir_bytes = dir_name.into_bytes_with_nul();
-        let ret = unsafe {
-            // Creating the directory isn't unsafe.  The fact that it modifies the guts of the path
-            // is also OK because it only overwrites the last 6 Xs added above.
-            libc::mkdtemp(dir_bytes.as_mut_ptr() as *mut libc::c_char)
-        };
-        if ret.is_null() {
-            return errno_result();
-        }
-        dir_bytes.pop(); // Remove the null becasue from_vec can't handle it.
-        Ok(TempDir {
-            path: Some(PathBuf::from(OsString::from_vec(dir_bytes))),
-        })
-    }
-
-    /// Removes the temporary directory.  Calling this is optional as dropping a `TempDir` object
-    /// will also remove the directory.  Calling remove explicitly allows for better error handling.
-    pub fn remove(mut self) -> Result<()> {
-        let path = self.path.take();
-        path.map_or(Ok(()), fs::remove_dir_all)?;
-        Ok(())
-    }
-
-    /// Returns the path to the tempdir if it is currently valid
-    pub fn as_path(&self) -> Option<&Path> {
-        self.path.as_ref().map(PathBuf::as_path)
-    }
-}
-
-impl Drop for TempDir {
-    fn drop(&mut self) {
-        if let Some(p) = &self.path {
-            // Nothing can be done here if this returns an error.
-            let _ = fs::remove_dir_all(p);
-        }
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn create_dir() {
-        let t = TempDir::new("/tmp/asdf").unwrap();
-        let path = t.as_path().unwrap();
-        assert!(path.exists());
-        assert!(path.is_dir());
-        assert!(path.starts_with("/tmp/"));
-    }
-
-    #[test]
-    fn remove_dir() {
-        let t = TempDir::new("/tmp/asdf").unwrap();
-        let path = t.as_path().unwrap().to_owned();
-        assert!(t.remove().is_ok());
-        assert!(!path.exists());
-    }
-}
diff --git a/sys_util/src/write_zeroes.rs b/sys_util/src/write_zeroes.rs
index 51f0dbe..e3b531e 100644
--- a/sys_util/src/write_zeroes.rs
+++ b/sys_util/src/write_zeroes.rs
@@ -59,14 +59,12 @@ mod tests {
     use super::*;
     use std::fs::OpenOptions;
     use std::io::{Read, Seek, SeekFrom};
-    use std::path::PathBuf;
-
-    use crate::TempDir;
+    use tempfile::TempDir;
 
     #[test]
     fn simple_test() {
-        let tempdir = TempDir::new("/tmp/write_zeroes_test").unwrap();
-        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        let tempdir = TempDir::new().unwrap();
+        let mut path = tempdir.path().to_owned();
         path.push("file");
         let mut f = OpenOptions::new()
             .read(true)
@@ -131,8 +129,8 @@ mod tests {
 
     #[test]
     fn large_write_zeroes() {
-        let tempdir = TempDir::new("/tmp/write_zeroes_test").unwrap();
-        let mut path = PathBuf::from(tempdir.as_path().unwrap());
+        let tempdir = TempDir::new().unwrap();
+        let mut path = tempdir.path().to_owned();
         path.push("file");
         let mut f = OpenOptions::new()
             .read(true)
diff --git a/tempfile/src/lib.rs b/tempfile/src/lib.rs
index 0b9e39b..92f766b 100644
--- a/tempfile/src/lib.rs
+++ b/tempfile/src/lib.rs
@@ -2,20 +2,39 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-//! Simplified tempfile which doesn't depend on the `rand` crate, instead using
-//! /dev/urandom as a source of entropy
+//! Simplified tempfile which doesn't depend on the `rand` crate.
+//!
+//! # Example
+//!
+//! ```
+//! use std::io::Result;
+//! use std::path::{Path, PathBuf};
+//! use tempfile::TempDir;
+//!
+//! fn main() -> Result<()> {
+//!     let t = TempDir::new()?;
+//!     assert!(t.path().exists());
+//!
+//!     Ok(())
+//! }
+//! ```
 
 use libc::mkdtemp;
 use std::env;
 use std::ffi::CString;
 use std::fs;
 use std::io::{Error, ErrorKind, Result};
+use std::mem::ManuallyDrop;
 use std::path::{Path, PathBuf};
+use std::ptr;
 
 pub struct Builder {
     prefix: String,
 }
 
+
+// Note: we implement a builder because the protoc-rust crate uses this API from
+// crates.io's tempfile. Our code mostly uses TempDir::new directly.
 impl Builder {
     pub fn new() -> Self {
         Builder {
@@ -31,10 +50,9 @@ impl Builder {
         self
     }
 
-    /// Tries to make a tempdir inside of `env::temp_dir()` with a specified
-    /// prefix. The directory and it's content is destroyed when TempDir is
-    /// dropped.
-    /// If the directory can not be created, `Err` is returned.
+    /// Creates a new temporary directory under libc's preferred system
+    /// temporary directory. The new directory will be removed when the returned
+    /// handle of type `TempDir` is dropped.
     pub fn tempdir(&self) -> Result<TempDir> {
         // mkdtemp() requires the template to end in 6 X chars, which will be replaced
         // with random characters to make the path unique.
@@ -68,17 +86,43 @@ impl Builder {
     }
 }
 
+/// Temporary directory. The directory will be removed when this object is
+/// dropped.
 pub struct TempDir {
     path: PathBuf,
+
+    // When adding new fields to TempDir: note that anything with a Drop impl
+    // will need to be dropped explicitly via ptr::read inside TempDir::remove
+    // or else it gets leaked (memory safe but not ideal).
 }
 
 impl TempDir {
+    pub fn new() -> Result<Self> {
+        Builder::new().tempdir()
+    }
+
     /// Accesses the tempdir's [`Path`].
     ///
     /// [`Path`]: http://doc.rust-lang.org/std/path/struct.Path.html
     pub fn path(&self) -> &Path {
         self.path.as_ref()
     }
+
+    /// Removes the temporary directory.
+    ///
+    /// Calling this is optional as dropping a TempDir object will also remove
+    /// the directory. Calling remove explicitly allows for any resulting error
+    /// to be handled.
+    pub fn remove(self) -> Result<()> {
+        // Place self inside ManuallyDrop so its Drop impl doesn't run, but nor
+        // does the path inside get dropped. Then use ptr::read to take out the
+        // PathBuf so that it *does* get dropped correctly at the bottom of this
+        // function.
+        let dont_drop = ManuallyDrop::new(self);
+        let path: PathBuf = unsafe { ptr::read(&dont_drop.path) };
+
+        fs::remove_dir_all(path)
+    }
 }
 
 impl Drop for TempDir {
@@ -86,3 +130,24 @@ impl Drop for TempDir {
         let _ = fs::remove_dir_all(&self.path);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::TempDir;
+
+    #[test]
+    fn create_dir() {
+        let t = TempDir::new().unwrap();
+        let path = t.path();
+        assert!(path.exists());
+        assert!(path.is_dir());
+    }
+
+    #[test]
+    fn remove_dir() {
+        let t = TempDir::new().unwrap();
+        let path = t.path().to_owned();
+        assert!(t.remove().is_ok());
+        assert!(!path.exists());
+    }
+}