Secure implementation of MakeTempDir

This commit is contained in:
T.E.A de Souza 2015-11-24 14:33:27 +08:00 коммит произвёл Tea
Родитель 4e9b449ecb
Коммит 34ee5df55d
1 изменённых файлов: 15 добавлений и 6 удалений

Просмотреть файл

@ -9,6 +9,10 @@
#include "caffe/common.hpp" #include "caffe/common.hpp"
#include "caffe/proto/caffe.pb.h" #include "caffe/proto/caffe.pb.h"
#ifndef CAFFE_TMP_DIR_RETRIES
#define CAFFE_TMP_DIR_RETRIES 100
#endif
namespace caffe { namespace caffe {
using ::google::protobuf::Message; using ::google::protobuf::Message;
@ -23,12 +27,17 @@ inline void MakeTempFilename(string* temp_filename) {
inline void MakeTempDir(string* temp_dirname) { inline void MakeTempDir(string* temp_dirname) {
temp_dirname->clear(); temp_dirname->clear();
const path& model = boost::filesystem::temp_directory_path() const path& model =
/"caffe_test.%%%%%%"; boost::filesystem::temp_directory_path()/"caffe_test.%%%%-%%%%";
const path& dir = boost::filesystem::unique_path(model).string(); for ( int i = 0; i < CAFFE_TMP_DIR_RETRIES; i++ ) {
bool directoryCreated = boost::filesystem::create_directory(dir); const path& dir = boost::filesystem::unique_path(model).string();
CHECK(directoryCreated); bool done = boost::filesystem::create_directory(dir);
*temp_dirname = dir.string(); if ( done ) {
*temp_dirname = dir.string();
return;
}
}
LOG(FATAL) << "Failed to create a temporary directory.";
} }
bool ReadProtoFromTextFile(const char* filename, Message* proto); bool ReadProtoFromTextFile(const char* filename, Message* proto);