Secure temporary file creation

jihochu · Nov 28, 2015 · 33905d5 · 33905d5
1 parent 34ee5df
commit 33905d5
Showing 1 changed file with 16 additions and 7 deletions.
diff --git a/include/caffe/util/io.hpp b/include/caffe/util/io.hpp
@@ -2,12 +2,15 @@
 #define CAFFE_UTIL_IO_H_
 
 #include <boost/filesystem.hpp>
+#include <iomanip>
+#include <iostream>  // NOLINT(readability/streams)
 #include <string>
 
 #include "google/protobuf/message.h"
 
 #include "caffe/common.hpp"
 #include "caffe/proto/caffe.pb.h"
+#include "caffe/util/format.hpp"
 
 #ifndef CAFFE_TMP_DIR_RETRIES
 #define CAFFE_TMP_DIR_RETRIES 100
@@ -18,13 +21,6 @@ namespace caffe {
 using ::google::protobuf::Message;
 using ::boost::filesystem::path;
 
-inline void MakeTempFilename(string* temp_filename) {
-  temp_filename->clear();
-  const path& model = boost::filesystem::temp_directory_path()
-    /"caffe_test.%%%%%%";
-  *temp_filename = boost::filesystem::unique_path(model).string();
-}
-
 inline void MakeTempDir(string* temp_dirname) {
   temp_dirname->clear();
   const path& model =
@@ -40,6 +36,19 @@ inline void MakeTempDir(string* temp_dirname) {
   LOG(FATAL) << "Failed to create a temporary directory.";
 }
 
+inline void MakeTempFilename(string* temp_filename) {
+  static path temp_files_subpath;
+  static uint64_t next_temp_file = 0;
+  temp_filename->clear();
+  if ( temp_files_subpath.empty() ) {
+    string path_string="";
+    MakeTempDir(&path_string);
+    temp_files_subpath = path_string;
+  }
+  *temp_filename =
+    (temp_files_subpath/caffe::format_int(next_temp_file++, 9)).string();
+}
+
 bool ReadProtoFromTextFile(const char* filename, Message* proto);
 
 inline bool ReadProtoFromTextFile(const string& filename, Message* proto) {