From c6897da7f0588adca8858488cb1c234457d3bd97 Mon Sep 17 00:00:00 2001
From: Knight <Knight42@mail.ustc.edu.cn>
Date: Sat, 18 Jun 2016 17:41:47 +0800
Subject: [PATCH] mktemp: make temp dir

---
 src/mktemp/mktemp.rs  | 31 ++++++++++++++++++++++++-------
 src/mktemp/tempdir.rs | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 7 deletions(-)
 create mode 100644 src/mktemp/tempdir.rs

diff --git a/src/mktemp/mktemp.rs b/src/mktemp/mktemp.rs
index a4df5adf5..22b8b014b 100644
--- a/src/mktemp/mktemp.rs
+++ b/src/mktemp/mktemp.rs
@@ -26,6 +26,8 @@ use std::iter;
 use rand::Rng;
 use tempfile::NamedTempFileOptions;
 
+mod tempdir;
+
 static NAME: &'static str = "mktemp";
 static VERSION: &'static str = env!("CARGO_PKG_VERSION");
 static DEFAULT_TEMPLATE: &'static str = "tmp.XXXXXXXXXX";
@@ -173,7 +175,18 @@ pub fn dry_exec(mut tmpdir: PathBuf, prefix: &str, rand: usize, suffix: &str) ->
 fn exec(tmpdir: PathBuf, prefix: &str, rand: usize, suffix: &str, make_dir: bool) -> i32 {
     // TODO: respect make_dir option
     if make_dir {
-        crash!(1, "Directory option is not supported yet. Sorry.");
+        match tempdir::new_in(&tmpdir, prefix, rand, suffix) {
+            Ok(ref f) => {
+                println!("{}", f);
+                return 0;
+            }
+            Err(e) => {
+                if !quiet {
+                    show_info!("{}", e);
+                }
+                return 1;
+            }
+        }
     }
 
     let tmpfile = NamedTempFileOptions::new()
@@ -184,13 +197,17 @@ fn exec(tmpdir: PathBuf, prefix: &str, rand: usize, suffix: &str, make_dir: bool
 
     let tmpfile = match tmpfile {
         Ok(f) => f,
-        Err(_) => crash!(1, "failed to create tempfile")
+        Err(e) => {
+            if !quiet {
+                show_info!("failed to create tempfile: {}", e);
+            }
+            return 1;
+        }
     };
-    
-    let tmpname = tmpfile
-        .path()
-        .to_string_lossy()
-        .to_string();
+
+    let tmpname = tmpfile.path()
+                         .to_string_lossy()
+                         .to_string();
 
     println!("{}", tmpname);
 
diff --git a/src/mktemp/tempdir.rs b/src/mktemp/tempdir.rs
new file mode 100644
index 000000000..186bcd8f5
--- /dev/null
+++ b/src/mktemp/tempdir.rs
@@ -0,0 +1,43 @@
+// Mainly taken from crate `tempdir`
+
+extern crate rand;
+use rand::{Rng, thread_rng};
+
+use std::io::Result as IOResult;
+use std::io::{Error, ErrorKind};
+use std::path::Path;
+
+// How many times should we (re)try finding an unused random name? It should be
+// enough that an attacker will run out of luck before we run out of patience.
+const NUM_RETRIES: u32 = 1 << 31;
+
+#[cfg(unix)]
+fn create_dir<P: AsRef<Path>>(path: P) -> IOResult<()> {
+    use std::fs::DirBuilder;
+    use std::os::unix::fs::DirBuilderExt;
+
+    DirBuilder::new().mode(0o700).create(path)
+}
+
+#[cfg(windows)]
+fn create_dir<P: AsRef<Path>>(path: P) -> IOResult<()> {
+    ::std::fs::create_dir(path)
+}
+
+pub fn new_in<P: AsRef<Path>>(tmpdir: P, prefix: &str, rand: usize, suffix: &str) -> IOResult<String> {
+
+    let mut rng = thread_rng();
+    for _ in 0..NUM_RETRIES {
+        let rand_chars: String = rng.gen_ascii_chars().take(rand).collect();
+        let leaf = format!("{}{}{}", prefix, rand_chars, suffix);
+        let path = tmpdir.as_ref().join(&leaf);
+        match create_dir(&path) {
+            Ok(_) => return Ok(path.to_string_lossy().into_owned()),
+            Err(ref e) if e.kind() == ErrorKind::AlreadyExists => {}
+            Err(e) => return Err(e),
+        }
+    }
+
+    Err(Error::new(ErrorKind::AlreadyExists,
+                   "too many temporary directories already exist"))
+}